This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new e4b3343ac4 [spark] Fix bucket metadata reads for chain tables (#10652)
e4b3343ac4 is described below
commit e4b3343ac4d968c06e219785b71e08feccf35eaa
Author: Ran Tao <[email protected]>
AuthorDate: Tue Oct 20 15:51:34 2026 +0800
[spark] Fix bucket metadata reads for chain tables (#10652)
---
.../org/apache/paimon/table/source/ChainSplit.java | 16 ++
.../apache/paimon/table/source/ChainSplitTest.java | 21 +-
.../apache/paimon/utils/ChainTableUtilsTest.java | 3 +-
.../paimon/spark/PaimonRecordReaderIterator.scala | 32 ++-
.../spark/util/SplitUtilsQueryAuthTest.scala | 76 +++++-
.../paimon/spark/SparkChainTableMetadataTest.java | 272 +++++++++++++++++++++
6 files changed, 401 insertions(+), 19 deletions(-)
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/ChainSplit.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/ChainSplit.java
index c74a8ac903..57003859d8 100644
--- a/paimon-core/src/main/java/org/apache/paimon/table/source/ChainSplit.java
+++ b/paimon-core/src/main/java/org/apache/paimon/table/source/ChainSplit.java
@@ -19,6 +19,7 @@
package org.apache.paimon.table.source;
import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.fs.Path;
import org.apache.paimon.io.DataFileMeta;
import org.apache.paimon.io.DataFileMetaSerializer;
import org.apache.paimon.io.DataFileMetaWriteColsLegacySerializer;
@@ -42,6 +43,9 @@ import java.util.Objects;
import java.util.Optional;
import java.util.OptionalLong;
+import static org.apache.paimon.utils.FileStorePathFactory.BUCKET_PATH_PREFIX;
+import static org.apache.paimon.utils.Preconditions.checkState;
+
/**
* A split describes chain table read scope. It follows DataSplit's custom
serialization pattern and
* extends it with branch + partition per data file.
@@ -78,6 +82,18 @@ public class ChainSplit implements Split {
return logicalPartition;
}
+ /** Returns the physical bucket shared by all files in this split. */
+ public int bucket() {
+ // Chain splits are built per bucket. Use the stored paths so older
serialized splits
+ // can also provide bucket metadata without changing the serialization
format.
+ checkState(!fileBucketPathMapping.isEmpty(), "No bucket path in chain
split.");
+ String bucketPath = fileBucketPathMapping.values().iterator().next();
+ String bucketName = new Path(bucketPath).getName();
+ checkState(
+ bucketName.startsWith(BUCKET_PATH_PREFIX), "Invalid bucket
path: %s", bucketPath);
+ return
Integer.parseInt(bucketName.substring(BUCKET_PATH_PREFIX.length()));
+ }
+
public List<DataFileMeta> dataFiles() {
return dataFiles;
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/ChainSplitTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/ChainSplitTest.java
index 402ac581bb..76901d9f9c 100644
---
a/paimon-core/src/test/java/org/apache/paimon/table/source/ChainSplitTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/table/source/ChainSplitTest.java
@@ -30,6 +30,7 @@ import org.junit.jupiter.api.Test;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
+import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -52,7 +53,7 @@ public class ChainSplitTest {
Map<String, String> fileBucketPathMapping = new HashMap<>();
Map<String, String> fileBranchMapping = new HashMap<>();
for (DataFileMeta dataFile : dataFiles) {
- fileBucketPathMapping.put(dataFile.fileName(),
"dt=20251202/bucket_0");
+ fileBucketPathMapping.put(dataFile.fileName(),
"dt=20251202/bucket-7");
fileBranchMapping.put(dataFile.fileName(), "branch_0");
}
ChainSplit split =
@@ -68,6 +69,7 @@ public class ChainSplitTest {
assertThat(fileBucketPathMapping).isEqualTo(newSplit.fileBucketPathMapping());
assertThat(fileBranchMapping).isEqualTo(newSplit.fileBranchMapping());
assertThat(newSplit.deletionFiles()).isEmpty();
+ assertThat(newSplit.bucket()).isEqualTo(7);
assertThat(newSplit).isEqualTo(split);
}
@@ -89,7 +91,7 @@ public class ChainSplitTest {
null,
new DeletionFile("dv3", 0, 50, 1L));
for (DataFileMeta dataFile : dataFiles) {
- fileBucketPathMapping.put(dataFile.fileName(),
"dt=20251202/bucket_0");
+ fileBucketPathMapping.put(dataFile.fileName(),
"dt=20251202/bucket-7");
fileBranchMapping.put(dataFile.fileName(), "branch_0");
}
ChainSplit split =
@@ -106,6 +108,21 @@ public class ChainSplitTest {
assertThat(fileBranchMapping).isEqualTo(newSplit.fileBranchMapping());
assertThat(newSplit.deletionFiles().isPresent());
assertThat(newSplit.deletionFiles().get()).isEqualTo(deletionFiles);
+ assertThat(newSplit.bucket()).isEqualTo(7);
assertThat(newSplit).isEqualTo(split);
}
+
+ @Test
+ public void testBucketFromDataSplit() {
+ DataFileMeta file =
DataFileTestDataGenerator.builder().build().next().meta;
+ DataSplit dataSplit =
+ DataSplit.builder()
+ .withPartition(BinaryRow.EMPTY_ROW)
+ .withBucket(3)
+ .withBucketPath(
+
"file:/warehouse/branch/branch-snapshot/dt=20251202/bucket-3")
+ .withDataFiles(Collections.singletonList(file))
+ .build();
+ assertThat(ChainSplit.from(dataSplit,
"snapshot").bucket()).isEqualTo(3);
+ }
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/utils/ChainTableUtilsTest.java
b/paimon-core/src/test/java/org/apache/paimon/utils/ChainTableUtilsTest.java
index 9797af2a91..4f510680e0 100644
--- a/paimon-core/src/test/java/org/apache/paimon/utils/ChainTableUtilsTest.java
+++ b/paimon-core/src/test/java/org/apache/paimon/utils/ChainTableUtilsTest.java
@@ -774,6 +774,7 @@ public class ChainTableUtilsTest {
1L);
assertThat(splits).hasSize(2);
+ assertThat(splits).allSatisfy(split ->
assertThat(split.bucket()).isZero());
List<Set<String>> groups =
splits.stream().map(ChainTableUtilsTest::fileNames).collect(Collectors.toList());
@@ -875,7 +876,7 @@ public class ChainTableUtilsTest {
return DataSplit.builder()
.withPartition(partition)
.withBucket(bucket)
- .withBucketPath("bucket_" + bucket)
+ .withBucketPath("bucket-" + bucket)
.withTotalBuckets(1)
.withDataFiles(files)
.build();
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 d2c758e93e..1c6dc81315 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
@@ -22,9 +22,9 @@ import org.apache.paimon.data.{BinaryString, GenericRow,
InternalRow => PaimonIn
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.utils.{CloseableIterator, Preconditions}
+import
org.apache.paimon.spark.schema.PaimonMetadataColumn.{PATH_AND_INDEX_META_COLUMNS,
VECTOR_SEARCH_META_COLUMN_NAMES}
+import org.apache.paimon.table.source.{ChainSplit, DataSplit, Split, Splits}
+import org.apache.paimon.utils.{CloseableIterator, IOUtils, Preconditions}
import org.apache.spark.sql.PaimonUtils
@@ -37,14 +37,32 @@ case class PaimonRecordReaderIterator(
split: Split)
extends CloseableIterator[PaimonInternalRow] {
+ // Keep partition metadata restricted to the original DataSplit. Unwrapping
authorization
+ // here would expose raw partition values after the reader has applied
column masks.
if (
- metadataColumns.exists(c =>
PARTITION_AND_BUCKET_META_COLUMNS.contains(c.name)) && !split
+ metadataColumns.exists(_.name == PaimonMetadataColumn.PARTITION_COLUMN) &&
!split
.isInstanceOf[DataSplit]
) {
- throw new RuntimeException(
- "There need be DataSplit when path and index metadata columns are
required")
+ IOUtils.closeQuietly(reader)
+ throw new UnsupportedOperationException(
+ "Partition metadata is only supported for DataSplit without query
authorization.")
}
+ private lazy val bucketMetadata =
+ if (metadataColumns.exists(_.name == PaimonMetadataColumn.BUCKET_COLUMN)) {
+ // Inspect the underlying split only for metadata. The reader still
receives the original
+ // split so fallback routing and query authorization are preserved.
+ Some(Splits.underlying(split) match {
+ case dataSplit: DataSplit => dataSplit.bucket()
+ case chainSplit: ChainSplit => chainSplit.bucket()
+ case _ =>
+ throw new RuntimeException(
+ "There need be DataSplit or ChainSplit when bucket metadata is
required")
+ })
+ } else {
+ None
+ }
+
private val needMetadata = metadataColumns.nonEmpty
private val needPathAndIndexMetadata =
metadataColumns.exists(c => PATH_AND_INDEX_META_COLUMNS.contains(c.name))
@@ -173,7 +191,7 @@ case class PaimonRecordReaderIterator(
case PaimonMetadataColumn.PARTITION_COLUMN =>
metadataRow.setField(index,
split.asInstanceOf[DataSplit].partition())
case PaimonMetadataColumn.BUCKET_COLUMN =>
- metadataRow.setField(index, split.asInstanceOf[DataSplit].bucket())
+ metadataRow.setField(index, bucketMetadata.get)
}
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/util/SplitUtilsQueryAuthTest.scala
b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/util/SplitUtilsQueryAuthTest.scala
index cada2b24a2..b02f43695e 100644
---
a/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/util/SplitUtilsQueryAuthTest.scala
+++
b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/util/SplitUtilsQueryAuthTest.scala
@@ -19,14 +19,17 @@
package org.apache.paimon.spark.util
import org.apache.paimon.catalog.TableQueryAuthResult
-import org.apache.paimon.data.BinaryRow
+import org.apache.paimon.data.{BinaryRow, GenericRow, InternalRow}
import org.apache.paimon.io.DataFileMeta
+import org.apache.paimon.reader.RecordReader
import org.apache.paimon.spark.PaimonRecordReaderIterator
import org.apache.paimon.spark.schema.PaimonMetadataColumn
import org.apache.paimon.stats.SimpleStats
-import org.apache.paimon.table.source.{DataSplit, QueryAuthSplit, Split}
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplitImpl
+import org.apache.paimon.table.source.{ChainSplit, DataSplit, QueryAuthSplit,
Split}
+import org.apache.paimon.utils.IteratorRecordReader
-import org.apache.spark.sql.types.BinaryType
+import org.apache.spark.sql.types.{StringType, StructField, StructType}
import org.scalatest.funsuite.AnyFunSuite
import java.util.Collections
@@ -72,12 +75,67 @@ class SplitUtilsQueryAuthTest extends AnyFunSuite {
assert(SplitUtils.dataFileCount(wrapped) == 1L, "file count collapsed
under query auth")
}
- test("partition metadata columns are refused on a query-authorized split") {
- val partitionColumn =
- PaimonMetadataColumn(-1, PaimonMetadataColumn.PARTITION_COLUMN,
BinaryType)
- val error = intercept[RuntimeException] {
- PaimonRecordReaderIterator(null, Seq(partitionColumn),
authorized(dataSplit()))
+ private val partitionColumn =
+ PaimonMetadataColumn.PARTITION(StructType(Seq(StructField("dt",
StringType))))
+
+ private class EmptyReader extends RecordReader[InternalRow] {
+ var closed = false
+
+ override def readBatch(): RecordReader.RecordIterator[InternalRow] = null
+
+ override def close(): Unit = {
+ closed = true
+ }
+ }
+
+ test("partition metadata remains supported for an unwrapped DataSplit") {
+ val iterator =
+ PaimonRecordReaderIterator(new EmptyReader, Seq(partitionColumn),
dataSplit())
+ try {
+ assert(!iterator.hasNext)
+ } finally {
+ iterator.close()
+ }
+ }
+
+ private def chainAndAuthorizedSplits(): Seq[Split] = {
+ val chainSplit = ChainSplit.from(dataSplit(), "snapshot")
+ Seq(
+ authorized(dataSplit()),
+ chainSplit,
+ authorized(chainSplit),
+ new FallbackSplitImpl(chainSplit, true),
+ new FallbackSplitImpl(authorized(chainSplit), true))
+ }
+
+ test("partition metadata remains unsupported for authorized and chain
splits") {
+ chainAndAuthorizedSplits().foreach {
+ split =>
+ val reader = new EmptyReader
+ val error = intercept[UnsupportedOperationException] {
+ PaimonRecordReaderIterator(reader, Seq(partitionColumn), split)
+ }
+ assert(error.getMessage.contains("Partition metadata is only supported
for DataSplit"))
+ assert(reader.closed)
+ }
+ }
+
+ test("bucket metadata supports authorized and chain splits") {
+ chainAndAuthorizedSplits().foreach {
+ split =>
+ val reader = new IteratorRecordReader[InternalRow](
+
Collections.singletonList[InternalRow](GenericRow.of(Int.box(42))).iterator())
+ val iterator =
+ PaimonRecordReaderIterator(reader, Seq(PaimonMetadataColumn.BUCKET),
split)
+ try {
+ assert(iterator.hasNext)
+ val row = iterator.next()
+ assert(row.getInt(0) == 42)
+ assert(row.getInt(1) == 0)
+ assert(!iterator.hasNext)
+ } finally {
+ iterator.close()
+ }
}
- assert(error.getMessage.contains("DataSplit"))
}
}
diff --git
a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkChainTableMetadataTest.java
b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkChainTableMetadataTest.java
new file mode 100644
index 0000000000..fa4da2b923
--- /dev/null
+++
b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkChainTableMetadataTest.java
@@ -0,0 +1,272 @@
+/*
+ * 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.CatalogLoader;
+import org.apache.paimon.catalog.DelegateCatalog;
+import org.apache.paimon.catalog.TableQueryAuthResult;
+import org.apache.paimon.data.BinaryString;
+import org.apache.paimon.options.Options;
+import org.apache.paimon.predicate.Predicate;
+import org.apache.paimon.predicate.PredicateBuilder;
+import org.apache.paimon.table.CatalogEnvironment;
+import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.FileStoreTableFactory;
+import org.apache.paimon.table.source.ChainSplit;
+import org.apache.paimon.table.source.QueryAuthSplit;
+import org.apache.paimon.table.source.Split;
+import org.apache.paimon.utils.JsonSerdeUtil;
+
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
+import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
+import org.apache.spark.sql.connector.catalog.Identifier;
+import org.apache.spark.sql.connector.catalog.Table;
+import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.nio.file.Path;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests bucket metadata when Spark reads chain splits, including query
authorization. */
+public class SparkChainTableMetadataTest {
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ public void testBucketMetadata(boolean queryAuth, @TempDir Path warehouse)
throws Exception {
+ try (SparkSession spark = createSparkSession(warehouse, queryAuth)) {
+ spark.sql("CREATE DATABASE spark_test.chain_metadata");
+ spark.sql("USE spark_test.chain_metadata");
+ spark.sql(
+ "CREATE TABLE T (id INT, payload STRING, seq BIGINT, dt
STRING) "
+ + "USING paimon PARTITIONED BY (dt) TBLPROPERTIES
("
+ + "'primary-key'='dt,id', 'bucket-key'='id',
'bucket'='2', "
+ + "'sequence.field'='seq',
'chain-table.enabled'='true', "
+ + "'partition.timestamp-pattern'='$dt', "
+ + "'partition.timestamp-formatter'='yyyyMMdd')");
+ spark.sql("CALL spark_test.sys.create_branch('chain_metadata.T',
'snapshot')");
+ spark.sql("CALL spark_test.sys.create_branch('chain_metadata.T',
'delta')");
+ for (String table : Arrays.asList("T", "T$branch_snapshot",
"T$branch_delta")) {
+ spark.sql(
+ "ALTER TABLE `"
+ + table
+ + "` SET TBLPROPERTIES ("
+ + "'scan.fallback-snapshot-branch'='snapshot',
"
+ + "'scan.fallback-delta-branch'='delta')");
+ }
+ spark.sql(
+ "INSERT INTO `T$branch_snapshot` "
+ + "SELECT CAST(id AS INT), 'old', 1, '20250810'
FROM range(16)");
+ spark.sql(
+ "INSERT INTO `T$branch_delta` "
+ + "SELECT CAST(id AS INT), 'new', 2, '20250811'
FROM range(4)");
+
+ List<Row> bucketRows =
+ spark.sql("SELECT id, __paimon_bucket FROM
`T$branch_snapshot`")
+ .collectAsList();
+ assertThat(bucketRows).hasSize(16);
+ assertThat(bucketRows.stream().map(row ->
row.getInt(1)).collect(Collectors.toSet()))
+ .containsExactlyInAnyOrder(0, 1);
+
+ if (queryAuth) {
+ // Filter on a non-key field after merging: updated rows must
not reappear with
+ // their old snapshot values, even when payload is absent from
the projection.
+ for (String name : Arrays.asList("T", "T$branch_snapshot",
"T$branch_delta")) {
+ spark.sql(
+ "ALTER TABLE `"
+ + name
+ + "` SET TBLPROPERTIES
('query-auth.enabled'='true')");
+ }
+ }
+
+ // Cover both snapshot-only ChainSplit.from and snapshot + delta
chain merging.
+ for (String dt : Arrays.asList("20250810", "20250811")) {
+ if (queryAuth) {
+ assertQueryAuthChainSplits(spark, dt);
+ }
+ List<Row> visibleBucketRows =
+ bucketRows.stream()
+ .filter(
+ row ->
+ !queryAuth
+ ||
dt.equals("20250810")
+ || row.getInt(0) >= 4)
+ .collect(Collectors.toList());
+ String predicate = " WHERE dt = '" + dt + "'";
+ assertThat(
+ spark.sql("SELECT id, __paimon_bucket FROM T"
+ predicate)
+ .collectAsList())
+
.containsExactlyInAnyOrderElementsOf(visibleBucketRows);
+ // Partition metadata on chain and authorized splits is a
separate concern.
+ for (String query :
+ Arrays.asList(
+ "SELECT id, __paimon_partition,
__paimon_bucket FROM T" + predicate,
+ "SELECT id FROM T"
+ + predicate
+ + " AND __paimon_partition.dt = '"
+ + dt
+ + "'")) {
+ assertThatThrownBy(() -> spark.sql(query).collectAsList())
+ .rootCause()
+ .isInstanceOf(UnsupportedOperationException.class)
+ .hasMessageContaining(
+ "Partition metadata is only supported for
DataSplit");
+ }
+
+ for (int bucket = 0; bucket < 2; bucket++) {
+ final int selectedBucket = bucket;
+ List<Row> expected =
+ visibleBucketRows.stream()
+ .filter(row -> row.getInt(1) ==
selectedBucket)
+ .map(
+ row ->
+ RowFactory.create(
+ row.getInt(0),
+
dt.equals("20250811")
+ &&
row.getInt(0) < 4
+ ? "new"
+ : "old"))
+ .collect(Collectors.toList());
+ assertThat(
+ spark.sql(
+ "SELECT id, payload FROM T"
+ + predicate
+ + " AND
__paimon_bucket = "
+ + bucket)
+ .collectAsList())
+ .containsExactlyInAnyOrderElementsOf(expected);
+ }
+ }
+ spark.sql("DROP TABLE T");
+ spark.sql("USE spark_test.default");
+ spark.sql("DROP DATABASE spark_test.chain_metadata");
+ }
+ }
+
+ private SparkSession createSparkSession(Path warehouse, boolean queryAuth)
{
+ return SparkSession.builder()
+ .master("local[2]")
+ .config(
+ "spark.sql.catalog.spark_test",
+ (queryAuth ? RowFilterSparkCatalog.class :
SparkCatalog.class).getName())
+ .config("spark.sql.catalog.spark_test.warehouse",
warehouse.toUri().toString())
+ .config("spark.sql.defaultCatalog", "spark_test")
+ .config(
+ "spark.sql.extensions",
+
"org.apache.paimon.spark.extensions.PaimonSparkSessionExtensions")
+ .getOrCreate();
+ }
+
+ private FileStoreTable loadTable(SparkSession spark) throws Exception {
+ SparkCatalog catalog =
+ (SparkCatalog)
spark.sessionState().catalogManager().currentCatalog();
+ SparkTable table =
+ (SparkTable) catalog.loadTable(Identifier.of(new String[]
{"chain_metadata"}, "T"));
+ return (FileStoreTable) table.table();
+ }
+
+ private void assertQueryAuthChainSplits(SparkSession spark, String dt)
throws Exception {
+ List<Split> splits =
+ loadTable(spark)
+ .newReadBuilder()
+ .withPartitionFilter(Collections.singletonMap("dt",
dt))
+ .newScan()
+ .plan()
+ .splits();
+ assertThat(splits)
+ .isNotEmpty()
+ .allSatisfy(
+ split -> {
+
assertThat(split).isInstanceOf(FallbackSplit.class);
+ Split wrapped = ((FallbackSplit) split).wrapped();
+
assertThat(wrapped).isInstanceOf(QueryAuthSplit.class);
+ QueryAuthSplit authSplit = (QueryAuthSplit)
wrapped;
+ assertThat(authSplit.authResult()).isNotNull();
+
assertThat(authSplit.authResult().hasRules()).isTrue();
+
assertThat(authSplit.split()).isInstanceOf(ChainSplit.class);
+ });
+ }
+
+ /** Supplies a row-filter policy while keeping the real filesystem catalog
and table reads. */
+ public static class RowFilterSparkCatalog extends SparkCatalog {
+
+ @Override
+ protected Table loadSparkTable(Identifier identifier, Map<String,
String> extraOptions)
+ throws NoSuchTableException {
+ SparkTable sparkTable = (SparkTable)
super.loadSparkTable(identifier, extraOptions);
+ FileStoreTable table = (FileStoreTable) sparkTable.table();
+ return table.coreOptions().queryAuthEnabled()
+ ? SparkTable$.MODULE$.of(withRowFilter(table))
+ : sparkTable;
+ }
+ }
+
+ private static FileStoreTable withRowFilter(FileStoreTable table) {
+ Predicate rowFilter =
+ new PredicateBuilder(table.rowType()).equal(1,
BinaryString.fromString("old"));
+ TableQueryAuthResult authResult =
+ new TableQueryAuthResult(
+
Collections.singletonList(JsonSerdeUtil.toFlatJson(rowFilter)), null);
+ CatalogEnvironment environment = table.catalogEnvironment();
+ CatalogLoader loader = environment.catalogLoader();
+ CatalogEnvironment authorizedEnvironment =
+ new CatalogEnvironment(
+ environment.identifier(),
+ environment.uuid(),
+ authCatalogLoader(loader, authResult),
+ environment.lockFactory(),
+ environment.lockContext(),
+ environment.catalogContext(),
+ false,
+ false);
+ return FileStoreTableFactory.create(
+ table.fileIO(),
+ table.location(),
+ table.schema(),
+ new Options(table.options()),
+ authorizedEnvironment);
+ }
+
+ private static CatalogLoader authCatalogLoader(
+ CatalogLoader loader, TableQueryAuthResult authResult) {
+ return () ->
+ new DelegateCatalog(loader.load()) {
+ @Override
+ public TableQueryAuthResult authTableQuery(
+ org.apache.paimon.catalog.Identifier identifier,
List<String> select) {
+ return authResult;
+ }
+
+ @Override
+ public CatalogLoader catalogLoader() {
+ return authCatalogLoader(loader, authResult);
+ }
+ };
+ }
+}
