[spark] Fix Fluss Spark read paths

fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/source/IcebergLakeSource.java

@@ -114,6 +114,9 @@ public Planner<IcebergSplit> createPlanner(PlannerContext context) throws IOExce
     public RecordReader createRecordReader(ReaderContext<IcebergSplit> context) throws IOException {
         Catalog catalog = IcebergCatalogUtils.createIcebergCatalog(icebergConfig);
         Table table = catalog.loadTable(toIceberg(tablePath));
+        if (context.requireSortedRecords()) {
+            return new IcebergSortedRecordReader(table, context.lakeSplit(), project);
+        }
         return new IcebergRecordReader(context.lakeSplit().fileScanTask(), table, project);
     }
 

fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/source/IcebergRecordReader.java

@@ -93,8 +93,7 @@ public static class IcebergRecordAsFlussRecordIterator implements CloseableItera
 
         private final org.apache.iceberg.io.CloseableIterator<Record> icebergRecordIterator;
 
-        private final ProjectedRow projectedRow;
-        private final IcebergRecordAsFlussRow icebergRecordAsFlussRow;
+        private final int[] projectedPositions;
 
         public IcebergRecordAsFlussRecordIterator(
                 CloseableIterable<Record> icebergRecordIterator, Types.StructType struct) {
@@ -107,8 +106,7 @@ public IcebergRecordAsFlussRecordIterator(
                     IcebergUtils.isLegacyTable(new Schema(struct.fields()))
                             ? struct.fields().size() - LEGACY_SYSTEM_COLUMNS.size()
                             : struct.fields().size();
-            projectedRow = ProjectedRow.from(IntStream.range(0, businessFieldCount).toArray());
-            icebergRecordAsFlussRow = new IcebergRecordAsFlussRow();
+            projectedPositions = IntStream.range(0, businessFieldCount).toArray();
         }
 
         @Override
@@ -136,8 +134,8 @@ public LogRecord next() {
                     NO_SYSTEM_COLUMN_VALUE,
                     NO_SYSTEM_COLUMN_VALUE,
                     ChangeType.INSERT,
-                    projectedRow.replaceRow(
-                            icebergRecordAsFlussRow.replaceIcebergRecord(icebergRecord)));
+                    ProjectedRow.from(projectedPositions)
+                            .replaceRow(new IcebergRecordAsFlussRow(icebergRecord)));
         }
     }
 }

fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/source/IcebergSortedRecordReader.java

@@ -0,0 +1,225 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.fluss.lake.iceberg.source;
+
+import org.apache.fluss.lake.source.SortedRecordReader;
+import org.apache.fluss.record.LogRecord;
+import org.apache.fluss.row.InternalRow;
+import org.apache.fluss.utils.CloseableIterator;
+
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.types.Type;
+import org.apache.iceberg.types.Types;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+
+import static org.apache.fluss.utils.Preconditions.checkState;
+
+/** Sorted Iceberg reader used by primary-key lake union reads. */
+public class IcebergSortedRecordReader implements SortedRecordReader {
+
+    private final @Nullable IcebergRecordReader delegate;
+    private final List<Types.NestedField> keyFields;
+    private final int[] keyPositions;
+
+    /**
+     * Creates an Iceberg sorted reader.
+     *
+     * @param table Iceberg table to read
+     * @param split file split, or null when only the comparator is required
+     * @param project top-level projection, or null for the full table projection
+     */
+    public IcebergSortedRecordReader(
+            Table table, @Nullable IcebergSplit split, @Nullable int[][] project) {
+        this.delegate =
+                split == null
+                        ? null
+                        : new IcebergRecordReader(split.fileScanTask(), table, project);
+        SortOrder sortOrder = createSortOrder(table.schema(), project);
+        this.keyFields = sortOrder.keyFields;
+        this.keyPositions = sortOrder.keyPositions;
+    }
+
+    @Override
+    public CloseableIterator<LogRecord> read() throws IOException {
+        if (delegate == null) {
+            return CloseableIterator.wrap(Collections.emptyIterator());
+        }
+
+        List<LogRecord> records = new ArrayList<>();
+        CloseableIterator<LogRecord> iterator = delegate.read();
+        try {
+            while (iterator.hasNext()) {
+                records.add(iterator.next());
+            }
+        } finally {
+            iterator.close();
+        }
+        records.sort((record1, record2) -> compareRows(record1.getRow(), record2.getRow()));
+        return CloseableIterator.wrap(records.iterator());
+    }
+
+    @Override
+    public Comparator<InternalRow> order() {
+        return this::compareKeyRows;
+    }
+
+    private int compareRows(InternalRow row1, InternalRow row2) {
+        for (int i = 0; i < keyFields.size(); i++) {
+            int position = keyPositions[i];
+            checkState(
+                    !row1.isNullAt(position) && !row2.isNullAt(position),
+                    "Iceberg identifier field at position %s must not be null.",
+                    position);
+            int result = compareValue(row1, row2, position, keyFields.get(i).type());
+            if (result != 0) {
+                return result;
+            }
+        }
+        return 0;
+    }
+
+    /**
+     * Compares rows containing only primary-key fields.
+     *
+     * <p>The client-side sort/merge reader passes primary-key projections to {@link #order()}, so
+     * their positions always start at zero regardless of the positions in the Iceberg table.
+     */
+    private int compareKeyRows(InternalRow row1, InternalRow row2) {
+        for (int i = 0; i < keyFields.size(); i++) {
+            checkState(
+                    !row1.isNullAt(i) && !row2.isNullAt(i),
+                    "Iceberg identifier field at key position %s must not be null.",
+                    i);
+            int result = compareValue(row1, row2, i, keyFields.get(i).type());
+            if (result != 0) {
+                return result;
+            }
+        }
+        return 0;
+    }
+
+    private int compareValue(InternalRow row1, InternalRow row2, int position, Type type) {
+        switch (type.typeId()) {
+            case BOOLEAN:
+                return Boolean.compare(row1.getBoolean(position), row2.getBoolean(position));
+            case INTEGER:
+            case DATE:
+            case TIME:
+                return Integer.compare(row1.getInt(position), row2.getInt(position));
+            case LONG:
+                return Long.compare(row1.getLong(position), row2.getLong(position));
+            case FLOAT:
+                return Float.compare(row1.getFloat(position), row2.getFloat(position));
+            case DOUBLE:
+                return Double.compare(row1.getDouble(position), row2.getDouble(position));
+            case STRING:
+                return row1.getString(position).compareTo(row2.getString(position));
+            case DECIMAL:
+                Types.DecimalType decimalType = (Types.DecimalType) type;
+                return row1.getDecimal(position, decimalType.precision(), decimalType.scale())
+                        .compareTo(
+                                row2.getDecimal(
+                                        position, decimalType.precision(), decimalType.scale()));
+            case TIMESTAMP:
+                return compareTimestamp(
+                        row1, row2, position, ((Types.TimestampType) type).shouldAdjustToUTC());
+            case TIMESTAMP_NANO:
+                return compareTimestamp(
+                        row1, row2, position, ((Types.TimestampNanoType) type).shouldAdjustToUTC());
+            case BINARY:
+            case FIXED:
+                return compareBytes(row1.getBytes(position), row2.getBytes(position));
+            default:
+                throw new UnsupportedOperationException(
+                        "Unsupported Iceberg identifier type: " + type.typeId());
+        }
+    }
+
+    private int compareTimestamp(
+            InternalRow row1, InternalRow row2, int position, boolean shouldAdjustToUTC) {
+        if (shouldAdjustToUTC) {
+            return row1.getTimestampLtz(position, 6).compareTo(row2.getTimestampLtz(position, 6));
+        }
+        return row1.getTimestampNtz(position, 6).compareTo(row2.getTimestampNtz(position, 6));
+    }
+
+    private int compareBytes(byte[] bytes1, byte[] bytes2) {
+        int length = Math.min(bytes1.length, bytes2.length);
+        for (int i = 0; i < length; i++) {
+            int result = Byte.compare(bytes1[i], bytes2[i]);
+            if (result != 0) {
+                return result;
+            }
+        }
+        return Integer.compare(bytes1.length, bytes2.length);
+    }
+
+    private static SortOrder createSortOrder(Schema schema, @Nullable int[][] project) {
+        List<Types.NestedField> keyFields = new ArrayList<>();
+        List<Integer> originalPositions = new ArrayList<>();
+        List<Types.NestedField> columns = schema.columns();
+        for (int i = 0; i < columns.size(); i++) {
+            Types.NestedField field = columns.get(i);
+            if (schema.identifierFieldIds().contains(field.fieldId())) {
+                keyFields.add(field);
+                originalPositions.add(i);
+            }
+        }
+
+        int[] keyPositions = new int[keyFields.size()];
+        for (int i = 0; i < originalPositions.size(); i++) {
+            int position = findProjectedPosition(originalPositions.get(i), project);
+            checkState(
+                    position >= 0,
+                    "Iceberg identifier field at position %s is missing from the projection.",
+                    originalPositions.get(i));
+            keyPositions[i] = position;
+        }
+        return new SortOrder(keyFields, keyPositions);
+    }
+
+    private static int findProjectedPosition(int originalPosition, @Nullable int[][] project) {
+        if (project == null) {
+            return originalPosition;
+        }
+        for (int i = 0; i < project.length; i++) {
+            if (project[i].length > 0 && project[i][0] == originalPosition) {
+                return i;
+            }
+        }
+        return -1;
+    }
+
+    private static final class SortOrder {
+        private final List<Types.NestedField> keyFields;
+        private final int[] keyPositions;
+
+        private SortOrder(List<Types.NestedField> keyFields, int[] keyPositions) {
+            this.keyFields = keyFields;
+            this.keyPositions = keyPositions;
+        }
+    }
+}

fluss-lake/fluss-lake-iceberg/src/test/java/org/apache/fluss/lake/iceberg/source/IcebergLakeSourceTest.java

@@ -21,11 +21,14 @@
 import org.apache.fluss.config.Configuration;
 import org.apache.fluss.lake.source.LakeSource;
 import org.apache.fluss.lake.source.RecordReader;
+import org.apache.fluss.lake.source.SortedRecordReader;
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.predicate.Predicate;
 import org.apache.fluss.predicate.PredicateBuilder;
 import org.apache.fluss.record.LogRecord;
 import org.apache.fluss.row.BinaryString;
+import org.apache.fluss.row.GenericRow;
+import org.apache.fluss.row.TimestampLtz;
 import org.apache.fluss.types.DataTypes;
 import org.apache.fluss.types.IntType;
 import org.apache.fluss.types.RowType;
@@ -89,6 +92,84 @@ void testEmptyFiltersDoNotAccessCatalog() {
         assertThat(result.remainingPredicates()).isEmpty();
     }
 
+    @Test
+    void testSortedReaderSupportsMissingLakeSplit() throws Exception {
+        TablePath tablePath = TablePath.of("fluss", "test_sorted_reader_without_split");
+        Schema schema =
+                new Schema(
+                        Arrays.asList(
+                                required(1, "event_time", Types.TimestampType.withZone()),
+                                optional(2, "name", Types.StringType.get())),
+                        Collections.singleton(1));
+        createTable(tablePath, schema, PartitionSpec.unpartitioned());
+
+        LakeSource<IcebergSplit> lakeSource = lakeStorage.createLakeSource(tablePath);
+        RecordReader recordReader =
+                lakeSource.createRecordReader(
+                        new LakeSource.ReaderContext<IcebergSplit>() {
+                            @Override
+                            public IcebergSplit lakeSplit() {
+                                return null;
+                            }
+
+                            @Override
+                            public boolean requireSortedRecords() {
+                                return true;
+                            }
+                        });
+
+        assertThat(recordReader).isInstanceOf(SortedRecordReader.class);
+        SortedRecordReader sortedRecordReader = (SortedRecordReader) recordReader;
+        assertThat(
+                        sortedRecordReader
+                                .order()
+                                .compare(
+                                        GenericRow.of(TimestampLtz.fromEpochMillis(1)),
+                                        GenericRow.of(TimestampLtz.fromEpochMillis(2))))
+                .isNegative();
+        try (CloseableIterator<LogRecord> iterator = sortedRecordReader.read()) {
+            assertThat(iterator.hasNext()).isFalse();
+        }
+    }
+
+    @Test
+    void testSortedReaderComparesPrimaryKeyProjection() throws Exception {
+        TablePath tablePath = TablePath.of("fluss", "test_sorted_reader_key_projection");
+        Schema schema =
+                new Schema(
+                        Arrays.asList(
+                                optional(1, "name", Types.StringType.get()),
+                                required(2, "event_time", Types.TimestampType.withZone()),
+                                optional(3, "value", Types.StringType.get())),
+                        Collections.singleton(2));
+        createTable(tablePath, schema, PartitionSpec.unpartitioned());
+
+        LakeSource<IcebergSplit> lakeSource = lakeStorage.createLakeSource(tablePath);
+        lakeSource.withProject(new int[][] {{0}, {1}});
+        SortedRecordReader sortedRecordReader =
+                (SortedRecordReader)
+                        lakeSource.createRecordReader(
+                                new LakeSource.ReaderContext<IcebergSplit>() {
+                                    @Override
+                                    public IcebergSplit lakeSplit() {
+                                        return null;
+                                    }
+
+                                    @Override
+                                    public boolean requireSortedRecords() {
+                                        return true;
+                                    }
+                                });
+
+        assertThat(
+                        sortedRecordReader
+                                .order()
+                                .compare(
+                                        GenericRow.of(TimestampLtz.fromEpochMillis(1)),
+                                        GenericRow.of(TimestampLtz.fromEpochMillis(2))))
+                .isNegative();
+    }
+
     @Test
     void testWithFilters() throws Exception {
         TablePath tablePath = TablePath.of("fluss", "test_filters");