[lake/hudi] Add Hudi Flink union read primary key table IT

fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LakeSnapshotAndLogSplitScanner.java

@@ -188,10 +188,10 @@ private void initializeLakeRecordIterators() throws IOException {
         List<RecordReader> recordReaders = new ArrayList<>();
         if (lakeSplits == null || lakeSplits.isEmpty()) {
             // pass null split to get rowComparator
-            recordReaders.add(lakeSource.createRecordReader(() -> null));
+            recordReaders.add(lakeSource.createRecordReader(sortedReaderContext(null)));
         } else {
             for (LakeSplit lakeSplit : lakeSplits) {
-                recordReaders.add(lakeSource.createRecordReader(() -> lakeSplit));
+                recordReaders.add(lakeSource.createRecordReader(sortedReaderContext(lakeSplit)));
             }
         }
         for (RecordReader reader : recordReaders) {
@@ -206,6 +206,21 @@ private void initializeLakeRecordIterators() throws IOException {
         lakeRecordIteratorsInitialized = true;
     }
 
+    private LakeSource.ReaderContext<LakeSplit> sortedReaderContext(@Nullable LakeSplit lakeSplit) {
+        return new LakeSource.ReaderContext<LakeSplit>() {
+            @Nullable
+            @Override
+            public LakeSplit lakeSplit() {
+                return lakeSplit;
+            }
+
+            @Override
+            public boolean requireSortedRecords() {
+                return true;
+            }
+        };
+    }
+
     private void pollLogRecords(Duration timeout) {
         ScanRecords scanRecords = logScanner.poll(timeout);
         for (ScanRecord scanRecord : scanRecords) {

fluss-common/src/main/java/org/apache/fluss/lake/source/LakeSource.java

@@ -21,6 +21,8 @@
 import org.apache.fluss.lake.serializer.SimpleVersionedSerializer;
 import org.apache.fluss.predicate.Predicate;
 
+import javax.annotation.Nullable;
+
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.List;
@@ -91,12 +93,28 @@ interface PlannerContext extends Serializable {
     }
 
     /**
-     * Context interface for record readers, providing access to the lake split being read.
+     * Context interface for record readers, providing access to the lake split being read and the
+     * required reading semantics.
      *
      * @param <Split> The type of lake split
      */
     interface ReaderContext<Split extends LakeSplit> extends Serializable {
+        /**
+         * Returns the lake split to read.
+         *
+         * <p>The split can be null when the caller only needs reader-level metadata, such as a sort
+         * comparator.
+         */
+        @Nullable
         Split lakeSplit();
+
+        /**
+         * Returns whether records produced by this reader must follow the order defined by {@link
+         * SortedRecordReader#order()}.
+         */
+        default boolean requireSortedRecords() {
+            return false;
+        }
     }
 
     /**

fluss-lake/fluss-lake-hudi/src/main/java/org/apache/fluss/lake/hudi/source/HudiLakeSource.java

@@ -27,6 +27,7 @@
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.predicate.Predicate;
 
+import org.apache.hudi.common.model.HoodieTableType;
 import org.apache.hudi.org.apache.avro.Schema;
 import org.apache.hudi.source.ExpressionPredicates;
 import org.apache.hudi.util.StreamerUtil;
@@ -102,6 +103,17 @@ public Planner<HudiSplit> createPlanner(PlannerContext context) throws IOExcepti
     @Override
     public RecordReader createRecordReader(ReaderContext<HudiSplit> context) throws IOException {
         try {
+            if (context.requireSortedRecords()) {
+                if (!shouldUseSortedRecordReader()) {
+                    throw new UnsupportedOperationException(
+                            "Hudi lake source can not provide sorted records for "
+                                    + tablePath
+                                    + ". Sorted records are only supported for MERGE_ON_READ tables"
+                                    + " whose projection contains all Hudi record key fields.");
+                }
+                return new HudiSortedRecordReader(
+                        hudiConfig, tablePath, context.lakeSplit(), project, predicates);
+            }
             return new HudiRecordReader(
                     hudiConfig, tablePath, context.lakeSplit(), project, predicates);
         } catch (Exception e) {
@@ -121,4 +133,11 @@ private Schema getHudiSchema() {
             throw new RuntimeException("Fail to get Hudi schema for " + tablePath + ".", e);
         }
     }
+
+    private boolean shouldUseSortedRecordReader() throws IOException {
+        try (HudiTableInfo hudiTableInfo = HudiTableInfo.create(tablePath, hudiConfig)) {
+            return hudiTableInfo.getTableType() == HoodieTableType.MERGE_ON_READ
+                    && HudiSortedRecordReader.canSortByRecordKey(hudiTableInfo, project);
+        }
+    }
 }