This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new fd1e6edb09 [flink] Preserve query authorization when narrowing splits
in the Flink source (#9949)
fd1e6edb09 is described below
commit fd1e6edb093d8aa12c8b2d0c68e6aad2fd4dbc4f
Author: Jiajia Li <[email protected]>
AuthorDate: Sun Sep 20 13:48:38 2026 +0800
[flink] Preserve query authorization when narrowing splits in the Flink
source (#9949)
---
.../apache/paimon/table/ChainGroupReadTable.java | 5 +-
.../paimon/table/ChainTableFileStoreTable.java | 4 +-
.../apache/paimon/table/source/QueryAuthSplit.java | 10 -
.../org/apache/paimon/table/source/Splits.java | 77 ++++++
.../org/apache/paimon/table/source/SplitsTest.java | 167 +++++++++++++
.../source/ContinuousFileSplitEnumerator.java | 14 +-
.../source/DynamicPartitionFilteringInfo.java | 29 +++
.../flink/source/FileStoreSourceSplitReader.java | 6 +-
.../source/StaticFileStoreSplitEnumerator.java | 4 +-
.../AlignedContinuousFileSplitEnumerator.java | 3 +-
.../source/assigners/AlignedSplitAssigner.java | 18 +-
.../assigners/DynamicPartitionPruningAssigner.java | 20 +-
.../source/assigners/PreAssignSplitAssigner.java | 31 +--
.../flink/source/operator/MonitorSource.java | 11 +-
.../paimon/flink/source/operator/ReadOperator.java | 16 +-
.../apache/paimon/flink/utils/TableScanUtils.java | 8 +-
.../source/ContinuousFileSplitEnumeratorTest.java | 53 +++++
.../source/FileStoreSourceSplitReaderTest.java | 48 +++-
.../flink/source/FlinkSourceBuilderTest.java | 221 +++++++++++++++++
...ContinuousFileSplitEnumeratorQueryAuthTest.java | 145 ++++++++++++
.../source/assigners/AlignedSplitAssignerTest.java | 245 +++++++++++++++++++
.../DynamicPartitionPruningAssignerTest.java | 262 +++++++++++++++++++++
.../assigners/PreAssignSplitAssignerTest.java | 219 +++++++++++++++++
.../operator/ReadOperatorSplitShapeTest.java | 230 ++++++++++++++++++
.../paimon/flink/utils/TableScanUtilsTest.java | 116 +++++++++
25 files changed, 1887 insertions(+), 75 deletions(-)
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java
b/paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java
index 5bfeec5129..5564260d8e 100644
--- a/paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java
+++ b/paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java
@@ -39,6 +39,7 @@ import org.apache.paimon.table.source.DataTableScan;
import org.apache.paimon.table.source.InnerTableRead;
import org.apache.paimon.table.source.QueryAuthSplit;
import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.table.source.TableRead;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.ChainPartitionProjector;
@@ -623,7 +624,7 @@ public class ChainGroupReadTable extends
FallbackReadFileStoreTable {
@Override
public RecordReader<InternalRow> createReader(Split split) throws
IOException {
// fallbackRead unwraps and applies the rules itself, so it gets
the wrapper untouched.
- Split inner = QueryAuthSplit.unwrap(split);
+ Split inner = Splits.underlying(split);
if (inner instanceof ChainSplit || inner instanceof DataSplit) {
return fallbackRead.createReader(split);
}
@@ -653,7 +654,7 @@ public class ChainGroupReadTable extends
FallbackReadFileStoreTable {
"Branch scans of the same chain table returned
different authorization rules.");
authResult = next;
}
- dataSplits.add((DataSplit) QueryAuthSplit.unwrap(split));
+ dataSplits.add((DataSplit) Splits.underlying(split));
}
return dataSplits;
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableFileStoreTable.java
b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableFileStoreTable.java
index 8daed52edc..dec16c5504 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableFileStoreTable.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableFileStoreTable.java
@@ -29,8 +29,8 @@ import org.apache.paimon.schema.TableSchema;
import org.apache.paimon.table.source.ChainSplit;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.InnerTableRead;
-import org.apache.paimon.table.source.QueryAuthSplit;
import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.table.source.StreamDataTableScan;
import org.apache.paimon.table.source.TableRead;
import org.apache.paimon.types.RowType;
@@ -219,7 +219,7 @@ public class ChainTableFileStoreTable extends
FallbackReadFileStoreTable {
return fallbackRead.createReader(split);
}
// Route on what the wrapper carries, so the wrapper still reaches
the read.
- Split inner = QueryAuthSplit.unwrap(split);
+ Split inner = Splits.underlying(split);
if (inner instanceof ChainSplit || inner instanceof DataSplit) {
return chainGroupRead.createReader(split);
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/QueryAuthSplit.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/QueryAuthSplit.java
index f38fbc500f..3f6106d015 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/QueryAuthSplit.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/QueryAuthSplit.java
@@ -54,16 +54,6 @@ public class QueryAuthSplit implements Split {
return split;
}
- /**
- * The split this wrapper carries, or {@code split} itself when it is not
one. Narrowing a
- * planned split to a concrete type must go through this, or the cast
fails once the table
- * carries a rule. A caller handing the split to a read passes the wrapper
on instead, or the
- * authorization never reaches the reader.
- */
- public static Split unwrap(Split split) {
- return split instanceof QueryAuthSplit ? ((QueryAuthSplit)
split).split() : split;
- }
-
@Nullable
public TableQueryAuthResult authResult() {
return authResult;
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/Splits.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/Splits.java
new file mode 100644
index 0000000000..5a11c481ae
--- /dev/null
+++ b/paimon-core/src/main/java/org/apache/paimon/table/source/Splits.java
@@ -0,0 +1,77 @@
+/*
+ * 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.paimon.table.source;
+
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit;
+
+import javax.annotation.Nullable;
+
+/** Utilities for looking at what a planned {@link Split} is built from. */
+public class Splits {
+
+ private Splits() {}
+
+ /**
+ * The split a planned one is built from, looked at through every wrapper.
One layer is not
+ * enough: an authorized fallback table hands out a {@link FallbackSplit}
holding a {@link
+ * QueryAuthSplit}.
+ *
+ * <p>Stops at a {@link DataSplit}. {@code FallbackDataSplit} is one and
wraps itself, so that
+ * condition ends the loop as much as it preserves what a plan carrying no
rules resolved to.
+ *
+ * <p>Only for reading what the split is made of. A caller handing it to a
read passes the
+ * planned split on as it is, or the authorization never reaches the
reader.
+ */
+ public static Split underlying(Split split) {
+ Split current = split;
+ while (!(current instanceof DataSplit)) {
+ if (current instanceof QueryAuthSplit) {
+ current = ((QueryAuthSplit) current).split();
+ } else if (current instanceof FallbackSplit) {
+ current = ((FallbackSplit) current).wrapped();
+ } else {
+ return current;
+ }
+ }
+ return current;
+ }
+
+ /**
+ * The authorization a planned split carries, or null when it carries
none. The companion of
+ * {@link #underlying(Split)}, which drops the wrapper the rules live on.
+ *
+ * <p>A {@code FallbackDataSplit} wraps itself, which ends the walk: it
carries no rules either.
+ */
+ @Nullable
+ public static TableQueryAuthResult authResult(Split split) {
+ Split current = split;
+ while (!(current instanceof QueryAuthSplit)) {
+ if (!(current instanceof FallbackSplit)) {
+ return null;
+ }
+ Split wrapped = ((FallbackSplit) current).wrapped();
+ if (wrapped == current) {
+ return null;
+ }
+ current = wrapped;
+ }
+ return ((QueryAuthSplit) current).authResult();
+ }
+}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/SplitsTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/SplitsTest.java
new file mode 100644
index 0000000000..3f403f5909
--- /dev/null
+++ b/paimon-core/src/test/java/org/apache/paimon/table/source/SplitsTest.java
@@ -0,0 +1,167 @@
+/*
+ * 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.paimon.table.source;
+
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.table.FallbackReadFileStoreTable;
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackDataSplit;
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit;
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplitImpl;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+
+import static org.apache.paimon.io.DataFileTestUtils.newFile;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Test for {@link Splits}. */
+public class SplitsTest {
+
+ @Test
+ public void testDataSplitIsItsOwnUnderlyingSplit() {
+ DataSplit split = dataSplit();
+
+ assertThat(Splits.underlying(split)).isSameAs(split);
+ }
+
+ @Test
+ public void testAuthorizationWrapperIsLookedThrough() {
+ DataSplit split = dataSplit();
+
+ assertThat(Splits.underlying(withQueryAuth(split))).isSameAs(split);
+ }
+
+ @Test
+ public void testFallbackWrapperAroundAuthorizationWrapperIsLookedThrough()
{
+ DataSplit split = dataSplit();
+ Split nested = toFallbackSplit(withQueryAuth(split));
+
+ assertThat(nested).isInstanceOf(FallbackSplitImpl.class);
+ assertThat(((FallbackSplit)
nested).wrapped()).isInstanceOf(QueryAuthSplit.class);
+ assertThat(Splits.underlying(nested)).isSameAs(split);
+ }
+
+ @Test
+ public void testFallbackDataSplitIsItsOwnUnderlyingSplit() {
+ DataSplit split = dataSplit();
+ Split fallback = toFallbackSplit(split);
+
+ assertThat(fallback).isInstanceOf(FallbackDataSplit.class);
+ assertThat(Splits.underlying(fallback)).isSameAs(fallback);
+ }
+
+ @Test
+ public void
testSplitThatIsNeitherADataSplitNorAWrapperIsItsOwnUnderlyingSplit() {
+ ChainSplit split = chainSplit();
+
+ assertThat(Splits.underlying(split)).isSameAs(split);
+ }
+
+ @Test
+ public void
testFallbackWrapperAroundASplitThatIsNotADataSplitIsLookedThrough() {
+ ChainSplit split = chainSplit();
+
+ assertThat(Splits.underlying(toFallbackSplit(split))).isSameAs(split);
+ }
+
+ @Test
+ public void
testAuthorizationWrapperAroundASplitThatIsNotADataSplitIsLookedThrough() {
+ ChainSplit split = chainSplit();
+
+ assertThat(Splits.underlying(withQueryAuth(split))).isSameAs(split);
+ }
+
+ @Test
+ public void
testAuthorizationWrapperAroundAFallbackDataSplitStopsAtTheFallbackDataSplit() {
+ Split fallback = toFallbackSplit(dataSplit());
+
+
assertThat(Splits.underlying(withQueryAuth(fallback))).isSameAs(fallback);
+ }
+
+ @Test
+ public void testSplitCarryingNoAuthorizationHasNoAuthResult() {
+ assertThat(Splits.authResult(dataSplit())).isNull();
+ assertThat(Splits.authResult(chainSplit())).isNull();
+ assertThat(Splits.authResult(toFallbackSplit(dataSplit()))).isNull();
+ assertThat(Splits.authResult(toFallbackSplit(chainSplit()))).isNull();
+ }
+
+ @Test
+ public void testAuthResultIsReadThroughTheAuthorizationWrapper() {
+ TableQueryAuthResult authResult = authResult();
+
+ assertThat(Splits.authResult(new QueryAuthSplit(dataSplit(),
authResult)))
+ .isSameAs(authResult);
+ }
+
+ @Test
+ public void testAuthResultIsReadThroughAFallbackWrapper() {
+ TableQueryAuthResult authResult = authResult();
+ Split nested = toFallbackSplit(new QueryAuthSplit(dataSplit(),
authResult));
+
+ assertThat(Splits.authResult(nested)).isSameAs(authResult);
+ }
+
+ @Test
+ public void testAuthResultIsReadFromAWrapperAroundAFallbackDataSplit() {
+ TableQueryAuthResult authResult = authResult();
+ Split split = new QueryAuthSplit(toFallbackSplit(dataSplit()),
authResult);
+
+ assertThat(Splits.authResult(split)).isSameAs(authResult);
+ }
+
+ @Test
+ public void testAuthorizationWrapperWithoutRulesHasNoAuthResult() {
+ assertThat(Splits.authResult(new QueryAuthSplit(dataSplit(),
null))).isNull();
+ }
+
+ private static TableQueryAuthResult authResult() {
+ return new TableQueryAuthResult(
+ null, Collections.singletonMap("name", "{\"name\":\"NULL\"}"));
+ }
+
+ private static DataSplit dataSplit() {
+ return DataSplit.builder()
+ .withSnapshot(1)
+ .withPartition(BinaryRow.EMPTY_ROW)
+ .withBucket(0)
+ .withBucketPath("bucket-0")
+ .withDataFiles(Collections.singletonList(newFile(0L, 1L)))
+ .build();
+ }
+
+ private static ChainSplit chainSplit() {
+ return new ChainSplit(
+ BinaryRow.EMPTY_ROW,
+ Collections.singletonList(newFile(0L, 1L)),
+ Collections.emptyMap(),
+ Collections.emptyMap(),
+ null);
+ }
+
+ private static Split withQueryAuth(Split split) {
+ return new QueryAuthSplit(split, new TableQueryAuthResult(null, null));
+ }
+
+ private static Split toFallbackSplit(Split split) {
+ return FallbackReadFileStoreTable.toFallbackSplit(split, true);
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumerator.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumerator.java
index 75c75cfddb..abbafd2f65 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumerator.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumerator.java
@@ -31,6 +31,8 @@ import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.EndOfScanException;
import org.apache.paimon.table.source.IncrementalSplit;
import org.apache.paimon.table.source.SnapshotNotExistPlan;
+import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.table.source.StreamTableScan;
import org.apache.paimon.table.source.TableScan;
@@ -329,12 +331,14 @@ public class ContinuousFileSplitEnumerator
protected int assignSuggestedTask(FileStoreSourceSplit split) {
int task;
- if (split.split() instanceof DataSplit) {
- task = assignSuggestedTask((DataSplit) split.split());
- } else if (split.split() instanceof ChainSplit) {
- task = assignSuggestedTask((ChainSplit) split.split());
+ // Reads the underlying split, assigns the original.
+ Split inner = Splits.underlying(split.split());
+ if (inner instanceof DataSplit) {
+ task = assignSuggestedTask((DataSplit) inner);
+ } else if (inner instanceof ChainSplit) {
+ task = assignSuggestedTask((ChainSplit) inner);
} else {
- task = assignSuggestedTask((IncrementalSplit) split.split());
+ task = assignSuggestedTask((IncrementalSplit) inner);
}
// Split assigners keep splits in a map keyed by task, but only ever
hand out splits for
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/DynamicPartitionFilteringInfo.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/DynamicPartitionFilteringInfo.java
index f70739e950..aa5374a27c 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/DynamicPartitionFilteringInfo.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/DynamicPartitionFilteringInfo.java
@@ -18,11 +18,20 @@
package org.apache.paimon.flink.source;
+import org.apache.paimon.catalog.TableQueryAuthResult;
import org.apache.paimon.codegen.CodeGenUtils;
import org.apache.paimon.codegen.Projection;
+import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.flink.FlinkRowData;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.types.RowType;
+import org.apache.flink.table.connector.source.DynamicFilteringData;
+
import java.io.Serializable;
+import java.util.Collections;
import java.util.List;
/** Manage dynamic partition filtering fields and table partition row. */
@@ -49,4 +58,24 @@ public class DynamicPartitionFilteringInfo implements
Serializable {
return partitionRowProjection;
}
+
+ /**
+ * Whether {@code split} may still hold a row {@code dynamicFilteringData}
asks for, answered
+ * from the partition the plan recorded rather than by opening the split.
+ *
+ * <p>A mask on a filtering field breaks that shortcut: the reader hands
out a value other than
+ * the recorded one, so such a split is kept for the join to filter on the
masked value.
+ */
+ public boolean mayMatch(DynamicFilteringData dynamicFilteringData, Split
split) {
+ TableQueryAuthResult authResult = Splits.authResult(split);
+ if (authResult != null
+ && !Collections.disjoint(
+ authResult.extractColumnMasking().keySet(),
+ dynamicPartitionFilteringFields)) {
+ return true;
+ }
+ BinaryRow partition = ((DataSplit)
Splits.underlying(split)).partition();
+ return dynamicFilteringData.contains(
+ new
FlinkRowData(getPartitionRowProjection().apply(partition)));
+ }
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/FileStoreSourceSplitReader.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/FileStoreSourceSplitReader.java
index 8568e2a9dd..1c230ab59b 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/FileStoreSourceSplitReader.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/FileStoreSourceSplitReader.java
@@ -26,6 +26,7 @@ import org.apache.paimon.reader.RecordReader;
import org.apache.paimon.reader.RecordReader.RecordIterator;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.table.source.TableRead;
import org.apache.paimon.types.BlobType;
import org.apache.paimon.types.RowType;
@@ -219,9 +220,10 @@ public class FileStoreSourceSplitReader
}
// update metric when split changes
- if (nextSplit.split() instanceof DataSplit) {
+ Split inner = Splits.underlying(nextSplit.split());
+ if (inner instanceof DataSplit) {
long eventTime =
- ((DataSplit) nextSplit.split())
+ ((DataSplit) inner)
.earliestFileCreationEpochMillis()
.orElse(FileStoreSourceReaderMetrics.UNDEFINED);
metrics.recordSnapshotUpdate(eventTime);
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/StaticFileStoreSplitEnumerator.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/StaticFileStoreSplitEnumerator.java
index 7bd8ff990f..c0ecc363ac 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/StaticFileStoreSplitEnumerator.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/StaticFileStoreSplitEnumerator.java
@@ -148,7 +148,7 @@ public class StaticFileStoreSplitEnumerator
this.splitAssigner =
((PreAssignSplitAssigner) splitAssigner)
.ofDynamicPartitionPruning(
-
dynamicPartitionFilteringInfo.getPartitionRowProjection(),
+ dynamicPartitionFilteringInfo,
((DynamicFilteringEvent)
sourceEvent).getData());
} else {
this.splitAssigner =
@@ -156,7 +156,7 @@ public class StaticFileStoreSplitEnumerator
.createDynamicPartitionPruningAssignerIfNeeded(
subtaskId,
splitAssigner,
-
dynamicPartitionFilteringInfo.getPartitionRowProjection(),
+ dynamicPartitionFilteringInfo,
sourceEvent,
LOG);
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumerator.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumerator.java
index c728134c07..90164b5d36 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumerator.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumerator.java
@@ -26,6 +26,7 @@ import org.apache.paimon.flink.source.assigners.SplitAssigner;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.EndOfScanException;
import org.apache.paimon.table.source.SnapshotNotExistPlan;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.table.source.StreamTableScan;
import org.apache.paimon.table.source.TableScan;
import org.apache.paimon.utils.Preconditions;
@@ -123,7 +124,7 @@ public class AlignedContinuousFileSplitEnumerator extends
ContinuousFileSplitEnu
Map<Long, List<FileStoreSourceSplit>> splitsBySnapshot = new
TreeMap<>();
for (FileStoreSourceSplit split : splits) {
- long snapshotId = ((DataSplit) split.split()).snapshotId();
+ long snapshotId = ((DataSplit)
Splits.underlying(split.split())).snapshotId();
splitsBySnapshot.computeIfAbsent(snapshotId, snapshot -> new
ArrayList<>()).add(split);
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/AlignedSplitAssigner.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/AlignedSplitAssigner.java
index 648758f838..5f5686093b 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/AlignedSplitAssigner.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/AlignedSplitAssigner.java
@@ -21,6 +21,7 @@ package org.apache.paimon.flink.source.assigners;
import org.apache.paimon.flink.source.FileStoreSourceSplit;
import org.apache.paimon.flink.source.align.PlaceholderSplit;
import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.utils.Preconditions;
import javax.annotation.Nullable;
@@ -66,9 +67,9 @@ public class AlignedSplitAssigner implements SplitAssigner {
@Override
public void addSplit(int subtask, FileStoreSourceSplit splits) {
- long snapshotId = ((DataSplit) splits.split()).snapshotId();
+ long snapshotId = dataSplit(splits).snapshotId();
PendingSnapshot last = pendingSplitAssignment.peekLast();
- boolean isPlaceholder = splits.split() instanceof PlaceholderSplit;
+ boolean isPlaceholder = dataSplit(splits) instanceof PlaceholderSplit;
if (last == null || last.snapshotId != snapshotId) {
last = new PendingSnapshot(snapshotId, isPlaceholder, new
HashMap<>());
last.add(subtask, splits);
@@ -85,8 +86,8 @@ public class AlignedSplitAssigner implements SplitAssigner {
return;
}
- long snapshotId = ((DataSplit) splits.get(0).split()).snapshotId();
- boolean isPlaceholder = splits.get(0).split() instanceof
PlaceholderSplit;
+ long snapshotId = dataSplit(splits.get(0)).snapshotId();
+ boolean isPlaceholder = dataSplit(splits.get(0)) instanceof
PlaceholderSplit;
PendingSnapshot head = pendingSplitAssignment.peek();
if (head == null || snapshotId != head.snapshotId) {
head = new PendingSnapshot(snapshotId, isPlaceholder, new
HashMap<>());
@@ -154,7 +155,7 @@ public class AlignedSplitAssigner implements SplitAssigner {
public void add(int subtask, FileStoreSourceSplit split) {
Preconditions.checkArgument(
- ((DataSplit) split.split()).snapshotId() == snapshotId,
+ dataSplit(split).snapshotId() == snapshotId,
"SnapshotId not equal. This is a bug, please file an
issue.");
subtaskSplits.computeIfAbsent(subtask, id -> new
ArrayList<>()).add(split);
}
@@ -166,7 +167,7 @@ public class AlignedSplitAssigner implements SplitAssigner {
splits.forEach(
split ->
Preconditions.checkArgument(
- ((DataSplit) split.split()).snapshotId()
== snapshotId,
+ dataSplit(split).snapshotId() ==
snapshotId,
"SnapshotId not equal"));
subtaskSplits.put(subtask, splits);
}
@@ -175,4 +176,9 @@ public class AlignedSplitAssigner implements SplitAssigner {
return subtaskSplits.isEmpty() || isPlaceHolder;
}
}
+
+ /** Reads the underlying split, assigns the original. */
+ private static DataSplit dataSplit(FileStoreSourceSplit split) {
+ return (DataSplit) Splits.underlying(split.split());
+ }
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssigner.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssigner.java
index 9221f1f27e..5e4475c8b1 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssigner.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssigner.java
@@ -18,11 +18,8 @@
package org.apache.paimon.flink.source.assigners;
-import org.apache.paimon.codegen.Projection;
-import org.apache.paimon.data.BinaryRow;
-import org.apache.paimon.flink.FlinkRowData;
+import org.apache.paimon.flink.source.DynamicPartitionFilteringInfo;
import org.apache.paimon.flink.source.FileStoreSourceSplit;
-import org.apache.paimon.table.source.DataSplit;
import org.apache.flink.api.connector.source.SourceEvent;
import org.apache.flink.table.connector.source.DynamicFilteringData;
@@ -41,15 +38,15 @@ import java.util.stream.Collectors;
public class DynamicPartitionPruningAssigner implements SplitAssigner {
private final SplitAssigner innerAssigner;
- private final Projection partitionRowProjection;
+ private final DynamicPartitionFilteringInfo dynamicPartitionFilteringInfo;
private final DynamicFilteringData dynamicFilteringData;
public DynamicPartitionPruningAssigner(
SplitAssigner innerAssigner,
- Projection partitionRowProjection,
+ DynamicPartitionFilteringInfo dynamicPartitionFilteringInfo,
DynamicFilteringData dynamicFilteringData) {
this.innerAssigner = innerAssigner;
- this.partitionRowProjection = partitionRowProjection;
+ this.dynamicPartitionFilteringInfo = dynamicPartitionFilteringInfo;
this.dynamicFilteringData = dynamicFilteringData;
}
@@ -90,7 +87,7 @@ public class DynamicPartitionPruningAssigner implements
SplitAssigner {
public static SplitAssigner createDynamicPartitionPruningAssignerIfNeeded(
int subtaskId,
SplitAssigner oriAssigner,
- Projection partitionRowProjection,
+ DynamicPartitionFilteringInfo dynamicPartitionFilteringInfo,
SourceEvent sourceEvent,
Logger logger) {
DynamicFilteringData dynamicFilteringData = ((DynamicFilteringEvent)
sourceEvent).getData();
@@ -100,7 +97,7 @@ public class DynamicPartitionPruningAssigner implements
SplitAssigner {
dynamicFilteringData.isFiltering());
return dynamicFilteringData.isFiltering()
? new DynamicPartitionPruningAssigner(
- oriAssigner, partitionRowProjection,
dynamicFilteringData)
+ oriAssigner, dynamicPartitionFilteringInfo,
dynamicFilteringData)
: oriAssigner;
}
@@ -115,9 +112,6 @@ public class DynamicPartitionPruningAssigner implements
SplitAssigner {
}
private boolean filter(FileStoreSourceSplit sourceSplit) {
- DataSplit dataSplit = (DataSplit) sourceSplit.split();
- BinaryRow partition = dataSplit.partition();
- FlinkRowData projected = new
FlinkRowData(partitionRowProjection.apply(partition));
- return dynamicFilteringData.contains(projected);
+ return dynamicPartitionFilteringInfo.mayMatch(dynamicFilteringData,
sourceSplit.split());
}
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java
index 24ac4a2911..c0077e89ee 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java
@@ -18,11 +18,8 @@
package org.apache.paimon.flink.source.assigners;
-import org.apache.paimon.codegen.Projection;
-import org.apache.paimon.data.BinaryRow;
-import org.apache.paimon.flink.FlinkRowData;
+import org.apache.paimon.flink.source.DynamicPartitionFilteringInfo;
import org.apache.paimon.flink.source.FileStoreSourceSplit;
-import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.utils.BinPacking;
import org.apache.paimon.utils.SerializableFunction;
@@ -92,13 +89,13 @@ public class PreAssignSplitAssigner implements
SplitAssigner {
int splitBatchSize,
int parallelism,
Collection<FileStoreSourceSplit> splits,
- Projection partitionRowProjection,
+ DynamicPartitionFilteringInfo dynamicPartitionFilteringInfo,
DynamicFilteringData dynamicFilteringData) {
this(
splitBatchSize,
parallelism,
splits,
- partitionRowProjection,
+ dynamicPartitionFilteringInfo,
dynamicFilteringData,
split -> split.split().rowCount());
}
@@ -107,14 +104,17 @@ public class PreAssignSplitAssigner implements
SplitAssigner {
int splitBatchSize,
int parallelism,
Collection<FileStoreSourceSplit> splits,
- Projection partitionRowProjection,
+ DynamicPartitionFilteringInfo dynamicPartitionFilteringInfo,
DynamicFilteringData dynamicFilteringData,
SerializableFunction<FileStoreSourceSplit, Long> weightFunc) {
this(
splitBatchSize,
parallelism,
splits.stream()
- .filter(s -> filter(partitionRowProjection,
dynamicFilteringData, s))
+ .filter(
+ s ->
+ dynamicPartitionFilteringInfo.mayMatch(
+ dynamicFilteringData,
s.split()))
.collect(Collectors.toList()),
weightFunc);
}
@@ -279,23 +279,14 @@ public class PreAssignSplitAssigner implements
SplitAssigner {
}
public SplitAssigner ofDynamicPartitionPruning(
- Projection partitionRowProjection, DynamicFilteringData
dynamicFilteringData) {
+ DynamicPartitionFilteringInfo dynamicPartitionFilteringInfo,
+ DynamicFilteringData dynamicFilteringData) {
return new PreAssignSplitAssigner(
splitBatchSize,
parallelism,
splits,
- partitionRowProjection,
+ dynamicPartitionFilteringInfo,
dynamicFilteringData,
weightFunc);
}
-
- private static boolean filter(
- Projection partitionRowProjection,
- DynamicFilteringData dynamicFilteringData,
- FileStoreSourceSplit sourceSplit) {
- DataSplit dataSplit = (DataSplit) sourceSplit.split();
- BinaryRow partition = dataSplit.partition();
- FlinkRowData projected = new
FlinkRowData(partitionRowProjection.apply(partition));
- return dynamicFilteringData.contains(projected);
- }
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/operator/MonitorSource.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/operator/MonitorSource.java
index 740cb59311..1b1ad2392f 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/operator/MonitorSource.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/operator/MonitorSource.java
@@ -32,8 +32,10 @@ import org.apache.paimon.table.Table;
import org.apache.paimon.table.sink.ChannelComputer;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.EndOfScanException;
+import org.apache.paimon.table.source.IncrementalSplit;
import org.apache.paimon.table.source.ReadBuilder;
import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.table.source.StreamTableScan;
import org.apache.paimon.table.source.TableScan;
import org.apache.paimon.types.RowType;
@@ -417,7 +419,14 @@ public class MonitorSource extends
AbstractNonCoordinatedSource<Split> {
return ChannelComputer.select(key.f1, numPartitions);
},
split -> {
- DataSplit dataSplit = (DataSplit) split;
+ // Keys on the underlying split, routes the original.
Either shape orders
+ // by partition and bucket.
+ Split inner = Splits.underlying(split);
+ if (inner instanceof IncrementalSplit) {
+ IncrementalSplit incrementalSplit = (IncrementalSplit)
inner;
+ return Tuple2.of(incrementalSplit.partition(),
incrementalSplit.bucket());
+ }
+ DataSplit dataSplit = (DataSplit) inner;
return Tuple2.of(dataSplit.partition(),
dataSplit.bucket());
});
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/operator/ReadOperator.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/operator/ReadOperator.java
index 95bf5606e7..5004058561 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/operator/ReadOperator.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/operator/ReadOperator.java
@@ -27,6 +27,7 @@ import org.apache.paimon.flink.source.RecordLimiter;
import org.apache.paimon.flink.source.metrics.FileStoreSourceReaderMetrics;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.table.source.TableRead;
import org.apache.paimon.types.BlobType;
import org.apache.paimon.types.RowType;
@@ -144,10 +145,14 @@ public class ReadOperator extends
AbstractStreamOperator<RowData>
Split split = record.getValue();
// update metric when reading a new split
+ // A shape this metric cannot read must not fail the read.
+ Split inner = Splits.underlying(split);
long eventTime =
- ((DataSplit) split)
- .earliestFileCreationEpochMillis()
- .orElse(FileStoreSourceReaderMetrics.UNDEFINED);
+ inner instanceof DataSplit
+ ? ((DataSplit) inner)
+ .earliestFileCreationEpochMillis()
+ .orElse(FileStoreSourceReaderMetrics.UNDEFINED)
+ : FileStoreSourceReaderMetrics.UNDEFINED;
sourceReaderMetrics.recordSnapshotUpdate(eventTime);
// update idleStartTime when reading a new split
idleStartTime = FileStoreSourceReaderMetrics.ACTIVE;
@@ -156,7 +161,10 @@ public class ReadOperator extends
AbstractStreamOperator<RowData>
try (CloseableIterator<InternalRow> iterator =
read.createReader(split).toCloseableIterator()) {
while (!reachLimit() && iterator.hasNext()) {
- emitEventTimeLag = System.currentTimeMillis() - eventTime;
+ emitEventTimeLag =
+ eventTime == FileStoreSourceReaderMetrics.UNDEFINED
+ ? FileStoreSourceReaderMetrics.UNDEFINED
+ : System.currentTimeMillis() - eventTime;
// each Split is already counted as one input record,
// so we don't need to count the first record
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/utils/TableScanUtils.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/utils/TableScanUtils.java
index ca901ad1bf..d756ab8bd2 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/utils/TableScanUtils.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/utils/TableScanUtils.java
@@ -23,6 +23,8 @@ import org.apache.paimon.flink.source.FileStoreSourceSplit;
import org.apache.paimon.options.Options;
import org.apache.paimon.table.Table;
import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.Splits;
import org.apache.paimon.table.source.TableScan;
import java.util.HashMap;
@@ -67,8 +69,10 @@ public class TableScanUtils {
/** Get snapshot id from {@link FileStoreSourceSplit}. */
public static Optional<Long> getSnapshotId(FileStoreSourceSplit split) {
- if (split.split() instanceof DataSplit) {
- return Optional.of(((DataSplit) split.split()).snapshotId());
+ // An empty answer does not fail here, it stalls consumer progress and
watermarks.
+ Split inner = Splits.underlying(split.split());
+ if (inner instanceof DataSplit) {
+ return Optional.of(((DataSplit) inner).snapshotId());
}
return Optional.empty();
}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumeratorTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumeratorTest.java
index 0a50038bb4..103ae68683 100644
---
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumeratorTest.java
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumeratorTest.java
@@ -18,11 +18,14 @@
package org.apache.paimon.flink.source;
+import org.apache.paimon.catalog.TableQueryAuthResult;
import org.apache.paimon.io.DataFileMeta;
import org.apache.paimon.table.BucketMode;
+import org.apache.paimon.table.FallbackReadFileStoreTable;
import org.apache.paimon.table.source.DataFilePlan;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.IncrementalSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
import org.apache.paimon.table.source.StreamTableScan;
import org.apache.paimon.table.source.TableScan;
@@ -858,6 +861,56 @@ public class ContinuousFileSplitEnumeratorTest
Assertions.assertThat(enumerator.nextSnapshotId).isEqualTo(3);
}
+ @Test
+ public void testQueryAuthSplitAssignedLikeTheSplitItWraps() {
+ int parallelism = 3;
+ ContinuousFileSplitEnumerator enumerator =
buildEnumerator(parallelism);
+
+ FileStoreSourceSplit plain = createSnapshotSplit(1, 2,
Collections.emptyList());
+ FileStoreSourceSplit wrapped =
+ new FileStoreSourceSplit(
+ plain.splitId(),
+ new QueryAuthSplit(plain.split(), new
TableQueryAuthResult(null, null)));
+
+ assertThat(enumerator.assignSuggestedTask(wrapped))
+ .isEqualTo(enumerator.assignSuggestedTask(plain));
+ }
+
+ @Test
+ public void testFallbackWrappedQueryAuthSplitAssignedLikeTheSplitItWraps()
{
+ int parallelism = 3;
+ ContinuousFileSplitEnumerator enumerator =
buildEnumerator(parallelism);
+
+ FileStoreSourceSplit plain = createSnapshotSplit(1, 2,
Collections.emptyList());
+ FileStoreSourceSplit nested =
+ new FileStoreSourceSplit(
+ plain.splitId(),
+ FallbackReadFileStoreTable.toFallbackSplit(
+ new QueryAuthSplit(
+ plain.split(), new
TableQueryAuthResult(null, null)),
+ true));
+
+ assertThat(nested.split()).isNotInstanceOf(DataSplit.class);
+ assertThat(enumerator.assignSuggestedTask(nested))
+ .isEqualTo(enumerator.assignSuggestedTask(plain));
+ }
+
+ @Test
+ public void testFallbackDataSplitAssignedLikeTheSplitItCopies() {
+ int parallelism = 3;
+ ContinuousFileSplitEnumerator enumerator =
buildEnumerator(parallelism);
+
+ FileStoreSourceSplit plain = createSnapshotSplit(1, 2,
Collections.emptyList());
+ FileStoreSourceSplit fallback =
+ new FileStoreSourceSplit(
+ plain.splitId(),
+
FallbackReadFileStoreTable.toFallbackSplit(plain.split(), true));
+
+ assertThat(fallback.split()).isInstanceOf(DataSplit.class);
+ assertThat(enumerator.assignSuggestedTask(fallback))
+ .isEqualTo(enumerator.assignSuggestedTask(plain));
+ }
+
@Test
public void testPostponeBucketIncrementalSplitAssignedByWriteId() {
int parallelism = 3;
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FileStoreSourceSplitReaderTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FileStoreSourceSplitReaderTest.java
index b3a9414093..42e3c1406e 100644
---
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FileStoreSourceSplitReaderTest.java
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FileStoreSourceSplitReaderTest.java
@@ -19,6 +19,7 @@
package org.apache.paimon.flink.source;
import org.apache.paimon.KeyValue;
+import org.apache.paimon.catalog.TableQueryAuthResult;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.disk.IOManager;
@@ -32,6 +33,7 @@ import org.apache.paimon.reader.RecordReader;
import org.apache.paimon.schema.FileSystemSchemaManager;
import org.apache.paimon.schema.Schema;
import org.apache.paimon.schema.SchemaManager;
+import org.apache.paimon.table.source.QueryAuthSplit;
import org.apache.paimon.table.source.Split;
import org.apache.paimon.table.source.TableRead;
import org.apache.paimon.utils.RecordWriter;
@@ -138,13 +140,49 @@ public class FileStoreSourceSplitReaderTest {
assertThat(thread.isAlive()).isFalse();
}
+ @Test
+ public void testWrappedSplitReportsTheFileCreationTimeOfTheSplitItWraps()
throws Exception {
+ TestChangelogDataReadWrite rw = new
TestChangelogDataReadWrite(tempDir.toString());
+ List<DataFileMeta> files = rw.writeFiles(row(1), 0, kvs());
+
+ FileStoreSourceSplit split = newSourceSplit("id1", row(1), 0, files);
+ FileStoreSourceSplit wrapped =
+ new FileStoreSourceSplit(
+ "id1",
+ new QueryAuthSplit(split.split(), new
TableQueryAuthResult(null, null)));
+
+ long plain = readAndReportFileCreationTime(rw.createReadWithKey(),
split);
+ long authorized =
readAndReportFileCreationTime(rw.createReadWithKey(), wrapped);
+
+ assertThat(plain).isNotEqualTo(FileStoreSourceReaderMetrics.UNDEFINED);
+ assertThat(authorized).isEqualTo(plain);
+ }
+
+ private long readAndReportFileCreationTime(TableRead tableRead,
FileStoreSourceSplit split)
+ throws Exception {
+ FileStoreSourceReaderMetrics metrics =
+ new FileStoreSourceReaderMetrics(new DummyMetricGroup());
+ FileStoreSourceSplitReader reader = createReader(tableRead, null,
metrics);
+ try {
+ assignSplit(reader, split);
+ readRecords(reader.fetch(), split.splitId(), 0);
+ assertThat(reader.fetch().finishedSplits())
+ .isEqualTo(Collections.singleton(split.splitId()));
+ } finally {
+ reader.close();
+ }
+ return metrics.getLatestFileCreationTime();
+ }
+
private FileStoreSourceSplitReader createReader(TableRead tableRead,
@Nullable Long limit) {
+ return createReader(
+ tableRead, limit, new FileStoreSourceReaderMetrics(new
DummyMetricGroup()));
+ }
+
+ private FileStoreSourceSplitReader createReader(
+ TableRead tableRead, @Nullable Long limit,
FileStoreSourceReaderMetrics metrics) {
return new FileStoreSourceSplitReader(
- tableRead,
- limit == null ? null : new RecordLimiter(limit),
- new FileStoreSourceReaderMetrics(new DummyMetricGroup()),
- null,
- false);
+ tableRead, limit == null ? null : new RecordLimiter(limit),
metrics, null, false);
}
private void innerTestOnce(int skip) throws Exception {
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FlinkSourceBuilderTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FlinkSourceBuilderTest.java
index 4f8a097e47..709b0e733b 100644
---
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FlinkSourceBuilderTest.java
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FlinkSourceBuilderTest.java
@@ -22,22 +22,36 @@ import org.apache.paimon.catalog.Catalog;
import org.apache.paimon.catalog.CatalogContext;
import org.apache.paimon.catalog.CatalogFactory;
import org.apache.paimon.catalog.Identifier;
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.flink.source.operator.MonitorSource;
import org.apache.paimon.schema.Schema;
+import org.apache.paimon.table.FallbackReadFileStoreTable;
import org.apache.paimon.table.Table;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.IncrementalSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
+import org.apache.paimon.table.source.Split;
import org.apache.paimon.types.DataTypes;
import org.apache.flink.api.dag.Transformation;
+import org.apache.flink.runtime.plugable.SerializationDelegate;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.transformations.PartitionTransformation;
import org.apache.flink.streaming.api.transformations.SourceTransformation;
+import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import java.nio.file.Path;
+import java.util.Collections;
import static org.apache.paimon.flink.LogicalTypeConversion.toLogicalType;
import static org.assertj.core.api.Assertions.assertThat;
@@ -204,4 +218,211 @@ public class FlinkSourceBuilderTest {
.getSource())
.isInstanceOf(PaimonDataStreamSource.class));
}
+
+ @ValueSource(booleans = {false, true})
+ @ParameterizedTest
+ public void testOrderedShuffleRoutesQueryAuthSplitLikeTheSplitItWraps(
+ boolean shuffleBucketWithPartition) throws Exception {
+ Table table = createTable("ordered_shuffle_" +
shuffleBucketWithPartition, false, 2, true);
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+
+ DataStream<RowData> dataStream =
+ MonitorSource.buildSource(
+ env,
+ "source",
+ InternalTypeInfo.of(toLogicalType(table.rowType())),
+ table.newReadBuilder(),
+ 10,
+ false,
+ shuffleBucketWithPartition,
+ false,
+ null,
+ true,
+ null);
+
+ Transformation<?> input =
dataStream.getTransformation().getInputs().get(0);
+ assertThat(input).isInstanceOf(PartitionTransformation.class);
+ @SuppressWarnings("unchecked")
+ StreamPartitioner<Split> partitioner =
+ ((PartitionTransformation<Split>) input).getPartitioner();
+ partitioner.setup(4);
+
+ DataSplit bucketOne = dataSplit(1);
+ DataSplit bucketThree = dataSplit(3);
+
+ assertThat(selectChannel(partitioner, withQueryAuth(bucketOne)))
+ .isEqualTo(selectChannel(partitioner, bucketOne));
+ assertThat(selectChannel(partitioner, withQueryAuth(bucketThree)))
+ .isEqualTo(selectChannel(partitioner, bucketThree));
+
+ assertThat(selectChannel(partitioner, withQueryAuth(bucketOne)))
+ .isNotEqualTo(selectChannel(partitioner,
withQueryAuth(bucketThree)));
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ public void testOrderedShuffleRoutesIncrementalSplitByPartitionAndBucket(
+ boolean shuffleBucketWithPartition) throws Exception {
+ Table table =
+ createTable("ordered_incremental_" +
shuffleBucketWithPartition, false, 2, true);
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+
+ DataStream<RowData> dataStream =
+ MonitorSource.buildSource(
+ env,
+ "source",
+ InternalTypeInfo.of(toLogicalType(table.rowType())),
+ table.newReadBuilder(),
+ 10,
+ false,
+ shuffleBucketWithPartition,
+ false,
+ null,
+ true,
+ null);
+
+ Transformation<?> input =
dataStream.getTransformation().getInputs().get(0);
+ @SuppressWarnings("unchecked")
+ StreamPartitioner<Split> partitioner =
+ ((PartitionTransformation<Split>) input).getPartitioner();
+ partitioner.setup(4);
+
+ Split bucketOne = incrementalSplit(1);
+ Split bucketThree = incrementalSplit(3);
+
+ assertThat(selectChannel(partitioner, bucketOne))
+ .isEqualTo(selectChannel(partitioner, dataSplit(1)));
+ assertThat(selectChannel(partitioner, withQueryAuth(bucketOne)))
+ .isEqualTo(selectChannel(partitioner, bucketOne));
+ assertThat(selectChannel(partitioner, withQueryAuth(bucketThree)))
+ .isEqualTo(selectChannel(partitioner, bucketThree));
+
+ assertThat(selectChannel(partitioner, withQueryAuth(bucketOne)))
+ .isNotEqualTo(selectChannel(partitioner,
withQueryAuth(bucketThree)));
+ }
+
+ @ValueSource(booleans = {false, true})
+ @ParameterizedTest
+ public void
testOrderedShuffleRoutesFallbackWrappedQueryAuthSplitLikeTheSplitItWraps(
+ boolean shuffleBucketWithPartition) throws Exception {
+ Table table = createTable("ordered_nested_" +
shuffleBucketWithPartition, false, 2, true);
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+
+ DataStream<RowData> dataStream =
+ MonitorSource.buildSource(
+ env,
+ "source",
+ InternalTypeInfo.of(toLogicalType(table.rowType())),
+ table.newReadBuilder(),
+ 10,
+ false,
+ shuffleBucketWithPartition,
+ false,
+ null,
+ true,
+ null);
+
+ Transformation<?> input =
dataStream.getTransformation().getInputs().get(0);
+ assertThat(input).isInstanceOf(PartitionTransformation.class);
+ @SuppressWarnings("unchecked")
+ StreamPartitioner<Split> partitioner =
+ ((PartitionTransformation<Split>) input).getPartitioner();
+ partitioner.setup(4);
+
+ DataSplit bucketOne = dataSplit(1);
+ DataSplit bucketThree = dataSplit(3);
+ Split nestedOne = withFallback(withQueryAuth(bucketOne));
+ Split nestedThree = withFallback(withQueryAuth(bucketThree));
+
+ assertThat(nestedOne).isNotInstanceOf(DataSplit.class);
+ assertThat(selectChannel(partitioner, nestedOne))
+ .isEqualTo(selectChannel(partitioner, bucketOne));
+ assertThat(selectChannel(partitioner, nestedThree))
+ .isEqualTo(selectChannel(partitioner, bucketThree));
+ assertThat(selectChannel(partitioner, nestedOne))
+ .isNotEqualTo(selectChannel(partitioner, nestedThree));
+
+ Split nestedIncrementalOne =
withFallback(withQueryAuth(incrementalSplit(1)));
+ assertThat(selectChannel(partitioner, nestedIncrementalOne))
+ .isEqualTo(selectChannel(partitioner, bucketOne));
+ }
+
+ @ValueSource(booleans = {false, true})
+ @ParameterizedTest
+ public void testOrderedShuffleRoutesFallbackDataSplitLikeTheSplitItCopies(
+ boolean shuffleBucketWithPartition) throws Exception {
+ Table table = createTable("ordered_fallback_" +
shuffleBucketWithPartition, false, 2, true);
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+
+ DataStream<RowData> dataStream =
+ MonitorSource.buildSource(
+ env,
+ "source",
+ InternalTypeInfo.of(toLogicalType(table.rowType())),
+ table.newReadBuilder(),
+ 10,
+ false,
+ shuffleBucketWithPartition,
+ false,
+ null,
+ true,
+ null);
+
+ Transformation<?> input =
dataStream.getTransformation().getInputs().get(0);
+ @SuppressWarnings("unchecked")
+ StreamPartitioner<Split> partitioner =
+ ((PartitionTransformation<Split>) input).getPartitioner();
+ partitioner.setup(4);
+
+ DataSplit bucketOne = dataSplit(1);
+ DataSplit bucketThree = dataSplit(3);
+ Split fallbackOne = withFallback(bucketOne);
+ Split fallbackThree = withFallback(bucketThree);
+
+ assertThat(fallbackOne).isInstanceOf(DataSplit.class);
+ assertThat(selectChannel(partitioner, fallbackOne))
+ .isEqualTo(selectChannel(partitioner, bucketOne));
+ assertThat(selectChannel(partitioner, fallbackThree))
+ .isEqualTo(selectChannel(partitioner, bucketThree));
+ assertThat(selectChannel(partitioner, fallbackOne))
+ .isNotEqualTo(selectChannel(partitioner, fallbackThree));
+ }
+
+ private static int selectChannel(StreamPartitioner<Split> partitioner,
Split split) {
+ SerializationDelegate<StreamRecord<Split>> delegate = new
SerializationDelegate<>(null);
+ delegate.setInstance(new StreamRecord<>(split));
+ return partitioner.selectChannel(delegate);
+ }
+
+ private static DataSplit dataSplit(int bucket) {
+ return DataSplit.builder()
+ .withSnapshot(1)
+ .withPartition(BinaryRow.EMPTY_ROW)
+ .withBucket(bucket)
+ .withDataFiles(Collections.emptyList())
+ .isStreaming(true)
+ .withBucketPath("/temp/xxx")
+ .build();
+ }
+
+ private static Split incrementalSplit(int bucket) {
+ return new IncrementalSplit(
+ 1L,
+ BinaryRow.EMPTY_ROW,
+ bucket,
+ 1,
+ Collections.emptyList(),
+ null,
+ Collections.emptyList(),
+ null,
+ true);
+ }
+
+ private static Split withQueryAuth(Split split) {
+ return new QueryAuthSplit(split, new TableQueryAuthResult(null, null));
+ }
+
+ private static Split withFallback(Split split) {
+ return FallbackReadFileStoreTable.toFallbackSplit(split, true);
+ }
}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumeratorQueryAuthTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumeratorQueryAuthTest.java
new file mode 100644
index 0000000000..e05fe061f4
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumeratorQueryAuthTest.java
@@ -0,0 +1,145 @@
+/*
+ * 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.paimon.flink.source.align;
+
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.flink.source.FileSplitEnumeratorTestBase;
+import org.apache.paimon.flink.source.FileStoreSourceSplit;
+import org.apache.paimon.io.DataFileMeta;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
+import org.apache.paimon.table.source.Split;
+
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+import
org.apache.flink.connector.testutils.source.reader.TestingSplitEnumeratorContext;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.TreeMap;
+import java.util.UUID;
+import java.util.stream.Collectors;
+
+import static org.apache.paimon.io.DataFileTestUtils.row;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests {@link AlignedContinuousFileSplitEnumerator} over {@link
QueryAuthSplit}. */
+public class AlignedContinuousFileSplitEnumeratorQueryAuthTest
+ extends FileSplitEnumeratorTestBase<FileStoreSourceSplit> {
+
+ @Test
+ public void testWrappedSplitsGroupedBySnapshotLikeTheSplitsTheyWrap()
throws Exception {
+ List<FileStoreSourceSplit> splits =
+ Arrays.asList(snapshotSplit(2, 0), snapshotSplit(1, 1),
snapshotSplit(1, 0));
+
+ List<FileStoreSourceSplit> plain = assignFirstSnapshot(splits);
+ List<FileStoreSourceSplit> authorized =
assignFirstSnapshot(withQueryAuth(splits));
+
+ assertThat(underlying(plain)).containsExactly(splits.get(2).split(),
splits.get(1).split());
+ assertThat(underlying(authorized)).isEqualTo(underlying(plain));
+ }
+
+ @Test
+ public void testWrappedSplitIsAssignedStillWrapped() throws Exception {
+ List<FileStoreSourceSplit> wrapped =
+ withQueryAuth(Collections.singletonList(snapshotSplit(1, 0)));
+
+ List<FileStoreSourceSplit> assigned = assignFirstSnapshot(wrapped);
+
+ assertThat(assigned).hasSize(1);
+ assertThat(assigned.get(0)).isSameAs(wrapped.get(0));
+ assertThat(assigned.get(0).split()).isInstanceOf(QueryAuthSplit.class);
+ }
+
+ private List<FileStoreSourceSplit>
assignFirstSnapshot(List<FileStoreSourceSplit> splits) {
+ TestingSplitEnumeratorContext<FileStoreSourceSplit> context =
getSplitEnumeratorContext(2);
+ AlignedContinuousFileSplitEnumerator enumerator = enumerator(context);
+ enumerator.addSplits(splits);
+ enumerator.handleSplitRequest(0, "test-host");
+ enumerator.handleSplitRequest(1, "test-host");
+
+ List<FileStoreSourceSplit> assigned = new ArrayList<>();
+ new TreeMap<>(context.getSplitAssignments())
+ .values()
+ .forEach(state -> assigned.addAll(state.getAssignedSplits()));
+ return assigned;
+ }
+
+ private static AlignedContinuousFileSplitEnumerator enumerator(
+ SplitEnumeratorContext<FileStoreSourceSplit> context) {
+ return new AlignedContinuousFileSplitEnumerator(
+ context,
+ Collections.emptyList(),
+ null,
+ Long.MAX_VALUE,
+ null,
+ false,
+ 30000L,
+ 10,
+ false,
+ -1,
+ 10);
+ }
+
+ private FileStoreSourceSplit snapshotSplit(int snapshotId, int bucket) {
+ return createSnapshotSplit(snapshotId, bucket,
Collections.emptyList());
+ }
+
+ private static FileStoreSourceSplit withQueryAuth(FileStoreSourceSplit
split) {
+ return new FileStoreSourceSplit(
+ split.splitId(),
+ new QueryAuthSplit(split.split(), new
TableQueryAuthResult(null, null)),
+ split.recordsToSkip());
+ }
+
+ private static List<FileStoreSourceSplit>
withQueryAuth(List<FileStoreSourceSplit> splits) {
+ return splits.stream()
+
.map(AlignedContinuousFileSplitEnumeratorQueryAuthTest::withQueryAuth)
+ .collect(Collectors.toList());
+ }
+
+ private static List<Split> underlying(Collection<FileStoreSourceSplit>
splits) {
+ List<Split> underlying = new ArrayList<>();
+ for (FileStoreSourceSplit split : splits) {
+ Split inner = split.split();
+ underlying.add(
+ inner instanceof QueryAuthSplit ? ((QueryAuthSplit)
inner).split() : inner);
+ }
+ return underlying;
+ }
+
+ @Override
+ protected FileStoreSourceSplit createSnapshotSplit(
+ int snapshotId, int bucket, List<DataFileMeta> files, int...
partitions) {
+ return new FileStoreSourceSplit(
+ UUID.randomUUID().toString(),
+ DataSplit.builder()
+ .withSnapshot(snapshotId)
+ .withPartition(row(partitions))
+ .withBucket(bucket)
+ .withDataFiles(files)
+ .isStreaming(true)
+ .withBucketPath("/temp/xxx")
+ .build(),
+ 0);
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/AlignedSplitAssignerTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/AlignedSplitAssignerTest.java
new file mode 100644
index 0000000000..a327451fba
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/AlignedSplitAssignerTest.java
@@ -0,0 +1,245 @@
+/*
+ * 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.paimon.flink.source.assigners;
+
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.flink.source.FileStoreSourceSplit;
+import org.apache.paimon.flink.source.align.PlaceholderSplit;
+import org.apache.paimon.table.FallbackReadFileStoreTable;
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplitImpl;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
+import org.apache.paimon.table.source.Split;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+import java.util.stream.Collectors;
+
+import static org.apache.paimon.io.DataFileTestUtils.row;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests {@link AlignedSplitAssigner} over {@link QueryAuthSplit}. */
+public class AlignedSplitAssignerTest {
+
+ @Test
+ public void testWrappedSplitsAssignedLikeTheSplitsTheyWrap() {
+ List<FileStoreSourceSplit> splits =
+ Arrays.asList(snapshotSplit(1, 0), snapshotSplit(1, 1),
snapshotSplit(2, 0));
+
+ AlignedSplitAssigner plain = new AlignedSplitAssigner();
+ AlignedSplitAssigner authorized = new AlignedSplitAssigner();
+ for (int i = 0; i < splits.size(); i++) {
+ plain.addSplit(i % 2, splits.get(i));
+ authorized.addSplit(i % 2, withQueryAuth(splits.get(i)));
+ }
+
+
assertThat(authorized.remainingSnapshots()).isEqualTo(plain.remainingSnapshots());
+
assertThat(authorized.numberOfRemainingSplits()).isEqualTo(plain.numberOfRemainingSplits());
+
assertThat(authorized.getNextSnapshotId(0)).isEqualTo(plain.getNextSnapshotId(0));
+ assertThat(authorized.isAligned()).isEqualTo(plain.isAligned());
+ assertThat(underlying(authorized.remainingSplits()))
+ .isEqualTo(underlying(plain.remainingSplits()));
+
+ assertThat(underlying(authorized.getNext(0, null)))
+ .isEqualTo(underlying(plain.getNext(0, null)));
+ assertThat(underlying(authorized.getNext(1, null)))
+ .isEqualTo(underlying(plain.getNext(1, null)));
+ assertThat(authorized.isAligned()).isEqualTo(plain.isAligned());
+
assertThat(authorized.numberOfRemainingSplits()).isEqualTo(plain.numberOfRemainingSplits());
+ }
+
+ @Test
+ public void testWrappedSplitIsAssignedStillWrapped() {
+ FileStoreSourceSplit wrapped = withQueryAuth(snapshotSplit(1, 0));
+
+ AlignedSplitAssigner assigner = new AlignedSplitAssigner();
+ assigner.addSplit(0, wrapped);
+
+ assertThat(assigner.remainingSplits()).containsExactly(wrapped);
+ List<FileStoreSourceSplit> assigned = assigner.getNext(0, null);
+ assertThat(assigned).hasSize(1);
+ assertThat(assigned.get(0)).isSameAs(wrapped);
+ assertThat(assigned.get(0).split()).isInstanceOf(QueryAuthSplit.class);
+ }
+
+ @Test
+ public void testWrappedPlaceholderStillCountsAsAligned() {
+ FileStoreSourceSplit placeholder = placeholderSplit(1);
+
+ AlignedSplitAssigner plain = new AlignedSplitAssigner();
+ plain.addSplit(0, placeholder);
+ AlignedSplitAssigner authorized = new AlignedSplitAssigner();
+ authorized.addSplit(0, withQueryAuth(placeholder));
+
+ assertThat(plain.isAligned()).isTrue();
+ assertThat(authorized.isAligned()).isEqualTo(plain.isAligned());
+ assertThat(authorized.getNext(0, null)).isEqualTo(plain.getNext(0,
null));
+ }
+
+ @Test
+ public void testWrappedSplitsAddedBackIntoTheSnapshotAlreadyPending() {
+ FileStoreSourceSplit pending = snapshotSplit(4, 0);
+ List<FileStoreSourceSplit> back = Arrays.asList(snapshotSplit(4, 1),
snapshotSplit(4, 2));
+
+ AlignedSplitAssigner plain = new AlignedSplitAssigner();
+ plain.addSplit(0, pending);
+ plain.addSplitsBack(1, back);
+ AlignedSplitAssigner authorized = new AlignedSplitAssigner();
+ authorized.addSplit(0, withQueryAuth(pending));
+ authorized.addSplitsBack(1, withQueryAuth(back));
+
+ assertThat(plain.remainingSnapshots()).isEqualTo(1);
+
assertThat(authorized.remainingSnapshots()).isEqualTo(plain.remainingSnapshots());
+
assertThat(authorized.numberOfRemainingSplits()).isEqualTo(plain.numberOfRemainingSplits());
+ assertThat(underlying(authorized.getNext(1, null)))
+ .isEqualTo(underlying(plain.getNext(1, null)));
+ }
+
+ @Test
+ public void testWrappedSplitsFromDifferentSnapshotsAreStillRejected() {
+ List<FileStoreSourceSplit> mixed =
+ withQueryAuth(Arrays.asList(snapshotSplit(1, 0),
snapshotSplit(2, 0)));
+
+ assertThatThrownBy(() -> new AlignedSplitAssigner().addSplitsBack(0,
mixed))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("SnapshotId not equal");
+ }
+
+ @Test
+ public void testWrappedPlaceholderAddedBackStillCountsAsAligned() {
+ FileStoreSourceSplit placeholder = placeholderSplit(5);
+
+ AlignedSplitAssigner plain = new AlignedSplitAssigner();
+ plain.addSplitsBack(0, Collections.singletonList(placeholder));
+ AlignedSplitAssigner authorized = new AlignedSplitAssigner();
+ authorized.addSplitsBack(0,
Collections.singletonList(withQueryAuth(placeholder)));
+
+ assertThat(plain.isAligned()).isTrue();
+ assertThat(authorized.isAligned()).isEqualTo(plain.isAligned());
+ }
+
+ @Test
+ public void testFallbackWrappedSplitsAssignedLikeTheSplitsTheyWrap() {
+ List<FileStoreSourceSplit> splits =
+ Arrays.asList(snapshotSplit(1, 0), snapshotSplit(1, 1),
snapshotSplit(2, 0));
+
+ AlignedSplitAssigner plain = new AlignedSplitAssigner();
+ AlignedSplitAssigner nested = new AlignedSplitAssigner();
+ for (int i = 0; i < splits.size(); i++) {
+ plain.addSplit(i % 2, splits.get(i));
+ nested.addSplit(i % 2, withFallback(withQueryAuth(splits.get(i))));
+ }
+
+
assertThat(nested.remainingSnapshots()).isEqualTo(plain.remainingSnapshots());
+
assertThat(nested.numberOfRemainingSplits()).isEqualTo(plain.numberOfRemainingSplits());
+
assertThat(nested.getNextSnapshotId(0)).isEqualTo(plain.getNextSnapshotId(0));
+ assertThat(underlying(nested.remainingSplits()))
+ .isEqualTo(underlying(plain.remainingSplits()));
+
+ assertThat(underlying(nested.getNext(0, null)))
+ .isEqualTo(underlying(plain.getNext(0, null)));
+ assertThat(underlying(nested.getNext(1, null)))
+ .isEqualTo(underlying(plain.getNext(1, null)));
+ }
+
+ @Test
+ public void testFallbackWrappedSplitIsAssignedStillFallbackWrapped() {
+ FileStoreSourceSplit nested =
withFallback(withQueryAuth(snapshotSplit(1, 0)));
+
+ AlignedSplitAssigner assigner = new AlignedSplitAssigner();
+ assigner.addSplit(0, nested);
+
+ assertThat(assigner.remainingSplits()).containsExactly(nested);
+ List<FileStoreSourceSplit> assigned = assigner.getNext(0, null);
+ assertThat(assigned).hasSize(1);
+ assertThat(assigned.get(0)).isSameAs(nested);
+
assertThat(assigned.get(0).split()).isInstanceOf(FallbackSplitImpl.class);
+ }
+
+ @Test
+ public void testFallbackWrappedPlaceholderStillCountsAsAligned() {
+ FileStoreSourceSplit placeholder = placeholderSplit(1);
+
+ AlignedSplitAssigner plain = new AlignedSplitAssigner();
+ plain.addSplit(0, placeholder);
+ AlignedSplitAssigner nested = new AlignedSplitAssigner();
+ nested.addSplit(0, withFallback(withQueryAuth(placeholder)));
+
+ assertThat(plain.isAligned()).isTrue();
+ assertThat(nested.isAligned()).isEqualTo(plain.isAligned());
+ }
+
+ private static FileStoreSourceSplit snapshotSplit(long snapshotId, int
bucket) {
+ return new FileStoreSourceSplit(
+ UUID.randomUUID().toString(),
+ DataSplit.builder()
+ .withSnapshot(snapshotId)
+ .withPartition(row(1))
+ .withBucket(bucket)
+ .withDataFiles(Collections.emptyList())
+ .isStreaming(true)
+ .withBucketPath("/temp/xxx") // not used
+ .build());
+ }
+
+ private static FileStoreSourceSplit placeholderSplit(long snapshotId) {
+ return new FileStoreSourceSplit(
+ UUID.randomUUID().toString(), new
PlaceholderSplit(snapshotId));
+ }
+
+ private static FileStoreSourceSplit withQueryAuth(FileStoreSourceSplit
split) {
+ return new FileStoreSourceSplit(
+ split.splitId(),
+ new QueryAuthSplit(split.split(), new
TableQueryAuthResult(null, null)),
+ split.recordsToSkip());
+ }
+
+ private static List<FileStoreSourceSplit>
withQueryAuth(List<FileStoreSourceSplit> splits) {
+ return splits.stream()
+ .map(AlignedSplitAssignerTest::withQueryAuth)
+ .collect(Collectors.toList());
+ }
+
+ private static FileStoreSourceSplit withFallback(FileStoreSourceSplit
split) {
+ return new FileStoreSourceSplit(
+ split.splitId(),
+ FallbackReadFileStoreTable.toFallbackSplit(split.split(),
true),
+ split.recordsToSkip());
+ }
+
+ private static List<Split> underlying(Collection<FileStoreSourceSplit>
splits) {
+ List<Split> underlying = new ArrayList<>();
+ for (FileStoreSourceSplit split : splits) {
+ Split inner = split.split();
+ if (inner instanceof FallbackSplitImpl) {
+ inner = ((FallbackSplitImpl) inner).wrapped();
+ }
+ underlying.add(
+ inner instanceof QueryAuthSplit ? ((QueryAuthSplit)
inner).split() : inner);
+ }
+ return underlying;
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssignerTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssignerTest.java
new file mode 100644
index 0000000000..00ebc9b986
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssignerTest.java
@@ -0,0 +1,262 @@
+/*
+ * 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.paimon.flink.source.assigners;
+
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.flink.source.DynamicPartitionFilteringInfo;
+import org.apache.paimon.flink.source.FileStoreSourceSplit;
+import org.apache.paimon.predicate.NullTransform;
+import org.apache.paimon.table.FallbackReadFileStoreTable;
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplitImpl;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
+import org.apache.paimon.table.source.Split;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+import org.apache.paimon.utils.JsonSerdeUtil;
+
+import org.apache.flink.api.java.typeutils.GenericTypeInfo;
+import org.apache.flink.table.connector.source.DynamicFilteringData;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.IntType;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+import java.util.stream.Collectors;
+
+import static org.apache.paimon.io.DataFileTestUtils.row;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests {@link DynamicPartitionPruningAssigner} over {@link QueryAuthSplit}.
*/
+public class DynamicPartitionPruningAssignerTest {
+
+ @Test
+ public void testWrappedSplitsPrunedLikeTheSplitsTheyWrap() {
+ FileStoreSourceSplit kept = partitionSplit(1);
+ FileStoreSourceSplit pruned = partitionSplit(2);
+ List<FileStoreSourceSplit> splits = Arrays.asList(kept, pruned);
+
+ SplitAssigner plain = pruningAssigner(splits);
+ SplitAssigner authorized = pruningAssigner(withQueryAuth(splits));
+
+
assertThat(underlying(plain.remainingSplits())).containsExactly(kept.split());
+ assertThat(underlying(authorized.remainingSplits()))
+ .isEqualTo(underlying(plain.remainingSplits()));
+
+ List<FileStoreSourceSplit> authorizedNext = authorized.getNext(0,
null);
+
assertThat(underlying(authorizedNext)).isEqualTo(underlying(plain.getNext(0,
null)));
+ assertThat(underlying(authorizedNext)).containsExactly(kept.split());
+
+ assertThat(authorized.getNext(0, null)).isEqualTo(plain.getNext(0,
null)).isEmpty();
+ }
+
+ @Test
+ public void testWrappedSplitIsKeptStillWrapped() {
+ FileStoreSourceSplit wrapped = withQueryAuth(partitionSplit(1));
+
+ SplitAssigner assigner = pruningAssigner(Collections.emptyList());
+ assigner.addSplit(0, wrapped);
+
+ assertThat(assigner.remainingSplits()).containsExactly(wrapped);
+ List<FileStoreSourceSplit> assigned = assigner.getNext(0, null);
+ assertThat(assigned).hasSize(1);
+ assertThat(assigned.get(0)).isSameAs(wrapped);
+ assertThat(assigned.get(0).split()).isInstanceOf(QueryAuthSplit.class);
+ }
+
+ @Test
+ public void testFallbackWrappedSplitsPrunedLikeTheSplitsTheyWrap() {
+ FileStoreSourceSplit kept = partitionSplit(1);
+ FileStoreSourceSplit pruned = partitionSplit(2);
+ List<FileStoreSourceSplit> splits = Arrays.asList(kept, pruned);
+
+ SplitAssigner plain = pruningAssigner(splits);
+ SplitAssigner nested =
pruningAssigner(withFallback(withQueryAuth(splits)));
+
+ assertThat(underlying(nested.remainingSplits()))
+ .isEqualTo(underlying(plain.remainingSplits()));
+
+ List<FileStoreSourceSplit> nestedNext = nested.getNext(0, null);
+
assertThat(underlying(nestedNext)).isEqualTo(underlying(plain.getNext(0,
null)));
+ assertThat(underlying(nestedNext)).containsExactly(kept.split());
+
+ assertThat(nested.getNext(0, null)).isEqualTo(plain.getNext(0,
null)).isEmpty();
+ }
+
+ @Test
+ public void testFallbackWrappedSplitIsKeptStillFallbackWrapped() {
+ FileStoreSourceSplit nested =
withFallback(withQueryAuth(partitionSplit(1)));
+
+ SplitAssigner assigner = pruningAssigner(Collections.emptyList());
+ assigner.addSplit(0, nested);
+
+ assertThat(assigner.remainingSplits()).containsExactly(nested);
+ List<FileStoreSourceSplit> assigned = assigner.getNext(0, null);
+ assertThat(assigned).hasSize(1);
+ assertThat(assigned.get(0)).isSameAs(nested);
+
assertThat(assigned.get(0).split()).isInstanceOf(FallbackSplitImpl.class);
+ assertThat(((FallbackSplitImpl) assigned.get(0).split()).wrapped())
+ .isInstanceOf(QueryAuthSplit.class);
+ }
+
+ @Test
+ public void testSplitIsKeptWhenAFilteringFieldIsMasked() {
+ FileStoreSourceSplit masked = withColumnMasking(partitionSplit(2),
"f0");
+
+ SplitAssigner assigner =
pruningAssigner(Collections.singletonList(masked));
+
+ assertThat(assigner.remainingSplits()).containsExactly(masked);
+ assertThat(assigner.getNext(0, null)).containsExactly(masked);
+ }
+
+ @Test
+ public void testSplitIsPrunedWhenTheMaskedColumnIsNotAFilteringField() {
+ FileStoreSourceSplit masked = withColumnMasking(partitionSplit(2),
"other");
+
+ SplitAssigner assigner =
pruningAssigner(Collections.singletonList(masked));
+
+ assertThat(assigner.remainingSplits()).isEmpty();
+ assertThat(assigner.getNext(0, null)).isEmpty();
+ }
+
+ @Test
+ public void testMaskedFilteringFieldKeepsTheSplitThroughAFallbackWrapper()
{
+ FileStoreSourceSplit masked =
withFallback(withColumnMasking(partitionSplit(2), "f0"));
+
+ SplitAssigner assigner =
pruningAssigner(Collections.singletonList(masked));
+
+ assertThat(assigner.remainingSplits()).containsExactly(masked);
+ }
+
+ private static SplitAssigner
pruningAssigner(Collection<FileStoreSourceSplit> splits) {
+ DynamicPartitionFilteringInfo filteringInfo =
+ new DynamicPartitionFilteringInfo(
+ RowType.of(DataTypes.INT()),
Collections.singletonList("f0"));
+ DynamicFilteringData filteringData =
+ new MockDynamicFilteringData(
+ org.apache.flink.table.types.logical.RowType.of(new
IntType()),
+ new RowData[] {GenericRowData.of(1)});
+ return new DynamicPartitionPruningAssigner(
+ new FIFOSplitAssigner(splits), filteringInfo, filteringData);
+ }
+
+ private static FileStoreSourceSplit partitionSplit(int partition) {
+ return new FileStoreSourceSplit(
+ UUID.randomUUID().toString(),
+ DataSplit.builder()
+ .withSnapshot(1)
+ .withPartition(row(partition))
+ .withBucket(0)
+ .withDataFiles(Collections.emptyList())
+ .isStreaming(true)
+ .withBucketPath("/temp/xxx") // not used
+ .build());
+ }
+
+ private static FileStoreSourceSplit withQueryAuth(FileStoreSourceSplit
split) {
+ return new FileStoreSourceSplit(
+ split.splitId(),
+ new QueryAuthSplit(split.split(), new
TableQueryAuthResult(null, null)),
+ split.recordsToSkip());
+ }
+
+ private static FileStoreSourceSplit withColumnMasking(
+ FileStoreSourceSplit split, String column) {
+ return new FileStoreSourceSplit(
+ split.splitId(),
+ new QueryAuthSplit(
+ split.split(),
+ new TableQueryAuthResult(
+ null,
+ Collections.singletonMap(
+ column,
JsonSerdeUtil.toJson(NullTransform.INSTANCE)))),
+ split.recordsToSkip());
+ }
+
+ private static List<FileStoreSourceSplit>
withQueryAuth(List<FileStoreSourceSplit> splits) {
+ return splits.stream()
+ .map(DynamicPartitionPruningAssignerTest::withQueryAuth)
+ .collect(Collectors.toList());
+ }
+
+ private static FileStoreSourceSplit withFallback(FileStoreSourceSplit
split) {
+ return new FileStoreSourceSplit(
+ split.splitId(),
+ FallbackReadFileStoreTable.toFallbackSplit(split.split(),
true),
+ split.recordsToSkip());
+ }
+
+ private static List<FileStoreSourceSplit>
withFallback(List<FileStoreSourceSplit> splits) {
+ return splits.stream()
+ .map(DynamicPartitionPruningAssignerTest::withFallback)
+ .collect(Collectors.toList());
+ }
+
+ private static List<Split> underlying(Collection<FileStoreSourceSplit>
splits) {
+ List<Split> underlying = new ArrayList<>();
+ for (FileStoreSourceSplit split : splits) {
+ Split inner = split.split();
+ if (inner instanceof FallbackSplitImpl) {
+ inner = ((FallbackSplitImpl) inner).wrapped();
+ }
+ underlying.add(
+ inner instanceof QueryAuthSplit ? ((QueryAuthSplit)
inner).split() : inner);
+ }
+ return underlying;
+ }
+
+ private static class MockDynamicFilteringData extends DynamicFilteringData
{
+
+ private final org.apache.flink.table.types.logical.RowType rowType;
+ private final RowData[] neededPartitions;
+
+ public MockDynamicFilteringData(
+ org.apache.flink.table.types.logical.RowType rowType,
RowData[] neededPartitions) {
+ super(new GenericTypeInfo<>(RowData.class), rowType,
Collections.emptyList(), true);
+ this.rowType = rowType;
+ this.neededPartitions = neededPartitions;
+ }
+
+ public boolean contains(RowData row) {
+ checkArgument(rowType.getFieldCount() == row.getArity());
+ for (RowData mayMatch : neededPartitions) {
+ if (matchRow(row, mayMatch)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ private boolean matchRow(RowData row, RowData mayMatch) {
+ for (int i = 0; i < rowType.getFieldCount(); i++) {
+ if (row.getInt(i) != mayMatch.getInt(i)) {
+ return false;
+ }
+ }
+ return true;
+ }
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssignerTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssignerTest.java
new file mode 100644
index 0000000000..ea1e54873b
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssignerTest.java
@@ -0,0 +1,219 @@
+/*
+ * 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.paimon.flink.source.assigners;
+
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.flink.source.DynamicPartitionFilteringInfo;
+import org.apache.paimon.flink.source.FileStoreSourceSplit;
+import org.apache.paimon.predicate.NullTransform;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
+import org.apache.paimon.table.source.Split;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+import org.apache.paimon.utils.JsonSerdeUtil;
+
+import org.apache.flink.api.java.typeutils.GenericTypeInfo;
+import org.apache.flink.table.connector.source.DynamicFilteringData;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.IntType;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+import java.util.stream.Collectors;
+
+import static org.apache.paimon.io.DataFileTestUtils.row;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests {@link PreAssignSplitAssigner} pruning over {@link QueryAuthSplit}.
*/
+public class PreAssignSplitAssignerTest {
+
+ @Test
+ public void testWrappedSplitsPrunedLikeTheSplitsTheyWrap() {
+ FileStoreSourceSplit keptOne = partitionSplit(1);
+ FileStoreSourceSplit pruned = partitionSplit(2);
+ FileStoreSourceSplit keptThree = partitionSplit(3);
+ List<FileStoreSourceSplit> splits = Arrays.asList(keptOne, pruned,
keptThree);
+
+ SplitAssigner plain = pruningAssigner(splits);
+ SplitAssigner authorized = pruningAssigner(withQueryAuth(splits));
+
+ assertThat(underlying(plain.remainingSplits()))
+ .containsExactly(keptOne.split(), keptThree.split());
+ assertThat(underlying(authorized.remainingSplits()))
+ .isEqualTo(underlying(plain.remainingSplits()));
+
assertThat(authorized.numberOfRemainingSplits()).isEqualTo(plain.numberOfRemainingSplits());
+
+ List<FileStoreSourceSplit> authorizedNext = authorized.getNext(0,
null);
+
assertThat(underlying(authorizedNext)).isEqualTo(underlying(plain.getNext(0,
null)));
+
assertThat(underlying(authorizedNext)).containsExactly(keptOne.split(),
keptThree.split());
+
+ assertThat(authorized.getNext(0, null)).isEqualTo(plain.getNext(0,
null)).isEmpty();
+
assertThat(authorized.numberOfRemainingSplits()).isEqualTo(plain.numberOfRemainingSplits());
+ }
+
+ @Test
+ public void testWrappedSplitIsKeptStillWrapped() {
+ FileStoreSourceSplit wrapped = withQueryAuth(partitionSplit(1));
+
+ SplitAssigner assigner =
pruningAssigner(Collections.singletonList(wrapped));
+
+ assertThat(assigner.remainingSplits()).containsExactly(wrapped);
+ List<FileStoreSourceSplit> assigned = assigner.getNext(0, null);
+ assertThat(assigned).hasSize(1);
+ assertThat(assigned.get(0)).isSameAs(wrapped);
+ assertThat(assigned.get(0).split()).isInstanceOf(QueryAuthSplit.class);
+ }
+
+ @Test
+ public void testWrappedSplitOutsideTheFilterIsNotKept() {
+ FileStoreSourceSplit pruned = partitionSplit(2);
+
+ SplitAssigner plain =
pruningAssigner(Collections.singletonList(pruned));
+ SplitAssigner authorized =
+
pruningAssigner(Collections.singletonList(withQueryAuth(pruned)));
+
+ assertThat(plain.remainingSplits()).isEmpty();
+ assertThat(authorized.remainingSplits()).isEmpty();
+
assertThat(authorized.numberOfRemainingSplits()).isEqualTo(plain.numberOfRemainingSplits());
+ assertThat(authorized.getNext(0, null)).isEqualTo(plain.getNext(0,
null)).isEmpty();
+ }
+
+ @Test
+ public void testSplitIsKeptWhenAFilteringFieldIsMasked() {
+ FileStoreSourceSplit masked = withColumnMasking(partitionSplit(2),
"f0");
+
+ SplitAssigner assigner =
pruningAssigner(Collections.singletonList(masked));
+
+ assertThat(assigner.remainingSplits()).containsExactly(masked);
+ assertThat(assigner.numberOfRemainingSplits()).isEqualTo(1);
+ assertThat(assigner.getNext(0, null)).containsExactly(masked);
+ }
+
+ @Test
+ public void testSplitIsPrunedWhenTheMaskedColumnIsNotAFilteringField() {
+ FileStoreSourceSplit masked = withColumnMasking(partitionSplit(2),
"other");
+
+ SplitAssigner assigner =
pruningAssigner(Collections.singletonList(masked));
+
+ assertThat(assigner.remainingSplits()).isEmpty();
+ assertThat(assigner.numberOfRemainingSplits()).isZero();
+ assertThat(assigner.getNext(0, null)).isEmpty();
+ }
+
+ private static SplitAssigner
pruningAssigner(Collection<FileStoreSourceSplit> splits) {
+ DynamicPartitionFilteringInfo filteringInfo =
+ new DynamicPartitionFilteringInfo(
+ RowType.of(DataTypes.INT()),
Collections.singletonList("f0"));
+ DynamicFilteringData filteringData =
+ new MockDynamicFilteringData(
+ org.apache.flink.table.types.logical.RowType.of(new
IntType()),
+ new RowData[] {GenericRowData.of(1),
GenericRowData.of(3)});
+ return new PreAssignSplitAssigner(10, 1, splits)
+ .ofDynamicPartitionPruning(filteringInfo, filteringData);
+ }
+
+ private static FileStoreSourceSplit partitionSplit(int partition) {
+ return new FileStoreSourceSplit(
+ UUID.randomUUID().toString(),
+ DataSplit.builder()
+ .withSnapshot(1)
+ .withPartition(row(partition))
+ .withBucket(0)
+ .withDataFiles(Collections.emptyList())
+ .isStreaming(false)
+ .withBucketPath("/temp/xxx")
+ .build());
+ }
+
+ private static FileStoreSourceSplit withQueryAuth(FileStoreSourceSplit
split) {
+ return new FileStoreSourceSplit(
+ split.splitId(),
+ new QueryAuthSplit(split.split(), new
TableQueryAuthResult(null, null)),
+ split.recordsToSkip());
+ }
+
+ private static FileStoreSourceSplit withColumnMasking(
+ FileStoreSourceSplit split, String column) {
+ return new FileStoreSourceSplit(
+ split.splitId(),
+ new QueryAuthSplit(
+ split.split(),
+ new TableQueryAuthResult(
+ null,
+ Collections.singletonMap(
+ column,
JsonSerdeUtil.toJson(NullTransform.INSTANCE)))),
+ split.recordsToSkip());
+ }
+
+ private static List<FileStoreSourceSplit>
withQueryAuth(List<FileStoreSourceSplit> splits) {
+ return splits.stream()
+ .map(PreAssignSplitAssignerTest::withQueryAuth)
+ .collect(Collectors.toList());
+ }
+
+ private static List<Split> underlying(Collection<FileStoreSourceSplit>
splits) {
+ List<Split> underlying = new ArrayList<>();
+ for (FileStoreSourceSplit split : splits) {
+ Split inner = split.split();
+ underlying.add(
+ inner instanceof QueryAuthSplit ? ((QueryAuthSplit)
inner).split() : inner);
+ }
+ return underlying;
+ }
+
+ private static class MockDynamicFilteringData extends DynamicFilteringData
{
+
+ private final org.apache.flink.table.types.logical.RowType rowType;
+ private final RowData[] neededPartitions;
+
+ public MockDynamicFilteringData(
+ org.apache.flink.table.types.logical.RowType rowType,
RowData[] neededPartitions) {
+ super(new GenericTypeInfo<>(RowData.class), rowType,
Collections.emptyList(), true);
+ this.rowType = rowType;
+ this.neededPartitions = neededPartitions;
+ }
+
+ public boolean contains(RowData row) {
+ checkArgument(rowType.getFieldCount() == row.getArity());
+ for (RowData mayMatch : neededPartitions) {
+ if (matchRow(row, mayMatch)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ private boolean matchRow(RowData row, RowData mayMatch) {
+ for (int i = 0; i < rowType.getFieldCount(); i++) {
+ if (row.getInt(i) != mayMatch.getInt(i)) {
+ return false;
+ }
+ }
+ return true;
+ }
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/operator/ReadOperatorSplitShapeTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/operator/ReadOperatorSplitShapeTest.java
new file mode 100644
index 0000000000..b2c6574748
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/operator/ReadOperatorSplitShapeTest.java
@@ -0,0 +1,230 @@
+/*
+ * 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.paimon.flink.source.operator;
+
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.GenericRow;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.disk.IOManager;
+import org.apache.paimon.flink.utils.TestingMetricUtils;
+import org.apache.paimon.metrics.MetricRegistry;
+import org.apache.paimon.reader.RecordReader;
+import org.apache.paimon.table.FallbackReadFileStoreTable;
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplitImpl;
+import org.apache.paimon.table.source.ChainSplit;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
+import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.TableRead;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowType;
+import org.apache.paimon.utils.IteratorRecordReader;
+
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.runtime.typeutils.InternalSerializers;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import static org.apache.paimon.flink.LogicalTypeConversion.toLogicalType;
+import static org.apache.paimon.io.DataFileTestUtils.newFile;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests {@link ReadOperator} over every shape of {@link Split} it can be
handed. */
+public class ReadOperatorSplitShapeTest {
+
+ private static final RowType READ_TYPE = RowType.of(DataTypes.INT());
+
+ private static final int VALUE = 42;
+
+ @Test
+ public void testWrappedDataSplitReadsTheRowsOfTheSplitItWraps() throws
Exception {
+ Split split = dataSplit();
+
+ ReadResult plain = read(split);
+ ReadResult authorized = read(withQueryAuth(split));
+
+ assertThat(plain.values).containsExactly(VALUE);
+ assertThat(authorized.values).isEqualTo(plain.values);
+ }
+
+ @Test
+ public void testWrappedDataSplitIsHandedToTheReadStillWrapped() throws
Exception {
+ Split wrapped = withQueryAuth(dataSplit());
+
+ ReadResult authorized = read(wrapped);
+
+ assertThat(authorized.readSplit).isSameAs(wrapped);
+ assertThat(authorized.readSplit).isInstanceOf(QueryAuthSplit.class);
+ }
+
+ @Test
+ public void testSplitThatIsNotADataSplitIsReadWrappedOrNot() throws
Exception {
+ ReadResult plain = read(chainSplit());
+ ReadResult authorized = read(withQueryAuth(chainSplit()));
+
+ assertThat(plain.values).containsExactly(VALUE);
+ assertThat(authorized.values).isEqualTo(plain.values);
+ }
+
+ @Test
+ public void
testFallbackWrappedQueryAuthSplitReadsTheRowsOfTheSplitItWraps() throws
Exception {
+ Split split = dataSplit();
+
+ ReadResult plain = read(split);
+ ReadResult nested = read(withFallback(withQueryAuth(split)));
+
+ assertThat(plain.values).containsExactly(VALUE);
+ assertThat(nested.values).isEqualTo(plain.values);
+ }
+
+ @Test
+ public void
testFallbackWrappedQueryAuthSplitIsHandedToTheReadStillWrapped() throws
Exception {
+ Split nested = withFallback(withQueryAuth(dataSplit()));
+
+ ReadResult result = read(nested);
+
+ assertThat(result.readSplit).isSameAs(nested);
+ assertThat(result.readSplit).isInstanceOf(FallbackSplitImpl.class);
+ }
+
+ @Test
+ public void
testFallbackWrappedQueryAuthSplitKeepsTheEventTimeOfTheSplitItWraps()
+ throws Exception {
+ Split split = dataSplit();
+
+ ReadResult plain = read(split);
+ ReadResult authorized = read(withQueryAuth(split));
+ ReadResult nested = read(withFallback(withQueryAuth(split)));
+ ReadResult fallback = read(withFallback(split));
+
+ assertThat(plain.fetchEventTimeLag).isGreaterThan(0L);
+ assertThat(authorized.fetchEventTimeLag).isGreaterThan(0L);
+ assertThat(nested.fetchEventTimeLag).isGreaterThan(0L);
+ assertThat(fallback.fetchEventTimeLag).isGreaterThan(0L);
+ }
+
+ @Test
+ public void testSplitWithoutAnEventTimeLeavesTheMetricUndefined() throws
Exception {
+ assertThat(read(chainSplit()).fetchEventTimeLag).isEqualTo(-1L);
+
assertThat(read(withFallback(withQueryAuth(chainSplit()))).fetchEventTimeLag)
+ .isEqualTo(-1L);
+ }
+
+ private static ReadResult read(Split split) throws Exception {
+ RecordingTableRead tableRead = new RecordingTableRead();
+ ReadOperator operator = new ReadOperator(() -> tableRead, null, null,
READ_TYPE, false);
+
+ OneInputStreamOperatorTestHarness<Split, RowData> harness =
+ new OneInputStreamOperatorTestHarness<>(operator);
+ harness.setup(InternalSerializers.create(toLogicalType(READ_TYPE)));
+ harness.open();
+ try {
+ harness.processElement(new StreamRecord<>(split));
+
+ List<Integer> values = new ArrayList<>();
+ for (Object record : harness.getOutput()) {
+ @SuppressWarnings("unchecked")
+ StreamRecord<RowData> streamRecord = (StreamRecord<RowData>)
record;
+ values.add(streamRecord.getValue().getInt(0));
+ }
+ long fetchEventTimeLag =
+ (Long)
+ TestingMetricUtils.getGauge(
+ operator.getMetricGroup(),
"currentFetchEventTimeLag")
+ .getValue();
+ return new ReadResult(values, tableRead.split, fetchEventTimeLag);
+ } finally {
+ harness.close();
+ }
+ }
+
+ private static Split dataSplit() {
+ return DataSplit.builder()
+ .withSnapshot(1)
+ .withPartition(BinaryRow.EMPTY_ROW)
+ .withBucket(0)
+ .withBucketPath("bucket-0")
+ .withDataFiles(Collections.singletonList(newFile(0L, 1L)))
+ .build();
+ }
+
+ private static Split chainSplit() {
+ return new ChainSplit(
+ BinaryRow.EMPTY_ROW,
+ Collections.singletonList(newFile(0L, 1L)),
+ Collections.emptyMap(),
+ Collections.emptyMap(),
+ null);
+ }
+
+ private static Split withQueryAuth(Split split) {
+ return new QueryAuthSplit(split, new TableQueryAuthResult(null, null));
+ }
+
+ private static Split withFallback(Split split) {
+ return FallbackReadFileStoreTable.toFallbackSplit(split, true);
+ }
+
+ private static class ReadResult {
+
+ private final List<Integer> values;
+ private final Split readSplit;
+ private final long fetchEventTimeLag;
+
+ private ReadResult(List<Integer> values, Split readSplit, long
fetchEventTimeLag) {
+ this.values = values;
+ this.readSplit = readSplit;
+ this.fetchEventTimeLag = fetchEventTimeLag;
+ }
+ }
+
+ private static class RecordingTableRead implements TableRead {
+
+ private Split split;
+
+ @Override
+ public TableRead withMetricRegistry(MetricRegistry registry) {
+ return this;
+ }
+
+ @Override
+ public TableRead executeFilter() {
+ return this;
+ }
+
+ @Override
+ public TableRead withIOManager(IOManager ioManager) {
+ return this;
+ }
+
+ @Override
+ public RecordReader<InternalRow> createReader(Split split) throws
IOException {
+ this.split = split;
+ InternalRow row = GenericRow.of(VALUE);
+ return new
IteratorRecordReader<>(Collections.singleton(row).iterator());
+ }
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/utils/TableScanUtilsTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/utils/TableScanUtilsTest.java
new file mode 100644
index 0000000000..c32d531263
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/utils/TableScanUtilsTest.java
@@ -0,0 +1,116 @@
+/*
+ * 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.paimon.flink.utils;
+
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.flink.source.FileStoreSourceSplit;
+import org.apache.paimon.table.FallbackReadFileStoreTable;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
+import org.apache.paimon.table.source.Split;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.Optional;
+import java.util.OptionalLong;
+import java.util.UUID;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link TableScanUtils}. */
+public class TableScanUtilsTest {
+
+ @Test
+ public void testSnapshotIdReadThroughTheAuthorizationWrapper() {
+ FileStoreSourceSplit plain = sourceSplit(dataSplit(7L));
+ FileStoreSourceSplit wrapped = sourceSplit(new
QueryAuthSplit(plain.split(), authResult()));
+
+ assertThat(TableScanUtils.getSnapshotId(plain)).hasValue(7L);
+ assertThat(TableScanUtils.getSnapshotId(wrapped))
+ .isEqualTo(TableScanUtils.getSnapshotId(plain));
+ }
+
+ @Test
+ public void testSnapshotIdOfASplitWithoutOneIsEmpty() {
+ Split noSnapshotId =
+ new Split() {
+ @Override
+ public long rowCount() {
+ return 0;
+ }
+
+ @Override
+ public OptionalLong mergedRowCount() {
+ return OptionalLong.empty();
+ }
+ };
+
+ assertThat(TableScanUtils.getSnapshotId(sourceSplit(noSnapshotId)))
+ .isEqualTo(Optional.empty());
+ assertThat(
+ TableScanUtils.getSnapshotId(
+ sourceSplit(new QueryAuthSplit(noSnapshotId,
authResult()))))
+ .isEqualTo(Optional.empty());
+ }
+
+ @Test
+ public void testSnapshotIdReadThroughTheFallbackAndAuthorizationWrappers()
{
+ FileStoreSourceSplit plain = sourceSplit(dataSplit(7L));
+ FileStoreSourceSplit nested =
+ sourceSplit(withFallback(new QueryAuthSplit(plain.split(),
authResult())));
+
+ assertThat(nested.split()).isNotInstanceOf(DataSplit.class);
+ assertThat(TableScanUtils.getSnapshotId(nested))
+ .isEqualTo(TableScanUtils.getSnapshotId(plain))
+ .hasValue(7L);
+ }
+
+ @Test
+ public void testSnapshotIdOfAFallbackDataSplitIsTheSnapshotItCopies() {
+ FileStoreSourceSplit fallback =
sourceSplit(withFallback(dataSplit(7L)));
+
+ assertThat(fallback.split()).isInstanceOf(DataSplit.class);
+ assertThat(TableScanUtils.getSnapshotId(fallback)).hasValue(7L);
+ }
+
+ private static TableQueryAuthResult authResult() {
+ return new TableQueryAuthResult(null, null);
+ }
+
+ private static Split withFallback(Split split) {
+ return FallbackReadFileStoreTable.toFallbackSplit(split, true);
+ }
+
+ private static DataSplit dataSplit(long snapshotId) {
+ return DataSplit.builder()
+ .withSnapshot(snapshotId)
+ .withPartition(BinaryRow.EMPTY_ROW)
+ .withBucket(0)
+ .withDataFiles(Collections.emptyList())
+ .isStreaming(true)
+ .withBucketPath("/temp/xxx")
+ .build();
+ }
+
+ private static FileStoreSourceSplit sourceSplit(Split split) {
+ return new FileStoreSourceSplit(UUID.randomUUID().toString(), split,
0);
+ }
+}
