From 68229cf9522e93ebc81f99f9df17681ca0977919 Mon Sep 17 00:00:00 2001 From: umi Date: Fri, 17 Jul 2026 12:06:49 +0800 Subject: [PATCH 01/31] [core] Handle query auth wrappers in split consumers --- .../paimon/crosspartition/IndexBootstrap.java | 10 ++- .../table/FallbackReadFileStoreTable.java | 13 ++-- .../paimon/table/source/QueryAuthSplit.java | 5 ++ .../crosspartition/IndexBootstrapTest.java | 15 +++++ .../table/FallbackReadFileStoreTableTest.java | 61 +++++++++++++++++++ 5 files changed, 95 insertions(+), 9 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java b/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java index 24166243b231..4c906ea80791 100644 --- a/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java +++ b/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java @@ -29,6 +29,7 @@ import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.DataTableScan; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.ReadBuilder; import org.apache.paimon.table.source.Split; import org.apache.paimon.types.DataField; @@ -119,7 +120,7 @@ public RecordReader bootstrap(int numAssigners, int assignId) throw options.pageSize(), options.crossPartitionUpsertBootstrapParallelism(), split -> { - DataSplit dataSplit = ((DataSplit) split); + DataSplit dataSplit = unwrapDataSplit(split); int bucket = dataSplit.bucket(); return partBucketConverter.toGenericRow( new JoinedRow(dataSplit.partition(), GenericRow.of(bucket))); @@ -129,7 +130,7 @@ public RecordReader bootstrap(int numAssigners, int assignId) throw @VisibleForTesting static boolean filterSplit(Split split, long indexTtl, long currentTime) { - List files = ((DataSplit) split).dataFiles(); + List files = unwrapDataSplit(split).dataFiles(); for (DataFileMeta file : files) { long fileTime = file.creationTimeEpochMillis(); if (currentTime <= fileTime + indexTtl) { @@ -139,6 +140,11 @@ static boolean filterSplit(Split split, long indexTtl, long currentTime) { return false; } + @VisibleForTesting + static DataSplit unwrapDataSplit(Split split) { + return (DataSplit) QueryAuthSplit.unwrap(split); + } + public static RowType bootstrapType(TableSchema schema) { List primaryKeys = schema.trimmedPrimaryKeys(); List partitionKeys = schema.partitionKeys(); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/FallbackReadFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/FallbackReadFileStoreTable.java index bc84bca88afe..ec559807f3da 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/FallbackReadFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/FallbackReadFileStoreTable.java @@ -564,8 +564,7 @@ public TableScan.Plan plan() { Set completePartitions = new HashSet<>(newPartitionListingScan(true).listPartitions()); for (Split split : mainScan.plan().splits()) { - DataSplit dataSplit = (DataSplit) split; - splits.add(toFallbackSplit(dataSplit, false)); + splits.add(toFallbackSplit(split, false)); } List remainingPartitions = @@ -691,9 +690,10 @@ public TableRead withIOManager(IOManager ioManager) { public RecordReader createReader(Split split) throws IOException { if (split instanceof FallbackSplit) { FallbackSplit fallbackSplit = (FallbackSplit) split; + Split wrappedSplit = fallbackSplit.wrapped(); if (fallbackSplit.isFallback()) { try { - return fallbackRead.createReader(fallbackSplit.wrapped()); + return fallbackRead.createReader(wrappedSplit); } catch (Exception e) { if (fallbackReadFailFast) { if (e instanceof IOException) { @@ -703,16 +703,15 @@ public RecordReader createReader(Split split) throws IOException { throw (RuntimeException) e; } throw new IOException( - "Failed to read fallback branch split: " - + fallbackSplit.wrapped(), - e); + "Failed to read fallback branch split: " + wrappedSplit, e); } LOG.error( "Reading from supplemental branch has problems: {}", - fallbackSplit.wrapped(), + wrappedSplit, e); } } + return mainRead.createReader(wrappedSplit); } return mainRead.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 84646181e11f..6bd03379f714 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 @@ -49,6 +49,11 @@ public Split split() { return split; } + /** Unwraps one query authorization layer to expose the wrapped split. */ + 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/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java b/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java index 526b6bd29f9f..3a8b897ec53e 100644 --- a/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java @@ -20,6 +20,7 @@ import org.apache.paimon.CoreOptions; import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; import org.apache.paimon.data.Timestamp; @@ -31,6 +32,7 @@ import org.apache.paimon.table.Table; import org.apache.paimon.table.TableTestBase; import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.Pair; @@ -47,6 +49,7 @@ import static org.apache.paimon.crosspartition.IndexBootstrap.BUCKET_FIELD; import static org.apache.paimon.crosspartition.IndexBootstrap.filterSplit; +import static org.apache.paimon.crosspartition.IndexBootstrap.unwrapDataSplit; import static org.apache.paimon.data.BinaryRow.EMPTY_ROW; import static org.apache.paimon.stats.SimpleStats.EMPTY_STATS; import static org.assertj.core.api.Assertions.assertThat; @@ -129,6 +132,18 @@ public void testFilterSplit() { assertThat(filterSplit(newSplit(newFile(100), newFile(200)), 200, 230)).isTrue(); } + @Test + public void testQueryAuthSplit() { + DataSplit dataSplit = newSplit(newFile(100), newFile(200)); + QueryAuthSplit authSplit = + new QueryAuthSplit( + dataSplit, new TableQueryAuthResult(Collections.emptyList(), null)); + + assertThat(unwrapDataSplit(authSplit)).isSameAs(dataSplit); + assertThat(filterSplit(authSplit, 50, 230)).isTrue(); + assertThat(filterSplit(authSplit, 50, 300)).isFalse(); + } + private DataSplit newSplit(DataFileMeta... files) { return DataSplit.builder() .withSnapshot(1) diff --git a/paimon-core/src/test/java/org/apache/paimon/table/FallbackReadFileStoreTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/FallbackReadFileStoreTableTest.java index e0841df4d0e4..c24fc71f8939 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/FallbackReadFileStoreTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/FallbackReadFileStoreTableTest.java @@ -20,6 +20,7 @@ import org.apache.paimon.CoreOptions; import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.BinaryRow; import org.apache.paimon.data.InternalRow; import org.apache.paimon.fs.FileIO; @@ -29,6 +30,8 @@ import org.apache.paimon.manifest.PartitionEntry; import org.apache.paimon.options.Options; import org.apache.paimon.partition.PartitionPredicate; +import org.apache.paimon.predicate.FieldRef; +import org.apache.paimon.predicate.FieldTransform; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.reader.RecordReader; @@ -40,11 +43,14 @@ import org.apache.paimon.table.sink.StreamTableWrite; 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.TableRead; import org.apache.paimon.types.DataType; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.InstantiationUtil; +import org.apache.paimon.utils.JsonSerdeUtil; import org.apache.paimon.utils.Pair; import org.apache.paimon.utils.TraceableFileIO; @@ -53,6 +59,7 @@ import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.AdditionalAnswers; import org.mockito.Mockito; import java.io.IOException; @@ -210,6 +217,60 @@ public void testPlanWithDataFilter(boolean wrappedFirst) throws Exception { .isTrue(); } + @Test + public void testPlanAndReadWithQueryAuthSplit() throws Exception { + FileStoreTable mainTable = createTable(); + writeDataIntoTable(mainTable, 0, rowData(1, 10)); + + mainTable.createBranch("bc"); + FileStoreTable branchTable = createTableFromBranch(mainTable, "bc"); + writeDataIntoTable(branchTable, 0, rowData(2, 20)); + + FallbackReadFileStoreTable table = + new FallbackReadFileStoreTable(mainTable, branchTable, true); + TableQueryAuthResult authResult = + new TableQueryAuthResult( + null, + Collections.singletonMap( + "a", + JsonSerdeUtil.toFlatJson( + new FieldTransform( + new FieldRef(0, "pt", DataTypes.INT()))))); + DataTableScan scan = + table.newFallbackScan( + fileStoreTable -> queryAuthScan(fileStoreTable.newScan(), authResult)); + + List splits = scan.plan().splits(); + assertThat(splits).hasSize(2); + assertThat(splits) + .allSatisfy( + split -> { + assertThat(split) + .isInstanceOf(FallbackReadFileStoreTable.FallbackSplit.class); + Split wrapped = + ((FallbackReadFileStoreTable.FallbackSplit) split).wrapped(); + assertThat(wrapped).isInstanceOf(QueryAuthSplit.class); + }); + + List> result = new ArrayList<>(); + for (Split split : splits) { + Split deserialized = + InstantiationUtil.deserializeObject( + InstantiationUtil.serializeObject(split), getClass().getClassLoader()); + RecordReader reader = table.newRead().createReader(deserialized); + reader.forEachRemaining(r -> result.add(Pair.of(r.getInt(0), r.getInt(1)))); + reader.close(); + } + assertThat(result).containsExactlyInAnyOrder(Pair.of(1, 1), Pair.of(2, 2)); + } + + private DataTableScan queryAuthScan(DataTableScan delegate, TableQueryAuthResult authResult) { + DataTableScan scan = + Mockito.mock(DataTableScan.class, AdditionalAnswers.delegatesTo(delegate)); + Mockito.doAnswer(ignored -> authResult.convertPlan(delegate.plan())).when(scan).plan(); + return scan; + } + @ParameterizedTest @ValueSource(booleans = {true, false}) public void testWriteGoesToWrapped(boolean wrappedFirst) throws Exception { From 37bf952c795ca225b75bd82318549a01150baa62 Mon Sep 17 00:00:00 2001 From: umi Date: Fri, 17 Jul 2026 17:35:53 +0800 Subject: [PATCH 02/31] [core] Bypass query auth for index bootstrap --- .../paimon/crosspartition/IndexBootstrap.java | 15 +- .../crosspartition/IndexBootstrapTest.java | 186 +++++++++++++++++- 2 files changed, 184 insertions(+), 17 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java b/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java index 4c906ea80791..8f9d99edd6dd 100644 --- a/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java +++ b/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java @@ -24,12 +24,12 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.data.JoinedRow; import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.options.Options; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.DataTableScan; -import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.ReadBuilder; import org.apache.paimon.table.source.Split; import org.apache.paimon.types.DataField; @@ -42,12 +42,12 @@ import java.io.Serializable; import java.time.Duration; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.function.Consumer; import java.util.stream.Collectors; import java.util.stream.Stream; +import static org.apache.paimon.CoreOptions.QUERY_AUTH_ENABLED; import static org.apache.paimon.CoreOptions.SCAN_MODE; import static org.apache.paimon.CoreOptions.StartupMode.LATEST; import static org.apache.paimon.io.SplitsParallelReadUtil.parallelExecute; @@ -81,11 +81,12 @@ public RecordReader bootstrap(int numAssigners, int assignId) throw .mapToInt(Integer::intValue) .toArray(); - // force using the latest scan mode + // Force using the latest scan mode and bypass query auth for this internal index read. + Options bootstrapOptions = new Options(); + bootstrapOptions.set(SCAN_MODE, LATEST); + bootstrapOptions.set(QUERY_AUTH_ENABLED, false); ReadBuilder readBuilder = - table.copy(Collections.singletonMap(SCAN_MODE.key(), LATEST.toString())) - .newReadBuilder() - .withProjection(keyProjection); + table.copy(bootstrapOptions.toMap()).newReadBuilder().withProjection(keyProjection); DataTableScan tableScan = (DataTableScan) readBuilder.newScan(); List splits = @@ -142,7 +143,7 @@ static boolean filterSplit(Split split, long indexTtl, long currentTime) { @VisibleForTesting static DataSplit unwrapDataSplit(Split split) { - return (DataSplit) QueryAuthSplit.unwrap(split); + return (DataSplit) split; } public static RowType bootstrapType(TableSchema schema) { diff --git a/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java b/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java index 3a8b897ec53e..0921ddf565ea 100644 --- a/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java @@ -27,17 +27,30 @@ import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.manifest.FileSource; import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.FieldRef; +import org.apache.paimon.predicate.FieldTransform; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.Schema; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.DelegatedFileStoreTable; +import org.apache.paimon.table.FallbackReadFileStoreTable; +import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.Table; import org.apache.paimon.table.TableTestBase; import org.apache.paimon.table.source.DataSplit; -import org.apache.paimon.table.source.QueryAuthSplit; +import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.Filter; +import org.apache.paimon.utils.JsonSerdeUtil; import org.apache.paimon.utils.Pair; import org.junit.jupiter.api.Test; +import org.mockito.AdditionalAnswers; +import org.mockito.Mockito; import java.time.Instant; import java.time.ZoneId; @@ -45,11 +58,11 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.function.Consumer; import static org.apache.paimon.crosspartition.IndexBootstrap.BUCKET_FIELD; import static org.apache.paimon.crosspartition.IndexBootstrap.filterSplit; -import static org.apache.paimon.crosspartition.IndexBootstrap.unwrapDataSplit; import static org.apache.paimon.data.BinaryRow.EMPTY_ROW; import static org.apache.paimon.stats.SimpleStats.EMPTY_STATS; import static org.assertj.core.api.Assertions.assertThat; @@ -100,7 +113,11 @@ public void testBoostrap() throws Exception { } private Table createTable() throws Exception { - Identifier identifier = identifier("T"); + return createTable("T"); + } + + private Table createTable(String tableName) throws Exception { + Identifier identifier = identifier(tableName); Options options = new Options(); options.set(CoreOptions.BUCKET, -1); Schema schema = @@ -133,15 +150,164 @@ public void testFilterSplit() { } @Test - public void testQueryAuthSplit() { + public void testFallbackDataSplit() { DataSplit dataSplit = newSplit(newFile(100), newFile(200)); - QueryAuthSplit authSplit = - new QueryAuthSplit( - dataSplit, new TableQueryAuthResult(Collections.emptyList(), null)); + FallbackSplit fallbackSplit = FallbackReadFileStoreTable.toFallbackSplit(dataSplit, true); + + assertThat(fallbackSplit).isInstanceOf(DataSplit.class); + assertThat(filterSplit(fallbackSplit, 50, 230)).isTrue(); + assertThat(filterSplit(fallbackSplit, 50, 300)).isFalse(); + assertThat(fallbackSplit.isFallback()).isTrue(); + } + + @Test + public void testBootstrapIgnoresQueryAuth() throws Exception { + FileStoreTable table = (FileStoreTable) createTable(); + write( + table, + row(1, 1, 1, 2), + row(1, 2, 2, 3), + row(1, 3, 3, 4), + row(2, 4, 4, 5), + row(2, 5, 5, 6), + row(3, 6, 6, 7), + row(3, 7, 7, 8)); + + TableQueryAuthResult authResult = queryAuthResult(table); + IndexBootstrap indexBootstrap = + new IndexBootstrap(new QueryAuthFileStoreTable(table, authResult)); + + List result = new ArrayList<>(); + try (RecordReader reader = indexBootstrap.bootstrap(1, 0)) { + reader.forEachRemaining( + row -> result.add(GenericRow.of(row.getInt(0), row.getInt(1), row.getInt(2)))); + } + + assertThat(result) + .containsExactlyInAnyOrder( + GenericRow.of(1, 1, 2), + GenericRow.of(2, 1, 3), + GenericRow.of(3, 1, 4), + GenericRow.of(4, 2, 5), + GenericRow.of(5, 2, 6), + GenericRow.of(6, 3, 7), + GenericRow.of(7, 3, 8)); + } + + @Test + public void testBootstrapIgnoresQueryAuthWithFallback() throws Exception { + FileStoreTable mainTable = (FileStoreTable) createTable("MAIN"); + FileStoreTable fallbackTable = (FileStoreTable) createTable("FALLBACK"); + write(mainTable, row(1, 10, 10, 2)); + write(fallbackTable, row(2, 20, 20, 3)); + + TableQueryAuthResult authResult = queryAuthResult(mainTable); + FallbackReadFileStoreTable table = + new FallbackReadFileStoreTable( + new QueryAuthFileStoreTable(mainTable, authResult), + new QueryAuthFileStoreTable(fallbackTable, authResult), + true); + + List result = new ArrayList<>(); + try (RecordReader reader = new IndexBootstrap(table).bootstrap(1, 0)) { + reader.forEachRemaining( + row -> result.add(GenericRow.of(row.getInt(0), row.getInt(1), row.getInt(2)))); + } + + assertThat(result) + .containsExactlyInAnyOrder(GenericRow.of(10, 1, 2), GenericRow.of(20, 2, 3)); + } + + private TableQueryAuthResult queryAuthResult(FileStoreTable table) { + Predicate filter = new PredicateBuilder(table.rowType()).equal(0, 1); + return new TableQueryAuthResult( + Collections.singletonList(JsonSerdeUtil.toFlatJson(filter)), + Collections.singletonMap( + "pk", + JsonSerdeUtil.toFlatJson( + new FieldTransform(new FieldRef(0, "pt", DataTypes.INT()))))); + } + + private static class QueryAuthFileStoreTable extends DelegatedFileStoreTable { + + private final TableQueryAuthResult authResult; + private final boolean queryAuthEnabled; + + private QueryAuthFileStoreTable(FileStoreTable wrapped, TableQueryAuthResult authResult) { + this(wrapped, authResult, true); + } + + private QueryAuthFileStoreTable( + FileStoreTable wrapped, TableQueryAuthResult authResult, boolean queryAuthEnabled) { + super(wrapped); + this.authResult = authResult; + this.queryAuthEnabled = queryAuthEnabled; + } + + @Override + public DataTableScan newScan() { + DataTableScan delegate = wrapped.newScan(); + if (!queryAuthEnabled) { + return delegate; + } + DataTableScan scan = + Mockito.mock(DataTableScan.class, AdditionalAnswers.delegatesTo(delegate)); + Mockito.doAnswer( + invocation -> { + Filter filter = invocation.getArgument(0); + delegate.withBucketFilter(filter); + return scan; + }) + .when(scan) + .withBucketFilter(Mockito.any()); + Mockito.doAnswer( + invocation -> { + Filter filter = invocation.getArgument(0); + delegate.withLevelFilter(filter); + return scan; + }) + .when(scan) + .withLevelFilter(Mockito.any()); + Mockito.doAnswer(ignored -> authResult.convertPlan(delegate.plan())).when(scan).plan(); + return scan; + } + + @Override + public FileStoreTable copy(Map dynamicOptions) { + return new QueryAuthFileStoreTable( + wrapped.copy(dynamicOptions), authResult, queryAuthEnabled(dynamicOptions)); + } + + @Override + public FileStoreTable copy(TableSchema newTableSchema) { + return new QueryAuthFileStoreTable( + wrapped.copy(newTableSchema), authResult, queryAuthEnabled); + } + + @Override + public FileStoreTable copyWithoutTimeTravel(Map dynamicOptions) { + return new QueryAuthFileStoreTable( + wrapped.copyWithoutTimeTravel(dynamicOptions), + authResult, + queryAuthEnabled(dynamicOptions)); + } + + @Override + public FileStoreTable copyWithLatestSchema() { + return new QueryAuthFileStoreTable( + wrapped.copyWithLatestSchema(), authResult, queryAuthEnabled); + } + + @Override + public FileStoreTable switchToBranch(String branchName) { + return new QueryAuthFileStoreTable( + wrapped.switchToBranch(branchName), authResult, queryAuthEnabled); + } - assertThat(unwrapDataSplit(authSplit)).isSameAs(dataSplit); - assertThat(filterSplit(authSplit, 50, 230)).isTrue(); - assertThat(filterSplit(authSplit, 50, 300)).isFalse(); + private boolean queryAuthEnabled(Map dynamicOptions) { + String value = dynamicOptions.get(CoreOptions.QUERY_AUTH_ENABLED.key()); + return value == null ? queryAuthEnabled : Boolean.parseBoolean(value); + } } private DataSplit newSplit(DataFileMeta... files) { From 478edc9469a954323c880f05abd3fc6e905f78e1 Mon Sep 17 00:00:00 2001 From: umi Date: Fri, 17 Jul 2026 17:45:39 +0800 Subject: [PATCH 03/31] [core] Preserve query auth in chain table scans # Conflicts: # paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java --- .../paimon/table/ChainGroupReadTable.java | 35 +++- .../paimon/table/ChainTableStreamScan.java | 37 ++-- .../paimon/table/source/QueryAuthSplit.java | 8 + .../paimon/table/ChainGroupReadTableTest.java | 184 ++++++++++++++++++ 4 files changed, 243 insertions(+), 21 deletions(-) create mode 100644 paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java 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 c1e67e8a9591..caf689b199f7 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 @@ -19,6 +19,7 @@ package org.apache.paimon.table; import org.apache.paimon.CoreOptions; +import org.apache.paimon.annotation.VisibleForTesting; import org.apache.paimon.codegen.CodeGenUtils; import org.apache.paimon.codegen.RecordComparator; import org.apache.paimon.data.BinaryRow; @@ -36,6 +37,7 @@ import org.apache.paimon.table.source.DataSplit; 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.TableRead; import org.apache.paimon.types.RowType; @@ -60,6 +62,7 @@ import java.util.stream.Collectors; import static org.apache.paimon.utils.Preconditions.checkArgument; +import static org.apache.paimon.utils.Preconditions.checkNotNull; /** * Chain table which mainly read from the snapshot branch. However, if the snapshot branch does not @@ -494,19 +497,21 @@ private Set preloadTargetSnapshotSplits(List splits) { } for (Split split : mainScan.plan().splits()) { - DataSplit dataSplit = (DataSplit) split; + DataSplit dataSplit = unwrapDataSplit(split); HashMap fileBucketPathMapping = new HashMap<>(); HashMap fileBranchMapping = new HashMap<>(); for (DataFileMeta file : dataSplit.dataFiles()) { - fileBucketPathMapping.put(file.fileName(), ((DataSplit) split).bucketPath()); + fileBucketPathMapping.put(file.fileName(), dataSplit.bucketPath()); fileBranchMapping.put(file.fileName(), options.scanFallbackSnapshotBranch()); } splits.add( - new ChainSplit( - dataSplit.partition(), - dataSplit.dataFiles(), - fileBranchMapping, - fileBucketPathMapping)); + QueryAuthSplit.retainAuth( + split, + new ChainSplit( + dataSplit.partition(), + dataSplit.dataFiles(), + fileBranchMapping, + fileBucketPathMapping))); } snapshotPartitions.addAll( @@ -515,6 +520,19 @@ private Set preloadTargetSnapshotSplits(List splits) { return snapshotPartitions; } + private static DataSplit unwrapDataSplit(Split split) { + return (DataSplit) QueryAuthSplit.unwrap(split); + } + + private static Split retainQueryAuth(List sourceSplits, Split replacement) { + for (Split sourceSplit : sourceSplits) { + if (sourceSplit instanceof QueryAuthSplit) { + return QueryAuthSplit.retainAuth(sourceSplit, replacement); + } + } + return replacement; + } + private DataTableScan newFilteredScan(boolean snapshot) { DataTableScan scan = snapshot @@ -582,7 +600,8 @@ public TableRead withIOManager(IOManager ioManager) { @Override public RecordReader createReader(Split split) throws IOException { - if (split instanceof ChainSplit || split instanceof DataSplit) { + Split wrappedSplit = QueryAuthSplit.unwrap(split); + if (wrappedSplit instanceof ChainSplit || wrappedSplit instanceof DataSplit) { return fallbackRead.createReader(split); } throw new IllegalArgumentException( diff --git a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java index 5812a2611bc7..ed913a630e35 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java @@ -33,6 +33,7 @@ import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.DataTableStreamScan; import org.apache.paimon.table.source.InnerTableScan; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.SnapshotNotExistPlan; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.StreamDataTableScan; @@ -238,7 +239,7 @@ private TableScan.Plan planStarting() { } // 1. Read delta branch data at the pinned snapshot, grouped by partition. - Map> deltaSplitsByPartition; + Map> deltaSplitsByPartition; if (deltaLatestId != null) { FileStoreTable pinnedDelta = deltaTable.copy(pinnedOptions(deltaLatestId)); DataTableScan pinnedDeltaScan = pinnedDelta.newScan(); @@ -273,7 +274,7 @@ private TableScan.Plan planStarting() { // 3. Scan file splits for latest snapshot partitions only, at the pinned snapshot. // Reuse the pinnedSnapshot from step 2 to avoid redundant copy operations. List latestPartitions = new ArrayList<>(latestChainPartitionPerGroup.values()); - Map> snapshotSplitsByPartition; + Map> snapshotSplitsByPartition; if (!latestPartitions.isEmpty() && pinnedSnapshot != null) { DataTableScan snapshotScan = pinnedSnapshot.newScan(); snapshotScan.withPartitionFilter(latestPartitions); @@ -332,13 +333,16 @@ private List buildLightweightStartingSplits( Map latestChainPartitionPerGroup) { List allSplits = new ArrayList<>(); - for (Map.Entry> entry : snapshotSplitsByPartition.entrySet()) { - for (DataSplit ds : entry.getValue()) { - allSplits.add(ChainSplit.from(ds, snapshotBranch)); + for (Map.Entry> entry : snapshotSplitsByPartition.entrySet()) { + for (Split split : entry.getValue()) { + DataSplit dataSplit = unwrapDataSplit(split); + allSplits.add( + QueryAuthSplit.retainAuth( + split, ChainSplit.from(dataSplit, snapshotBranch))); } } - for (Map.Entry> entry : deltaSplitsByPartition.entrySet()) { + for (Map.Entry> entry : deltaSplitsByPartition.entrySet()) { BinaryRow partition = entry.getKey(); Object groupKey = toGroupKey(partition); BinaryRow latestPartition = latestChainPartitionPerGroup.get(groupKey); @@ -350,8 +354,11 @@ private List buildLightweightStartingSplits( partitionProjector.extractChainPartition(partition), partitionProjector.extractChainPartition(latestPartition)) > 0) { - for (DataSplit ds : entry.getValue()) { - allSplits.add(ChainSplit.from(ds, deltaBranch)); + for (Split split : entry.getValue()) { + DataSplit dataSplit = unwrapDataSplit(split); + allSplits.add( + QueryAuthSplit.retainAuth( + split, ChainSplit.from(dataSplit, deltaBranch))); } } } @@ -477,15 +484,19 @@ private Long captureDeltaPosition(FileStoreTable deltaTable) { } /** Plans a scan and groups the resulting splits by partition. */ - private static Map> groupByPartition(DataTableScan scan) { - Map> grouped = new LinkedHashMap<>(); - for (Split s : scan.plan().splits()) { - DataSplit ds = (DataSplit) s; - grouped.computeIfAbsent(ds.partition(), k -> new ArrayList<>()).add(ds); + private static Map> groupByPartition(DataTableScan scan) { + Map> grouped = new LinkedHashMap<>(); + for (Split split : scan.plan().splits()) { + DataSplit dataSplit = unwrapDataSplit(split); + grouped.computeIfAbsent(dataSplit.partition(), k -> new ArrayList<>()).add(split); } return grouped; } + private static DataSplit unwrapDataSplit(Split split) { + return (DataSplit) QueryAuthSplit.unwrap(split); + } + /** * Extracts a stable group key from a full partition row. When there is no group partition (all * fields are chain keys), returns a shared singleton to avoid zero-field {@link BinaryRow} 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 6bd03379f714..a354fafd0af6 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,6 +54,14 @@ public static Split unwrap(Split split) { return split instanceof QueryAuthSplit ? ((QueryAuthSplit) split).split() : split; } + /** Retains the source split's query authorization result on a replacement split. */ + public static Split retainAuth(Split source, Split replacement) { + if (source instanceof QueryAuthSplit) { + return new QueryAuthSplit(replacement, ((QueryAuthSplit) source).authResult()); + } + return replacement; + } + @Nullable public TableQueryAuthResult authResult() { return authResult; diff --git a/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java new file mode 100644 index 000000000000..b975c093893d --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java @@ -0,0 +1,184 @@ +/* + * 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; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.catalog.TableQueryAuthResult; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.io.DataFileTestDataGenerator; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.source.ChainSplit; +import org.apache.paimon.table.source.DataFilePlan; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.DataTableScan; +import org.apache.paimon.table.source.DataTableStreamScan; +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.types.DataField; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.utils.SnapshotManager; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Tests for query authorization handling in chain table scans. */ +public class ChainGroupReadTableTest { + + @Test + public void testBatchScanRetainsQueryAuth() throws Exception { + TableSchema schema = tableSchema(); + PrimaryKeyFileStoreTable snapshotTable = table(schema, "snapshot"); + PrimaryKeyFileStoreTable deltaTable = table(schema, "delta"); + ChainGroupReadTable chainTable = new ChainGroupReadTable(snapshotTable, deltaTable); + + BinaryRow partition = BinaryRow.singleColumn(1); + QueryAuthSplit sourceSplit = queryAuthSplit(dataSplit(partition)); + DataTableScan snapshotScan = mock(DataTableScan.class); + when(snapshotScan.plan()) + .thenReturn(new DataFilePlan<>(Collections.singletonList(sourceSplit))); + when(snapshotScan.listPartitions()).thenReturn(Collections.singletonList(partition)); + + DataTableScan deltaScan = mock(DataTableScan.class); + when(deltaScan.listPartitions()).thenReturn(Collections.emptyList()); + + ChainGroupReadTable.ChainTableBatchScan scan = + new ChainGroupReadTable.ChainTableBatchScan( + schema, + chainTable, + table -> table == snapshotTable ? snapshotScan : deltaScan); + List splits = scan.plan().splits(); + + assertThat(splits).hasSize(1); + QueryAuthSplit result = (QueryAuthSplit) splits.get(0); + assertThat(result.authResult()).isSameAs(sourceSplit.authResult()); + assertThat(result.split()).isInstanceOf(ChainSplit.class); + + InnerTableRead snapshotRead = mock(InnerTableRead.class); + InnerTableRead deltaRead = mock(InnerTableRead.class); + RecordReader reader = mock(RecordReader.class); + when(snapshotTable.newRead()).thenReturn(snapshotRead); + when(deltaTable.newRead()).thenReturn(deltaRead); + when(deltaRead.createReader(result)).thenReturn(reader); + + assertThat(chainTable.newRead().createReader(result)).isSameAs(reader); + verify(deltaRead).createReader(result); + } + + @Test + public void testStreamStartingPlanRetainsQueryAuth() { + TableSchema schema = tableSchema(); + PrimaryKeyFileStoreTable snapshotTable = table(schema, "snapshot"); + PrimaryKeyFileStoreTable deltaTable = table(schema, "delta"); + + DataTableScan unusedScan = mock(DataTableScan.class); + when(snapshotTable.newScan()).thenReturn(unusedScan); + when(deltaTable.newScan()).thenReturn(unusedScan); + + DataTableStreamScan deltaStreamScan = mock(DataTableStreamScan.class); + when(deltaTable.newStreamScan()).thenReturn(deltaStreamScan); + + SnapshotManager snapshotManager = mock(SnapshotManager.class); + when(snapshotManager.latestSnapshotId()).thenReturn(null); + when(snapshotTable.snapshotManager()).thenReturn(snapshotManager); + + SnapshotManager deltaSnapshotManager = mock(SnapshotManager.class); + when(deltaSnapshotManager.latestSnapshotId()).thenReturn(1L); + when(deltaTable.snapshotManager()).thenReturn(deltaSnapshotManager); + + BinaryRow partition = BinaryRow.singleColumn(1); + QueryAuthSplit sourceSplit = queryAuthSplit(dataSplit(partition)); + DataTableScan pinnedDeltaScan = mock(DataTableScan.class); + when(pinnedDeltaScan.plan()) + .thenReturn(new DataFilePlan<>(Collections.singletonList(sourceSplit))); + PrimaryKeyFileStoreTable pinnedDeltaTable = table(schema, "delta"); + when(pinnedDeltaTable.newScan()).thenReturn(pinnedDeltaScan); + when(deltaTable.copy(anyMap())).thenReturn(pinnedDeltaTable); + + ChainTableStreamScan scan = + new ChainTableStreamScan(new ChainGroupReadTable(snapshotTable, deltaTable)); + List splits = scan.plan().splits(); + + assertThat(splits).hasSize(1); + QueryAuthSplit result = (QueryAuthSplit) splits.get(0); + assertThat(result.authResult()).isSameAs(sourceSplit.authResult()); + assertThat(result.split()).isInstanceOf(ChainSplit.class); + verify(deltaStreamScan).restore(2L); + } + + private static PrimaryKeyFileStoreTable table(TableSchema schema, String branch) { + PrimaryKeyFileStoreTable table = mock(PrimaryKeyFileStoreTable.class); + when(table.schema()).thenReturn(schema); + Map options = new HashMap<>(schema.options()); + options.put(CoreOptions.BRANCH.key(), branch); + when(table.coreOptions()).thenReturn(CoreOptions.fromMap(options)); + return table; + } + + private static TableSchema tableSchema() { + List fields = + Arrays.asList( + new DataField(0, "pt", DataTypes.INT()), + new DataField(1, "k", DataTypes.INT()), + new DataField(2, "v", DataTypes.INT())); + Map options = new HashMap<>(); + options.put(CoreOptions.SCAN_FALLBACK_SNAPSHOT_BRANCH.key(), "snapshot"); + options.put(CoreOptions.SCAN_FALLBACK_DELTA_BRANCH.key(), "delta"); + return new TableSchema( + 0, + fields, + 2, + Collections.singletonList("pt"), + Arrays.asList("pt", "k"), + options, + ""); + } + + private static DataSplit dataSplit(BinaryRow partition) { + return DataSplit.builder() + .withSnapshot(1L) + .withPartition(partition) + .withBucket(0) + .withBucketPath("pt=1/bucket-0") + .withTotalBuckets(1) + .withDataFiles( + Collections.singletonList( + DataFileTestDataGenerator.builder().build().next().meta)) + .build(); + } + + private static QueryAuthSplit queryAuthSplit(DataSplit split) { + TableQueryAuthResult authResult = + new TableQueryAuthResult(null, Collections.singletonMap("v", "mask")); + return new QueryAuthSplit(split, authResult); + } +} From b1e6b868b32f7b07bc2b2692b0d351a8f96bc964 Mon Sep 17 00:00:00 2001 From: umi Date: Fri, 17 Jul 2026 17:47:31 +0800 Subject: [PATCH 04/31] [core] Preserve query auth in data evolution scans --- .../globalindex/DataEvolutionBatchScan.java | 18 ++++++--- .../DataEvolutionBatchScanTest.java | 38 +++++++++++++++++++ 2 files changed, 50 insertions(+), 6 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/globalindex/DataEvolutionBatchScan.java b/paimon-core/src/main/java/org/apache/paimon/globalindex/DataEvolutionBatchScan.java index ebdfb0894f4b..36fc575066de 100644 --- a/paimon-core/src/main/java/org/apache/paimon/globalindex/DataEvolutionBatchScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/globalindex/DataEvolutionBatchScan.java @@ -37,6 +37,7 @@ import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.InnerTableScan; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.Split; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.Filter; @@ -332,12 +333,17 @@ private Optional evalGlobalIndex() { public static Plan wrapToIndexSplits( List splits, RowRangeIndex rowRangeIndex, ScoreGetter scoreGetter) { List indexedSplits = new ArrayList<>(); - Function> process = - split -> - Collections.singletonList( - split instanceof IndexedSplit - ? (IndexedSplit) split - : wrap((DataSplit) split, rowRangeIndex, scoreGetter)); + Function> process = + split -> { + Split wrappedSplit = QueryAuthSplit.unwrap(split); + if (wrappedSplit instanceof IndexedSplit) { + return Collections.singletonList(split); + } + IndexedSplit indexedSplit = + wrap((DataSplit) wrappedSplit, rowRangeIndex, scoreGetter); + return Collections.singletonList( + QueryAuthSplit.retainAuth(split, indexedSplit)); + }; randomlyExecuteSequentialReturn(process, splits, null).forEachRemaining(indexedSplits::add); return () -> indexedSplits; } diff --git a/paimon-core/src/test/java/org/apache/paimon/globalindex/DataEvolutionBatchScanTest.java b/paimon-core/src/test/java/org/apache/paimon/globalindex/DataEvolutionBatchScanTest.java index 78771719d068..5c149cc13433 100644 --- a/paimon-core/src/test/java/org/apache/paimon/globalindex/DataEvolutionBatchScanTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/globalindex/DataEvolutionBatchScanTest.java @@ -18,13 +18,16 @@ package org.apache.paimon.globalindex; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.BinaryRow; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.table.source.AppendBatchTableScan; +import org.apache.paimon.table.source.DataFilePlan; import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.DataTableScan; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.snapshot.SnapshotReader; import org.apache.paimon.types.DataField; @@ -112,6 +115,41 @@ public void testWithShardKeepsDataEvolutionWrapper() { verify(batchScan).withShard(0, 2); } + @Test + public void testPlanRetainsQueryAuthWhenWrappingIndexedSplit() { + DataSplit dataSplit = + DataSplit.builder() + .withSnapshot(1L) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withDataFiles(Collections.singletonList(newAppendFile(10L, 10L, "file"))) + .build(); + TableQueryAuthResult authResult = + new TableQueryAuthResult(null, Collections.singletonMap("f0", "mask")); + QueryAuthSplit queryAuthSplit = new QueryAuthSplit(dataSplit, authResult); + RowRangeIndex rowRangeIndex = + RowRangeIndex.create(Collections.singletonList(new Range(12L, 15L))); + + AppendBatchTableScan batchScan = mock(AppendBatchTableScan.class); + when(batchScan.withRowRangeIndex(rowRangeIndex)).thenReturn(batchScan); + when(batchScan.plan()) + .thenReturn(new DataFilePlan<>(Collections.singletonList(queryAuthSplit))); + + List splits = + new DataEvolutionBatchScan(null, batchScan) + .withRowRangeIndex(rowRangeIndex) + .plan() + .splits(); + + assertThat(splits).hasSize(1); + QueryAuthSplit result = (QueryAuthSplit) splits.get(0); + assertThat(result.authResult()).isSameAs(authResult); + IndexedSplit indexedSplit = (IndexedSplit) result.split(); + assertThat(indexedSplit.dataSplit()).isSameAs(dataSplit); + assertThat(indexedSplit.rowRanges()).containsExactly(new Range(12L, 15L)); + } + @Test public void testWrapToIndexSplitsRandomly() { Random random = new Random(); From 070a17cbbdec4bfef5f757d41b8485427ac9cded Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 11:53:59 +0800 Subject: [PATCH 05/31] fix --- .../apache/paimon/crosspartition/IndexBootstrap.java | 9 ++------- .../org/apache/paimon/table/ChainGroupReadTable.java | 10 ++++------ .../org/apache/paimon/table/ChainTableStreamScan.java | 10 +++------- .../org/apache/paimon/table/source/QueryAuthSplit.java | 5 +++++ .../apache/paimon/table/source/QueryAuthSplitTest.java | 9 +++++++++ 5 files changed, 23 insertions(+), 20 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java b/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java index 8f9d99edd6dd..2c5249fd45fb 100644 --- a/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java +++ b/paimon-core/src/main/java/org/apache/paimon/crosspartition/IndexBootstrap.java @@ -121,7 +121,7 @@ public RecordReader bootstrap(int numAssigners, int assignId) throw options.pageSize(), options.crossPartitionUpsertBootstrapParallelism(), split -> { - DataSplit dataSplit = unwrapDataSplit(split); + DataSplit dataSplit = (DataSplit) split; int bucket = dataSplit.bucket(); return partBucketConverter.toGenericRow( new JoinedRow(dataSplit.partition(), GenericRow.of(bucket))); @@ -131,7 +131,7 @@ public RecordReader bootstrap(int numAssigners, int assignId) throw @VisibleForTesting static boolean filterSplit(Split split, long indexTtl, long currentTime) { - List files = unwrapDataSplit(split).dataFiles(); + List files = ((DataSplit) split).dataFiles(); for (DataFileMeta file : files) { long fileTime = file.creationTimeEpochMillis(); if (currentTime <= fileTime + indexTtl) { @@ -141,11 +141,6 @@ static boolean filterSplit(Split split, long indexTtl, long currentTime) { return false; } - @VisibleForTesting - static DataSplit unwrapDataSplit(Split split) { - return (DataSplit) split; - } - public static RowType bootstrapType(TableSchema schema) { List primaryKeys = schema.trimmedPrimaryKeys(); List partitionKeys = schema.partitionKeys(); 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 caf689b199f7..256e3418c536 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 @@ -222,7 +222,8 @@ public ChainTableBatchScan( this(tableSchema, chainGroupReadTable, FileStoreTable::newScan); } - private ChainTableBatchScan( + @VisibleForTesting + ChainTableBatchScan( TableSchema tableSchema, ChainGroupReadTable chainGroupReadTable, Function scanCreator) { @@ -362,6 +363,7 @@ public Plan plan() { for (List deltaPartitionsInGroup : groupedDeltaPartitions.values()) { // Sort delta by chain dimension ascending. + // chainPartitionForCompare avoids copying BinaryRow in the comparator hot path. deltaPartitionsInGroup.sort( (a, b) -> chainPartitionComparator.compare( @@ -497,7 +499,7 @@ private Set preloadTargetSnapshotSplits(List splits) { } for (Split split : mainScan.plan().splits()) { - DataSplit dataSplit = unwrapDataSplit(split); + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); HashMap fileBucketPathMapping = new HashMap<>(); HashMap fileBranchMapping = new HashMap<>(); for (DataFileMeta file : dataSplit.dataFiles()) { @@ -520,10 +522,6 @@ private Set preloadTargetSnapshotSplits(List splits) { return snapshotPartitions; } - private static DataSplit unwrapDataSplit(Split split) { - return (DataSplit) QueryAuthSplit.unwrap(split); - } - private static Split retainQueryAuth(List sourceSplits, Split replacement) { for (Split sourceSplit : sourceSplits) { if (sourceSplit instanceof QueryAuthSplit) { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java index ed913a630e35..8f711238cf63 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java @@ -335,7 +335,7 @@ private List buildLightweightStartingSplits( for (Map.Entry> entry : snapshotSplitsByPartition.entrySet()) { for (Split split : entry.getValue()) { - DataSplit dataSplit = unwrapDataSplit(split); + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); allSplits.add( QueryAuthSplit.retainAuth( split, ChainSplit.from(dataSplit, snapshotBranch))); @@ -355,7 +355,7 @@ private List buildLightweightStartingSplits( partitionProjector.extractChainPartition(latestPartition)) > 0) { for (Split split : entry.getValue()) { - DataSplit dataSplit = unwrapDataSplit(split); + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); allSplits.add( QueryAuthSplit.retainAuth( split, ChainSplit.from(dataSplit, deltaBranch))); @@ -487,16 +487,12 @@ private Long captureDeltaPosition(FileStoreTable deltaTable) { private static Map> groupByPartition(DataTableScan scan) { Map> grouped = new LinkedHashMap<>(); for (Split split : scan.plan().splits()) { - DataSplit dataSplit = unwrapDataSplit(split); + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); grouped.computeIfAbsent(dataSplit.partition(), k -> new ArrayList<>()).add(split); } return grouped; } - private static DataSplit unwrapDataSplit(Split split) { - return (DataSplit) QueryAuthSplit.unwrap(split); - } - /** * Extracts a stable group key from a full partition row. When there is no group partition (all * fields are chain keys), returns a shared singleton to avoid zero-field {@link BinaryRow} 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 a354fafd0af6..c6842de8eb74 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,6 +54,11 @@ public static Split unwrap(Split split) { return split instanceof QueryAuthSplit ? ((QueryAuthSplit) split).split() : split; } + /** Unwraps one query authorization layer and returns the wrapped data split. */ + public static DataSplit unwrapDataSplit(Split split) { + return (DataSplit) unwrap(split); + } + /** Retains the source split's query authorization result on a replacement split. */ public static Split retainAuth(Split source, Split replacement) { if (source instanceof QueryAuthSplit) { diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/QueryAuthSplitTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/QueryAuthSplitTest.java index 2d40be47e459..97c24f94b6e0 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/QueryAuthSplitTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/QueryAuthSplitTest.java @@ -37,6 +37,15 @@ /** Test for {@link QueryAuthSplit}. */ public class QueryAuthSplitTest { + @Test + public void testUnwrapDataSplit() { + DataSplit dataSplit = dataSplit(); + + assertThat(QueryAuthSplit.unwrapDataSplit(dataSplit)).isSameAs(dataSplit); + assertThat(QueryAuthSplit.unwrapDataSplit(new QueryAuthSplit(dataSplit, authResult()))) + .isSameAs(dataSplit); + } + @Test public void testSerializeAndDeserialize() throws Exception { QueryAuthSplit split = new QueryAuthSplit(dataSplit(), authResult()); From e7563043cc8b6c400121c327132e5712df0b7401 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 12:17:27 +0800 Subject: [PATCH 06/31] rebaseMaster --- .../paimon/table/ChainGroupReadTable.java | 22 ++------ .../paimon/table/ChainTableStreamScan.java | 23 +++++---- .../apache/paimon/utils/ChainTableUtils.java | 51 ++++++++++++------- .../paimon/table/ChainGroupReadTableTest.java | 18 ++++++- 4 files changed, 66 insertions(+), 48 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 256e3418c536..7e113a3d70c2 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 @@ -62,7 +62,6 @@ import java.util.stream.Collectors; import static org.apache.paimon.utils.Preconditions.checkArgument; -import static org.apache.paimon.utils.Preconditions.checkNotNull; /** * Chain table which mainly read from the snapshot branch. However, if the snapshot branch does not @@ -435,18 +434,12 @@ public Plan plan() { deltaScan.withPartitionFilter(selectedDeltaPartitions); } - List deltaSubSplits = - deltaScan.plan().splits().stream() - .map(s -> (DataSplit) s) - .collect(Collectors.toList()); - List snapshotSubSplits = new ArrayList<>(); + List deltaSubSplits = deltaScan.plan().splits(); + List snapshotSubSplits = Collections.emptyList(); if (partitionPairs.getValue() != null) { snapshotScan.withPartitionFilter( Collections.singletonList(partitionPairs.getValue())); - snapshotSubSplits = - snapshotScan.plan().splits().stream() - .map(s -> (DataSplit) s) - .collect(Collectors.toList()); + snapshotSubSplits = snapshotScan.plan().splits(); } splits.addAll( ChainTableUtils.buildChainSplits( @@ -522,15 +515,6 @@ private Set preloadTargetSnapshotSplits(List splits) { return snapshotPartitions; } - private static Split retainQueryAuth(List sourceSplits, Split replacement) { - for (Split sourceSplit : sourceSplits) { - if (sourceSplit instanceof QueryAuthSplit) { - return QueryAuthSplit.retainAuth(sourceSplit, replacement); - } - } - return replacement; - } - private DataTableScan newFilteredScan(boolean snapshot) { DataTableScan scan = snapshot diff --git a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java index 8f711238cf63..02c751a8b4de 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java @@ -328,8 +328,8 @@ private TableScan.Plan planStarting() { private List buildLightweightStartingSplits( String snapshotBranch, String deltaBranch, - Map> snapshotSplitsByPartition, - Map> deltaSplitsByPartition, + Map> snapshotSplitsByPartition, + Map> deltaSplitsByPartition, Map latestChainPartitionPerGroup) { List allSplits = new ArrayList<>(); @@ -375,15 +375,15 @@ private List buildLightweightStartingSplits( private List buildMergedStartingSplits( String snapshotBranch, String deltaBranch, - Map> snapshotSplitsByPartition, - Map> deltaSplitsByPartition, + Map> snapshotSplitsByPartition, + Map> deltaSplitsByPartition, Map latestChainPartitionPerGroup) { List allSplits = new ArrayList<>(); // Pre-group delta splits and find the latest delta partition per group. - Map> deltaSplitsByGroup = new HashMap<>(); + Map> deltaSplitsByGroup = new HashMap<>(); Map latestDeltaPartitionPerGroup = new HashMap<>(); - for (Map.Entry> e : deltaSplitsByPartition.entrySet()) { + for (Map.Entry> e : deltaSplitsByPartition.entrySet()) { BinaryRow deltaPartition = e.getKey(); Object groupKey = toGroupKey(deltaPartition); deltaSplitsByGroup @@ -404,7 +404,7 @@ private List buildMergedStartingSplits( for (Map.Entry entry : latestChainPartitionPerGroup.entrySet()) { Object groupKey = entry.getKey(); BinaryRow snapshotPartition = entry.getValue(); - List snapshotSplits = + List snapshotSplits = snapshotSplitsByPartition.getOrDefault( snapshotPartition, Collections.emptyList()); @@ -418,15 +418,16 @@ private List buildMergedStartingSplits( snapshotPartition)) > 0; - List selectedDeltaSplits = new ArrayList<>(); + List selectedDeltaSplits = new ArrayList<>(); if (hasDeltaAfterSnapshot) { - for (DataSplit dataSplit : deltaSplitsByGroup.get(groupKey)) { + for (Split split : deltaSplitsByGroup.get(groupKey)) { + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); BinaryRow deltaPartition = dataSplit.partition(); if (chainPartitionComparator.compare( partitionProjector.extractChainPartition(deltaPartition), partitionProjector.extractChainPartition(snapshotPartition)) > 0) { - selectedDeltaSplits.add(dataSplit); + selectedDeltaSplits.add(split); } } } @@ -444,7 +445,7 @@ private List buildMergedStartingSplits( // Delta-only groups: there is no snapshot anchor, so merge all delta partitions in the // group into the latest delta partition. - for (Map.Entry> entry : deltaSplitsByGroup.entrySet()) { + for (Map.Entry> entry : deltaSplitsByGroup.entrySet()) { Object groupKey = entry.getKey(); if (!latestChainPartitionPerGroup.containsKey(groupKey)) { BinaryRow logicalPartition = latestDeltaPartitionPerGroup.get(groupKey); diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/ChainTableUtils.java b/paimon-core/src/main/java/org/apache/paimon/utils/ChainTableUtils.java index 498cc4dacc2c..0cfcc0f26af1 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/ChainTableUtils.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/ChainTableUtils.java @@ -31,6 +31,8 @@ import org.apache.paimon.table.FileStoreTable; 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.types.RowType; import java.time.LocalDateTime; @@ -375,40 +377,43 @@ public static Predicate createGroupChainPredicate( /** * Builds per-bucket {@link ChainSplit}s from the given snapshot and delta splits. Files that * originate from the snapshot splits are tagged with {@code snapshotBranch}; all other files - * are tagged with {@code deltaBranch}. + * are tagged with {@code deltaBranch}. If an input split carries query authorization, the + * corresponding output split retains it. * * @param logicalPartition the logical partition for the resulting ChainSplits * @param snapshotSplits splits from the snapshot branch * @param deltaSplits splits from the delta branch * @param snapshotBranch name of the snapshot branch * @param deltaBranch name of the delta branch - * @return one ChainSplit per bucket + * @return one split per bucket, optionally wrapped with query authorization */ - public static List buildChainSplits( + public static List buildChainSplits( BinaryRow logicalPartition, - List snapshotSplits, - List deltaSplits, + List snapshotSplits, + List deltaSplits, String snapshotBranch, String deltaBranch) { Set snapshotFileNames = snapshotSplits.stream() + .map(QueryAuthSplit::unwrapDataSplit) .flatMap(s -> s.dataFiles().stream().map(DataFileMeta::fileName)) .collect(Collectors.toSet()); - Map> bucketSplits = new LinkedHashMap<>(); + Map> bucketSplits = new LinkedHashMap<>(); Integer bucketInAll = null; - for (DataSplit ds : snapshotSplits) { - bucketInAll = addToBucketMap(ds, bucketSplits, bucketInAll); + for (Split split : snapshotSplits) { + bucketInAll = addToBucketMap(split, bucketSplits, bucketInAll); } - for (DataSplit ds : deltaSplits) { - bucketInAll = addToBucketMap(ds, bucketSplits, bucketInAll); + for (Split split : deltaSplits) { + bucketInAll = addToBucketMap(split, bucketSplits, bucketInAll); } - List result = new ArrayList<>(); - for (Map.Entry> entry : bucketSplits.entrySet()) { + List result = new ArrayList<>(); + for (Map.Entry> entry : bucketSplits.entrySet()) { Map fileBranchMapping = new HashMap<>(); Map fileBucketPathMapping = new HashMap<>(); - for (DataSplit ds : entry.getValue()) { + for (Split split : entry.getValue()) { + DataSplit ds = QueryAuthSplit.unwrapDataSplit(split); for (DataFileMeta file : ds.dataFiles()) { fileBucketPathMapping.put(file.fileName(), ds.bucketPath()); String branch = @@ -418,30 +423,42 @@ public static List buildChainSplits( fileBranchMapping.put(file.fileName(), branch); } } - result.add( + ChainSplit chainSplit = new ChainSplit( logicalPartition, entry.getValue().stream() + .map(QueryAuthSplit::unwrapDataSplit) .flatMap(ds -> ds.dataFiles().stream()) .collect(Collectors.toList()), fileBranchMapping, - fileBucketPathMapping)); + fileBucketPathMapping); + result.add(retainQueryAuth(entry.getValue(), chainSplit)); } return result; } private static Integer addToBucketMap( - DataSplit ds, Map> bucketSplits, Integer bucketInAll) { + Split split, Map> bucketSplits, Integer bucketInAll) { + DataSplit ds = QueryAuthSplit.unwrapDataSplit(split); Integer totalBuckets = ds.totalBuckets(); checkNotNull(totalBuckets, "totalBuckets should not be null"); if (bucketInAll != null) { checkArgument( totalBuckets.equals(bucketInAll), "Inconsistent bucket num " + ds.bucket()); } - bucketSplits.computeIfAbsent(ds.bucket(), k -> new ArrayList<>()).add(ds); + bucketSplits.computeIfAbsent(ds.bucket(), k -> new ArrayList<>()).add(split); return totalBuckets; } + private static Split retainQueryAuth(List sourceSplits, Split replacement) { + for (Split sourceSplit : sourceSplits) { + if (sourceSplit instanceof QueryAuthSplit) { + return QueryAuthSplit.retainAuth(sourceSplit, replacement); + } + } + return replacement; + } + /** * Validates that the chain table configuration is compatible with incremental read paths * (streaming read and lookup join). All validation rules for incremental reads should be diff --git a/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java index b975c093893d..ec56c930fac6 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java @@ -96,7 +96,16 @@ public void testBatchScanRetainsQueryAuth() throws Exception { @Test public void testStreamStartingPlanRetainsQueryAuth() { - TableSchema schema = tableSchema(); + testStreamStartingPlanRetainsQueryAuth(false); + } + + @Test + public void testStreamMergedStartingPlanRetainsQueryAuth() { + testStreamStartingPlanRetainsQueryAuth(true); + } + + private void testStreamStartingPlanRetainsQueryAuth(boolean mergeSnapshot) { + TableSchema schema = tableSchema(mergeSnapshot); PrimaryKeyFileStoreTable snapshotTable = table(schema, "snapshot"); PrimaryKeyFileStoreTable deltaTable = table(schema, "delta"); @@ -145,6 +154,10 @@ private static PrimaryKeyFileStoreTable table(TableSchema schema, String branch) } private static TableSchema tableSchema() { + return tableSchema(false); + } + + private static TableSchema tableSchema(boolean mergeSnapshot) { List fields = Arrays.asList( new DataField(0, "pt", DataTypes.INT()), @@ -153,6 +166,9 @@ private static TableSchema tableSchema() { Map options = new HashMap<>(); options.put(CoreOptions.SCAN_FALLBACK_SNAPSHOT_BRANCH.key(), "snapshot"); options.put(CoreOptions.SCAN_FALLBACK_DELTA_BRANCH.key(), "delta"); + options.put( + CoreOptions.CHAIN_TABLE_STREAMING_MERGE_SNAPSHOT.key(), + String.valueOf(mergeSnapshot)); return new TableSchema( 0, fields, From 60243a3c9fd499b0bd95671d26b6664b9d94a688 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 13:06:00 +0800 Subject: [PATCH 07/31] fix --- .../table/ChainTableFileStoreTable.java | 4 +- .../table/ChainTableFileStoreTableTest.java | 102 ++++++++++++++++++ 2 files changed, 105 insertions(+), 1 deletion(-) create mode 100644 paimon-core/src/test/java/org/apache/paimon/table/ChainTableFileStoreTableTest.java 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 863cd969bc89..3e4a70653e61 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 @@ -28,6 +28,7 @@ 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.StreamDataTableScan; import org.apache.paimon.table.source.TableRead; @@ -209,7 +210,8 @@ public RecordReader createReader(Split split) throws IOException { // FallbackSplit (including FallbackDataSplit): use inherited fallback read logic return fallbackRead.createReader(split); } - if (split instanceof ChainSplit || split instanceof DataSplit) { + Split wrappedSplit = QueryAuthSplit.unwrap(split); + if (wrappedSplit instanceof ChainSplit || wrappedSplit instanceof DataSplit) { return chainGroupRead.createReader(split); } // Other split types: use inherited fallback read logic diff --git a/paimon-core/src/test/java/org/apache/paimon/table/ChainTableFileStoreTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/ChainTableFileStoreTableTest.java new file mode 100644 index 000000000000..eaaa6a0d3d1b --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/ChainTableFileStoreTableTest.java @@ -0,0 +1,102 @@ +/* + * 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; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.catalog.TableQueryAuthResult; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.reader.RecordReader; +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.InnerTableRead; +import org.apache.paimon.table.source.QueryAuthSplit; +import org.apache.paimon.table.source.Split; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Tests for split routing in {@link ChainTableFileStoreTable}. */ +public class ChainTableFileStoreTableTest { + + @Test + public void testQueryAuthChainSplitRoutedToChainGroupRead() throws Exception { + assertQueryAuthSplitRoutedToChainGroupRead(mock(ChainSplit.class)); + } + + @Test + public void testQueryAuthDataSplitRoutedToChainGroupRead() throws Exception { + assertQueryAuthSplitRoutedToChainGroupRead(mock(DataSplit.class)); + } + + @Test + public void testFallbackSplitStillUsesFallbackRead() throws Exception { + ReadFixture fixture = new ReadFixture(); + QueryAuthSplit wrappedSplit = queryAuthSplit(mock(DataSplit.class)); + FallbackSplitImpl fallbackSplit = new FallbackSplitImpl(wrappedSplit, false); + RecordReader reader = mock(RecordReader.class); + when(fixture.wrappedRead.createReader(wrappedSplit)).thenReturn(reader); + + assertThat(fixture.table.newRead().createReader(fallbackSplit)).isSameAs(reader); + + verify(fixture.wrappedRead).createReader(wrappedSplit); + } + + private void assertQueryAuthSplitRoutedToChainGroupRead(Split wrappedSplit) throws IOException { + ReadFixture fixture = new ReadFixture(); + QueryAuthSplit split = queryAuthSplit(wrappedSplit); + RecordReader reader = mock(RecordReader.class); + when(fixture.chainGroupRead.createReader(split)).thenReturn(reader); + + assertThat(fixture.table.newRead().createReader(split)).isSameAs(reader); + + verify(fixture.chainGroupRead).createReader(split); + verify(fixture.wrappedRead, never()).createReader(any(Split.class)); + } + + private static QueryAuthSplit queryAuthSplit(Split split) { + return new QueryAuthSplit( + split, new TableQueryAuthResult(null, Collections.singletonMap("v", "mask"))); + } + + private static class ReadFixture { + + private final FileStoreTable wrapped = mock(FileStoreTable.class); + private final ChainGroupReadTable chainGroup = mock(ChainGroupReadTable.class); + private final InnerTableRead wrappedRead = mock(InnerTableRead.class); + private final InnerTableRead chainGroupRead = mock(InnerTableRead.class); + private final ChainTableFileStoreTable table; + + private ReadFixture() { + when(wrapped.coreOptions()).thenReturn(CoreOptions.fromMap(Collections.emptyMap())); + when(wrapped.newRead()).thenReturn(wrappedRead); + when(chainGroup.newRead()).thenReturn(chainGroupRead); + table = new ChainTableFileStoreTable(wrapped, chainGroup); + } + } +} From 359e30d6afebf5d81f9cf4676a76add3c8ffc67e Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 14:27:03 +0800 Subject: [PATCH 08/31] [core] Apply query auth after chain split assembly --- .../paimon/table/ChainGroupReadTable.java | 61 ++++++- .../paimon/table/ChainTableStreamScan.java | 44 ++++- .../apache/paimon/utils/ChainTableUtils.java | 17 +- .../paimon/table/ChainGroupReadTableTest.java | 170 ++++++++++++++++-- 4 files changed, 244 insertions(+), 48 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 7e113a3d70c2..464af1bf9575 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 @@ -20,6 +20,7 @@ import org.apache.paimon.CoreOptions; import org.apache.paimon.annotation.VisibleForTesting; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.codegen.CodeGenUtils; import org.apache.paimon.codegen.RecordComparator; import org.apache.paimon.data.BinaryRow; @@ -39,6 +40,7 @@ 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.TableQueryAuth; import org.apache.paimon.table.source.TableRead; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.ChainPartitionProjector; @@ -47,6 +49,8 @@ import org.apache.paimon.utils.Pair; import org.apache.paimon.utils.RowDataToObjectArrayConverter; +import javax.annotation.Nullable; + import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; @@ -121,6 +125,13 @@ private DataTableScan newDeltaScan(Function scanC return scanCreator.apply(other()); } + static FileStoreTable withoutQueryAuth(FileStoreTable table) { + if (!table.coreOptions().queryAuthEnabled()) { + return table; + } + return table.copy(Collections.singletonMap(CoreOptions.QUERY_AUTH_ENABLED.key(), "false")); + } + @Override public FileStoreTable copy(Map dynamicOptions) { Map wrappedOptions = @@ -212,6 +223,8 @@ public static class ChainTableBatchScan extends FallbackReadScan { private final ChainGroupReadTable chainGroupReadTable; private final RecordComparator chainPartitionComparator; private final ChainPartitionProjector partitionProjector; + private final TableQueryAuth queryAuth; + @Nullable private RowType readType; private Predicate dataPredicate; private Filter bucketFilter; protected boolean preloadTargetSnapshot = true; @@ -230,9 +243,10 @@ public ChainTableBatchScan( chainGroupReadTable.wrapped, chainGroupReadTable.other(), tableSchema, - scanCreator); + physicalScanCreator(scanCreator)); this.options = CoreOptions.fromMap(tableSchema.options()); this.chainGroupReadTable = chainGroupReadTable; + this.queryAuth = chainGroupReadTable.catalogEnvironment().tableQueryAuth(this.options); this.partitionConverter = new RowDataToObjectArrayConverter(tableSchema.logicalPartitionType()); this.partitionComparator = @@ -251,6 +265,19 @@ public ChainTableBatchScan( partitionProjector.chainPartitionType().getFieldTypes()); } + private static Function physicalScanCreator( + Function scanCreator) { + return table -> scanCreator.apply(ChainGroupReadTable.withoutQueryAuth(table)); + } + + @Override + public ChainTableBatchScan withReadType(@Nullable RowType readType) { + this.readType = readType; + mainScan.withReadType(readType); + fallbackScan.withReadType(readType); + return this; + } + @Override public ChainTableBatchScan withFilter(Predicate predicate) { super.withFilter(predicate); @@ -268,6 +295,18 @@ public ChainTableBatchScan withFilter(Predicate predicate) { return this; } + @Override + public ChainTableBatchScan withLimit(int limit) { + // Query authorization is applied only after the physical chain inputs have been + // assembled. Pushing a limit into those raw inputs could exhaust the limit with rows + // which are later removed by authorization. ReadBuilderImpl applies the logical limit + // after query authorization instead. + if (!options.queryAuthEnabled()) { + super.withLimit(limit); + } + return this; + } + @Override public ChainTableBatchScan withPartitionFilter(Map partitionSpec) { super.withPartitionFilter(partitionSpec); @@ -335,6 +374,8 @@ public FallbackReadScan skipPreloadTargetSnapshot() { */ @Override public Plan plan() { + TableQueryAuthResult queryAuthResult = + queryAuth.auth(readType == null ? null : readType.getFieldNames()); List splits = new ArrayList<>(); PredicateBuilder builder = new PredicateBuilder(tableSchema.logicalPartitionType()); Set snapshotPartitions = preloadTargetSnapshotSplits(splits); @@ -451,7 +492,8 @@ public Plan plan() { } } } - return new DataFilePlan<>(splits); + Plan plan = new DataFilePlan<>(splits); + return queryAuthResult == null ? plan : queryAuthResult.convertPlan(plan); } @Override @@ -500,13 +542,11 @@ private Set preloadTargetSnapshotSplits(List splits) { fileBranchMapping.put(file.fileName(), options.scanFallbackSnapshotBranch()); } splits.add( - QueryAuthSplit.retainAuth( - split, - new ChainSplit( - dataSplit.partition(), - dataSplit.dataFiles(), - fileBranchMapping, - fileBucketPathMapping))); + new ChainSplit( + dataSplit.partition(), + dataSplit.dataFiles(), + fileBranchMapping, + fileBucketPathMapping)); } snapshotPartitions.addAll( @@ -526,6 +566,9 @@ private DataTableScan newFilteredScan(boolean snapshot) { if (bucketFilter != null) { scan.withBucketFilter(bucketFilter); } + if (readType != null) { + scan.withReadType(readType); + } return scan; } } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java index 02c751a8b4de..35a37c4cbf64 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/ChainTableStreamScan.java @@ -20,6 +20,7 @@ import org.apache.paimon.CoreOptions; import org.apache.paimon.Snapshot; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.codegen.CodeGenUtils; import org.apache.paimon.codegen.RecordComparator; import org.apache.paimon.data.BinaryRow; @@ -37,8 +38,10 @@ import org.apache.paimon.table.source.SnapshotNotExistPlan; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.StreamDataTableScan; +import org.apache.paimon.table.source.TableQueryAuth; import org.apache.paimon.table.source.TableScan; import org.apache.paimon.table.source.snapshot.StartingContext; +import org.apache.paimon.types.RowType; import org.apache.paimon.utils.ChainPartitionProjector; import org.apache.paimon.utils.ChainTableUtils; import org.apache.paimon.utils.Filter; @@ -90,6 +93,11 @@ public class ChainTableStreamScan implements StreamDataTableScan { /** Projector for splitting full partition into group and chain parts. */ private final ChainPartitionProjector partitionProjector; + /** Query authorization for the logical chain table, applied after physical split assembly. */ + private final TableQueryAuth queryAuth; + + @Nullable private RowType readType; + /** Comparator for chain partition keys only. */ private final RecordComparator chainPartitionComparator; @@ -132,7 +140,14 @@ public ChainTableStreamScan(ChainGroupReadTable chainGroupReadTable) { this.batchScan = new ChainGroupReadTable.ChainTableBatchScan( chainGroupReadTable.schema(), chainGroupReadTable); - this.deltaStreamScan = (DataTableStreamScan) chainGroupReadTable.other().newStreamScan(); + this.deltaStreamScan = + (DataTableStreamScan) + ChainGroupReadTable.withoutQueryAuth(chainGroupReadTable.other()) + .newStreamScan(); + this.queryAuth = + chainGroupReadTable + .catalogEnvironment() + .tableQueryAuth(chainGroupReadTable.coreOptions()); ChainTableUtils.validateChainTableForIncrementalRead(chainGroupReadTable); @@ -161,6 +176,13 @@ public StartingContext startingContext() { @Override public TableScan.Plan plan() { + TableQueryAuthResult queryAuthResult = + queryAuth.auth(readType == null ? null : readType.getFieldNames()); + TableScan.Plan plan = planWithoutQueryAuth(); + return queryAuthResult == null ? plan : queryAuthResult.convertPlan(plan); + } + + private TableScan.Plan planWithoutQueryAuth() { if (!startingDone) { return planStarting(); } @@ -336,9 +358,7 @@ private List buildLightweightStartingSplits( for (Map.Entry> entry : snapshotSplitsByPartition.entrySet()) { for (Split split : entry.getValue()) { DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); - allSplits.add( - QueryAuthSplit.retainAuth( - split, ChainSplit.from(dataSplit, snapshotBranch))); + allSplits.add(ChainSplit.from(dataSplit, snapshotBranch)); } } @@ -356,9 +376,7 @@ private List buildLightweightStartingSplits( > 0) { for (Split split : entry.getValue()) { DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); - allSplits.add( - QueryAuthSplit.retainAuth( - split, ChainSplit.from(dataSplit, deltaBranch))); + allSplits.add(ChainSplit.from(dataSplit, deltaBranch)); } } } @@ -531,6 +549,14 @@ public InnerTableScan withFilter(Predicate predicate) { return this; } + @Override + public InnerTableScan withReadType(@Nullable RowType readType) { + this.readType = readType; + batchScan.withReadType(readType); + deltaStreamScan.withReadType(readType); + return this; + } + @Override public InnerTableScan withPartitionFilter(Map partitionSpec) { throw new UnsupportedOperationException( @@ -587,6 +613,7 @@ private static Map pinnedOptions(long snapshotId) { Map options = new HashMap<>(); options.put(CoreOptions.SCAN_SNAPSHOT_ID.key(), String.valueOf(snapshotId)); options.put(CoreOptions.SCAN_MODE.key(), CoreOptions.StartupMode.FROM_SNAPSHOT.toString()); + options.put(CoreOptions.QUERY_AUTH_ENABLED.key(), "false"); return options; } @@ -604,6 +631,9 @@ private void applyPredicatesShardAndBucket(DataTableScan scan) { if (bucketFilter != null) { scan.withBucketFilter(bucketFilter); } + if (readType != null) { + scan.withReadType(readType); + } } @Nullable diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/ChainTableUtils.java b/paimon-core/src/main/java/org/apache/paimon/utils/ChainTableUtils.java index 0cfcc0f26af1..36492a5dc234 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/ChainTableUtils.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/ChainTableUtils.java @@ -377,15 +377,15 @@ public static Predicate createGroupChainPredicate( /** * Builds per-bucket {@link ChainSplit}s from the given snapshot and delta splits. Files that * originate from the snapshot splits are tagged with {@code snapshotBranch}; all other files - * are tagged with {@code deltaBranch}. If an input split carries query authorization, the - * corresponding output split retains it. + * are tagged with {@code deltaBranch}. Query authorization belongs to the logical chain-table + * plan and must be applied by the caller after all physical splits have been assembled. * * @param logicalPartition the logical partition for the resulting ChainSplits * @param snapshotSplits splits from the snapshot branch * @param deltaSplits splits from the delta branch * @param snapshotBranch name of the snapshot branch * @param deltaBranch name of the delta branch - * @return one split per bucket, optionally wrapped with query authorization + * @return one raw chain split per bucket */ public static List buildChainSplits( BinaryRow logicalPartition, @@ -432,7 +432,7 @@ public static List buildChainSplits( .collect(Collectors.toList()), fileBranchMapping, fileBucketPathMapping); - result.add(retainQueryAuth(entry.getValue(), chainSplit)); + result.add(chainSplit); } return result; } @@ -450,15 +450,6 @@ private static Integer addToBucketMap( return totalBuckets; } - private static Split retainQueryAuth(List sourceSplits, Split replacement) { - for (Split sourceSplit : sourceSplits) { - if (sourceSplit instanceof QueryAuthSplit) { - return QueryAuthSplit.retainAuth(sourceSplit, replacement); - } - } - return replacement; - } - /** * Validates that the chain table configuration is compatible with incremental read paths * (streaming read and lookup join). All validation rules for incremental reads should be diff --git a/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java index ec56c930fac6..dffa7de42bf7 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java @@ -23,6 +23,7 @@ import org.apache.paimon.data.BinaryRow; import org.apache.paimon.data.InternalRow; import org.apache.paimon.io.DataFileTestDataGenerator; +import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.source.ChainSplit; @@ -35,6 +36,8 @@ import org.apache.paimon.table.source.Split; import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.JsonSerdeUtil; import org.apache.paimon.utils.SnapshotManager; import org.junit.jupiter.api.Test; @@ -46,8 +49,13 @@ import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -55,14 +63,22 @@ public class ChainGroupReadTableTest { @Test - public void testBatchScanRetainsQueryAuth() throws Exception { + public void testBatchScanUsesLogicalQueryAuth() throws Exception { TableSchema schema = tableSchema(); - PrimaryKeyFileStoreTable snapshotTable = table(schema, "snapshot"); - PrimaryKeyFileStoreTable deltaTable = table(schema, "delta"); + TableQueryAuthResult logicalAuth = maskingAuth("logical"); + CatalogEnvironment catalogEnvironment = catalogEnvironment(logicalAuth); + PrimaryKeyFileStoreTable snapshotTable = + table(schema, "snapshot", true, catalogEnvironment); + PrimaryKeyFileStoreTable deltaTable = table(schema, "delta", true, catalogEnvironment); + PrimaryKeyFileStoreTable rawSnapshotTable = + table(schema, "snapshot", false, catalogEnvironment); + PrimaryKeyFileStoreTable rawDeltaTable = table(schema, "delta", false, catalogEnvironment); + when(snapshotTable.copy(anyMap())).thenReturn(rawSnapshotTable); + when(deltaTable.copy(anyMap())).thenReturn(rawDeltaTable); ChainGroupReadTable chainTable = new ChainGroupReadTable(snapshotTable, deltaTable); BinaryRow partition = BinaryRow.singleColumn(1); - QueryAuthSplit sourceSplit = queryAuthSplit(dataSplit(partition)); + QueryAuthSplit sourceSplit = queryAuthSplit(dataSplit(partition), maskingAuth("physical")); DataTableScan snapshotScan = mock(DataTableScan.class); when(snapshotScan.plan()) .thenReturn(new DataFilePlan<>(Collections.singletonList(sourceSplit))); @@ -75,13 +91,16 @@ public void testBatchScanRetainsQueryAuth() throws Exception { new ChainGroupReadTable.ChainTableBatchScan( schema, chainTable, - table -> table == snapshotTable ? snapshotScan : deltaScan); + table -> table == rawSnapshotTable ? snapshotScan : deltaScan); List splits = scan.plan().splits(); assertThat(splits).hasSize(1); QueryAuthSplit result = (QueryAuthSplit) splits.get(0); - assertThat(result.authResult()).isSameAs(sourceSplit.authResult()); + assertThat(result.authResult()).isSameAs(logicalAuth); + assertThat(result.authResult()).isNotSameAs(sourceSplit.authResult()); assertThat(result.split()).isInstanceOf(ChainSplit.class); + verifyQueryAuthDisabled(snapshotTable); + verifyQueryAuthDisabled(deltaTable); InnerTableRead snapshotRead = mock(InnerTableRead.class); InnerTableRead deltaRead = mock(InnerTableRead.class); @@ -94,6 +113,71 @@ public void testBatchScanRetainsQueryAuth() throws Exception { verify(deltaRead).createReader(result); } + @Test + public void testBatchScanKeepsHistoricalAnchorOutsideAuthorizedPartition() { + TableSchema schema = tableSchema(); + RowType rowType = schema.logicalRowType(); + TableQueryAuthResult logicalAuth = + new TableQueryAuthResult( + Collections.singletonList( + JsonSerdeUtil.toFlatJson( + new PredicateBuilder(rowType).equal(0, 20260720))), + null); + CatalogEnvironment catalogEnvironment = catalogEnvironment(logicalAuth); + PrimaryKeyFileStoreTable snapshotTable = + table(schema, "snapshot", true, catalogEnvironment); + PrimaryKeyFileStoreTable deltaTable = table(schema, "delta", true, catalogEnvironment); + PrimaryKeyFileStoreTable rawSnapshotTable = + table(schema, "snapshot", false, catalogEnvironment); + PrimaryKeyFileStoreTable rawDeltaTable = table(schema, "delta", false, catalogEnvironment); + when(snapshotTable.copy(anyMap())).thenReturn(rawSnapshotTable); + when(deltaTable.copy(anyMap())).thenReturn(rawDeltaTable); + + BinaryRow historicalPartition = BinaryRow.singleColumn(20260719); + BinaryRow authorizedPartition = BinaryRow.singleColumn(20260720); + DataTableScan snapshotScan = mock(DataTableScan.class); + when(snapshotScan.listPartitions()) + .thenReturn(Collections.singletonList(historicalPartition)); + when(snapshotScan.plan()) + .thenReturn( + new DataFilePlan<>( + Collections.singletonList( + queryAuthSplit( + dataSplit(historicalPartition), + maskingAuth("snapshot"))))); + DataTableScan deltaScan = mock(DataTableScan.class); + when(deltaScan.listPartitions()).thenReturn(Collections.singletonList(authorizedPartition)); + when(deltaScan.plan()) + .thenReturn( + new DataFilePlan<>( + Collections.singletonList( + queryAuthSplit( + dataSplit(authorizedPartition), + maskingAuth("delta"))))); + + ChainGroupReadTable chainTable = new ChainGroupReadTable(snapshotTable, deltaTable); + ChainGroupReadTable.ChainTableBatchScan scan = + new ChainGroupReadTable.ChainTableBatchScan( + schema, + chainTable, + table -> table == rawSnapshotTable ? snapshotScan : deltaScan); + scan.skipPreloadTargetSnapshot(); + scan.withLimit(1); + + List splits = scan.plan().splits(); + + assertThat(splits).hasSize(1); + QueryAuthSplit result = (QueryAuthSplit) splits.get(0); + assertThat(result.authResult()).isSameAs(logicalAuth); + ChainSplit chainSplit = (ChainSplit) result.split(); + assertThat(chainSplit.logicalPartition()).isEqualTo(authorizedPartition); + assertThat(chainSplit.dataFiles()).hasSize(2); + verify(snapshotScan, never()).withLimit(anyInt()); + verify(deltaScan, never()).withLimit(anyInt()); + verifyQueryAuthDisabled(snapshotTable); + verifyQueryAuthDisabled(deltaTable); + } + @Test public void testStreamStartingPlanRetainsQueryAuth() { testStreamStartingPlanRetainsQueryAuth(false); @@ -106,15 +190,21 @@ public void testStreamMergedStartingPlanRetainsQueryAuth() { private void testStreamStartingPlanRetainsQueryAuth(boolean mergeSnapshot) { TableSchema schema = tableSchema(mergeSnapshot); - PrimaryKeyFileStoreTable snapshotTable = table(schema, "snapshot"); - PrimaryKeyFileStoreTable deltaTable = table(schema, "delta"); + TableQueryAuthResult logicalAuth = maskingAuth("logical"); + CatalogEnvironment catalogEnvironment = catalogEnvironment(logicalAuth); + PrimaryKeyFileStoreTable snapshotTable = + table(schema, "snapshot", true, catalogEnvironment); + PrimaryKeyFileStoreTable deltaTable = table(schema, "delta", true, catalogEnvironment); + PrimaryKeyFileStoreTable rawSnapshotTable = + table(schema, "snapshot", false, catalogEnvironment); + PrimaryKeyFileStoreTable rawDeltaTable = table(schema, "delta", false, catalogEnvironment); DataTableScan unusedScan = mock(DataTableScan.class); - when(snapshotTable.newScan()).thenReturn(unusedScan); - when(deltaTable.newScan()).thenReturn(unusedScan); + when(rawSnapshotTable.newScan()).thenReturn(unusedScan); + when(rawDeltaTable.newScan()).thenReturn(unusedScan); DataTableStreamScan deltaStreamScan = mock(DataTableStreamScan.class); - when(deltaTable.newStreamScan()).thenReturn(deltaStreamScan); + when(rawDeltaTable.newStreamScan()).thenReturn(deltaStreamScan); SnapshotManager snapshotManager = mock(SnapshotManager.class); when(snapshotManager.latestSnapshotId()).thenReturn(null); @@ -125,13 +215,22 @@ private void testStreamStartingPlanRetainsQueryAuth(boolean mergeSnapshot) { when(deltaTable.snapshotManager()).thenReturn(deltaSnapshotManager); BinaryRow partition = BinaryRow.singleColumn(1); - QueryAuthSplit sourceSplit = queryAuthSplit(dataSplit(partition)); + QueryAuthSplit sourceSplit = queryAuthSplit(dataSplit(partition), maskingAuth("physical")); DataTableScan pinnedDeltaScan = mock(DataTableScan.class); when(pinnedDeltaScan.plan()) .thenReturn(new DataFilePlan<>(Collections.singletonList(sourceSplit))); - PrimaryKeyFileStoreTable pinnedDeltaTable = table(schema, "delta"); + PrimaryKeyFileStoreTable pinnedDeltaTable = + table(schema, "delta", false, catalogEnvironment); when(pinnedDeltaTable.newScan()).thenReturn(pinnedDeltaScan); - when(deltaTable.copy(anyMap())).thenReturn(pinnedDeltaTable); + when(snapshotTable.copy(anyMap())).thenReturn(rawSnapshotTable); + when(deltaTable.copy(anyMap())) + .thenAnswer( + invocation -> { + Map options = invocation.getArgument(0); + return options.containsKey(CoreOptions.SCAN_SNAPSHOT_ID.key()) + ? pinnedDeltaTable + : rawDeltaTable; + }); ChainTableStreamScan scan = new ChainTableStreamScan(new ChainGroupReadTable(snapshotTable, deltaTable)); @@ -139,20 +238,48 @@ private void testStreamStartingPlanRetainsQueryAuth(boolean mergeSnapshot) { assertThat(splits).hasSize(1); QueryAuthSplit result = (QueryAuthSplit) splits.get(0); - assertThat(result.authResult()).isSameAs(sourceSplit.authResult()); + assertThat(result.authResult()).isSameAs(logicalAuth); + assertThat(result.authResult()).isNotSameAs(sourceSplit.authResult()); assertThat(result.split()).isInstanceOf(ChainSplit.class); verify(deltaStreamScan).restore(2L); + verifyQueryAuthDisabled(snapshotTable); + verifyQueryAuthDisabled(deltaTable); } - private static PrimaryKeyFileStoreTable table(TableSchema schema, String branch) { + private static PrimaryKeyFileStoreTable table( + TableSchema schema, + String branch, + boolean queryAuthEnabled, + CatalogEnvironment catalogEnvironment) { PrimaryKeyFileStoreTable table = mock(PrimaryKeyFileStoreTable.class); when(table.schema()).thenReturn(schema); Map options = new HashMap<>(schema.options()); options.put(CoreOptions.BRANCH.key(), branch); + options.put(CoreOptions.QUERY_AUTH_ENABLED.key(), String.valueOf(queryAuthEnabled)); when(table.coreOptions()).thenReturn(CoreOptions.fromMap(options)); + when(table.catalogEnvironment()).thenReturn(catalogEnvironment); return table; } + private static CatalogEnvironment catalogEnvironment(TableQueryAuthResult authResult) { + CatalogEnvironment catalogEnvironment = mock(CatalogEnvironment.class); + when(catalogEnvironment.tableQueryAuth(any(CoreOptions.class))) + .thenReturn(select -> authResult); + return catalogEnvironment; + } + + private static void verifyQueryAuthDisabled(PrimaryKeyFileStoreTable table) { + verify(table, atLeastOnce()) + .copy( + argThat( + (Map options) -> + "false" + .equals( + options.get( + CoreOptions.QUERY_AUTH_ENABLED + .key())))); + } + private static TableSchema tableSchema() { return tableSchema(false); } @@ -166,6 +293,9 @@ private static TableSchema tableSchema(boolean mergeSnapshot) { Map options = new HashMap<>(); options.put(CoreOptions.SCAN_FALLBACK_SNAPSHOT_BRANCH.key(), "snapshot"); options.put(CoreOptions.SCAN_FALLBACK_DELTA_BRANCH.key(), "delta"); + options.put(CoreOptions.QUERY_AUTH_ENABLED.key(), "true"); + options.put(CoreOptions.PARTITION_TIMESTAMP_PATTERN.key(), "$pt"); + options.put(CoreOptions.PARTITION_TIMESTAMP_FORMATTER.key(), "yyyyMMdd"); options.put( CoreOptions.CHAIN_TABLE_STREAMING_MERGE_SNAPSHOT.key(), String.valueOf(mergeSnapshot)); @@ -192,9 +322,11 @@ private static DataSplit dataSplit(BinaryRow partition) { .build(); } - private static QueryAuthSplit queryAuthSplit(DataSplit split) { - TableQueryAuthResult authResult = - new TableQueryAuthResult(null, Collections.singletonMap("v", "mask")); + private static QueryAuthSplit queryAuthSplit(DataSplit split, TableQueryAuthResult authResult) { return new QueryAuthSplit(split, authResult); } + + private static TableQueryAuthResult maskingAuth(String value) { + return new TableQueryAuthResult(null, Collections.singletonMap("v", value)); + } } From cdad524c1f8c492a951e3524f1171c721ab4e43d Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 14:59:08 +0800 Subject: [PATCH 09/31] [flink] Unwrap query auth split for task assignment --- .../source/ContinuousFileSplitEnumerator.java | 13 ++++++++----- .../ContinuousFileSplitEnumeratorTest.java | 17 +++++++++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) 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 75c75cfddb71..6c2a9b5f4dee 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 @@ -30,7 +30,9 @@ 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.QueryAuthSplit; import org.apache.paimon.table.source.SnapshotNotExistPlan; +import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.StreamTableScan; import org.apache.paimon.table.source.TableScan; @@ -328,13 +330,14 @@ protected synchronized void assignSplits() { } protected int assignSuggestedTask(FileStoreSourceSplit split) { + Split unwrappedSplit = QueryAuthSplit.unwrap(split.split()); int task; - if (split.split() instanceof DataSplit) { - task = assignSuggestedTask((DataSplit) split.split()); - } else if (split.split() instanceof ChainSplit) { - task = assignSuggestedTask((ChainSplit) split.split()); + if (unwrappedSplit instanceof DataSplit) { + task = assignSuggestedTask((DataSplit) unwrappedSplit); + } else if (unwrappedSplit instanceof ChainSplit) { + task = assignSuggestedTask((ChainSplit) unwrappedSplit); } else { - task = assignSuggestedTask((IncrementalSplit) split.split()); + task = assignSuggestedTask((IncrementalSplit) unwrappedSplit); } // 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/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 0a50038bb405..bb807fd123a6 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,13 @@ 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.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; @@ -904,6 +906,21 @@ public void testIncrementalSplitWithRealBucketAssignedByBucket() { assertThat(enumerator.assignSuggestedTask(split)).isEqualTo(4 % parallelism); } + @Test + public void testQueryAuthSplitAssignedByWrappedDataSplit() { + int parallelism = 3; + ContinuousFileSplitEnumerator enumerator = buildEnumerator(parallelism); + DataSplit dataSplit = createDataSplit(1L, 4, Collections.emptyList()); + QueryAuthSplit queryAuthSplit = + new QueryAuthSplit( + dataSplit, + new TableQueryAuthResult(null, Collections.singletonMap("f0", "mask"))); + FileStoreSourceSplit sourceSplit = new FileStoreSourceSplit("split", queryAuthSplit); + + assertThat(enumerator.assignSuggestedTask(sourceSplit)).isEqualTo(4 % parallelism); + assertThat(sourceSplit.split()).isSameAs(queryAuthSplit); + } + private ContinuousFileSplitEnumerator buildEnumerator(int parallelism) { return new Builder() .setSplitEnumeratorContext(getSplitEnumeratorContext(parallelism)) From 9555d89125aca9cea7e2a662db881cd18e861e97 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:00:43 +0800 Subject: [PATCH 10/31] [flink] Read snapshot id from query auth split --- .../paimon/flink/utils/TableScanUtils.java | 7 ++- .../flink/utils/TableScanUtilsTest.java | 56 +++++++++++++++++++ 2 files changed, 61 insertions(+), 2 deletions(-) create mode 100644 paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/utils/TableScanUtilsTest.java 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 ca901ad1bfdb..ceb407bb69f4 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.options.Options; import org.apache.paimon.table.Table; 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.TableScan; import java.util.HashMap; @@ -67,8 +69,9 @@ public static void streamingReadingValidate(Table table) { /** Get snapshot id from {@link FileStoreSourceSplit}. */ public static Optional getSnapshotId(FileStoreSourceSplit split) { - if (split.split() instanceof DataSplit) { - return Optional.of(((DataSplit) split.split()).snapshotId()); + Split unwrappedSplit = QueryAuthSplit.unwrap(split.split()); + if (unwrappedSplit instanceof DataSplit) { + return Optional.of(((DataSplit) unwrappedSplit).snapshotId()); } return Optional.empty(); } 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 000000000000..484fe354c97d --- /dev/null +++ b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/utils/TableScanUtilsTest.java @@ -0,0 +1,56 @@ +/* + * 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.flink.source.FileStoreSourceSplit; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.QueryAuthSplit; + +import org.junit.jupiter.api.Test; + +import java.util.Collections; + +import static org.apache.paimon.io.DataFileTestUtils.row; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link TableScanUtils}. */ +public class TableScanUtilsTest { + + @Test + public void testGetSnapshotIdFromQueryAuthSplit() { + DataSplit dataSplit = + DataSplit.builder() + .withSnapshot(5L) + .withPartition(row(1)) + .withBucket(0) + .withDataFiles(Collections.emptyList()) + .isStreaming(true) + .withBucketPath("") + .build(); + QueryAuthSplit queryAuthSplit = + new QueryAuthSplit( + dataSplit, + new TableQueryAuthResult(null, Collections.singletonMap("f0", "mask"))); + FileStoreSourceSplit sourceSplit = new FileStoreSourceSplit("split", queryAuthSplit); + + assertThat(TableScanUtils.getSnapshotId(sourceSplit)).contains(5L); + assertThat(sourceSplit.split()).isSameAs(queryAuthSplit); + } +} From 138542975edebf6b5d1019e3e36e36474d920b48 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:03:40 +0800 Subject: [PATCH 11/31] [flink] Support query auth splits in aligned mode --- .../AlignedContinuousFileSplitEnumerator.java | 4 +- .../assigners/AlignedSplitAssigner.java | 16 +++++--- ...gnedContinuousFileSplitEnumeratorTest.java | 40 +++++++++++++++++++ 3 files changed, 52 insertions(+), 8 deletions(-) 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 c728134c07ed..76d693f11195 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 @@ -23,8 +23,8 @@ import org.apache.paimon.flink.source.PendingSplitsCheckpoint; import org.apache.paimon.flink.source.assigners.AlignedSplitAssigner; 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.QueryAuthSplit; import org.apache.paimon.table.source.SnapshotNotExistPlan; import org.apache.paimon.table.source.StreamTableScan; import org.apache.paimon.table.source.TableScan; @@ -123,7 +123,7 @@ protected void addSplits(Collection splits) { Map> splitsBySnapshot = new TreeMap<>(); for (FileStoreSourceSplit split : splits) { - long snapshotId = ((DataSplit) split.split()).snapshotId(); + long snapshotId = QueryAuthSplit.unwrapDataSplit(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 648758f83846..e8984c01a5f8 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 @@ 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.QueryAuthSplit; import org.apache.paimon.utils.Preconditions; import javax.annotation.Nullable; @@ -66,9 +67,10 @@ public List getNext(int subtask, @Nullable String hostname @Override public void addSplit(int subtask, FileStoreSourceSplit splits) { - long snapshotId = ((DataSplit) splits.split()).snapshotId(); + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(splits.split()); + long snapshotId = dataSplit.snapshotId(); PendingSnapshot last = pendingSplitAssignment.peekLast(); - boolean isPlaceholder = splits.split() instanceof PlaceholderSplit; + boolean isPlaceholder = dataSplit instanceof PlaceholderSplit; if (last == null || last.snapshotId != snapshotId) { last = new PendingSnapshot(snapshotId, isPlaceholder, new HashMap<>()); last.add(subtask, splits); @@ -85,8 +87,9 @@ public void addSplitsBack(int suggestedTask, List splits) return; } - long snapshotId = ((DataSplit) splits.get(0).split()).snapshotId(); - boolean isPlaceholder = splits.get(0).split() instanceof PlaceholderSplit; + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(splits.get(0).split()); + long snapshotId = dataSplit.snapshotId(); + boolean isPlaceholder = dataSplit instanceof PlaceholderSplit; PendingSnapshot head = pendingSplitAssignment.peek(); if (head == null || snapshotId != head.snapshotId) { head = new PendingSnapshot(snapshotId, isPlaceholder, new HashMap<>()); @@ -154,7 +157,7 @@ public List remove(int subtask) { public void add(int subtask, FileStoreSourceSplit split) { Preconditions.checkArgument( - ((DataSplit) split.split()).snapshotId() == snapshotId, + QueryAuthSplit.unwrapDataSplit(split.split()).snapshotId() == snapshotId, "SnapshotId not equal. This is a bug, please file an issue."); subtaskSplits.computeIfAbsent(subtask, id -> new ArrayList<>()).add(split); } @@ -166,7 +169,8 @@ public void addAll(int subtask, List splits) { splits.forEach( split -> Preconditions.checkArgument( - ((DataSplit) split.split()).snapshotId() == snapshotId, + QueryAuthSplit.unwrapDataSplit(split.split()).snapshotId() + == snapshotId, "SnapshotId not equal")); subtaskSplits.put(subtask, splits); } diff --git a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumeratorTest.java b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumeratorTest.java index ad71fa81ca12..82ec9357af39 100644 --- a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumeratorTest.java +++ b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/align/AlignedContinuousFileSplitEnumeratorTest.java @@ -19,6 +19,7 @@ package org.apache.paimon.flink.source.align; import org.apache.paimon.CoreOptions; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.consumer.ConsumerManager; import org.apache.paimon.flink.source.FileSplitEnumeratorTestBase; import org.apache.paimon.flink.source.FileStoreSourceSplit; @@ -34,6 +35,7 @@ import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.FileStoreTableFactory; import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.StreamTableScan; import org.apache.paimon.types.DataType; import org.apache.paimon.types.DataTypes; @@ -139,6 +141,35 @@ public void testSplitsAssignedBySnapshot() throws Exception { assertThat(assignments.get(1).getAssignedSplits()).containsExactly(expectedSplits.get(2)); } + @Test + public void testQueryAuthSplitsAssignedBySnapshot() { + final TestingSplitEnumeratorContext context = + getSplitEnumeratorContext(1); + FileStoreSourceSplit first = + queryAuthSplit(createSnapshotSplit(1, 0, Collections.emptyList())); + FileStoreSourceSplit second = + queryAuthSplit(createSnapshotSplit(1, 0, Collections.emptyList())); + + final AlignedContinuousFileSplitEnumerator enumerator = + new Builder() + .setSplitEnumeratorContext(context) + .setInitialSplits(Arrays.asList(first, second)) + .build(); + + enumerator.handleSplitRequest(0, "test-host"); + List assignedSplits = + context.getSplitAssignments().get(0).getAssignedSplits(); + assertThat(assignedSplits).containsExactly(first, second); + assertThat(assignedSplits) + .allSatisfy(split -> assertThat(split.split()).isInstanceOf(QueryAuthSplit.class)); + + context.getSplitAssignments().clear(); + enumerator.addSplitsBack(assignedSplits, 0); + enumerator.handleSplitRequest(0, "test-host"); + assertThat(context.getSplitAssignments().get(0).getAssignedSplits()) + .containsExactly(first, second); + } + @Test public void testEnumeratorSnapshotState() throws Exception { final TestingSplitEnumeratorContext context = @@ -257,6 +288,15 @@ public AlignedContinuousFileSplitEnumerator build() { } } + private static FileStoreSourceSplit queryAuthSplit(FileStoreSourceSplit sourceSplit) { + return new FileStoreSourceSplit( + sourceSplit.splitId(), + new QueryAuthSplit( + sourceSplit.split(), + new TableQueryAuthResult(null, Collections.singletonMap("f0", "mask"))), + sourceSplit.recordsToSkip()); + } + @Override protected FileStoreSourceSplit createSnapshotSplit( int snapshotId, int bucket, List files, int... partitions) { From d316fa3643b579e20a002736b5cf8ef27187efab Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:09:18 +0800 Subject: [PATCH 12/31] [flink] Support query auth in dynamic partition pruning --- .../DynamicPartitionPruningAssigner.java | 3 +- .../assigners/PreAssignSplitAssigner.java | 3 +- .../DynamicPartitionPruningAssignerTest.java | 113 ++++++++++++++++++ 3 files changed, 117 insertions(+), 2 deletions(-) create mode 100644 paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssignerTest.java 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 9221f1f27eb6..463e68803bbe 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 @@ -23,6 +23,7 @@ import org.apache.paimon.flink.FlinkRowData; import org.apache.paimon.flink.source.FileStoreSourceSplit; import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.flink.api.connector.source.SourceEvent; import org.apache.flink.table.connector.source.DynamicFilteringData; @@ -115,7 +116,7 @@ public int numberOfRemainingSplits() { } private boolean filter(FileStoreSourceSplit sourceSplit) { - DataSplit dataSplit = (DataSplit) sourceSplit.split(); + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(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/assigners/PreAssignSplitAssigner.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java index 24ac4a291164..281e52cb6728 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 @@ -23,6 +23,7 @@ import org.apache.paimon.flink.FlinkRowData; import org.apache.paimon.flink.source.FileStoreSourceSplit; import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.utils.BinPacking; import org.apache.paimon.utils.SerializableFunction; @@ -293,7 +294,7 @@ private static boolean filter( Projection partitionRowProjection, DynamicFilteringData dynamicFilteringData, FileStoreSourceSplit sourceSplit) { - DataSplit dataSplit = (DataSplit) sourceSplit.split(); + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(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/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 000000000000..fc3eb886b03a --- /dev/null +++ b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/assigners/DynamicPartitionPruningAssignerTest.java @@ -0,0 +1,113 @@ +/* + * 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.codegen.Projection; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.flink.source.FileStoreSourceSplit; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.QueryAuthSplit; + +import org.apache.flink.api.java.typeutils.GenericTypeInfo; +import org.apache.flink.table.connector.source.DynamicFilteringData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.types.logical.IntType; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; + +import static org.apache.paimon.io.DataFileTestUtils.row; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for dynamic partition pruning assigners. */ +public class DynamicPartitionPruningAssignerTest { + + private static final Projection IDENTITY_PROJECTION = input -> (BinaryRow) input; + + @Test + public void testPreAssignQueryAuthSplits() { + FileStoreSourceSplit selected = queryAuthSplit("selected", 1); + FileStoreSourceSplit filtered = queryAuthSplit("filtered", 2); + + PreAssignSplitAssigner assigner = + new PreAssignSplitAssigner( + 10, + 1, + Arrays.asList(selected, filtered), + IDENTITY_PROJECTION, + new TestDynamicFilteringData(1), + split -> 1L); + + assertThat(assigner.getNext(0, null)).containsExactly(selected); + assertThat(selected.split()).isInstanceOf(QueryAuthSplit.class); + } + + @Test + public void testPreemptiveQueryAuthSplits() { + FileStoreSourceSplit filtered = queryAuthSplit("filtered", 2); + FileStoreSourceSplit selected = queryAuthSplit("selected", 1); + DynamicPartitionPruningAssigner assigner = + new DynamicPartitionPruningAssigner( + new FIFOSplitAssigner(Arrays.asList(filtered, selected)), + IDENTITY_PROJECTION, + new TestDynamicFilteringData(1)); + + assertThat(assigner.remainingSplits()).containsExactly(selected); + assertThat(assigner.getNext(0, null)).containsExactly(selected); + assertThat(selected.split()).isInstanceOf(QueryAuthSplit.class); + } + + private static FileStoreSourceSplit queryAuthSplit(String id, int partition) { + DataSplit dataSplit = + DataSplit.builder() + .withSnapshot(1L) + .withPartition(row(partition)) + .withBucket(0) + .withDataFiles(Collections.emptyList()) + .isStreaming(false) + .withBucketPath("") + .build(); + return new FileStoreSourceSplit( + id, + new QueryAuthSplit( + dataSplit, + new TableQueryAuthResult(null, Collections.singletonMap("f0", "mask")))); + } + + private static class TestDynamicFilteringData extends DynamicFilteringData { + + private final int selectedPartition; + + private TestDynamicFilteringData(int selectedPartition) { + super( + new GenericTypeInfo<>(RowData.class), + org.apache.flink.table.types.logical.RowType.of(new IntType()), + Collections.emptyList(), + true); + this.selectedPartition = selectedPartition; + } + + @Override + public boolean contains(RowData row) { + return row.getInt(0) == selectedPartition; + } + } +} From 35b4d3de94c595071b926efeef58ac9992e965ed Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:05:44 +0800 Subject: [PATCH 13/31] [flink] Handle query auth splits in dedicated source --- .../flink/source/operator/MonitorSource.java | 12 ++-- .../flink/source/operator/ReadOperator.java | 4 +- .../source/operator/OperatorSourceTest.java | 68 +++++++++++++++++++ 3 files changed, 79 insertions(+), 5 deletions(-) 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 b6de64472b5b..f100a712ca2c 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 @@ -18,6 +18,7 @@ package org.apache.paimon.flink.source.operator; +import org.apache.paimon.data.BinaryRow; import org.apache.paimon.flink.NestedProjectedRowData; import org.apache.paimon.flink.source.AbstractNonCoordinatedSource; import org.apache.paimon.flink.source.AbstractNonCoordinatedSourceReader; @@ -30,6 +31,7 @@ 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.QueryAuthSplit; import org.apache.paimon.table.source.ReadBuilder; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.StreamTableScan; @@ -314,9 +316,11 @@ private static DataStream shuffleOrdered( } return ChannelComputer.select(key.f1, numPartitions); }, - split -> { - DataSplit dataSplit = (DataSplit) split; - return Tuple2.of(dataSplit.partition(), dataSplit.bucket()); - }); + MonitorSource::splitKey); + } + + static Tuple2 splitKey(Split split) { + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); + 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 30b2ab0d62e4..7981c97af35a 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 @@ -25,6 +25,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.QueryAuthSplit; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.TableRead; import org.apache.paimon.utils.CloseableIterator; @@ -114,8 +115,9 @@ public void processElement(StreamRecord record) throws Exception { Split split = record.getValue(); // update metric when reading a new split + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); long eventTime = - ((DataSplit) split) + dataSplit .earliestFileCreationEpochMillis() .orElse(FileStoreSourceReaderMetrics.UNDEFINED); sourceReaderMetrics.recordSnapshotUpdate(eventTime); diff --git a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/operator/OperatorSourceTest.java b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/operator/OperatorSourceTest.java index 9c57f27b866d..c9d2cb4a6435 100644 --- a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/operator/OperatorSourceTest.java +++ b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/operator/OperatorSourceTest.java @@ -22,13 +22,20 @@ import org.apache.paimon.catalog.CatalogContext; import org.apache.paimon.catalog.CatalogFactory; import org.apache.paimon.catalog.Identifier; +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.schema.Schema; import org.apache.paimon.table.Table; import org.apache.paimon.table.sink.BatchTableCommit; import org.apache.paimon.table.sink.BatchTableWrite; import org.apache.paimon.table.sink.BatchWriteBuilder; +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; @@ -62,6 +69,7 @@ import java.nio.file.Path; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.TimeUnit; @@ -218,6 +226,66 @@ public void testReadOperator() throws Exception { new StreamRecord<>(GenericRowData.of(2, 2, 2))); } + @Test + public void testQueryAuthSplitInDedicatedSourcePath() throws Exception { + DataSplit dataSplit = + DataSplit.builder() + .withSnapshot(1L) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withDataFiles(Collections.emptyList()) + .build(); + QueryAuthSplit queryAuthSplit = new QueryAuthSplit(dataSplit, null); + assertThat(MonitorSource.splitKey(queryAuthSplit)) + .isEqualTo(MonitorSource.splitKey(dataSplit)); + + AtomicReference readerSplit = new AtomicReference<>(); + TableRead capturingRead = + new TableRead() { + @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 createReader(Split split) { + readerSplit.set(split); + return new RecordReader() { + @Override + public RecordIterator readBatch() { + return null; + } + + @Override + public void close() {} + }; + } + }; + + ReadOperator readOperator = new ReadOperator(() -> capturingRead, null, null); + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>(readOperator); + harness.setup( + InternalSerializers.create( + RowType.of(new IntType(), new IntType(), new IntType()))); + harness.open(); + harness.processElement(new StreamRecord<>(queryAuthSplit)); + + assertThat(readerSplit.get()).isSameAs(queryAuthSplit); + assertThat(harness.getOutput()).isEmpty(); + } + @Test public void testReadOperatorWithLimit() throws Exception { ReadOperator readOperator = From 29e9d5977df0edaee580360c6f6d67c90641cdd9 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:11:01 +0800 Subject: [PATCH 14/31] [flink-cdc] Handle query auth splits in enumerator --- .../enumerator/CDCSourceEnumerator.java | 10 ++- .../enumerator/CDCSourceEnumeratorTest.java | 82 +++++++++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) diff --git a/paimon-flink/paimon-flink-cdc/src/main/java/org/apache/paimon/flink/pipeline/cdc/source/enumerator/CDCSourceEnumerator.java b/paimon-flink/paimon-flink-cdc/src/main/java/org/apache/paimon/flink/pipeline/cdc/source/enumerator/CDCSourceEnumerator.java index 5e62aeda291a..ebc8fe6a23b1 100644 --- a/paimon-flink/paimon-flink-cdc/src/main/java/org/apache/paimon/flink/pipeline/cdc/source/enumerator/CDCSourceEnumerator.java +++ b/paimon-flink/paimon-flink-cdc/src/main/java/org/apache/paimon/flink/pipeline/cdc/source/enumerator/CDCSourceEnumerator.java @@ -30,6 +30,7 @@ import org.apache.paimon.table.Table; import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.EndOfScanException; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.SnapshotNotExistPlan; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.StreamDataTableScan; @@ -315,8 +316,13 @@ protected TableAwareFileStoreSourceSplit toTableAwareSplit( FileStoreTable table, Identifier identifier, @Nullable Long lastSchemaId) { - Preconditions.checkState(split instanceof DataSplit); - long snapshotId = ((DataSplit) split).snapshotId(); + Split unwrappedSplit = QueryAuthSplit.unwrap(split); + Preconditions.checkState( + unwrappedSplit instanceof DataSplit, + "CDC source expects DataSplit, but got %s.", + unwrappedSplit == null ? "null" : unwrappedSplit.getClass().getName()); + DataSplit dataSplit = (DataSplit) unwrappedSplit; + long snapshotId = dataSplit.snapshotId(); long schemaId = table.snapshot(snapshotId).schemaId(); return new TableAwareFileStoreSourceSplit( splitId, split, 0, identifier, lastSchemaId, schemaId); diff --git a/paimon-flink/paimon-flink-cdc/src/test/java/org/apache/paimon/flink/pipeline/cdc/source/enumerator/CDCSourceEnumeratorTest.java b/paimon-flink/paimon-flink-cdc/src/test/java/org/apache/paimon/flink/pipeline/cdc/source/enumerator/CDCSourceEnumeratorTest.java index 8d5f227ef32e..e539956dd6e0 100644 --- a/paimon-flink/paimon-flink-cdc/src/test/java/org/apache/paimon/flink/pipeline/cdc/source/enumerator/CDCSourceEnumeratorTest.java +++ b/paimon-flink/paimon-flink-cdc/src/test/java/org/apache/paimon/flink/pipeline/cdc/source/enumerator/CDCSourceEnumeratorTest.java @@ -19,6 +19,7 @@ package org.apache.paimon.flink.pipeline.cdc.source.enumerator; import org.apache.paimon.CoreOptions; +import org.apache.paimon.Snapshot; import org.apache.paimon.catalog.Catalog; import org.apache.paimon.catalog.CatalogContext; import org.apache.paimon.catalog.CatalogFactory; @@ -30,9 +31,12 @@ import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.options.Options; import org.apache.paimon.schema.Schema; +import org.apache.paimon.table.FallbackReadFileStoreTable; import org.apache.paimon.table.FileStoreTable; 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.Split; import org.apache.paimon.table.source.StreamDataTableScan; import org.apache.paimon.table.source.TableScan; @@ -61,6 +65,7 @@ import static org.apache.paimon.io.DataFileTestUtils.row; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link CDCSourceEnumerator}. */ public class CDCSourceEnumeratorTest @@ -321,6 +326,83 @@ public void testSnapshotEnumerator() { .containsExactly(splits.get(0), splits.get(1)); } + @Test + public void testQueryAuthSplitKeepsWrapperAndRejectsNonDataSplit() throws Exception { + Options options = new Options(); + options.setString("warehouse", warehouseFolder.toAbsolutePath().toString()); + Configuration cdcConfig = new Configuration(); + cdcConfig.set(toCDCOption(CDCOptions.DATABASE), DATABASE); + CDCSourceEnumerator enumerator = + new CDCSourceEnumerator( + getSplitEnumeratorContext(1), + new org.apache.flink.configuration.Configuration(), + 1L, + CatalogContext.create(options), + cdcConfig, + null); + + Identifier identifier = Identifier.create(DATABASE, TABLE + 0); + FileStoreTable wrappedTable = (FileStoreTable) catalog.getTable(identifier); + FileStoreTable table = + new FallbackReadFileStoreTable(wrappedTable, wrappedTable, true) { + @Override + public Snapshot snapshot(long snapshotId) { + return new Snapshot( + 0, + snapshotId, + 7L, + null, + null, + null, + null, + null, + null, + null, + "user", + 0L, + Snapshot.CommitKind.APPEND, + 0L, + 0L, + 0L, + null, + null, + null, + null, + null, + null); + } + }; + DataSplit dataSplit = createDataSplit(5L, 0, Collections.emptyList()); + QueryAuthSplit queryAuthSplit = new QueryAuthSplit(dataSplit, null); + + TableAwareFileStoreSourceSplit result = + enumerator.toTableAwareSplit("split-1", queryAuthSplit, table, identifier, null); + + assertThat(result.split()).isSameAs(queryAuthSplit); + assertThat(result.getSchemaId()).isEqualTo(7L); + + IncrementalSplit incrementalSplit = + new IncrementalSplit( + 5L, + row(0), + 0, + 1, + Collections.emptyList(), + null, + Collections.emptyList(), + null, + true); + assertThatThrownBy( + () -> + enumerator.toTableAwareSplit( + "split-2", incrementalSplit, table, identifier, null)) + .isInstanceOf(IllegalStateException.class) + .hasMessage( + "CDC source expects DataSplit, but got " + + IncrementalSplit.class.getName() + + "."); + } + @Test public void testAssignSameTableSplitToSameTask() { final TestingSplitEnumeratorContext context = From efa166b5627462c6e3fcc8b14b5eb489627ab2ce Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:16:23 +0800 Subject: [PATCH 15/31] [spark] Read metadata from query auth splits --- .../spark/PaimonRecordReaderIterator.scala | 13 ++++---- .../spark/SparkCatalogWithRestTest.java | 30 +++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala index 93256df2c00f..770874ad84f3 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala @@ -23,7 +23,7 @@ import org.apache.paimon.fs.Path import org.apache.paimon.reader.{FileRecordIterator, RecordReader, ScoreRecordIterator, ScoreRecordReader} import org.apache.paimon.spark.schema.PaimonMetadataColumn import org.apache.paimon.spark.schema.PaimonMetadataColumn.{PARTITION_AND_BUCKET_META_COLUMNS, PATH_AND_INDEX_META_COLUMNS, VECTOR_SEARCH_META_COLUMN_NAMES} -import org.apache.paimon.table.source.{DataSplit, Split} +import org.apache.paimon.table.source.{DataSplit, QueryAuthSplit, Split} import org.apache.paimon.utils.{CloseableIterator, Preconditions} import org.apache.spark.sql.PaimonUtils @@ -37,12 +37,15 @@ case class PaimonRecordReaderIterator( split: Split) extends CloseableIterator[PaimonInternalRow] { + private val metadataSplit = QueryAuthSplit.unwrap(split) + if ( - metadataColumns.exists(c => PARTITION_AND_BUCKET_META_COLUMNS.contains(c.name)) && !split + metadataColumns.exists( + c => PARTITION_AND_BUCKET_META_COLUMNS.contains(c.name)) && !metadataSplit .isInstanceOf[DataSplit] ) { throw new RuntimeException( - "There need be DataSplit when path and index metadata columns are required") + "A DataSplit is required when partition or bucket metadata columns are requested") } private val needMetadata = metadataColumns.nonEmpty @@ -169,9 +172,9 @@ case class PaimonRecordReaderIterator( case PaimonMetadataColumn.FILE_PATH_COLUMN => metadataRow.setField(index, BinaryString.fromString(lastFilePath.toString)) case PaimonMetadataColumn.PARTITION_COLUMN => - metadataRow.setField(index, split.asInstanceOf[DataSplit].partition()) + metadataRow.setField(index, metadataSplit.asInstanceOf[DataSplit].partition()) case PaimonMetadataColumn.BUCKET_COLUMN => - metadataRow.setField(index, split.asInstanceOf[DataSplit].bucket()) + metadataRow.setField(index, metadataSplit.asInstanceOf[DataSplit].bucket()) } } } diff --git a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java index fcde06202ed0..f30989569e32 100644 --- a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java +++ b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java @@ -733,6 +733,36 @@ public void testRowFilterBucketKeyTable() { .isEqualTo("[[3,c]]"); } + @Test + public void testRowFilterWithPartitionAndBucketMetadata() { + spark.sql( + "CREATE TABLE t_auth_metadata (id INT, name STRING, pt STRING) " + + "PARTITIONED BY (pt) TBLPROPERTIES " + + "('bucket'='2', 'bucket-key'='id', 'query-auth.enabled'='true')"); + spark.sql( + "INSERT INTO t_auth_metadata VALUES " + + "(1, 'blocked', 'p1'), (2, 'allowed', 'p2')"); + + Predicate idEq2Predicate = + LeafPredicate.of( + new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), + Equal.INSTANCE, + Collections.singletonList(2)); + restCatalogServer.setRowFilterAuth( + Identifier.create("db2", "t_auth_metadata"), + Collections.singletonList(idEq2Predicate)); + + List rows = + spark.sql( + "SELECT id, __paimon_partition, __paimon_bucket " + + "FROM t_auth_metadata ORDER BY id") + .collectAsList(); + assertThat(rows).hasSize(1); + assertThat(rows.get(0).getInt(0)).isEqualTo(2); + assertThat(rows.get(0).getStruct(1).getString(0)).isEqualTo("p2"); + assertThat(rows.get(0).getInt(2)).isBetween(0, 1); + } + @Test public void testRowFilterDeletionVectorsTable() { // Deletion-vectors table: deleted rows excluded via the deletion vector, then filtered. From bc5e71d3069216dbb2e215a46846c976fe4ed474 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:02:08 +0800 Subject: [PATCH 16/31] [spark] Preserve query auth in streaming splits --- .../spark/sources/PaimonReadLimits.scala | 2 +- .../paimon/spark/sources/StreamHelper.scala | 20 +++++---- .../spark/SparkCatalogWithRestTest.java | 41 +++++++++++++++++++ 3 files changed, 54 insertions(+), 9 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonReadLimits.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonReadLimits.scala index 241aeaef54db..66df521e61c0 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonReadLimits.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonReadLimits.scala @@ -110,7 +110,7 @@ case class PaimonReadLimitGuard(limits: Array[ReadLimit], lastTriggerMillis: Lon private def getBytesAndRows(indexedDataSplit: IndexedDataSplit): (Long, Long) = { var rows = 0L var bytes = 0L - indexedDataSplit.entry.dataFiles().asScala.foreach { + indexedDataSplit.dataSplit.dataFiles().asScala.foreach { file => rows += file.rowCount() bytes += file.fileSize() diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/StreamHelper.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/StreamHelper.scala index 7e61d71ac183..cfa19a81aeaf 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/StreamHelper.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/StreamHelper.scala @@ -22,7 +22,7 @@ import org.apache.paimon.CoreOptions import org.apache.paimon.data.BinaryRow import org.apache.paimon.spark.SparkTypeUtils import org.apache.paimon.table.DataTable -import org.apache.paimon.table.source.{DataSplit, StreamDataTableScan} +import org.apache.paimon.table.source.{DataSplit, QueryAuthSplit, Split, StreamDataTableScan} import org.apache.paimon.table.source.TableScan.Plan import org.apache.paimon.table.source.snapshot.StartingContext import org.apache.paimon.utils.{InternalRowPartitionComputer, TypeUtils} @@ -34,7 +34,7 @@ import org.apache.spark.sql.types.StructType import scala.collection.JavaConverters._ import scala.collection.mutable -case class IndexedDataSplit(snapshotId: Long, index: Long, entry: DataSplit) +case class IndexedDataSplit(snapshotId: Long, index: Long, entry: Split, dataSplit: DataSplit) private[spark] trait StreamHelper { @@ -122,16 +122,20 @@ private[spark] trait StreamHelper { /** Sort the [[DataSplit]] list and index them. */ private def convertPlanToIndexedSplits(plan: Plan): Array[IndexedDataSplit] = { - val dataSplits = - plan.splits().asScala.collect { case dataSplit: DataSplit => dataSplit }.toArray - val snapshotId = dataSplits.head.snapshotId() + val dataSplits = plan + .splits() + .asScala + .map(split => (split, QueryAuthSplit.unwrap(split))) + .collect { case (entry, dataSplit: DataSplit) => (entry, dataSplit) } + .toArray + val snapshotId = dataSplits.head._2.snapshotId() dataSplits - .sortWith((ds1, ds2) => compareByPartitionAndBucket(ds1, ds2) < 0) + .sortWith((ds1, ds2) => compareByPartitionAndBucket(ds1._2, ds2._2) < 0) .zipWithIndex .map { - case (split, idx) => - IndexedDataSplit(snapshotId, idx, split) + case ((entry, dataSplit), idx) => + IndexedDataSplit(snapshotId, idx, entry, dataSplit) } } diff --git a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java index f30989569e32..ce4395d5dea1 100644 --- a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java +++ b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java @@ -56,6 +56,7 @@ import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; import org.apache.spark.sql.connector.catalog.CatalogManager; +import org.apache.spark.sql.streaming.StreamingQuery; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -763,6 +764,46 @@ public void testRowFilterWithPartitionAndBucketMetadata() { assertThat(rows.get(0).getInt(2)).isBetween(0, 1); } + @Test + public void testStreamingRowFilter() throws Exception { + spark.sql( + "CREATE TABLE t_stream_auth (id INT, name STRING) TBLPROPERTIES " + + "('query-auth.enabled'='true')"); + spark.sql("INSERT INTO t_stream_auth VALUES (1, 'blocked'), (2, 'allowed')"); + + Predicate idEq2Predicate = + LeafPredicate.of( + new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), + Equal.INSTANCE, + Collections.singletonList(2)); + restCatalogServer.setRowFilterAuth( + Identifier.create("db2", "t_stream_auth"), + Collections.singletonList(idEq2Predicate)); + + StreamingQuery query = + spark.readStream() + .format("paimon") + .table("t_stream_auth") + .writeStream() + .format("memory") + .option( + "checkpointLocation", + tempFile.resolve("stream-auth-checkpoint").toString()) + .queryName("stream_auth_result") + .outputMode("append") + .start(); + try { + query.processAllAvailable(); + assertThat( + spark.sql("SELECT * FROM stream_auth_result ORDER BY id") + .collectAsList() + .toString()) + .isEqualTo("[[2,allowed]]"); + } finally { + query.stop(); + } + } + @Test public void testRowFilterDeletionVectorsTable() { // Deletion-vectors table: deleted rows excluded via the deletion vector, then filtered. From e5b5e7066214304acc39f396fd89ba0e1e72516c Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:04:06 +0800 Subject: [PATCH 17/31] [spark] Handle query auth in primary key vector search --- .../spark/execution/PaimonStrategy.scala | 4 +- .../spark/SparkCatalogWithRestTest.java | 41 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala index 74c0e6c47dbd..aefe14254fd7 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala @@ -32,7 +32,7 @@ import org.apache.paimon.spark.data.SparkInternalRow import org.apache.paimon.spark.read.VectorSearchResultUtils import org.apache.paimon.spark.schema.PaimonMetadataColumn import org.apache.paimon.table.{InnerTable, SpecialFields, Table} -import org.apache.paimon.table.source.{BatchVectorSearchBuilder, DataSplit, InnerTableScan, PrimaryKeyScoredResult, PrimaryKeySearchPosition, PrimaryKeyVectorResult, ReadBuilder, VectorScan} +import org.apache.paimon.table.source.{BatchVectorSearchBuilder, DataSplit, InnerTableScan, PrimaryKeyScoredResult, PrimaryKeySearchPosition, PrimaryKeyVectorResult, QueryAuthSplit, ReadBuilder, VectorScan} import org.apache.paimon.types.RowType import org.apache.paimon.utils.RoaringNavigableMap64 @@ -567,7 +567,7 @@ case class LateralVectorSearchExec( scan.plan().splits().asScala.iterator.flatMap { split => - val indexedSplit = split.asInstanceOf[IndexedSplit] + val indexedSplit = QueryAuthSplit.unwrap(split).asInstanceOf[IndexedSplit] val dataSplit = indexedSplit.dataSplit() val file = LateralVectorSearchPhysicalFile.from(dataSplit) val reader = diff --git a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java index ce4395d5dea1..d9f3546ba465 100644 --- a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java +++ b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java @@ -804,6 +804,47 @@ public void testStreamingRowFilter() throws Exception { } } + @Test + public void testLateralPrimaryKeyVectorSearchWithRowFilter() { + spark.sql( + "CREATE TABLE t_auth_vector (" + + "id INT, embedding ARRAY, query_embedding ARRAY) " + + "TBLPROPERTIES (" + + "'primary-key'='id', " + + "'bucket'='1', " + + "'deletion-vectors.enabled'='true', " + + "'vector-field'='embedding', " + + "'field.embedding.vector-dim'='2', " + + "'pk-vector.index.columns'='embedding', " + + "'fields.embedding.pk-vector.index.type'='test-vector-ann', " + + "'fields.embedding.pk-vector.distance.metric'='l2', " + + "'test.vector.dimension'='2', " + + "'test.vector.metric'='l2', " + + "'query-auth.enabled'='true')"); + spark.sql( + "INSERT INTO t_auth_vector VALUES " + + "(1, array(5.0f, 0.0f), array(0.0f, 0.0f)), " + + "(2, array(1.0f, 0.0f), array(0.0f, 0.0f))"); + + Predicate idEq1Predicate = + LeafPredicate.of( + new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), + Equal.INSTANCE, + Collections.singletonList(1)); + restCatalogServer.setRowFilterAuth( + Identifier.create("db2", "t_auth_vector"), + Collections.singletonList(idEq1Predicate)); + + List rows = + spark.sql( + "SELECT q.id AS query_id, r.id AS result_id " + + "FROM t_auth_vector AS q, " + + "LATERAL (SELECT id FROM vector_search(" + + "'t_auth_vector', 'embedding', q.query_embedding, 1)) AS r") + .collectAsList(); + assertThat(rows).isEmpty(); + } + @Test public void testRowFilterDeletionVectorsTable() { // Deletion-vectors table: deleted rows excluded via the deletion vector, then filtered. From ab5d2cd636a637c5fa7213cf455978b9cfc4c84e Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:21:43 +0800 Subject: [PATCH 18/31] [flink] Use full cache lookup with query auth --- .../flink/lookup/FileStoreLookupFunction.java | 3 +- .../lookup/PrimaryKeyPartialLookupTable.java | 4 + .../lookup/FileStoreLookupFunctionTest.java | 106 +++++++++++++++++- 3 files changed, 106 insertions(+), 7 deletions(-) diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FileStoreLookupFunction.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FileStoreLookupFunction.java index 77d74526a6bc..07ae23160f56 100644 --- a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FileStoreLookupFunction.java +++ b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FileStoreLookupFunction.java @@ -213,7 +213,8 @@ private void open() throws Exception { table instanceof FallbackReadFileStoreTable && ((FallbackReadFileStoreTable) table).other() instanceof ChainGroupReadTable; - if (!isChainTable + if (!table.coreOptions().queryAuthEnabled() + && !isChainTable && options.get(LOOKUP_CACHE_MODE) == LookupCacheMode.AUTO && new HashSet<>(table.primaryKeys()).equals(new HashSet<>(joinKeys))) { if (isRemoteServiceAvailable(table)) { diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyPartialLookupTable.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyPartialLookupTable.java index 9964a8fa6c7f..f7a008bbffa1 100644 --- a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyPartialLookupTable.java +++ b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyPartialLookupTable.java @@ -83,6 +83,10 @@ private PrimaryKeyPartialLookupTable( } CoreOptions coreOptions = CoreOptions.fromMap(table.options()); + if (coreOptions.queryAuthEnabled()) { + throw new UnsupportedOperationException( + "Primary key partial lookup does not support query authorization."); + } if (!coreOptions.needLookup() && coreOptions.mergeEngine() != CoreOptions.MergeEngine.DEDUPLICATE) { diff --git a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java index ad315237b9e6..d4ff5f3010f1 100644 --- a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java +++ b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java @@ -19,6 +19,9 @@ package org.apache.paimon.flink.lookup; import org.apache.paimon.CoreOptions; +import org.apache.paimon.catalog.FileSystemCatalog; +import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; @@ -29,10 +32,13 @@ import org.apache.paimon.flink.lookup.PrimaryKeyPartialLookupTable.RemoteQueryExecutor; import org.apache.paimon.lookup.rocksdb.RocksDBOptions; import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.schema.Schema; import org.apache.paimon.schema.SchemaManager; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.service.ServiceManager; +import org.apache.paimon.table.CatalogEnvironment; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.FileStoreTableFactory; import org.apache.paimon.table.sink.CommitMessage; @@ -42,6 +48,7 @@ import org.apache.paimon.types.DataType; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.JsonSerdeUtil; import org.apache.paimon.utils.TraceableFileIO; import org.apache.flink.table.data.RowData; @@ -137,6 +144,23 @@ private FileStoreTable createFileStoreTable( boolean refreshAsync, Integer fullLoadThreshold) throws Exception { + return createFileStoreTable( + isPartition, + dynamicPartition, + refreshAsync, + fullLoadThreshold, + false, + CatalogEnvironment.empty()); + } + + private FileStoreTable createFileStoreTable( + boolean isPartition, + boolean dynamicPartition, + boolean refreshAsync, + Integer fullLoadThreshold, + boolean queryAuthEnabled, + CatalogEnvironment catalogEnvironment) + throws Exception { SchemaManager schemaManager = new SchemaManager(fileIO, tablePath); Options conf = new Options(); conf.set(FlinkConnectorOptions.LOOKUP_REFRESH_ASYNC, refreshAsync); @@ -151,11 +175,11 @@ private FileStoreTable createFileStoreTable( if (fullLoadThreshold != null) { conf.set(FlinkConnectorOptions.LOOKUP_REFRESH_FULL_LOAD_THRESHOLD, fullLoadThreshold); } + if (queryAuthEnabled) { + conf.set(CoreOptions.QUERY_AUTH_ENABLED, true); + } - RowType rowType = - RowType.of( - new DataType[] {DataTypes.INT(), DataTypes.INT(), DataTypes.BIGINT()}, - new String[] {"pt", "k", "v"}); + RowType rowType = tableRowType(); Schema schema = new Schema( rowType.getFields(), @@ -164,8 +188,7 @@ private FileStoreTable createFileStoreTable( conf.toMap(), ""); TableSchema tableSchema = schemaManager.createTable(schema); - return FileStoreTableFactory.create( - fileIO, new org.apache.paimon.fs.Path(tempDir.toString()), tableSchema); + return FileStoreTableFactory.create(fileIO, tablePath, tableSchema, catalogEnvironment); } @AfterEach @@ -219,6 +242,77 @@ public void testDefaultRemotePartial(boolean refreshAsync) throws Exception { assertThat(queryExecutor).isInstanceOf(RemoteQueryExecutor.class); } + @Test + public void testQueryAuthUsesFullCacheAndAppliesRowFilter() throws Exception { + Predicate authFilter = new PredicateBuilder(tableRowType()).equal(2, 20L); + TableQueryAuthResult authResult = + new TableQueryAuthResult( + Collections.singletonList(JsonSerdeUtil.toFlatJson(authFilter)), null); + Identifier identifier = Identifier.create("default", "table"); + CatalogEnvironment catalogEnvironment = + new CatalogEnvironment( + identifier, + null, + () -> + new FileSystemCatalog(fileIO, tablePath) { + @Override + public TableQueryAuthResult authTableQuery( + Identifier ignored, List select) { + return authResult; + } + }, + null, + null, + null, + false, + false); + table = createFileStoreTable(false, false, false, null, true, catalogEnvironment); + + StreamTableWrite writer = table.newStreamWriteBuilder().newWrite(); + writer.write(GenericRow.of(1, 1, 10L)); + writer.write(GenericRow.of(2, 2, 20L)); + commit(writer.prepareCommit(true, 1)); + writer.close(); + + ServiceManager serviceManager = new ServiceManager(fileIO, tablePath); + serviceManager.resetService( + PRIMARY_KEY_LOOKUP, new InetSocketAddress[] {new InetSocketAddress(1)}); + lookupFunction = createLookupFunction(table, true); + lookupFunction.open(tempDir.toString()); + + assertThat(lookupFunction.lookupTable()).isInstanceOf(FullCacheLookupTable.class); + assertThat(lookupFunction.lookup(new FlinkRowData(GenericRow.of(1, 1)))).isEmpty(); + assertThat(lookupFunction.lookup(new FlinkRowData(GenericRow.of(2, 2)))).hasSize(1); + } + + @Test + public void testPartialLookupRejectsQueryAuth() throws Exception { + table = createFileStoreTable(false, false, false, null, true, CatalogEnvironment.empty()); + + assertThatThrownBy( + () -> + PrimaryKeyPartialLookupTable.createLocalTable( + table, + new int[] {0, 1}, + tempDir.resolve("local").toFile(), + Arrays.asList("pt", "k"), + Collections.emptySet())) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("does not support query authorization"); + assertThatThrownBy( + () -> + PrimaryKeyPartialLookupTable.createRemoteTable( + table, new int[] {0, 1}, Arrays.asList("pt", "k"))) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("does not support query authorization"); + } + + private RowType tableRowType() { + return RowType.of( + new DataType[] {DataTypes.INT(), DataTypes.INT(), DataTypes.BIGINT()}, + new String[] {"pt", "k", "v"}); + } + @Test public void testLookupScanLeak() throws Exception { createLookupFunction(false); From a3f092ac0d3513c31c80cb9ce7c965de27090613 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:22:50 +0800 Subject: [PATCH 19/31] [core] Bypass query auth for compact bucket scans --- .../table/system/CompactBucketsTable.java | 9 +++- .../apache/paimon/rest/RESTCatalogTest.java | 45 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/table/system/CompactBucketsTable.java b/paimon-core/src/main/java/org/apache/paimon/table/system/CompactBucketsTable.java index f18d184b8437..d98dd7f00d7f 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/system/CompactBucketsTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/system/CompactBucketsTable.java @@ -108,10 +108,15 @@ public CompactBucketsTable(FileStoreTable wrapped, boolean isContinuous) { // if need to specify the database of a table, use this method public CompactBucketsTable(FileStoreTable wrapped, boolean isContinuous, String databaseName) { - this.wrapped = wrapped; + this.wrapped = + wrapped.coreOptions().queryAuthEnabled() + ? wrapped.copy( + Collections.singletonMap( + CoreOptions.QUERY_AUTH_ENABLED.key(), "false")) + : wrapped; this.isContinuous = isContinuous; this.databaseName = databaseName; - this.baseSchemaId = wrapped.schema().id(); + this.baseSchemaId = this.wrapped.schema().id(); } @Override diff --git a/paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java b/paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java index 504dc6fab6ec..fcaa3a7b7b66 100644 --- a/paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java @@ -88,10 +88,13 @@ import org.apache.paimon.table.sink.StreamTableCommit; import org.apache.paimon.table.sink.StreamTableWrite; import org.apache.paimon.table.sink.TableWriteImpl; +import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.InnerTableScan; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.ReadBuilder; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.TableRead; +import org.apache.paimon.table.system.CompactBucketsTable; import org.apache.paimon.table.system.SystemTableLoader; import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; @@ -2820,6 +2823,48 @@ void testTableAuth() throws Exception { table.newReadBuilder().withProjection(new int[] {1}).newScan().plan(); } + @Test + void testCompactBucketsTableBypassesQueryAuth() throws Exception { + Identifier identifier = Identifier.create("test_table_db", "auth_compact_buckets_table"); + catalog.createDatabase(identifier.getDatabaseName(), true); + catalog.createTable( + identifier, + new Schema( + singletonList(new DataField(0, "id", DataTypes.INT())), + emptyList(), + emptyList(), + singletonMap(QUERY_AUTH_ENABLED.key(), "true"), + ""), + true); + + FileStoreTable table = (FileStoreTable) catalog.getTable(identifier); + batchWrite(table, singletonList(1)); + setRowFilter( + identifier, + singletonList( + LeafPredicate.of( + new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), + GreaterThan.INSTANCE, + singletonList(0)))); + + assertThat(table.newScan().plan().splits()) + .isNotEmpty() + .allSatisfy(split -> assertThat(split).isInstanceOf(QueryAuthSplit.class)); + + CompactBucketsTable bucketsTable = new CompactBucketsTable(table, false); + List splits = bucketsTable.newScan().plan().splits(); + assertThat(splits) + .isNotEmpty() + .allSatisfy(split -> assertThat(split).isInstanceOf(DataSplit.class)); + for (Split split : splits) { + try (RecordReader reader = bucketsTable.newRead().createReader(split)) { + List rows = new ArrayList<>(); + reader.forEachRemaining(rows::add); + assertThat(rows).hasSize(1); + } + } + } + @Test void testSnapshotMethods() throws Exception { Identifier identifier = Identifier.create("test_table_db", "snapshots_table"); From 712b4615a7b61013787fb58517850211a9cbd16b Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 15:03:42 +0800 Subject: [PATCH 20/31] [core] Bypass query auth for file monitor scans --- .../paimon/table/system/FileMonitorTable.java | 2 + .../apache/paimon/rest/RESTCatalogTest.java | 42 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/paimon-core/src/main/java/org/apache/paimon/table/system/FileMonitorTable.java b/paimon-core/src/main/java/org/apache/paimon/table/system/FileMonitorTable.java index 97903d881ea0..7522b74dd729 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/system/FileMonitorTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/system/FileMonitorTable.java @@ -66,6 +66,7 @@ import java.util.Optional; import static java.util.Collections.emptyList; +import static org.apache.paimon.CoreOptions.QUERY_AUTH_ENABLED; import static org.apache.paimon.CoreOptions.SCAN_BOUNDED_WATERMARK; import static org.apache.paimon.CoreOptions.STREAM_SCAN_MODE; import static org.apache.paimon.CoreOptions.StreamScanMode.FILE_MONITOR; @@ -98,6 +99,7 @@ public FileMonitorTable(FileStoreTable wrapped) { Map dynamicOptions = new HashMap<>(); dynamicOptions.put(STREAM_SCAN_MODE.key(), FILE_MONITOR.getValue()); dynamicOptions.put(SCAN_BOUNDED_WATERMARK.key(), null); + dynamicOptions.put(QUERY_AUTH_ENABLED.key(), "false"); this.wrapped = wrapped.copy(dynamicOptions); } diff --git a/paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java b/paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java index fcaa3a7b7b66..7dcf6f1b8e06 100644 --- a/paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/rest/RESTCatalogTest.java @@ -93,8 +93,10 @@ import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.ReadBuilder; import org.apache.paimon.table.source.Split; +import org.apache.paimon.table.source.StreamTableScan; import org.apache.paimon.table.source.TableRead; import org.apache.paimon.table.system.CompactBucketsTable; +import org.apache.paimon.table.system.FileMonitorTable; import org.apache.paimon.table.system.SystemTableLoader; import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; @@ -2865,6 +2867,46 @@ void testCompactBucketsTableBypassesQueryAuth() throws Exception { } } + @Test + void testFileMonitorTableBypassesQueryAuth() throws Exception { + Identifier identifier = Identifier.create("test_table_db", "auth_file_monitor_table"); + catalog.createDatabase(identifier.getDatabaseName(), true); + catalog.createTable( + identifier, + new Schema( + singletonList(new DataField(0, "id", DataTypes.INT())), + emptyList(), + emptyList(), + singletonMap(QUERY_AUTH_ENABLED.key(), "true"), + ""), + true); + + FileStoreTable table = (FileStoreTable) catalog.getTable(identifier); + batchWrite(table, singletonList(1)); + setRowFilter( + identifier, + singletonList( + LeafPredicate.of( + new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), + GreaterThan.INSTANCE, + singletonList(0)))); + + FileMonitorTable monitorTable = new FileMonitorTable(table); + ReadBuilder readBuilder = monitorTable.newReadBuilder(); + StreamTableScan scan = readBuilder.newStreamScan(); + List splits = scan.plan().splits(); + assertThat(splits) + .isNotEmpty() + .allSatisfy(split -> assertThat(split).isNotInstanceOf(QueryAuthSplit.class)); + for (Split split : splits) { + try (RecordReader reader = readBuilder.newRead().createReader(split)) { + List rows = new ArrayList<>(); + reader.forEachRemaining(rows::add); + assertThat(rows).hasSize(1); + } + } + } + @Test void testSnapshotMethods() throws Exception { Identifier identifier = Identifier.create("test_table_db", "snapshots_table"); From 1175f29e5a4793f50da2927c33f0625c7f4f06ff Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 18:23:26 +0800 Subject: [PATCH 21/31] fix --- .../paimon/table/source/QueryAuthSplit.java | 12 +++- .../table/source/QueryAuthSplitTest.java | 9 +++ .../PaimonRecordReaderIteratorTest.scala | 60 +++++++++++++++++++ 3 files changed, 79 insertions(+), 2 deletions(-) create mode 100644 paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonRecordReaderIteratorTest.scala 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 c6842de8eb74..1697430fecf6 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 @@ -23,6 +23,7 @@ import org.apache.paimon.io.DataInputViewStreamWrapper; import org.apache.paimon.io.DataOutputView; import org.apache.paimon.io.DataOutputViewStreamWrapper; +import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit; import javax.annotation.Nullable; @@ -49,12 +50,19 @@ public Split split() { return split; } - /** Unwraps one query authorization layer to expose the wrapped split. */ + /** Unwraps one fallback layer followed by one query authorization layer. */ public static Split unwrap(Split split) { + if (split instanceof FallbackSplit) { + Split wrapped = ((FallbackSplit) split).wrapped(); + // FallbackDataSplit is already a DataSplit and deliberately wraps itself. + if (wrapped != split) { + split = wrapped; + } + } return split instanceof QueryAuthSplit ? ((QueryAuthSplit) split).split() : split; } - /** Unwraps one query authorization layer and returns the wrapped data split. */ + /** Unwraps one fallback and query authorization layer and returns the data split. */ public static DataSplit unwrapDataSplit(Split split) { return (DataSplit) unwrap(split); } diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/QueryAuthSplitTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/QueryAuthSplitTest.java index 97c24f94b6e0..835ff5958928 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/QueryAuthSplitTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/QueryAuthSplitTest.java @@ -23,6 +23,7 @@ import org.apache.paimon.io.DataFileTestDataGenerator; import org.apache.paimon.io.DataInputDeserializer; import org.apache.paimon.io.DataOutputViewStreamWrapper; +import org.apache.paimon.table.FallbackReadFileStoreTable; import org.apache.paimon.utils.InstantiationUtil; import org.junit.jupiter.api.Test; @@ -44,6 +45,14 @@ public void testUnwrapDataSplit() { assertThat(QueryAuthSplit.unwrapDataSplit(dataSplit)).isSameAs(dataSplit); assertThat(QueryAuthSplit.unwrapDataSplit(new QueryAuthSplit(dataSplit, authResult()))) .isSameAs(dataSplit); + + Split fallbackDataSplit = FallbackReadFileStoreTable.toFallbackSplit(dataSplit, false); + assertThat(QueryAuthSplit.unwrapDataSplit(fallbackDataSplit)).isSameAs(fallbackDataSplit); + + Split fallbackQueryAuthSplit = + new FallbackReadFileStoreTable.FallbackSplitImpl( + new QueryAuthSplit(dataSplit, authResult()), false); + assertThat(QueryAuthSplit.unwrapDataSplit(fallbackQueryAuthSplit)).isSameAs(dataSplit); } @Test diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonRecordReaderIteratorTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonRecordReaderIteratorTest.scala new file mode 100644 index 000000000000..4351fdfc78ee --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonRecordReaderIteratorTest.scala @@ -0,0 +1,60 @@ +/* + * 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.spark + +import org.apache.paimon.catalog.TableQueryAuthResult +import org.apache.paimon.data.{BinaryRow, InternalRow => PaimonInternalRow} +import org.apache.paimon.reader.RecordReader +import org.apache.paimon.spark.schema.PaimonMetadataColumn +import org.apache.paimon.table.FallbackReadFileStoreTable +import org.apache.paimon.table.source.{DataSplit, QueryAuthSplit} + +import org.junit.jupiter.api.Assertions.assertDoesNotThrow + +import java.util.Collections + +/** Tests for [[PaimonRecordReaderIterator]]. */ +class PaimonRecordReaderIteratorTest extends PaimonSparkTestBase { + + test("Paimon: fallback query auth split exposes metadata") { + val dataSplit = DataSplit + .builder() + .withSnapshot(1L) + .withPartition(BinaryRow.singleColumn(1)) + .withBucket(2) + .withDataFiles(Collections.emptyList()) + .isStreaming(false) + .withBucketPath("") + .build() + val queryAuthSplit = new QueryAuthSplit( + dataSplit, + new TableQueryAuthResult(null, Collections.singletonMap("f0", "mask"))) + val fallbackSplit = + new FallbackReadFileStoreTable.FallbackSplitImpl(queryAuthSplit, false) + val emptyReader = new RecordReader[PaimonInternalRow] { + override def readBatch(): RecordReader.RecordIterator[PaimonInternalRow] = null + + override def close(): Unit = {} + } + + assertDoesNotThrow( + () => + PaimonRecordReaderIterator(emptyReader, Seq(PaimonMetadataColumn.BUCKET), fallbackSplit)) + } +} From d96d196c3f413febe245f71cdca65d90e7d404fc Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 19:07:53 +0800 Subject: [PATCH 22/31] [common] Preserve score readers through query auth --- .../paimon/reader/ScoreRecordReader.java | 39 ++++++++ .../paimon/reader/ScoreRecordReaderTest.java | 91 +++++++++++++++++++ .../spark/SparkCatalogWithRestTest.java | 42 +++++++++ 3 files changed, 172 insertions(+) create mode 100644 paimon-common/src/test/java/org/apache/paimon/reader/ScoreRecordReaderTest.java diff --git a/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordReader.java b/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordReader.java index 69b946332595..4c30a9af1c3a 100644 --- a/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordReader.java +++ b/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordReader.java @@ -18,9 +18,12 @@ package org.apache.paimon.reader; +import org.apache.paimon.utils.Filter; + import javax.annotation.Nullable; import java.io.IOException; +import java.util.function.Function; /** A {@link RecordReader} whose records expose vector-search scores and row identifiers. */ public interface ScoreRecordReader extends RecordReader { @@ -28,4 +31,40 @@ public interface ScoreRecordReader extends RecordReader { @Nullable @Override ScoreRecordIterator readBatch() throws IOException; + + @Override + default ScoreRecordReader transform(Function function) { + ScoreRecordReader thisReader = this; + return new ScoreRecordReader() { + @Nullable + @Override + public ScoreRecordIterator readBatch() throws IOException { + ScoreRecordIterator iterator = thisReader.readBatch(); + return iterator == null ? null : iterator.transform(function); + } + + @Override + public void close() throws IOException { + thisReader.close(); + } + }; + } + + @Override + default ScoreRecordReader filter(Filter filter) { + ScoreRecordReader thisReader = this; + return new ScoreRecordReader() { + @Nullable + @Override + public ScoreRecordIterator readBatch() throws IOException { + ScoreRecordIterator iterator = thisReader.readBatch(); + return iterator == null ? null : iterator.filter(filter); + } + + @Override + public void close() throws IOException { + thisReader.close(); + } + }; + } } diff --git a/paimon-common/src/test/java/org/apache/paimon/reader/ScoreRecordReaderTest.java b/paimon-common/src/test/java/org/apache/paimon/reader/ScoreRecordReaderTest.java new file mode 100644 index 000000000000..4b48485465a4 --- /dev/null +++ b/paimon-common/src/test/java/org/apache/paimon/reader/ScoreRecordReaderTest.java @@ -0,0 +1,91 @@ +/* + * 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.reader; + +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import java.io.IOException; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link ScoreRecordReader}. */ +public class ScoreRecordReaderTest { + + @Test + public void testFilterAndTransformPreserveScoreMetadata() throws IOException { + ScoreRecordReader reader = + new ScoreRecordReader() { + private boolean emitted; + + @Nullable + @Override + public ScoreRecordIterator readBatch() { + if (emitted) { + return null; + } + emitted = true; + return new TestingScoreRecordIterator(); + } + + @Override + public void close() {} + }; + + ScoreRecordReader transformed = + reader.filter(value -> value == 1).transform(String::valueOf); + ScoreRecordIterator iterator = transformed.readBatch(); + + assertThat(iterator).isNotNull(); + assertThat(iterator.next()).isEqualTo("1"); + assertThat(iterator.returnedRowId()).isEqualTo(11L); + assertThat(iterator.returnedScore()).isEqualTo(0.75f); + assertThat(iterator.next()).isNull(); + assertThat(transformed.readBatch()).isNull(); + } + + private static class TestingScoreRecordIterator implements ScoreRecordIterator { + + private final int[] values = {0, 1}; + private final long[] rowIds = {10L, 11L}; + private final float[] scores = {0.25f, 0.75f}; + private int index = -1; + + @Nullable + @Override + public Integer next() { + index++; + return index < values.length ? values[index] : null; + } + + @Override + public float returnedScore() { + return scores[index]; + } + + @Override + public long returnedRowId() { + return rowIds[index]; + } + + @Override + public void releaseBatch() {} + } +} diff --git a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java index d9f3546ba465..e7509f9ac29d 100644 --- a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java +++ b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java @@ -845,6 +845,48 @@ public void testLateralPrimaryKeyVectorSearchWithRowFilter() { assertThat(rows).isEmpty(); } + @Test + public void testLateralPrimaryKeyVectorSearchWithAllowedRowFilter() { + spark.sql( + "CREATE TABLE t_auth_vector_allowed (" + + "id INT, embedding ARRAY, query_embedding ARRAY) " + + "TBLPROPERTIES (" + + "'primary-key'='id', " + + "'bucket'='1', " + + "'deletion-vectors.enabled'='true', " + + "'vector-field'='embedding', " + + "'field.embedding.vector-dim'='2', " + + "'pk-vector.index.columns'='embedding', " + + "'fields.embedding.pk-vector.index.type'='test-vector-ann', " + + "'fields.embedding.pk-vector.distance.metric'='l2', " + + "'test.vector.dimension'='2', " + + "'test.vector.metric'='l2', " + + "'query-auth.enabled'='true')"); + spark.sql( + "INSERT INTO t_auth_vector_allowed VALUES " + + "(1, array(1.0f, 0.0f), array(0.0f, 0.0f)), " + + "(2, array(5.0f, 0.0f), array(0.0f, 0.0f))"); + + Predicate idEq1Predicate = + LeafPredicate.of( + new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), + Equal.INSTANCE, + Collections.singletonList(1)); + restCatalogServer.setRowFilterAuth( + Identifier.create("db2", "t_auth_vector_allowed"), + Collections.singletonList(idEq1Predicate)); + + List rows = + spark.sql( + "SELECT q.id AS query_id, r.id AS result_id " + + "FROM t_auth_vector_allowed AS q, " + + "LATERAL (SELECT id FROM vector_search(" + + "'t_auth_vector_allowed', 'embedding', " + + "q.query_embedding, 1)) AS r") + .collectAsList(); + assertThat(rows.toString()).isEqualTo("[[1,1]]"); + } + @Test public void testRowFilterDeletionVectorsTable() { // Deletion-vectors table: deleted rows excluded via the deletion vector, then filtered. From be05bcfab0a00587dd729a4417bbe4114c78e3b4 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 19:15:27 +0800 Subject: [PATCH 23/31] [flink] Unwrap query auth splits for lag metrics --- .../source/reader/CDCSourceSplitReader.java | 6 ++- .../reader/CDCSourceSplitReaderTest.java | 51 +++++++++++++++++-- .../source/FileStoreSourceSplitReader.java | 6 ++- .../FileStoreSourceSplitReaderTest.java | 46 +++++++++++++++-- 4 files changed, 97 insertions(+), 12 deletions(-) diff --git a/paimon-flink/paimon-flink-cdc/src/main/java/org/apache/paimon/flink/pipeline/cdc/source/reader/CDCSourceSplitReader.java b/paimon-flink/paimon-flink-cdc/src/main/java/org/apache/paimon/flink/pipeline/cdc/source/reader/CDCSourceSplitReader.java index 960665bf365e..8e5b608d60bc 100644 --- a/paimon-flink/paimon-flink-cdc/src/main/java/org/apache/paimon/flink/pipeline/cdc/source/reader/CDCSourceSplitReader.java +++ b/paimon-flink/paimon-flink-cdc/src/main/java/org/apache/paimon/flink/pipeline/cdc/source/reader/CDCSourceSplitReader.java @@ -30,6 +30,7 @@ import org.apache.paimon.reader.RecordReader.RecordIterator; import org.apache.paimon.schema.TableSchema; 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.utils.Pool; @@ -213,9 +214,10 @@ private void checkSplitOrStartNext() throws IOException { } // update metric when split changes - if (nextSplit.split() instanceof DataSplit) { + Split split = QueryAuthSplit.unwrap(nextSplit.split()); + if (split instanceof DataSplit) { long eventTime = - ((DataSplit) nextSplit.split()) + ((DataSplit) split) .earliestFileCreationEpochMillis() .orElse(FileStoreSourceReaderMetrics.UNDEFINED); metrics.recordSnapshotUpdate(eventTime); diff --git a/paimon-flink/paimon-flink-cdc/src/test/java/org/apache/paimon/flink/pipeline/cdc/source/reader/CDCSourceSplitReaderTest.java b/paimon-flink/paimon-flink-cdc/src/test/java/org/apache/paimon/flink/pipeline/cdc/source/reader/CDCSourceSplitReaderTest.java index 27ae504f21ec..1d85ea6ebb2a 100644 --- a/paimon-flink/paimon-flink-cdc/src/test/java/org/apache/paimon/flink/pipeline/cdc/source/reader/CDCSourceSplitReaderTest.java +++ b/paimon-flink/paimon-flink-cdc/src/test/java/org/apache/paimon/flink/pipeline/cdc/source/reader/CDCSourceSplitReaderTest.java @@ -23,6 +23,7 @@ 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.data.GenericRow; import org.apache.paimon.data.InternalRow; @@ -39,6 +40,7 @@ import org.apache.paimon.schema.Schema; import org.apache.paimon.schema.TableSchema; 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.utils.InstantiationUtil; @@ -132,6 +134,42 @@ public void testPrimaryKeySkip() throws Exception { innerTestOnce(4); } + @Test + public void testQueryAuthSplitUpdatesFetchLagMetric() throws Exception { + TestChangelogDataReadWrite rw = new TestChangelogDataReadWrite(tablePath); + FileStoreSourceReaderMetrics metrics = + new FileStoreSourceReaderMetrics(new DummyMetricGroup()); + CDCSourceSplitReader reader = + createReader( + new TestingTableRead( + new SingleBatchRecordReader(new TrackingRecordIterator())), + Collections.emptyList(), + metrics); + + List files = rw.writeFiles(row(1), 0, kvs()); + TableAwareFileStoreSourceSplit split = newSourceSplit("id1", row(1), 0, files); + assignSplit( + reader, + new TableAwareFileStoreSourceSplit( + split.splitId(), + new QueryAuthSplit(split.split(), new TableQueryAuthResult(null, null)), + split.recordsToSkip(), + split.getIdentifier(), + split.getLastSchemaId(), + split.getSchemaId(), + split.schemaChangeEventsToSkip())); + + RecordsWithSplitIds> records = reader.fetch(); + assertThat(metrics.getLatestFileCreationTime()) + .isEqualTo( + files.stream() + .mapToLong(DataFileMeta::creationTimeEpochMillis) + .min() + .orElse(FileStoreSourceReaderMetrics.UNDEFINED)); + records.recycle(); + reader.close(); + } + @Test public void testSplitReaderWakeupAble() throws Exception { TestChangelogDataReadWrite rw = new TestChangelogDataReadWrite(tablePath); @@ -167,10 +205,17 @@ private CDCSourceSplitReader createReader(TableRead tableRead) { private CDCSourceSplitReader createReader( TableRead tableRead, List schemaChangeEvents) { - return new TestCDCSourceSplitReader( - new FileStoreSourceReaderMetrics(new DummyMetricGroup()), + return createReader( tableRead, - schemaChangeEvents); + schemaChangeEvents, + new FileStoreSourceReaderMetrics(new DummyMetricGroup())); + } + + private CDCSourceSplitReader createReader( + TableRead tableRead, + List schemaChangeEvents, + FileStoreSourceReaderMetrics metrics) { + return new TestCDCSourceSplitReader(metrics, tableRead, schemaChangeEvents); } private void innerTestOnce(int skip) throws Exception { 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 8568e2a9dd23..db7832c2dc5e 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 @@ -25,6 +25,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.QueryAuthSplit; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.TableRead; import org.apache.paimon.types.BlobType; @@ -219,9 +220,10 @@ private void checkSplitOrStartNext() throws IOException { } // update metric when split changes - if (nextSplit.split() instanceof DataSplit) { + Split split = QueryAuthSplit.unwrap(nextSplit.split()); + if (split instanceof DataSplit) { long eventTime = - ((DataSplit) nextSplit.split()) + ((DataSplit) split) .earliestFileCreationEpochMillis() .orElse(FileStoreSourceReaderMetrics.UNDEFINED); metrics.recordSnapshotUpdate(eventTime); 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 135cb38dab0a..070af97460d7 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; @@ -31,6 +32,7 @@ import org.apache.paimon.reader.RecordReader; 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; @@ -108,6 +110,38 @@ public void testPrimaryKeySkip() throws Exception { innerTestOnce(4); } + @Test + public void testQueryAuthSplitUpdatesFetchLagMetric() throws Exception { + TestChangelogDataReadWrite rw = new TestChangelogDataReadWrite(tempDir.toString()); + FileStoreSourceReaderMetrics metrics = + new FileStoreSourceReaderMetrics(new DummyMetricGroup()); + FileStoreSourceSplitReader reader = + createReader( + new TestingTableRead( + new SingleBatchRecordReader(new TrackingRecordIterator())), + null, + metrics); + + List files = rw.writeFiles(row(1), 0, kvs()); + FileStoreSourceSplit split = newSourceSplit("id1", row(1), 0, files); + assignSplit( + reader, + new FileStoreSourceSplit( + split.splitId(), + new QueryAuthSplit(split.split(), new TableQueryAuthResult(null, null)), + split.recordsToSkip())); + + RecordsWithSplitIds> records = reader.fetch(); + assertThat(metrics.getLatestFileCreationTime()) + .isEqualTo( + files.stream() + .mapToLong(DataFileMeta::creationTimeEpochMillis) + .min() + .orElse(FileStoreSourceReaderMetrics.UNDEFINED)); + records.recycle(); + reader.close(); + } + @Test public void testSplitReaderWakeupAble() throws Exception { TestChangelogDataReadWrite rw = new TestChangelogDataReadWrite(tempDir.toString()); @@ -138,12 +172,14 @@ public void testSplitReaderWakeupAble() throws Exception { } 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 { From 1eec687c6c19aea78f64b79c0679ca7bed1a71d7 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 19:34:28 +0800 Subject: [PATCH 24/31] [spark] Apply query auth to metadata-only vector search --- .../spark/execution/PaimonStrategy.scala | 6 ++- .../spark/SparkCatalogWithRestTest.java | 42 +++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala index aefe14254fd7..3de8cf141297 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala @@ -315,6 +315,7 @@ case class LateralVectorSearchExec( private def createSearchContext( rightProjection: UnsafeProjection, readerTracker: LateralVectorSearchReaderTracker): LateralVectorSearchContext = { + val coreOptions = new CoreOptions(innerTable.options()) val rowType = innerTable.rowType() val readFieldNames = vectorSearchOutput .filterNot( @@ -356,7 +357,7 @@ case class LateralVectorSearchExec( val vectorPlan = vectorSearchBuilder.newVectorScan().scan() val batchSize = - Math.max(1, new CoreOptions(innerTable.options()).vectorSearchLateralJoinBatchSize()) + Math.max(1, coreOptions.vectorSearchLateralJoinBatchSize()) LateralVectorSearchContext( readBuilder, @@ -367,7 +368,8 @@ case class LateralVectorSearchExec( sparkRow, rowIdOrdinal = resultRowType.getFieldIndex(SpecialFields.ROW_ID.name()), metaColumnsOnly = - VectorSearchResultUtils.isVectorSearchMetaOnly(vectorSearchOutput.map(_.name)), + VectorSearchResultUtils.isVectorSearchMetaOnly(vectorSearchOutput.map(_.name)) && + !coreOptions.queryAuthEnabled(), projectionInputOrdinals = vectorSearchOutput.map { attr => if (attr.name == PaimonMetadataColumn.SEARCH_SCORE_COLUMN) { diff --git a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java index e7509f9ac29d..9c01d94169cc 100644 --- a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java +++ b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java @@ -845,6 +845,48 @@ public void testLateralPrimaryKeyVectorSearchWithRowFilter() { assertThat(rows).isEmpty(); } + @Test + public void testLateralPrimaryKeyVectorSearchMetadataOnlyWithRowFilter() { + spark.sql( + "CREATE TABLE t_auth_vector_metadata (" + + "id INT, embedding ARRAY, query_embedding ARRAY) " + + "TBLPROPERTIES (" + + "'primary-key'='id', " + + "'bucket'='1', " + + "'deletion-vectors.enabled'='true', " + + "'vector-field'='embedding', " + + "'field.embedding.vector-dim'='2', " + + "'pk-vector.index.columns'='embedding', " + + "'fields.embedding.pk-vector.index.type'='test-vector-ann', " + + "'fields.embedding.pk-vector.distance.metric'='l2', " + + "'test.vector.dimension'='2', " + + "'test.vector.metric'='l2', " + + "'query-auth.enabled'='true')"); + spark.sql( + "INSERT INTO t_auth_vector_metadata VALUES " + + "(1, array(5.0f, 0.0f), array(0.0f, 0.0f)), " + + "(2, array(1.0f, 0.0f), array(0.0f, 0.0f))"); + + Predicate idEq1Predicate = + LeafPredicate.of( + new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), + Equal.INSTANCE, + Collections.singletonList(1)); + restCatalogServer.setRowFilterAuth( + Identifier.create("db2", "t_auth_vector_metadata"), + Collections.singletonList(idEq1Predicate)); + + List rows = + spark.sql( + "SELECT q.id AS query_id, r._row_id AS row_id " + + "FROM t_auth_vector_metadata AS q, " + + "LATERAL (SELECT _row_id FROM vector_search(" + + "'t_auth_vector_metadata', 'embedding', " + + "q.query_embedding, 1)) AS r") + .collectAsList(); + assertThat(rows).isEmpty(); + } + @Test public void testLateralPrimaryKeyVectorSearchWithAllowedRowFilter() { spark.sql( From 63cd8e1f54ad3a0aff9fbff5bc19f1c2e289e083 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 19:38:42 +0800 Subject: [PATCH 25/31] [core] Avoid TopN pruning before chain query auth --- .../paimon/table/ChainGroupReadTable.java | 12 +++++++ .../paimon/table/ChainGroupReadTableTest.java | 31 +++++++++++++++++++ 2 files changed, 43 insertions(+) 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 464af1bf9575..40c7211e1418 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 @@ -31,6 +31,7 @@ import org.apache.paimon.partition.PartitionPredicate; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.predicate.TopN; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.source.ChainSplit; @@ -307,6 +308,17 @@ public ChainTableBatchScan withLimit(int limit) { return this; } + @Override + public ChainTableBatchScan withTopN(TopN topN) { + // Query authorization is applied only after the physical chain inputs have been + // assembled. Pushing TopN into those raw inputs could prune authorized rows when + // unauthorized rows rank higher. The engine applies the logical TopN after auth. + if (!options.queryAuthEnabled()) { + super.withTopN(topN); + } + return this; + } + @Override public ChainTableBatchScan withPartitionFilter(Map partitionSpec) { super.withPartitionFilter(partitionSpec); diff --git a/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java index dffa7de42bf7..df24c5036a11 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/ChainGroupReadTableTest.java @@ -24,6 +24,7 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.io.DataFileTestDataGenerator; import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.predicate.TopN; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.source.ChainSplit; @@ -178,6 +179,36 @@ public void testBatchScanKeepsHistoricalAnchorOutsideAuthorizedPartition() { verifyQueryAuthDisabled(deltaTable); } + @Test + public void testBatchScanDoesNotPushTopNBeforeQueryAuth() { + TableSchema schema = tableSchema(); + TableQueryAuthResult logicalAuth = maskingAuth("logical"); + CatalogEnvironment catalogEnvironment = catalogEnvironment(logicalAuth); + PrimaryKeyFileStoreTable snapshotTable = + table(schema, "snapshot", true, catalogEnvironment); + PrimaryKeyFileStoreTable deltaTable = table(schema, "delta", true, catalogEnvironment); + PrimaryKeyFileStoreTable rawSnapshotTable = + table(schema, "snapshot", false, catalogEnvironment); + PrimaryKeyFileStoreTable rawDeltaTable = table(schema, "delta", false, catalogEnvironment); + when(snapshotTable.copy(anyMap())).thenReturn(rawSnapshotTable); + when(deltaTable.copy(anyMap())).thenReturn(rawDeltaTable); + + DataTableScan snapshotScan = mock(DataTableScan.class); + DataTableScan deltaScan = mock(DataTableScan.class); + ChainGroupReadTable chainTable = new ChainGroupReadTable(snapshotTable, deltaTable); + ChainGroupReadTable.ChainTableBatchScan scan = + new ChainGroupReadTable.ChainTableBatchScan( + schema, + chainTable, + table -> table == rawSnapshotTable ? snapshotScan : deltaScan); + + TopN topN = mock(TopN.class); + scan.withTopN(topN); + + verify(snapshotScan, never()).withTopN(topN); + verify(deltaScan, never()).withTopN(topN); + } + @Test public void testStreamStartingPlanRetainsQueryAuth() { testStreamStartingPlanRetainsQueryAuth(false); From 33fb8c65890797fa6614a8a0b9d0c10402555e72 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 20:20:08 +0800 Subject: [PATCH 26/31] [spark] Apply query auth before lateral vector limit --- .../spark/execution/PaimonStrategy.scala | 21 ++++++++++++++++++- .../spark/SparkCatalogWithRestTest.java | 4 ++-- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala index 3de8cf141297..9a7fc98ad740 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala @@ -19,6 +19,7 @@ package org.apache.paimon.spark.execution import org.apache.paimon.CoreOptions +import org.apache.paimon.catalog.TableQueryAuthResult import org.apache.paimon.data.BinaryRow import org.apache.paimon.globalindex.{GlobalIndexResult, IndexedSplit, ScoredGlobalIndexResult} import org.apache.paimon.partition.PartitionPredicate @@ -31,7 +32,7 @@ import org.apache.paimon.spark.catalyst.plans.logical.{CopyIntoLocationCommand, import org.apache.paimon.spark.data.SparkInternalRow import org.apache.paimon.spark.read.VectorSearchResultUtils import org.apache.paimon.spark.schema.PaimonMetadataColumn -import org.apache.paimon.table.{InnerTable, SpecialFields, Table} +import org.apache.paimon.table.{FileStoreTable, InnerTable, SpecialFields, Table} import org.apache.paimon.table.source.{BatchVectorSearchBuilder, DataSplit, InnerTableScan, PrimaryKeyScoredResult, PrimaryKeySearchPosition, PrimaryKeyVectorResult, QueryAuthSplit, ReadBuilder, VectorScan} import org.apache.paimon.types.RowType import org.apache.paimon.utils.RoaringNavigableMap64 @@ -354,6 +355,7 @@ case class LateralVectorSearchExec( .withLimit(limit) .withOptions(options.asJava) pushSearchFilters(Seq(readBuilder, physicalReadBuilder), vectorSearchBuilder) + pushQueryAuthFilter(coreOptions, vectorSearchBuilder) val vectorPlan = vectorSearchBuilder.newVectorScan().scan() val batchSize = @@ -406,6 +408,23 @@ case class LateralVectorSearchExec( } } + private def pushQueryAuthFilter( + coreOptions: CoreOptions, + vectorSearchBuilder: BatchVectorSearchBuilder): Unit = { + if (!coreOptions.queryAuthEnabled()) { + return + } + + // Vector search applies its limit before physical readers enforce QueryAuthSplit. Push the row + // filter into candidate selection so unauthorized rows cannot consume the limit. + val table = innerTable.asInstanceOf[FileStoreTable] + val authResult = table.catalogEnvironment().tableQueryAuth(coreOptions).auth(null) + Option(authResult) + .flatMap(result => Option(result.extractPredicate())) + .flatMap(predicate => Option(TableQueryAuthResult.remapPredicate(predicate, table.rowType()))) + .foreach(vectorSearchBuilder.withFilter) + } + private def convertSearchFilters(): Seq[Predicate] = { if (searchFilters.isEmpty) { Seq.empty diff --git a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java index 9c01d94169cc..4f893a410f50 100644 --- a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java +++ b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java @@ -842,7 +842,7 @@ public void testLateralPrimaryKeyVectorSearchWithRowFilter() { + "LATERAL (SELECT id FROM vector_search(" + "'t_auth_vector', 'embedding', q.query_embedding, 1)) AS r") .collectAsList(); - assertThat(rows).isEmpty(); + assertThat(rows.toString()).isEqualTo("[[1,1]]"); } @Test @@ -884,7 +884,7 @@ public void testLateralPrimaryKeyVectorSearchMetadataOnlyWithRowFilter() { + "'t_auth_vector_metadata', 'embedding', " + "q.query_embedding, 1)) AS r") .collectAsList(); - assertThat(rows).isEmpty(); + assertThat(rows.toString()).isEqualTo("[[1,0]]"); } @Test From 705616d0ac6bb2111ccc498fc9ed456c067acd6c Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 20:53:31 +0800 Subject: [PATCH 27/31] [flink] Reject lookup joins with query auth --- .../flink/lookup/FileStoreLookupFunction.java | 8 ++- .../lookup/FileStoreLookupFunctionTest.java | 50 ++----------------- 2 files changed, 11 insertions(+), 47 deletions(-) diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FileStoreLookupFunction.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FileStoreLookupFunction.java index 07ae23160f56..99639aac735a 100644 --- a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FileStoreLookupFunction.java +++ b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FileStoreLookupFunction.java @@ -124,6 +124,11 @@ public FileStoreLookupFunction( int[] joinKeyIndex, @Nullable Predicate predicate, @Nullable ShuffleStrategy strategy) { + if (table.coreOptions().queryAuthEnabled()) { + throw new UnsupportedOperationException( + "Lookup join does not support query authorization."); + } + if (!TableScanUtils.supportCompactDiffStreamingReading(table)) { TableScanUtils.streamingReadingValidate(table); } @@ -213,8 +218,7 @@ private void open() throws Exception { table instanceof FallbackReadFileStoreTable && ((FallbackReadFileStoreTable) table).other() instanceof ChainGroupReadTable; - if (!table.coreOptions().queryAuthEnabled() - && !isChainTable + if (!isChainTable && options.get(LOOKUP_CACHE_MODE) == LookupCacheMode.AUTO && new HashSet<>(table.primaryKeys()).equals(new HashSet<>(joinKeys))) { if (isRemoteServiceAvailable(table)) { diff --git a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java index d4ff5f3010f1..ce98fcbf138d 100644 --- a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java +++ b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java @@ -19,9 +19,6 @@ package org.apache.paimon.flink.lookup; import org.apache.paimon.CoreOptions; -import org.apache.paimon.catalog.FileSystemCatalog; -import org.apache.paimon.catalog.Identifier; -import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; @@ -32,8 +29,6 @@ import org.apache.paimon.flink.lookup.PrimaryKeyPartialLookupTable.RemoteQueryExecutor; import org.apache.paimon.lookup.rocksdb.RocksDBOptions; import org.apache.paimon.options.Options; -import org.apache.paimon.predicate.Predicate; -import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.schema.Schema; import org.apache.paimon.schema.SchemaManager; import org.apache.paimon.schema.TableSchema; @@ -48,7 +43,6 @@ import org.apache.paimon.types.DataType; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; -import org.apache.paimon.utils.JsonSerdeUtil; import org.apache.paimon.utils.TraceableFileIO; import org.apache.flink.table.data.RowData; @@ -243,46 +237,12 @@ public void testDefaultRemotePartial(boolean refreshAsync) throws Exception { } @Test - public void testQueryAuthUsesFullCacheAndAppliesRowFilter() throws Exception { - Predicate authFilter = new PredicateBuilder(tableRowType()).equal(2, 20L); - TableQueryAuthResult authResult = - new TableQueryAuthResult( - Collections.singletonList(JsonSerdeUtil.toFlatJson(authFilter)), null); - Identifier identifier = Identifier.create("default", "table"); - CatalogEnvironment catalogEnvironment = - new CatalogEnvironment( - identifier, - null, - () -> - new FileSystemCatalog(fileIO, tablePath) { - @Override - public TableQueryAuthResult authTableQuery( - Identifier ignored, List select) { - return authResult; - } - }, - null, - null, - null, - false, - false); - table = createFileStoreTable(false, false, false, null, true, catalogEnvironment); - - StreamTableWrite writer = table.newStreamWriteBuilder().newWrite(); - writer.write(GenericRow.of(1, 1, 10L)); - writer.write(GenericRow.of(2, 2, 20L)); - commit(writer.prepareCommit(true, 1)); - writer.close(); - - ServiceManager serviceManager = new ServiceManager(fileIO, tablePath); - serviceManager.resetService( - PRIMARY_KEY_LOOKUP, new InetSocketAddress[] {new InetSocketAddress(1)}); - lookupFunction = createLookupFunction(table, true); - lookupFunction.open(tempDir.toString()); + public void testLookupRejectsQueryAuth() throws Exception { + table = createFileStoreTable(false, false, false, null, true, CatalogEnvironment.empty()); - assertThat(lookupFunction.lookupTable()).isInstanceOf(FullCacheLookupTable.class); - assertThat(lookupFunction.lookup(new FlinkRowData(GenericRow.of(1, 1)))).isEmpty(); - assertThat(lookupFunction.lookup(new FlinkRowData(GenericRow.of(2, 2)))).hasSize(1); + assertThatThrownBy(() -> createLookupFunction(table, true)) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("Lookup join does not support query authorization"); } @Test From 57df19d3e535702d0251cbea6671c3ea053330e9 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 21:09:26 +0800 Subject: [PATCH 28/31] [spark] Fix lateral vector search column auth --- .../spark/execution/PaimonStrategy.scala | 8 ++++-- .../spark/SparkCatalogWithRestTest.java | 28 ++++++++++--------- 2 files changed, 21 insertions(+), 15 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala index 9a7fc98ad740..e9a84a361166 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala @@ -355,7 +355,7 @@ case class LateralVectorSearchExec( .withLimit(limit) .withOptions(options.asJava) pushSearchFilters(Seq(readBuilder, physicalReadBuilder), vectorSearchBuilder) - pushQueryAuthFilter(coreOptions, vectorSearchBuilder) + pushQueryAuthFilter(coreOptions, (readFieldNames :+ columnName).distinct, vectorSearchBuilder) val vectorPlan = vectorSearchBuilder.newVectorScan().scan() val batchSize = @@ -410,6 +410,7 @@ case class LateralVectorSearchExec( private def pushQueryAuthFilter( coreOptions: CoreOptions, + authFieldNames: Seq[String], vectorSearchBuilder: BatchVectorSearchBuilder): Unit = { if (!coreOptions.queryAuthEnabled()) { return @@ -418,7 +419,10 @@ case class LateralVectorSearchExec( // Vector search applies its limit before physical readers enforce QueryAuthSplit. Push the row // filter into candidate selection so unauthorized rows cannot consume the limit. val table = innerTable.asInstanceOf[FileStoreTable] - val authResult = table.catalogEnvironment().tableQueryAuth(coreOptions).auth(null) + val authResult = table + .catalogEnvironment() + .tableQueryAuth(coreOptions) + .auth(authFieldNames.asJava) Option(authResult) .flatMap(result => Option(result.extractPredicate())) .flatMap(predicate => Option(TableQueryAuthResult.remapPredicate(predicate, table.rowType()))) diff --git a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java index 4f893a410f50..3869b45c5aac 100644 --- a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java +++ b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkCatalogWithRestTest.java @@ -805,10 +805,10 @@ public void testStreamingRowFilter() throws Exception { } @Test - public void testLateralPrimaryKeyVectorSearchWithRowFilter() { + public void testLateralPrimaryKeyVectorSearchWithRowAndColumnAuth() { spark.sql( "CREATE TABLE t_auth_vector (" - + "id INT, embedding ARRAY, query_embedding ARRAY) " + + "id INT, embedding ARRAY, query_embedding ARRAY, secret STRING) " + "TBLPROPERTIES (" + "'primary-key'='id', " + "'bucket'='1', " @@ -823,17 +823,18 @@ public void testLateralPrimaryKeyVectorSearchWithRowFilter() { + "'query-auth.enabled'='true')"); spark.sql( "INSERT INTO t_auth_vector VALUES " - + "(1, array(5.0f, 0.0f), array(0.0f, 0.0f)), " - + "(2, array(1.0f, 0.0f), array(0.0f, 0.0f))"); + + "(1, array(5.0f, 0.0f), array(0.0f, 0.0f), 's1'), " + + "(2, array(1.0f, 0.0f), array(0.0f, 0.0f), 's2')"); Predicate idEq1Predicate = LeafPredicate.of( new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), Equal.INSTANCE, Collections.singletonList(1)); - restCatalogServer.setRowFilterAuth( - Identifier.create("db2", "t_auth_vector"), - Collections.singletonList(idEq1Predicate)); + Identifier identifier = Identifier.create("db2", "t_auth_vector"); + restCatalogServer.setRowFilterAuth(identifier, Collections.singletonList(idEq1Predicate)); + restCatalogServer.addTableColumnAuth( + identifier, Arrays.asList("id", "embedding", "query_embedding")); List rows = spark.sql( @@ -849,7 +850,7 @@ public void testLateralPrimaryKeyVectorSearchWithRowFilter() { public void testLateralPrimaryKeyVectorSearchMetadataOnlyWithRowFilter() { spark.sql( "CREATE TABLE t_auth_vector_metadata (" - + "id INT, embedding ARRAY, query_embedding ARRAY) " + + "id INT, embedding ARRAY, query_embedding ARRAY, secret STRING) " + "TBLPROPERTIES (" + "'primary-key'='id', " + "'bucket'='1', " @@ -864,17 +865,18 @@ public void testLateralPrimaryKeyVectorSearchMetadataOnlyWithRowFilter() { + "'query-auth.enabled'='true')"); spark.sql( "INSERT INTO t_auth_vector_metadata VALUES " - + "(1, array(5.0f, 0.0f), array(0.0f, 0.0f)), " - + "(2, array(1.0f, 0.0f), array(0.0f, 0.0f))"); + + "(1, array(5.0f, 0.0f), array(0.0f, 0.0f), 's1'), " + + "(2, array(1.0f, 0.0f), array(0.0f, 0.0f), 's2')"); Predicate idEq1Predicate = LeafPredicate.of( new FieldTransform(new FieldRef(0, "id", DataTypes.INT())), Equal.INSTANCE, Collections.singletonList(1)); - restCatalogServer.setRowFilterAuth( - Identifier.create("db2", "t_auth_vector_metadata"), - Collections.singletonList(idEq1Predicate)); + Identifier identifier = Identifier.create("db2", "t_auth_vector_metadata"); + restCatalogServer.setRowFilterAuth(identifier, Collections.singletonList(idEq1Predicate)); + restCatalogServer.addTableColumnAuth( + identifier, Arrays.asList("id", "embedding", "query_embedding")); List rows = spark.sql( From 66eaaa58ba12efc7aed65944c655fcfc19488828 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 21:11:37 +0800 Subject: [PATCH 29/31] [core] Stabilize index bootstrap reader cleanup --- .../crosspartition/IndexBootstrapTest.java | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java b/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java index 0921ddf565ea..7f0435f10c41 100644 --- a/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java @@ -44,14 +44,17 @@ import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.CommonTestUtils; import org.apache.paimon.utils.Filter; import org.apache.paimon.utils.JsonSerdeUtil; import org.apache.paimon.utils.Pair; +import org.apache.paimon.utils.TraceableFileIO; import org.junit.jupiter.api.Test; import org.mockito.AdditionalAnswers; import org.mockito.Mockito; +import java.time.Duration; import java.time.Instant; import java.time.ZoneId; import java.util.ArrayList; @@ -106,10 +109,7 @@ public void testBoostrap() throws Exception { GenericRow.of(2, 1, 3), GenericRow.of(4, 2, 5), GenericRow.of(6, 3, 7)); result.clear(); - // In ParallelExecution, latch.countDown first, then close the reader, it may not be closed - // here, (this is good, beneficial for query speed) but TableTestBase.after will check leak - // streams. So sleep here to avoid unstable. - Thread.sleep(1000); + waitForParallelReadersClosed(); } private Table createTable() throws Exception { @@ -182,6 +182,7 @@ public void testBootstrapIgnoresQueryAuth() throws Exception { reader.forEachRemaining( row -> result.add(GenericRow.of(row.getInt(0), row.getInt(1), row.getInt(2)))); } + waitForParallelReadersClosed(); assertThat(result) .containsExactlyInAnyOrder( @@ -213,6 +214,7 @@ public void testBootstrapIgnoresQueryAuthWithFallback() throws Exception { reader.forEachRemaining( row -> result.add(GenericRow.of(row.getInt(0), row.getInt(1), row.getInt(2)))); } + waitForParallelReadersClosed(); assertThat(result) .containsExactlyInAnyOrder(GenericRow.of(10, 1, 2), GenericRow.of(20, 2, 3)); @@ -228,6 +230,16 @@ private TableQueryAuthResult queryAuthResult(FileStoreTable table) { new FieldTransform(new FieldRef(0, "pt", DataTypes.INT()))))); } + private void waitForParallelReadersClosed() throws Exception { + CommonTestUtils.waitUtil( + () -> + TraceableFileIO.openInputStreams( + path -> path.toString().contains(tempPath.toString())) + .isEmpty(), + Duration.ofSeconds(10), + Duration.ofMillis(10)); + } + private static class QueryAuthFileStoreTable extends DelegatedFileStoreTable { private final TableQueryAuthResult authResult; From 45e78aae361d76c1697d56ba3c97556ba9b25ce8 Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 22:38:39 +0800 Subject: [PATCH 30/31] [core] Handle query auth in maintenance and system tables --- .../paimon/operation/ListUnexistingFiles.java | 11 ++++-- .../paimon/table/system/BinlogTable.java | 3 +- .../table/system/FileKeyRangesTable.java | 8 ++-- .../paimon/table/system/FilesTable.java | 8 ++-- .../operation/ListUnexistingFilesTest.java | 37 ++++++++++++++++++ .../paimon/table/system/BinlogTableTest.java | 39 +++++++++++++++++-- .../table/system/FileKeyRangesTableTest.java | 28 +++++++++++++ .../paimon/table/system/FilesTableTest.java | 24 ++++++++++++ 8 files changed, 145 insertions(+), 13 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/ListUnexistingFiles.java b/paimon-core/src/main/java/org/apache/paimon/operation/ListUnexistingFiles.java index dc3d029fb6a7..d4c836d652d1 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/ListUnexistingFiles.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/ListUnexistingFiles.java @@ -40,6 +40,8 @@ import java.util.Map; import java.util.concurrent.ThreadPoolExecutor; +import static org.apache.paimon.CoreOptions.QUERY_AUTH_ENABLED; + /** List what data files recorded in manifests are missing from the filesystem. */ public class ListUnexistingFiles { @@ -48,11 +50,14 @@ public class ListUnexistingFiles { private final ThreadPoolExecutor executor; public ListUnexistingFiles(FileStoreTable table) { - this.table = table; - this.pathFactory = table.store().pathFactory(); + this.table = + table.coreOptions().queryAuthEnabled() + ? table.copy(Collections.singletonMap(QUERY_AUTH_ENABLED.key(), "false")) + : table; + this.pathFactory = this.table.store().pathFactory(); this.executor = FileOperationThreadPool.getExecutorService( - table.coreOptions().fileOperationThreadNum()); + this.table.coreOptions().fileOperationThreadNum()); } public Map> list(BinaryRow partition) throws Exception { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/system/BinlogTable.java b/paimon-core/src/main/java/org/apache/paimon/table/system/BinlogTable.java index 5c1268f96304..22603f32ab37 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/system/BinlogTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/system/BinlogTable.java @@ -28,6 +28,7 @@ import org.apache.paimon.table.Table; 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.types.ArrayType; import org.apache.paimon.types.DataField; @@ -118,7 +119,7 @@ public InnerTableRead withReadType(RowType readType) { @Override public RecordReader createReader(Split split) throws IOException { - DataSplit dataSplit = (DataSplit) split; + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(split); // When sequence number is enabled, the underlying data layout is: // [_SEQUENCE_NUMBER, pk, pt, col1, ...] // We need to offset the field index to skip the sequence number field. diff --git a/paimon-core/src/main/java/org/apache/paimon/table/system/FileKeyRangesTable.java b/paimon-core/src/main/java/org/apache/paimon/table/system/FileKeyRangesTable.java index 39e88f6e9afc..a223cdaf2824 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/system/FileKeyRangesTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/system/FileKeyRangesTable.java @@ -40,6 +40,7 @@ import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.InnerTableRead; import org.apache.paimon.table.source.InnerTableScan; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.ReadOnceTableScan; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.TableRead; @@ -274,13 +275,14 @@ public RowDataToObjectArrayConverter apply(Long schemaId) { }; List> iteratorList = new ArrayList<>(); - for (Split dataSplit : splits) { + for (Split plannedSplit : splits) { + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(plannedSplit); iteratorList.add( Iterators.transform( - ((DataSplit) dataSplit).dataFiles().iterator(), + dataSplit.dataFiles().iterator(), file -> toRow( - (DataSplit) dataSplit, + dataSplit, partitionCastExecutor, keyConverters, file))); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/system/FilesTable.java b/paimon-core/src/main/java/org/apache/paimon/table/system/FilesTable.java index 3840b4dc54cb..daffc8f2e0f1 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/system/FilesTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/system/FilesTable.java @@ -47,6 +47,7 @@ import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.InnerTableRead; import org.apache.paimon.table.source.InnerTableScan; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.ReadOnceTableScan; import org.apache.paimon.table.source.SingletonSplit; import org.apache.paimon.table.source.Split; @@ -375,13 +376,14 @@ public RowDataToObjectArrayConverter apply(Long schemaId) { }); } }; - for (Split dataSplit : splits) { + for (Split plannedSplit : splits) { + DataSplit dataSplit = QueryAuthSplit.unwrapDataSplit(plannedSplit); iteratorList.add( Iterators.transform( - ((DataSplit) dataSplit).dataFiles().iterator(), + dataSplit.dataFiles().iterator(), file -> toRow( - (DataSplit) dataSplit, + dataSplit, partitionCastExecutor, keyConverters, file, diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/ListUnexistingFilesTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/ListUnexistingFilesTest.java index cc2f0c16967a..99abb67a5c9f 100644 --- a/paimon-core/src/test/java/org/apache/paimon/operation/ListUnexistingFilesTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/operation/ListUnexistingFilesTest.java @@ -21,6 +21,7 @@ import org.apache.paimon.CoreOptions; import org.apache.paimon.catalog.FileSystemCatalog; import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.BinaryRow; import org.apache.paimon.data.BinaryRowWriter; import org.apache.paimon.data.GenericRow; @@ -32,14 +33,19 @@ import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.sink.TableCommitImpl; import org.apache.paimon.table.sink.TableWriteImpl; +import org.apache.paimon.table.source.DataTableScan; +import org.apache.paimon.table.source.Split; import org.apache.paimon.types.DataType; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowType; +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.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.AdditionalAnswers; +import org.mockito.Mockito; import java.util.ArrayList; import java.util.Arrays; @@ -92,6 +98,37 @@ public void testListFiles(int bucket, boolean dvEnabled) throws Exception { } } + @Test + public void testListFilesDisablesQueryAuth() throws Exception { + int[] numDeletes = new int[1]; + FileStoreTable table = + prepareRandomlyDeletedTable( + tempDir.toString(), "mydb", "auth_t", -1, 1, numDeletes, false); + List rawSplits = table.newScan().plan().splits(); + table = table.copy(Collections.singletonMap(CoreOptions.QUERY_AUTH_ENABLED.key(), "true")); + + TableQueryAuthResult authResult = + new TableQueryAuthResult(Collections.emptyList(), Collections.emptyMap()); + DataTableScan authScan = Mockito.mock(DataTableScan.class); + Mockito.when(authScan.withLevelFilter(Mockito.any())).thenReturn(authScan); + Mockito.when(authScan.withPartitionFilter(Mockito.>any())) + .thenReturn(authScan); + Mockito.when(authScan.plan()).thenReturn(authResult.convertPlan(() -> rawSplits)); + + FileStoreTable authTable = + Mockito.mock(FileStoreTable.class, AdditionalAnswers.delegatesTo(table)); + Mockito.when(authTable.newScan()).thenReturn(authScan); + + BinaryRow partition = new BinaryRow(1); + BinaryRowWriter writer = new BinaryRowWriter(partition); + writer.writeInt(0, 0); + writer.complete(); + + Map> result = + new ListUnexistingFiles(authTable).list(partition); + assertThat(result.values().stream().mapToInt(Map::size).sum()).isEqualTo(numDeletes[0]); + } + public static FileStoreTable prepareRandomlyDeletedTable( String warehouse, String databaseName, diff --git a/paimon-core/src/test/java/org/apache/paimon/table/system/BinlogTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/system/BinlogTableTest.java index 6f187fbffd98..107dab5cb46c 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/system/BinlogTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/system/BinlogTableTest.java @@ -20,6 +20,7 @@ import org.apache.paimon.CoreOptions; import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericArray; import org.apache.paimon.data.GenericRow; @@ -27,6 +28,8 @@ import org.apache.paimon.fs.FileIO; import org.apache.paimon.fs.Path; import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.Schema; import org.apache.paimon.schema.SchemaChange; import org.apache.paimon.schema.SchemaManager; @@ -35,12 +38,18 @@ import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.FileStoreTableFactory; import org.apache.paimon.table.TableTestBase; +import org.apache.paimon.table.source.QueryAuthSplit; +import org.apache.paimon.table.source.ReadBuilder; +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.RowKind; +import org.apache.paimon.utils.JsonSerdeUtil; import org.junit.jupiter.api.Test; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import static org.apache.paimon.catalog.Identifier.SYSTEM_TABLE_SPLITTER; @@ -49,6 +58,8 @@ /** Unit tests for {@link BinlogTable}. */ public class BinlogTableTest extends TableTestBase { + private FileStoreTable dataTable; + @Test public void testReadBinlogFromLatest() throws Exception { BinlogTable binlogTable = createBinlogTable("binlog_table", false); @@ -60,6 +71,29 @@ public void testReadBinlogFromLatest() throws Exception { assertThat(result).containsExactlyInAnyOrderElementsOf(expectRow); } + @Test + public void testReadQueryAuthSplits() throws Exception { + BinlogTable binlogTable = createBinlogTable("binlog_table_with_auth", false); + PredicateBuilder builder = new PredicateBuilder(dataTable.rowType()); + TableQueryAuthResult authResult = + new TableQueryAuthResult( + Collections.singletonList(JsonSerdeUtil.toFlatJson(builder.equal(1, 2))), + null); + + ReadBuilder readBuilder = binlogTable.newReadBuilder(); + TableRead read = readBuilder.newRead(); + List result = new ArrayList<>(); + for (Split split : readBuilder.newScan().plan().splits()) { + try (RecordReader reader = + read.createReader(new QueryAuthSplit(split, authResult))) { + reader.forEachRemaining( + row -> result.add(row.getString(0) + "-" + row.getArray(2).getInt(0))); + } + } + + assertThat(result).containsExactly(RowKind.UPDATE_AFTER.shortString() + "-2"); + } + @Test public void testReadSequenceNumberWithTableOption() throws Exception { BinlogTable binlogTable = createBinlogTable("binlog_table_with_seq", true); @@ -121,10 +155,9 @@ private BinlogTable createBinlogTable(String tableName, boolean enableSequenceNu TableSchema tableSchema = SchemaUtils.forceCommit( new SchemaManager(fileIO, tablePath), schemaBuilder.build()); - FileStoreTable table = - FileStoreTableFactory.create(LocalFileIO.create(), tablePath, tableSchema); + dataTable = FileStoreTableFactory.create(LocalFileIO.create(), tablePath, tableSchema); - writeTestData(table); + writeTestData(dataTable); Identifier binlogTableId = identifier(tableName + SYSTEM_TABLE_SPLITTER + BinlogTable.BINLOG); diff --git a/paimon-core/src/test/java/org/apache/paimon/table/system/FileKeyRangesTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/system/FileKeyRangesTableTest.java index de00e33b3062..6342cbc6e3b2 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/system/FileKeyRangesTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/system/FileKeyRangesTableTest.java @@ -19,26 +19,34 @@ package org.apache.paimon.table.system; import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; import org.apache.paimon.predicate.In; import org.apache.paimon.predicate.LeafPredicate; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.Schema; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.TableTestBase; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.ReadBuilder; +import org.apache.paimon.table.source.Split; import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; +import java.util.stream.Collectors; import static org.apache.paimon.catalog.Identifier.SYSTEM_TABLE_SPLITTER; import static org.apache.paimon.io.DataFileTestUtils.row; @@ -96,6 +104,26 @@ public void testReadBasic() throws Exception { }); } + @Test + public void testReadQueryAuthSplits() throws Exception { + TableQueryAuthResult authResult = + new TableQueryAuthResult(Collections.emptyList(), Collections.emptyMap()); + List splits = + table.newScan().plan().splits().stream() + .map(split -> new QueryAuthSplit(split, authResult)) + .collect(Collectors.toList()); + FilesTable.FilesSplit filesSplit = Mockito.mock(FilesTable.FilesSplit.class); + Mockito.when(filesSplit.splits(Mockito.any(FileStoreTable.class))).thenReturn(splits); + + List paths = new ArrayList<>(); + try (RecordReader reader = + fileKeyRangesTable.newRead().createReader(filesSplit)) { + reader.forEachRemaining(row -> paths.add(row.getString(2).toString())); + } + + assertThat(paths).hasSameSizeAs(read(fileKeyRangesTable)); + } + @Test public void testPartitionFilter() throws Exception { PredicateBuilder builder = new PredicateBuilder(FileKeyRangesTable.TABLE_TYPE); diff --git a/paimon-core/src/test/java/org/apache/paimon/table/system/FilesTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/system/FilesTableTest.java index 4c4e3c88598a..cf46675b9340 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/system/FilesTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/system/FilesTableTest.java @@ -20,6 +20,7 @@ import org.apache.paimon.CoreOptions; import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.catalog.TableQueryAuthResult; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; @@ -34,6 +35,7 @@ import org.apache.paimon.predicate.LeafPredicate; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.Schema; import org.apache.paimon.schema.SchemaChange; import org.apache.paimon.schema.SchemaManager; @@ -47,13 +49,16 @@ import org.apache.paimon.table.sink.BatchWriteBuilder; import org.apache.paimon.table.sink.CommitMessage; import org.apache.paimon.table.sink.CommitMessageImpl; +import org.apache.paimon.table.source.QueryAuthSplit; import org.apache.paimon.table.source.ReadBuilder; +import org.apache.paimon.table.source.Split; import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; import org.apache.paimon.utils.SnapshotManager; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; import java.io.IOException; import java.util.ArrayList; @@ -127,6 +132,25 @@ public void testReadWithFilter() throws Exception { .containsExactlyInAnyOrder("{2, 20}-0-5"); } + @Test + public void testReadQueryAuthSplits() throws Exception { + TableQueryAuthResult authResult = + new TableQueryAuthResult(Collections.emptyList(), Collections.emptyMap()); + List splits = + table.newScan().plan().splits().stream() + .map(split -> new QueryAuthSplit(split, authResult)) + .collect(Collectors.toList()); + FilesTable.FilesSplit filesSplit = Mockito.mock(FilesTable.FilesSplit.class); + Mockito.when(filesSplit.splits(Mockito.any(FileStoreTable.class))).thenReturn(splits); + + List paths = new ArrayList<>(); + try (RecordReader reader = filesTable.newRead().createReader(filesSplit)) { + reader.forEachRemaining(row -> paths.add(row.getString(2).toString())); + } + + assertThat(paths).hasSameSizeAs(read(filesTable)); + } + private List readPartBucketLevel(Predicate predicate) throws IOException { ReadBuilder readBuilder = filesTable.newReadBuilder().withFilter(predicate); List rows = new ArrayList<>(); From 6761e1b9da1d9b25db559068668f9b524c8d6eee Mon Sep 17 00:00:00 2001 From: umi Date: Mon, 20 Jul 2026 22:39:05 +0800 Subject: [PATCH 31/31] [spark] Preserve query auth through split planning --- .../org/apache/paimon/spark/PaimonScan.scala | 20 +++--- .../org/apache/paimon/spark/PaimonScan.scala | 24 +++---- .../paimon/spark/read/BinPackingSplits.scala | 66 ++++++++++++------- .../apache/paimon/spark/util/SplitUtils.scala | 22 +++++-- .../paimon/spark/BinPackingSplitsTest.scala | 31 ++++++++- 5 files changed, 115 insertions(+), 48 deletions(-) diff --git a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala index 9cc7e409c7fe..e9acf4dbd7a4 100644 --- a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala +++ b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala @@ -21,8 +21,9 @@ package org.apache.paimon.spark import org.apache.paimon.partition.PartitionPredicate import org.apache.paimon.predicate.{FullTextSearch, HybridSearch, Predicate, TopN, VectorSearch} import org.apache.paimon.spark.read.VariantExtractionInfo +import org.apache.paimon.spark.util.SplitUtils import org.apache.paimon.table.{BucketMode, FileStoreTable, InnerTable} -import org.apache.paimon.table.source.{DataSplit, Split} +import org.apache.paimon.table.source.Split import org.apache.spark.sql.connector.expressions._ import org.apache.spark.sql.connector.read.SupportsReportPartitioning @@ -82,12 +83,12 @@ case class PaimonScan( /** Extract the bucket number from the splits only if all splits have the same totalBuckets number. */ private def extractBucketNumber(): Option[Int] = { - val splits = inputSplits - if (splits.exists(!_.isInstanceOf[DataSplit])) { + val dataSplits = inputSplits.map(SplitUtils.dataSplit) + if (dataSplits.exists(_.isEmpty)) { None } else { val deduplicated = - splits.map(s => Option(s.asInstanceOf[DataSplit].totalBuckets())).toSeq.distinct + dataSplits.map(s => Option(s.get.totalBuckets())).toSeq.distinct deduplicated match { case Seq(Some(num)) => Some(num) @@ -108,16 +109,17 @@ case class PaimonScan( } override def getInputPartitions(splits: Array[Split]): Seq[PaimonInputPartition] = { - if (!shouldDoBucketedScan || splits.exists(!_.isInstanceOf[DataSplit])) { + val splitsWithMetadata = splits.map(split => (split, SplitUtils.dataSplit(split))) + if (!shouldDoBucketedScan || splitsWithMetadata.exists(_._2.isEmpty)) { return super.getInputPartitions(splits) } - splits - .map(_.asInstanceOf[DataSplit]) - .groupBy(_.bucket()) + splitsWithMetadata + .map { case (split, dataSplit) => (split, dataSplit.get) } + .groupBy(_._2.bucket()) .map { case (bucket, groupedSplits) => - PaimonBucketedInputPartition(groupedSplits, bucket) + PaimonBucketedInputPartition(groupedSplits.map(_._1), bucket) } .toSeq } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala index 1f0b1b75b038..338bf7344df8 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala @@ -23,8 +23,9 @@ import org.apache.paimon.partition.PartitionPredicate import org.apache.paimon.predicate.{FullTextSearch, HybridSearch, Predicate, TopN, VectorSearch} import org.apache.paimon.spark.commands.BucketExpression.quote import org.apache.paimon.spark.read.VariantExtractionInfo +import org.apache.paimon.spark.util.SplitUtils import org.apache.paimon.table.{BucketMode, FileStoreTable, InnerTable} -import org.apache.paimon.table.source.{DataSplit, Split} +import org.apache.paimon.table.source.Split import org.apache.spark.sql.PaimonUtils.fieldReference import org.apache.spark.sql.connector.expressions._ @@ -93,12 +94,12 @@ case class PaimonScan( /** Extract the bucket number from the splits only if all splits have the same totalBuckets number. */ private def extractBucketNumber(): Option[Int] = { - val splits = inputSplits - if (splits.exists(!_.isInstanceOf[DataSplit])) { + val dataSplits = inputSplits.map(SplitUtils.dataSplit) + if (dataSplits.exists(_.isEmpty)) { None } else { val deduplicated = - splits.map(s => Option(s.asInstanceOf[DataSplit].totalBuckets())).toSeq.distinct + dataSplits.map(s => Option(s.get.totalBuckets())).toSeq.distinct deduplicated match { case Seq(Some(num)) => Some(num) @@ -136,14 +137,14 @@ case class PaimonScan( val allSplitsKeepOrdering = inputPartitions.toSeq .map(_.asInstanceOf[PaimonBucketedInputPartition]) - .map(_.splits.asInstanceOf[Seq[DataSplit]]) + .map(_.splits.map(SplitUtils.dataSplit)) .forall { splits => // Only support report ordering if all matches: // - one `Split` per InputPartition (TODO: Re-construct splits using minKey/maxKey) // - `Split` is not rawConvertible so that the merge read can happen // - `Split` only contains one data file so it always sorted even without merge read - splits.size < 2 && splits.forall { + splits.forall(_.isDefined) && splits.size < 2 && splits.flatten.forall { split => !split.rawConvertible() || split.dataFiles().size() < 2 } } @@ -164,16 +165,17 @@ case class PaimonScan( } override protected def getInputPartitions(splits: Array[Split]): Seq[PaimonInputPartition] = { - if (!shouldDoBucketedScan || splits.exists(!_.isInstanceOf[DataSplit])) { + val splitsWithMetadata = splits.map(split => (split, SplitUtils.dataSplit(split))) + if (!shouldDoBucketedScan || splitsWithMetadata.exists(_._2.isEmpty)) { return super.getInputPartitions(splits) } - splits - .map(_.asInstanceOf[DataSplit]) - .groupBy(_.bucket()) + splitsWithMetadata + .map { case (split, dataSplit) => (split, dataSplit.get) } + .groupBy(_._2.bucket()) .map { case (bucket, groupedSplits) => - PaimonBucketedInputPartition(groupedSplits, bucket) + PaimonBucketedInputPartition(groupedSplits.map(_._1), bucket) } .toSeq } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BinPackingSplits.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BinPackingSplits.scala index f27af93cf604..91ff526e74a8 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BinPackingSplits.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BinPackingSplits.scala @@ -25,7 +25,7 @@ import org.apache.paimon.spark.PaimonInputPartition import org.apache.paimon.spark.util.SplitUtils import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit import org.apache.paimon.table.format.FormatDataSplit -import org.apache.paimon.table.source.{DataSplit, DeletionFile, Split} +import org.apache.paimon.table.source.{DataSplit, DeletionFile, QueryAuthSplit, Split} import org.apache.spark.internal.Logging import org.apache.spark.sql.PaimonSparkSession @@ -78,18 +78,22 @@ case class BinPackingSplits(coreOptions: CoreOptions, readRowSizeRatio: Double = def pack(splits: Array[Split]): Seq[PaimonInputPartition] = { val (toReshuffle, reserved) = splits.partition { case _: FallbackSplit => false - case split: DataSplit => split.rawConvertible() || coreOptions.dataEvolutionEnabled() - // FormatDataSplit is already packed with multiple files by target size in the core scan, - // so each split maps directly to one input partition here. - case _: FormatDataSplit => false - case _ => false + case split => + SplitUtils.unwrapQueryAuth(split) match { + case dataSplit: DataSplit => + dataSplit.rawConvertible() || coreOptions.dataEvolutionEnabled() + // FormatDataSplit is already packed with multiple files by target size in the core scan, + // so each split maps directly to one input partition here. + case _: FormatDataSplit => false + case _ => false + } } if (toReshuffle.nonEmpty) { val startTS = System.currentTimeMillis() val reshuffled = if (coreOptions.dataEvolutionEnabled()) { - packDataEvolutionSplit(toReshuffle.collect { case ds: DataSplit => ds }) + packDataEvolutionSplit(toReshuffle) } else { - packDataSplit(toReshuffle.collect { case ds: DataSplit => ds }) + packDataSplit(toReshuffle) } val all = reserved.map(PaimonInputPartition.apply) ++ reshuffled val duration = System.currentTimeMillis() - startTS @@ -102,22 +106,23 @@ case class BinPackingSplits(coreOptions: CoreOptions, readRowSizeRatio: Double = } } - private def packDataSplit(splits: Array[DataSplit]): Array[PaimonInputPartition] = { - val maxSplitBytes = computeMaxSplitBytes(splits) + private def packDataSplit(splits: Array[Split]): Array[PaimonInputPartition] = { + val maxSplitBytes = computeMaxSplitBytes(splits.map(SplitUtils.dataSplit(_).get)) var currentSize = 0L - val currentSplits = new ArrayBuffer[DataSplit] + val currentSplits = new ArrayBuffer[Split] val partitions = new ArrayBuffer[PaimonInputPartition] - var currentSplit: Option[DataSplit] = None + var currentSplit: Option[(Split, DataSplit)] = None val currentDataFiles = new ArrayBuffer[DataFileMeta] val currentDeletionFiles = new ArrayBuffer[DeletionFile] def closeDataSplit(): Unit = { if (currentSplit.nonEmpty && currentDataFiles.nonEmpty) { - val newSplit = - copyDataSplit(currentSplit.get, currentDataFiles.toSeq, currentDeletionFiles.toSeq) - currentSplits += newSplit + val (sourceSplit, dataSplit) = currentSplit.get + val replacement = + copyDataSplit(dataSplit, currentDataFiles.toSeq, currentDeletionFiles.toSeq) + currentSplits += QueryAuthSplit.retainAuth(sourceSplit, replacement) } currentDataFiles.clear() currentDeletionFiles.clear() @@ -134,13 +139,20 @@ case class BinPackingSplits(coreOptions: CoreOptions, readRowSizeRatio: Double = splits.foreach { split => - if (!currentSplit.exists(withSamePartitionAndBucket(_, split))) { + val dataSplit = SplitUtils.dataSplit(split).get + if ( + !currentSplit.exists { + case (source, currentDataSplit) => + withSamePartitionAndBucket(currentDataSplit, dataSplit) && + withSameQueryAuth(source, split) + } + ) { // close and open another data split closeDataSplit() - currentSplit = Some(split) + currentSplit = Some((split, dataSplit)) } - val ddFiles = dataFileAndDeletionFiles(split) + val ddFiles = dataFileAndDeletionFiles(dataSplit) ddFiles.foreach { case (dataFile, deletionFile) => val size = @@ -161,11 +173,11 @@ case class BinPackingSplits(coreOptions: CoreOptions, readRowSizeRatio: Double = partitions.toArray } - private def packDataEvolutionSplit(splits: Array[DataSplit]): Array[PaimonInputPartition] = { - val maxSplitBytes = computeMaxSplitBytes(splits) + private def packDataEvolutionSplit(splits: Array[Split]): Array[PaimonInputPartition] = { + val maxSplitBytes = computeMaxSplitBytes(splits.map(SplitUtils.dataSplit(_).get)) var currentSize = 0L - val currentSplits = new ArrayBuffer[DataSplit] + val currentSplits = new ArrayBuffer[Split] val partitions = new ArrayBuffer[PaimonInputPartition] def closeInputPartition(): Unit = { @@ -178,7 +190,7 @@ case class BinPackingSplits(coreOptions: CoreOptions, readRowSizeRatio: Double = splits.foreach { split => - val ddFiles = dataFileAndDeletionFiles(split) + val ddFiles = dataFileAndDeletionFiles(SplitUtils.dataSplit(split).get) val size = ddFiles.map { case (dataFile, deletionFile) => (dataFile.fileSize() * readRowSizeRatio).toLong + openCostInBytes + Option(deletionFile) @@ -219,6 +231,16 @@ case class BinPackingSplits(coreOptions: CoreOptions, readRowSizeRatio: Double = split1.partition().equals(split2.partition()) && split1.bucket() == split2.bucket() } + private def withSameQueryAuth(split1: Split, split2: Split): Boolean = { + (split1, split2) match { + case (auth1: QueryAuthSplit, auth2: QueryAuthSplit) => + Option(auth1.authResult()).map(r => (r.filter(), r.columnMasking())) == + Option(auth2.authResult()).map(r => (r.filter(), r.columnMasking())) + case (_: QueryAuthSplit, _) | (_, _: QueryAuthSplit) => false + case _ => true + } + } + private def dataFileAndDeletionFiles(split: DataSplit): Array[(DataFileMeta, DeletionFile)] = { if (deletionVectors && split.deletionFiles().isPresent) { val deletionFiles = split.deletionFiles().get().asScala diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/SplitUtils.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/SplitUtils.scala index c485fe4da345..c4924ceff7b5 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/SplitUtils.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/SplitUtils.scala @@ -19,7 +19,7 @@ package org.apache.paimon.spark.util import org.apache.paimon.table.format.FormatDataSplit -import org.apache.paimon.table.source.{DataSplit, Split} +import org.apache.paimon.table.source.{DataSplit, QueryAuthSplit, Split} import java.util.{Collections => JCollections} @@ -27,8 +27,22 @@ import scala.collection.JavaConverters._ object SplitUtils { - def splitSize(split: Split): Long = { + private[spark] def unwrapQueryAuth(split: Split): Split = { split match { + case _: QueryAuthSplit => QueryAuthSplit.unwrap(split) + case _ => split + } + } + + private[spark] def dataSplit(split: Split): Option[DataSplit] = { + unwrapQueryAuth(split) match { + case ds: DataSplit => Some(ds) + case _ => None + } + } + + def splitSize(split: Split): Long = { + unwrapQueryAuth(split) match { case ds: DataSplit => ds.dataFiles().asScala.map(_.fileSize).sum case fs: FormatDataSplit => @@ -40,7 +54,7 @@ object SplitUtils { def fileCount(split: Split): Long = dataFileCount(split) + deleteFileCount(split) def dataFileCount(split: Split): Long = { - split match { + unwrapQueryAuth(split) match { case ds: DataSplit => ds.dataFiles().size() case fs: FormatDataSplit => fs.fileCount() case _ => 0 @@ -48,7 +62,7 @@ object SplitUtils { } def deleteFileCount(split: Split): Long = { - split match { + unwrapQueryAuth(split) match { case ds: DataSplit => ds.deletionFiles() .orElse(JCollections.emptyList()) diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/BinPackingSplitsTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/BinPackingSplitsTest.scala index 9173681953b4..5b48f7aa8bdb 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/BinPackingSplitsTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/BinPackingSplitsTest.scala @@ -19,15 +19,17 @@ package org.apache.paimon.spark import org.apache.paimon.CoreOptions +import org.apache.paimon.catalog.TableQueryAuthResult import org.apache.paimon.data.BinaryRow import org.apache.paimon.io.DataFileMeta import org.apache.paimon.manifest.FileSource import org.apache.paimon.spark.read.BinPackingSplits -import org.apache.paimon.table.source.{DataSplit, DeletionFile, Split} +import org.apache.paimon.spark.util.SplitUtils +import org.apache.paimon.table.source.{DataSplit, DeletionFile, QueryAuthSplit, Split} import org.junit.jupiter.api.Assertions -import java.util.{HashMap => JHashMap} +import java.util.{Collections => JCollections, HashMap => JHashMap} import scala.collection.JavaConverters._ import scala.collection.mutable @@ -68,6 +70,31 @@ class BinPackingSplitsTest extends PaimonSparkTestBase { Assertions.assertEquals(1, reshuffled.length) } + test("Paimon: reshuffle query auth splits and retain auth") { + withSparkSQLConf("spark.sql.files.minPartitionNum" -> "1") { + val authResult = + new TableQueryAuthResult(JCollections.emptyList(), JCollections.emptyMap()) + val splits = (0 until 5).map { + index => + val dataSplit = newDataSplit(s"auth-$index", Seq(10L), rawConvertible = true) + new QueryAuthSplit(dataSplit, authResult): Split + }.toArray + val binPacking = BinPackingSplits( + CoreOptions.fromMap( + Map("source.split.open-file-cost" -> "0 B", "source.split.target-size" -> "1 MB").asJava)) + + val reshuffled = binPacking.pack(splits) + + Assertions.assertEquals(1, reshuffled.length) + Assertions.assertEquals(1, reshuffled.head.splits.length) + val packed = reshuffled.head.splits.head + Assertions.assertTrue(packed.isInstanceOf[QueryAuthSplit]) + Assertions.assertSame(authResult, packed.asInstanceOf[QueryAuthSplit].authResult()) + Assertions.assertEquals(5, SplitUtils.dataFileCount(packed)) + Assertions.assertEquals(50L, SplitUtils.splitSize(packed)) + } + } + test("Paimon: pack data evolution splits by split granularity") { withSparkSQLConf("spark.sql.files.minPartitionNum" -> "1") { val split1 = newDataSplit("split1", Seq(40L, 40L), deletionFileLength = Some(5L))