From ac6f71e93056e5160ebb23b9078aa0f4b3b3ee2d Mon Sep 17 00:00:00 2001 From: LsomeYeah Date: Tue, 4 Aug 2026 15:55:18 +0800 Subject: [PATCH 1/4] [spark] Expose written columns for streaming micro-batches --- .../paimon/table/source/AllColumns.java | 27 ++++ .../table/source/KnownWrittenColumns.java | 67 ++++++++ .../paimon/table/source/WrittenColumns.java | 33 ++++ .../paimon/utils/DataEvolutionUtils.java | 45 ++++++ .../paimon/utils/DataEvolutionUtilsTest.java | 132 ++++++++++++++++ .../sql/paimon/shims/MinorVersionShim.scala | 8 +- .../sql/paimon/shims/MinorVersionShim.scala | 8 +- .../paimon/spark/SparkConnectorOptions.java | 8 + .../paimon/spark/PaimonInputPartition.scala | 15 +- .../spark/PaimonSparkMicroBatchMetadata.scala | 94 ++++++++++++ .../sources/PaimonMicroBatchStream.scala | 61 +++++++- .../spark/sql/paimon/shims/SparkShim.scala | 6 +- .../paimon/spark/PaimonSourceTest.scala | 143 +++++++++++++++++- .../sql/paimon/shims/MinorVersionShim.scala | 8 +- .../spark/sql/paimon/shims/Spark3Shim.scala | 6 +- .../spark/sql/paimon/shims/Spark4Shim.scala | 8 +- 16 files changed, 652 insertions(+), 17 deletions(-) create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/AllColumns.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/KnownWrittenColumns.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/source/WrittenColumns.java create mode 100644 paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/AllColumns.java b/paimon-core/src/main/java/org/apache/paimon/table/source/AllColumns.java new file mode 100644 index 000000000000..e22a2171888c --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/AllColumns.java @@ -0,0 +1,27 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.table.source; + +import org.apache.paimon.annotation.Experimental; + +/** A conservative marker requiring consumers to assume that every column may be written. */ +@Experimental +public enum AllColumns implements WrittenColumns { + INSTANCE +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/KnownWrittenColumns.java b/paimon-core/src/main/java/org/apache/paimon/table/source/KnownWrittenColumns.java new file mode 100644 index 000000000000..514c72730e7a --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/KnownWrittenColumns.java @@ -0,0 +1,67 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.table.source; + +import org.apache.paimon.annotation.Experimental; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.TreeSet; + +/** A complete, immutable set of written field ids, ordered by field id. */ +@Experimental +public final class KnownWrittenColumns implements WrittenColumns { + + private static final long serialVersionUID = 1L; + + private final List fieldIds; + + public KnownWrittenColumns(Collection fieldIds) { + this.fieldIds = Collections.unmodifiableList(new ArrayList<>(new TreeSet<>(fieldIds))); + } + + public List fieldIds() { + return fieldIds; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof KnownWrittenColumns)) { + return false; + } + KnownWrittenColumns that = (KnownWrittenColumns) o; + return fieldIds.equals(that.fieldIds); + } + + @Override + public int hashCode() { + return Objects.hash(fieldIds); + } + + @Override + public String toString() { + return "KnownWrittenColumns{" + "fieldIds=" + fieldIds + '}'; + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/WrittenColumns.java b/paimon-core/src/main/java/org/apache/paimon/table/source/WrittenColumns.java new file mode 100644 index 000000000000..ca1070270db3 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/WrittenColumns.java @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.table.source; + +import org.apache.paimon.annotation.Experimental; + +import java.io.Serializable; + +/** + * Columns written by the data files selected for a scan. + * + *

The result is either {@link KnownWrittenColumns} or {@link AllColumns}. Consumers must treat + * {@link AllColumns} conservatively and must not interpret an empty {@link KnownWrittenColumns} as + * unknown. + */ +@Experimental +public interface WrittenColumns extends Serializable {} diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java b/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java index c1294f39462f..bc72672ccef2 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java @@ -20,13 +20,20 @@ import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.source.AllColumns; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.KnownWrittenColumns; +import org.apache.paimon.table.source.WrittenColumns; import org.apache.paimon.types.DataField; import java.util.Collection; import java.util.Comparator; +import java.util.HashMap; import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Set; +import java.util.TreeSet; import java.util.function.Function; import java.util.stream.Collectors; @@ -38,6 +45,44 @@ /** Util class for data evolution. */ public class DataEvolutionUtils { + /** Collect written field ids from data files in the selected splits. */ + public static WrittenColumns collectWrittenColumns( + Collection splits, Function schemaLoader) { + Set fieldIds = new TreeSet<>(); + Map>, Set> fieldIdsCache = new HashMap<>(); + for (DataSplit split : splits) { + for (DataFileMeta file : split.dataFiles()) { + try { + Pair> cacheKey = Pair.of(file.schemaId(), file.writeCols()); + Set fileFieldIds = fieldIdsCache.get(cacheKey); + if (fileFieldIds == null) { + fileFieldIds = computeFileFieldIds(schemaLoader, file); + fieldIdsCache.put(cacheKey, fileFieldIds); + } + fieldIds.addAll(fileFieldIds); + } catch (RuntimeException e) { + return AllColumns.INSTANCE; + } + } + } + return new KnownWrittenColumns(fieldIds); + } + + /** Resolve a data file's physical columns through the schema the file was written with. */ + public static Set computeFileFieldIds( + Function schemaLoader, DataFileMeta file) { + TableSchema fileSchema = schemaLoader.apply(file.schemaId()); + if (fileSchema == null) { + throw new IllegalArgumentException("Cannot find schema " + file.schemaId()); + } + + Set fieldIds = new TreeSet<>(); + for (DataField field : fileSchema.project(file.writeCols()).fields()) { + fieldIds.add(field.id()); + } + return fieldIds; + } + /** * Table field ids physically present in a file, resolved through the schema used to write it. */ diff --git a/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java b/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java index 33feb9d850e1..beb5ffd935ec 100644 --- a/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java @@ -18,11 +18,18 @@ package org.apache.paimon.utils; +import org.apache.paimon.data.BinaryRow; import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.schema.Schema; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.stats.SimpleStats; import org.apache.paimon.table.SpecialFields; +import org.apache.paimon.table.source.AllColumns; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.KnownWrittenColumns; +import org.apache.paimon.table.source.WrittenColumns; import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.IntType; import org.junit.jupiter.api.Test; @@ -31,10 +38,14 @@ import java.util.Collections; import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; /** Test for {@link DataEvolutionUtils}. */ public class DataEvolutionUtilsTest { @@ -109,6 +120,99 @@ public void testFileFieldIdsHandlesFullEmptyAndUnrelatedWrites() { .containsExactly(2); } + @Test + public void testCollectWrittenColumnsByFieldIdAcrossSchemas() { + Map schemas = new HashMap<>(); + schemas.put( + 0L, + tableSchema( + 0L, + new DataField(1, "a", DataTypes.INT()), + new DataField(2, "old_name", DataTypes.STRING()))); + schemas.put( + 1L, + tableSchema( + 1L, + new DataField(2, "new_name", DataTypes.STRING()), + new DataField(3, "c", DataTypes.BIGINT()))); + + DataFileMeta oldSchemaFile = dataFile(0L, Arrays.asList("a", "old_name")); + DataFileMeta newSchemaFile = dataFile(1L, Arrays.asList("new_name", "c")); + + WrittenColumns result = + DataEvolutionUtils.collectWrittenColumns( + Collections.singletonList(dataSplit(oldSchemaFile, newSchemaFile)), + schemas::get); + + assertThat(result).isInstanceOf(KnownWrittenColumns.class); + assertThat(((KnownWrittenColumns) result).fieldIds()).containsExactly(1, 2, 3); + } + + @Test + public void testCollectWrittenColumnsFallsBackWhenSchemaIsUnknown() { + DataFileMeta unknownSchemaFile = dataFile(99L, Collections.singletonList("a")); + + WrittenColumns result = + DataEvolutionUtils.collectWrittenColumns( + Collections.singletonList(dataSplit(unknownSchemaFile)), ignored -> null); + + assertThat(result).isSameAs(AllColumns.INSTANCE); + } + + @Test + public void testCollectWrittenColumnsFallsBackWhenSchemaResolutionFails() { + DataFileMeta file = dataFile(1L, Collections.singletonList("missing")); + + WrittenColumns result = + DataEvolutionUtils.collectWrittenColumns( + Collections.singletonList(dataSplit(file)), + ignored -> { + throw new IllegalArgumentException("schema cannot be resolved"); + }); + + assertThat(result).isSameAs(AllColumns.INSTANCE); + } + + @Test + public void testCollectWrittenColumnsCachesFileSchemaProjection() { + TableSchema schema = + tableSchema( + 1L, + new DataField(1, "a", DataTypes.INT()), + new DataField(2, "b", DataTypes.STRING())); + DataFileMeta first = dataFile(1L, Collections.singletonList("b")); + DataFileMeta second = dataFile(1L, Collections.singletonList("b")); + AtomicInteger schemaLoads = new AtomicInteger(); + + WrittenColumns result = + DataEvolutionUtils.collectWrittenColumns( + Collections.singletonList(dataSplit(first, second)), + ignored -> { + schemaLoads.incrementAndGet(); + return schema; + }); + + assertThat(((KnownWrittenColumns) result).fieldIds()).containsExactly(2); + assertThat(schemaLoads).hasValue(1); + } + + @Test + public void testCollectWrittenColumnsExpandsLegacyFileSchema() { + TableSchema schema = + tableSchema( + 1L, + new DataField(1, "a", DataTypes.INT()), + new DataField(2, "b", DataTypes.STRING())); + DataFileMeta legacyFile = dataFile(1L, null); + + WrittenColumns result = + DataEvolutionUtils.collectWrittenColumns( + Collections.singletonList(dataSplit(legacyFile)), ignored -> schema); + + assertThat(result).isInstanceOf(KnownWrittenColumns.class); + assertThat(((KnownWrittenColumns) result).fieldIds()).containsExactly(1, 2); + } + @Test public void testRetrieveAnchorFileSkipsSpecialFiles() { DataFileMeta blobFile = dataFile("blob-file.blob", 1); @@ -171,4 +275,32 @@ private static DataFileMeta dataFile( 0L, writeCols); } + + private static DataFileMeta dataFile(long schemaId, java.util.List writeCols) { + DataFileMeta file = mock(DataFileMeta.class); + when(file.schemaId()).thenReturn(schemaId); + when(file.writeCols()).thenReturn(writeCols); + return file; + } + + private static DataSplit dataSplit(DataFileMeta... files) { + return DataSplit.builder() + .withSnapshot(1L) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withDataFiles(Arrays.asList(files)) + .build(); + } + + private static TableSchema tableSchema(long id, DataField... fields) { + return TableSchema.create( + id, + new Schema( + Arrays.asList(fields), + Collections.emptyList(), + Collections.emptyList(), + Collections.emptyMap(), + null)); + } } diff --git a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala index ffcdee388c01..3cf8c41f2ecd 100644 --- a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala +++ b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala @@ -18,14 +18,18 @@ package org.apache.spark.sql.paimon.shims +import org.apache.spark.Partition import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression} import org.apache.spark.sql.catalyst.plans.logical.{CTERelationRef, LogicalPlan, MergeAction, MergeIntoTable} import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution} -import org.apache.spark.sql.connector.read.Scan -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.connector.read.{InputPartition, Scan} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceRDDPartition, DataSourceV2ScanRelation} object MinorVersionShim { + def dataSourceInputPartitions(partition: Partition): Seq[InputPartition] = + Seq(partition.asInstanceOf[DataSourceRDDPartition].inputPartition) + def createCTERelationRef( cteId: Long, resolved: Boolean, diff --git a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala index b66081d06567..aec98a1df5e4 100644 --- a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala +++ b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala @@ -18,14 +18,18 @@ package org.apache.spark.sql.paimon.shims +import org.apache.spark.Partition import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression} import org.apache.spark.sql.catalyst.plans.logical.{CTERelationRef, LogicalPlan, MergeAction, MergeIntoTable} import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution} -import org.apache.spark.sql.connector.read.Scan -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.connector.read.{InputPartition, Scan} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceRDDPartition, DataSourceV2ScanRelation} object MinorVersionShim { + def dataSourceInputPartitions(partition: Partition): Seq[InputPartition] = + partition.asInstanceOf[DataSourceRDDPartition].inputPartitions + def createCTERelationRef( cteId: Long, resolved: Boolean, diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java index 2f315b8df0f5..b60f04c6f29a 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java @@ -131,6 +131,14 @@ public class SparkConnectorOptions { .withDescription( "The maximum delay between two adjacent batches, which used to create MinRowsReadLimit with read.stream.minRowsPerTrigger together."); + public static final ConfigOption BATCH_WRITTEN_COLUMNS_ENABLED = + key("read.stream.batch-written-columns.enabled") + .booleanType() + .defaultValue(false) + .withDescription( + "Whether to expose the written field ids of an admitted micro-batch " + + "through PaimonSparkMicroBatchMetadata."); + public static final ConfigOption READ_CHANGELOG = key("read.changelog") .booleanType() diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala index 7e3dbf893b22..500acab8cf73 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala @@ -18,7 +18,7 @@ package org.apache.paimon.spark -import org.apache.paimon.table.source.Split +import org.apache.paimon.table.source.{Split, WrittenColumns} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.GenericInternalRow @@ -36,6 +36,19 @@ trait PaimonInputPartition extends InputPartition { } case class SimplePaimonInputPartition(splits: Seq[Split]) extends PaimonInputPartition + +private[spark] case class PaimonMicroBatchMetadata( + sourceId: String, + startOffset: String, + endOffset: String, + splitCount: Int, + writtenColumns: WrittenColumns) + +private[spark] case class PaimonMicroBatchInputPartition( + splits: Seq[Split], + metadata: PaimonMicroBatchMetadata) + extends PaimonInputPartition + object PaimonInputPartition { def apply(split: Split): PaimonInputPartition = { SimplePaimonInputPartition(Seq(split)) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala new file mode 100644 index 000000000000..5a952ed00baa --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala @@ -0,0 +1,94 @@ +/* + * 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.annotation.Experimental +import org.apache.paimon.table.source.WrittenColumns + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.Dataset +import org.apache.spark.sql.execution.datasources.v2.DataSourceRDD +import org.apache.spark.sql.paimon.shims.SparkShimLoader + +import java.util.{IdentityHashMap, Optional} + +import scala.collection.mutable +import scala.util.control.NonFatal + +/** Driver-side access to metadata planned for a Paimon streaming micro-batch. */ +@Experimental +final class PaimonSparkMicroBatchMetadata private () + +object PaimonSparkMicroBatchMetadata { + + /** + * Returns written columns for a raw foreachBatch Dataset with exactly one Paimon streaming + * source. This method only inspects driver-side RDD planning metadata and does not run a Spark + * job. The result is empty when metadata collection was not enabled, the Dataset is not backed by + * a Paimon source, the lineage is incomplete, or multiple Paimon sources make the result + * ambiguous. + */ + def writtenColumns(batch: Dataset[_]): Optional[WrittenColumns] = { + try { + extractWrittenColumns(batch) + } catch { + case NonFatal(_) => Optional.empty() + case _: LinkageError => Optional.empty() + } + } + + private def extractWrittenColumns(batch: Dataset[_]): Optional[WrittenColumns] = { + val visited = new IdentityHashMap[RDD[_], java.lang.Boolean]() + val metadata = mutable.ArrayBuffer.empty[PaimonMicroBatchMetadata] + var incompletePaimonSource = false + + def visit(rdd: RDD[_]): Unit = { + if (!visited.containsKey(rdd)) { + visited.put(rdd, java.lang.Boolean.TRUE) + rdd match { + case dataSourceRDD: DataSourceRDD => + dataSourceRDD.partitions.foreach { + partition => + SparkShimLoader.shim.dataSourceInputPartitions(partition).foreach { + case input: PaimonMicroBatchInputPartition => metadata += input.metadata + case _: PaimonInputPartition => incompletePaimonSource = true + case _ => + } + } + case _ => + } + rdd.dependencies.foreach(dependency => visit(dependency.rdd)) + } + } + + visit(batch.queryExecution.toRdd) + + val distinct = metadata.distinct + if (incompletePaimonSource || distinct.size != 1) { + Optional.empty() + } else { + val only = distinct.head + if (metadata.size != only.splitCount) { + Optional.empty() + } else { + Optional.of(only.writtenColumns) + } + } + } +} diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala index c3d2dfc8812d..7d5fe86b1fd5 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala @@ -20,15 +20,26 @@ package org.apache.paimon.spark.sources import org.apache.paimon.CoreOptions import org.apache.paimon.options.Options -import org.apache.paimon.spark.{PaimonImplicits, PaimonInputPartition, PaimonPartitionReaderFactory, SparkConnectorOptions} +import org.apache.paimon.schema.TableSchema +import org.apache.paimon.spark.{PaimonImplicits, PaimonInputPartition, PaimonMicroBatchInputPartition, PaimonMicroBatchMetadata, PaimonPartitionReaderFactory, SparkConnectorOptions} import org.apache.paimon.table.DataTable -import org.apache.paimon.table.source.ReadBuilder +import org.apache.paimon.table.source.{AllColumns, ReadBuilder} +import org.apache.paimon.utils.DataEvolutionUtils import org.apache.spark.internal.Logging import org.apache.spark.sql.connector.read.{InputPartition, PartitionReaderFactory} import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset, ReadLimit, SupportsTriggerAvailableNow} +import java.lang.{Long => JLong} +import java.util.function.Function + +import scala.collection.JavaConverters._ import scala.collection.mutable +import scala.util.control.NonFatal + +private[spark] case class PlannedMicroBatch( + admittedSplits: Array[IndexedDataSplit], + metadata: PaimonMicroBatchMetadata) class PaimonMicroBatchStream( originTable: DataTable, @@ -93,6 +104,9 @@ class PaimonMicroBatchStream( private lazy val blobAsDescriptor: Boolean = options.get(CoreOptions.BLOB_AS_DESCRIPTOR) + private lazy val batchWrittenColumnsEnabled: Boolean = + options.get(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED) + override def getDefaultReadLimit: ReadLimit = defaultReadLimit override def prepareForTriggerAvailableNow(): Unit = { @@ -134,9 +148,46 @@ class PaimonMicroBatchStream( } val endOffset = PaimonSourceOffset(end) - getBatch(startOffset, Some(endOffset), None) - .map(ids => PaimonInputPartition(ids.entry)) - .toArray[InputPartition] + val admittedSplits = getBatch(startOffset, Some(endOffset), None) + if (!batchWrittenColumnsEnabled) { + admittedSplits + .map(ids => PaimonInputPartition(ids.entry)) + .toArray[InputPartition] + } else { + val plannedBatch = createPlannedMicroBatch(startOffset, endOffset, admittedSplits) + plannedBatch.admittedSplits + .map(ids => PaimonMicroBatchInputPartition(Seq(ids.entry), plannedBatch.metadata)) + .toArray[InputPartition] + } + } + + private def createPlannedMicroBatch( + startOffset: PaimonSourceOffset, + endOffset: PaimonSourceOffset, + admittedSplits: Array[IndexedDataSplit]): PlannedMicroBatch = { + val writtenColumns = + try { + val schemaManager = table.schemaManager() + DataEvolutionUtils.collectWrittenColumns( + admittedSplits.map(_.entry).toSeq.asJava, + new Function[JLong, TableSchema] { + override def apply(schemaId: JLong): TableSchema = + schemaManager.schema(schemaId.longValue()) + } + ) + } catch { + case NonFatal(e) => + logWarning("Failed to collect written columns for a micro-batch; using all columns.", e) + AllColumns.INSTANCE + } + + val metadata = PaimonMicroBatchMetadata( + checkpointLocation, + startOffset.json(), + endOffset.json(), + admittedSplits.length, + writtenColumns) + PlannedMicroBatch(admittedSplits, metadata) } override def createReaderFactory(): PartitionReaderFactory = { diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/paimon/shims/SparkShim.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/paimon/shims/SparkShim.scala index df21168bd343..bf2eb27b49ea 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/paimon/shims/SparkShim.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/paimon/shims/SparkShim.scala @@ -26,6 +26,7 @@ import org.apache.paimon.spark.rowops.PaimonCopyOnWriteScan import org.apache.paimon.table.{FileStoreTable, FormatTable} import org.apache.paimon.types.{DataType, RowType} +import org.apache.spark.Partition import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.FunctionIdentifier import org.apache.spark.sql.catalyst.InternalRow @@ -38,7 +39,7 @@ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.ArrayData import org.apache.spark.sql.connector.catalog.{Column, Identifier, StagingTableCatalog, Table, TableCatalog} import org.apache.spark.sql.connector.expressions.Transform -import org.apache.spark.sql.connector.read.Scan +import org.apache.spark.sql.connector.read.{InputPartition, Scan} import org.apache.spark.sql.connector.write.BatchWrite import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation} @@ -55,6 +56,9 @@ trait SparkShim { def classicApi: ClassicApi + /** Returns the data source input partitions represented by a Spark RDD partition. */ + def dataSourceInputPartitions(partition: Partition): Seq[InputPartition] + def createSparkParser(delegate: ParserInterface): ParserInterface def createCustomResolution(spark: SparkSession): Rule[LogicalPlan] diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala index e8b685664c94..57f975af4be2 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala @@ -19,13 +19,16 @@ package org.apache.paimon.spark import org.apache.paimon.spark.sources.PaimonSourceOffset +import org.apache.paimon.table.source.{KnownWrittenColumns, WrittenColumns} -import org.apache.spark.sql.Row +import org.apache.spark.sql.{Dataset, Row} import org.apache.spark.sql.streaming.{StreamingQueryException, StreamTest, Trigger} import org.junit.jupiter.api.Assertions import java.util.concurrent.TimeUnit +import scala.collection.JavaConverters._ + class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { import testImplicits._ @@ -48,6 +51,144 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { } } + test("Paimon Source: expose written columns to raw foreachBatch") { + withTempDir { + checkpointDir => + val TableSnapshotState(_, location, _, _, _) = + prepareTableAndGetLocation(1, hasPk = true) + val expectedFieldIds = + loadTable("T").schema().fields().asScala.map(field => Integer.valueOf(field.id())).sorted + @volatile var writtenColumns: WrittenColumns = null + @volatile var metadataLookupStartedNoSparkJob = false + + val query = spark.readStream + .format("paimon") + .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) + .load(location) + .select("a") + .writeStream + .option("checkpointLocation", checkpointDir.getCanonicalPath) + .foreachBatch { + (batch: Dataset[Row], _: Long) => + val jobGroup = s"written-columns-metadata-${System.nanoTime()}" + val previousJobGroup = spark.sparkContext.getLocalProperty("spark.jobGroup.id") + spark.sparkContext.setLocalProperty("spark.jobGroup.id", jobGroup) + val metadata = + try { + PaimonSparkMicroBatchMetadata.writtenColumns(batch) + } finally { + metadataLookupStartedNoSparkJob = + spark.sparkContext.statusTracker.getJobIdsForGroup(jobGroup).isEmpty + spark.sparkContext.setLocalProperty("spark.jobGroup.id", previousJobGroup) + } + if (metadata.isPresent) { + writtenColumns = metadata.get() + } + batch.count() + () + } + .start() + + try { + query.processAllAvailable() + assert(writtenColumns.isInstanceOf[KnownWrittenColumns]) + assert( + writtenColumns.asInstanceOf[KnownWrittenColumns].fieldIds() == expectedFieldIds.asJava) + assert(metadataLookupStartedNoSparkJob) + } finally { + query.stop() + } + } + } + + test("Paimon Source: written columns metadata is disabled by default") { + withTempDir { + checkpointDir => + val TableSnapshotState(_, location, snapshotData, _, _) = + prepareTableAndGetLocation(1, hasPk = true) + @volatile var metadataAvailable = false + @volatile var rowCount = 0L + + val query = spark.readStream + .format("paimon") + .load(location) + .writeStream + .option("checkpointLocation", checkpointDir.getCanonicalPath) + .foreachBatch { + (batch: Dataset[Row], _: Long) => + metadataAvailable = PaimonSparkMicroBatchMetadata.writtenColumns(batch).isPresent + rowCount += batch.count() + () + } + .start() + + try { + query.processAllAvailable() + assert(!metadataAvailable) + assert(rowCount == snapshotData.size) + } finally { + query.stop() + } + } + } + + test("Paimon Source: expose partial data evolution written columns") { + withSparkSQLConf("spark.paimon.write.use-v2-write" -> "false") { + withTable("T") { + withTempDir { + checkpointDir => + spark.sql( + "CREATE TABLE T (id INT, b INT, c INT) " + + "TBLPROPERTIES ('row-tracking.enabled' = 'true', " + + "'data-evolution.enabled' = 'true')") + spark.sql("INSERT INTO T VALUES (1, 10, 100), (2, 20, 200)") + val fieldIds = + loadTable("T") + .schema() + .fields() + .asScala + .map(field => field.name() -> field.id()) + .toMap + @volatile var nonEmptyBatchColumns = Seq.empty[WrittenColumns] + + val query = spark.readStream + .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) + .option(SparkConnectorOptions.MAX_FILES_PER_TRIGGER.key(), 1) + .option("scan.mode", "latest") + .table("`T$row_tracking`") + .writeStream + .option("checkpointLocation", checkpointDir.getCanonicalPath) + .foreachBatch { + (batch: Dataset[Row], _: Long) => + val metadata = PaimonSparkMicroBatchMetadata.writtenColumns(batch) + if (batch.count() > 0 && metadata.isPresent) { + nonEmptyBatchColumns = nonEmptyBatchColumns :+ metadata.get() + } + () + } + .start() + + try { + query.processAllAvailable() + spark.sql("UPDATE T SET b = 22 WHERE id = 2") + spark.sql("UPDATE T SET c = NULL WHERE id = 1") + query.processAllAvailable() + + assert(nonEmptyBatchColumns.size >= 2) + assert(nonEmptyBatchColumns.forall(_.isInstanceOf[KnownWrittenColumns])) + val partialBatchColumns = nonEmptyBatchColumns.takeRight(2) + assert( + partialBatchColumns.map(_.asInstanceOf[KnownWrittenColumns].fieldIds()) == Seq( + Seq(Integer.valueOf(fieldIds("b"))).asJava, + Seq(Integer.valueOf(fieldIds("c"))).asJava)) + } finally { + query.stop() + } + } + } + } + } + test("Paimon Source: default scan mode") { withTempDir { checkpointDir => diff --git a/paimon-spark/paimon-spark3-common/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala b/paimon-spark/paimon-spark3-common/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala index 6c9e991d9904..33ff17bdf0ad 100644 --- a/paimon-spark/paimon-spark3-common/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala +++ b/paimon-spark/paimon-spark3-common/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala @@ -18,14 +18,18 @@ package org.apache.spark.sql.paimon.shims +import org.apache.spark.Partition import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression} import org.apache.spark.sql.catalyst.plans.logical.{CTERelationRef, LogicalPlan, MergeAction, MergeIntoTable} import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution} -import org.apache.spark.sql.connector.read.Scan -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.connector.read.{InputPartition, Scan} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceRDDPartition, DataSourceV2ScanRelation} object MinorVersionShim { + def dataSourceInputPartitions(partition: Partition): Seq[InputPartition] = + partition.asInstanceOf[DataSourceRDDPartition].inputPartitions + def createCTERelationRef( cteId: Long, resolved: Boolean, diff --git a/paimon-spark/paimon-spark3-common/src/main/scala/org/apache/spark/sql/paimon/shims/Spark3Shim.scala b/paimon-spark/paimon-spark3-common/src/main/scala/org/apache/spark/sql/paimon/shims/Spark3Shim.scala index 5ffbf6a14530..06a322de3381 100644 --- a/paimon-spark/paimon-spark3-common/src/main/scala/org/apache/spark/sql/paimon/shims/Spark3Shim.scala +++ b/paimon-spark/paimon-spark3-common/src/main/scala/org/apache/spark/sql/paimon/shims/Spark3Shim.scala @@ -31,6 +31,7 @@ import org.apache.paimon.types.{DataType, RowType} import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path +import org.apache.spark.Partition import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.{CTESubstitution, SubstituteUnresolvedOrdinals} @@ -48,7 +49,7 @@ import org.apache.spark.sql.catalyst.util.{ArrayData, GeneratedColumn, ResolveDe import org.apache.spark.sql.connector.catalog.{Column, Identifier, StagingTableCatalog, Table, TableCatalog} import org.apache.spark.sql.connector.catalog.CatalogV2Util.structTypeToV2Columns import org.apache.spark.sql.connector.expressions.Transform -import org.apache.spark.sql.connector.read.Scan +import org.apache.spark.sql.connector.read.{InputPartition, Scan} import org.apache.spark.sql.connector.write.BatchWrite import org.apache.spark.sql.execution.{SparkFormatTable, SparkPlan} import org.apache.spark.sql.execution.datasources.{PartitioningAwareFileIndex, PartitionSpec} @@ -64,6 +65,9 @@ class Spark3Shim extends SparkShim { override def classicApi: ClassicApi = new Classic3Api + override def dataSourceInputPartitions(partition: Partition): Seq[InputPartition] = + MinorVersionShim.dataSourceInputPartitions(partition) + override def createSparkParser(delegate: ParserInterface): ParserInterface = { new PaimonSpark3SqlExtensionsParser(delegate) } diff --git a/paimon-spark/paimon-spark4-common/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala b/paimon-spark/paimon-spark4-common/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala index 00e0b1ae4ff0..9ef987f568bf 100644 --- a/paimon-spark/paimon-spark4-common/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala +++ b/paimon-spark/paimon-spark4-common/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala @@ -31,6 +31,7 @@ import org.apache.paimon.types.{DataType, RowType} import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path +import org.apache.spark.Partition import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.CTESubstitution @@ -44,12 +45,12 @@ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.{ArrayData, GeneratedColumn, IdentityColumn, ResolveDefaultColumns} import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Column, Identifier, StagingTableCatalog, Table, TableCatalog} import org.apache.spark.sql.connector.expressions.Transform -import org.apache.spark.sql.connector.read.Scan +import org.apache.spark.sql.connector.read.{InputPartition, Scan} import org.apache.spark.sql.connector.write.BatchWrite import org.apache.spark.sql.execution.{SparkFormatTable, SparkPlan} import org.apache.spark.sql.execution.datasources.{PartitioningAwareFileIndex, PartitionSpec} import org.apache.spark.sql.execution.datasources.v2.{AtomicReplaceTableAsSelectExec, AtomicReplaceTableExec, ReplaceTableAsSelectExec, ReplaceTableExec} -import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceRDDPartition, DataSourceV2Relation, DataSourceV2ScanRelation} import org.apache.spark.sql.execution.streaming.runtime.MetadataLogFileIndex import org.apache.spark.sql.execution.streaming.sinks.FileStreamSink import org.apache.spark.sql.internal.SQLConf @@ -62,6 +63,9 @@ class Spark4Shim extends SparkShim { override def classicApi: ClassicApi = new Classic4Api + override def dataSourceInputPartitions(partition: Partition): Seq[InputPartition] = + partition.asInstanceOf[DataSourceRDDPartition].inputPartitions + override def createSparkParser(delegate: ParserInterface): ParserInterface = { new PaimonSpark4SqlExtensionsParser(delegate) } From c8d20fdc6f3d61cb43c5f759aaa73120676e1873 Mon Sep 17 00:00:00 2001 From: LsomeYeah Date: Tue, 4 Aug 2026 17:28:44 +0800 Subject: [PATCH 2/4] [spark] Address micro-batch metadata review feedback --- .../paimon/utils/DataEvolutionUtils.java | 5 +- .../paimon/utils/DataEvolutionUtilsTest.java | 23 +++++++ .../spark/sql/paimon/shims/Spark4Shim.scala | 8 ++- .../paimon/spark/PaimonInputPartition.scala | 2 +- .../spark/PaimonSparkMicroBatchMetadata.scala | 43 ++++++++++++- .../paimon/spark/PaimonSourceTest.scala | 64 +++++++++++++++++++ 6 files changed, 140 insertions(+), 5 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java b/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java index bc72672ccef2..5ada67ab0390 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java @@ -49,6 +49,9 @@ public class DataEvolutionUtils { public static WrittenColumns collectWrittenColumns( Collection splits, Function schemaLoader) { Set fieldIds = new TreeSet<>(); + Map schemaCache = new HashMap<>(); + Function cachedSchemaLoader = + schemaId -> schemaCache.computeIfAbsent(schemaId, schemaLoader); Map>, Set> fieldIdsCache = new HashMap<>(); for (DataSplit split : splits) { for (DataFileMeta file : split.dataFiles()) { @@ -56,7 +59,7 @@ public static WrittenColumns collectWrittenColumns( Pair> cacheKey = Pair.of(file.schemaId(), file.writeCols()); Set fileFieldIds = fieldIdsCache.get(cacheKey); if (fileFieldIds == null) { - fileFieldIds = computeFileFieldIds(schemaLoader, file); + fileFieldIds = computeFileFieldIds(cachedSchemaLoader, file); fieldIdsCache.put(cacheKey, fileFieldIds); } fieldIds.addAll(fileFieldIds); diff --git a/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java b/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java index beb5ffd935ec..9d2b2fb86f60 100644 --- a/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java @@ -196,6 +196,29 @@ public void testCollectWrittenColumnsCachesFileSchemaProjection() { assertThat(schemaLoads).hasValue(1); } + @Test + public void testCollectWrittenColumnsCachesSchemaAcrossProjections() { + TableSchema schema = + tableSchema( + 1L, + new DataField(1, "a", DataTypes.INT()), + new DataField(2, "b", DataTypes.STRING())); + DataFileMeta first = dataFile(1L, Collections.singletonList("a")); + DataFileMeta second = dataFile(1L, Collections.singletonList("b")); + AtomicInteger schemaLoads = new AtomicInteger(); + + WrittenColumns result = + DataEvolutionUtils.collectWrittenColumns( + Collections.singletonList(dataSplit(first, second)), + ignored -> { + schemaLoads.incrementAndGet(); + return schema; + }); + + assertThat(((KnownWrittenColumns) result).fieldIds()).containsExactly(1, 2); + assertThat(schemaLoads).hasValue(1); + } + @Test public void testCollectWrittenColumnsExpandsLegacyFileSchema() { TableSchema schema = diff --git a/paimon-spark/paimon-spark-4.0/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala b/paimon-spark/paimon-spark-4.0/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala index f9e7ba1c0b8d..b5be2b0f0662 100644 --- a/paimon-spark/paimon-spark-4.0/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala +++ b/paimon-spark/paimon-spark-4.0/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala @@ -31,6 +31,7 @@ import org.apache.paimon.types.{DataType, RowType} import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path +import org.apache.spark.Partition import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.{CTESubstitution, SubstituteUnresolvedOrdinals} @@ -44,12 +45,12 @@ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.{ArrayData, GeneratedColumn, IdentityColumn, ResolveDefaultColumns} import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Column, Identifier, StagingTableCatalog, Table, TableCatalog} import org.apache.spark.sql.connector.expressions.Transform -import org.apache.spark.sql.connector.read.Scan +import org.apache.spark.sql.connector.read.{InputPartition, Scan} import org.apache.spark.sql.connector.write.BatchWrite import org.apache.spark.sql.execution.{SparkFormatTable, SparkPlan} import org.apache.spark.sql.execution.datasources.{PartitioningAwareFileIndex, PartitionSpec} import org.apache.spark.sql.execution.datasources.v2.{AtomicReplaceTableAsSelectExec, AtomicReplaceTableExec, ReplaceTableAsSelectExec, ReplaceTableExec} -import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceRDDPartition, DataSourceV2Relation, DataSourceV2ScanRelation} import org.apache.spark.sql.execution.streaming.{FileStreamSink, MetadataLogFileIndex} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructType, VariantType} @@ -78,6 +79,9 @@ class Spark4Shim extends SparkShim { override def classicApi: ClassicApi = new Classic4Api + override def dataSourceInputPartitions(partition: Partition): Seq[InputPartition] = + partition.asInstanceOf[DataSourceRDDPartition].inputPartitions + override def createSparkParser(delegate: ParserInterface): ParserInterface = { new PaimonSpark4SqlExtensionsParser(delegate) } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala index 500acab8cf73..9596974d09de 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala @@ -46,7 +46,7 @@ private[spark] case class PaimonMicroBatchMetadata( private[spark] case class PaimonMicroBatchInputPartition( splits: Seq[Split], - metadata: PaimonMicroBatchMetadata) + @transient metadata: PaimonMicroBatchMetadata) extends PaimonInputPartition object PaimonInputPartition { diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala index 5a952ed00baa..b03c77ad471e 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala @@ -19,6 +19,7 @@ package org.apache.paimon.spark import org.apache.paimon.annotation.Experimental +import org.apache.paimon.spark.sources.PaimonMicroBatchStream import org.apache.paimon.table.source.WrittenColumns import org.apache.spark.rdd.RDD @@ -26,7 +27,7 @@ import org.apache.spark.sql.Dataset import org.apache.spark.sql.execution.datasources.v2.DataSourceRDD import org.apache.spark.sql.paimon.shims.SparkShimLoader -import java.util.{IdentityHashMap, Optional} +import java.util.{IdentityHashMap, Map => JMap, Optional, UUID} import scala.collection.mutable import scala.util.control.NonFatal @@ -37,6 +38,8 @@ final class PaimonSparkMicroBatchMetadata private () object PaimonSparkMicroBatchMetadata { + private val StreamingQueryIdKey = "sql.streaming.queryId" + /** * Returns written columns for a raw foreachBatch Dataset with exactly one Paimon streaming * source. This method only inspects driver-side RDD planning metadata and does not run a Spark @@ -54,6 +57,10 @@ object PaimonSparkMicroBatchMetadata { } private def extractWrittenColumns(batch: Dataset[_]): Optional[WrittenColumns] = { + if (!hasExactlyOnePaimonSource(batch)) { + return Optional.empty() + } + val visited = new IdentityHashMap[RDD[_], java.lang.Boolean]() val metadata = mutable.ArrayBuffer.empty[PaimonMicroBatchMetadata] var incompletePaimonSource = false @@ -91,4 +98,38 @@ object PaimonSparkMicroBatchMetadata { } } } + + private def hasExactlyOnePaimonSource(batch: Dataset[_]): Boolean = { + val queryId = batch.sparkSession.sparkContext.getLocalProperty(StreamingQueryIdKey) + if (queryId == null) { + return false + } + + val sharedState = + batch.sparkSession.getClass.getMethod("sharedState").invoke(batch.sparkSession) + val activeQueries = + sharedState.getClass + .getMethod("activeStreamingQueries") + .invoke(sharedState) + .asInstanceOf[JMap[UUID, AnyRef]] + val execution = activeQueries.get(UUID.fromString(queryId)) + if (execution == null) { + return false + } + + // Spark replaces sources without new offsets with LocalRelation before foreachBatch. Their + // RDD lineage therefore contains no InputPartition to inspect. The active StreamExecution is + // the only per-query structure which still retains every source. Keep this Spark-internal + // access isolated here and fail closed if a Spark version changes it. + val sources = + execution.getClass.getMethod("sources").invoke(execution).asInstanceOf[Seq[AnyRef]] + val distinctSources = new IdentityHashMap[AnyRef, java.lang.Boolean]() + sources.foreach(source => distinctSources.put(source, java.lang.Boolean.TRUE)) + + if (distinctSources.size() != 1) { + false + } else { + distinctSources.keySet().iterator().next().isInstanceOf[PaimonMicroBatchStream] + } + } } diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala index 57f975af4be2..90235e9c8d6b 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala @@ -20,6 +20,7 @@ package org.apache.paimon.spark import org.apache.paimon.spark.sources.PaimonSourceOffset import org.apache.paimon.table.source.{KnownWrittenColumns, WrittenColumns} +import org.apache.paimon.utils.InstantiationUtil import org.apache.spark.sql.{Dataset, Row} import org.apache.spark.sql.streaming.{StreamingQueryException, StreamTest, Trigger} @@ -51,6 +52,22 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { } } + test("Paimon Source: keep micro-batch metadata on the driver") { + val metadata = + PaimonMicroBatchMetadata( + "source", + "start", + "end", + 0, + new KnownWrittenColumns(Seq(Integer.valueOf(1)).asJava)) + val partition = PaimonMicroBatchInputPartition(Seq.empty, metadata) + + val restored = InstantiationUtil.clone(partition) + + assert(restored.splits.isEmpty) + assert(restored.metadata == null) + } + test("Paimon Source: expose written columns to raw foreachBatch") { withTempDir { checkpointDir => @@ -132,6 +149,53 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { } } + test("Paimon Source: written columns metadata is ambiguous with an empty second source") { + withTable("written_columns_source_1", "written_columns_source_2") { + withTempDir { + checkpointDir => + spark.sql("CREATE TABLE written_columns_source_1 (id INT)") + spark.sql("CREATE TABLE written_columns_source_2 (id INT)") + spark.sql("INSERT INTO written_columns_source_1 VALUES (1)") + spark.sql("INSERT INTO written_columns_source_2 VALUES (2)") + + val source1 = spark.readStream + .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) + .table("written_columns_source_1") + val source2 = spark.readStream + .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) + .table("written_columns_source_2") + @volatile var nonEmptyBatchMetadataPresent = Seq.empty[Boolean] + + val query = source1 + .union(source2) + .writeStream + .option("checkpointLocation", checkpointDir.getCanonicalPath) + .foreachBatch { + (batch: Dataset[Row], _: Long) => + val metadataPresent = + PaimonSparkMicroBatchMetadata.writtenColumns(batch).isPresent + if (batch.count() > 0) { + nonEmptyBatchMetadataPresent = nonEmptyBatchMetadataPresent :+ metadataPresent + } + () + } + .start() + + try { + query.processAllAvailable() + nonEmptyBatchMetadataPresent = Seq.empty + + spark.sql("INSERT INTO written_columns_source_1 VALUES (3)") + query.processAllAvailable() + + assert(nonEmptyBatchMetadataPresent == Seq(false)) + } finally { + query.stop() + } + } + } + } + test("Paimon Source: expose partial data evolution written columns") { withSparkSQLConf("spark.paimon.write.use-v2-write" -> "false") { withTable("T") { From c6df54ab43dee5639a772cad2a1aa5e79aa7721c Mon Sep 17 00:00:00 2001 From: LsomeYeah Date: Wed, 5 Aug 2026 16:38:16 +0800 Subject: [PATCH 3/4] [spark] Harden micro-batch written columns metadata --- docs/docs/spark/structured-streaming.md | 33 +++++++ .../spark_connector_configuration.html | 6 ++ .../paimon/utils/DataEvolutionUtils.java | 52 ++++++++--- .../paimon/utils/DataEvolutionUtilsTest.java | 43 ++++++++- .../spark/PaimonSparkMicroBatchMetadata.scala | 91 +++++++++++++------ .../sources/PaimonMicroBatchStream.scala | 20 +++- .../paimon/spark/PaimonSourceTest.scala | 66 +++++++++++++- 7 files changed, 261 insertions(+), 50 deletions(-) diff --git a/docs/docs/spark/structured-streaming.md b/docs/docs/spark/structured-streaming.md index 91801f3a4d92..bc601ed2ada6 100644 --- a/docs/docs/spark/structured-streaming.md +++ b/docs/docs/spark/structured-streaming.md @@ -198,6 +198,39 @@ val query = spark.readStream .start() ``` +### Written Columns of a Micro-Batch + +`foreachBatch` consumers can inspect which Paimon field IDs were written by the data files admitted to the current micro-batch. This experimental metadata collection is disabled by default. Enable `read.stream.batch-written-columns.enabled` on the Paimon streaming source, then call `PaimonSparkMicroBatchMetadata.writtenColumns` with the raw `Dataset` passed to `foreachBatch`. + +```scala +import org.apache.paimon.spark.PaimonSparkMicroBatchMetadata +import org.apache.paimon.table.source.{AllColumns, KnownWrittenColumns} +import org.apache.spark.sql.{Dataset, Row} + +val query = spark.readStream + .format("paimon") + .option("read.stream.batch-written-columns.enabled", "true") + .table("table_name") + .writeStream + .option("checkpointLocation", "/path/to/checkpoint") + .foreachBatch { (batch: Dataset[Row], _: Long) => + val writtenColumns = PaimonSparkMicroBatchMetadata.writtenColumns(batch) + if (!writtenColumns.isPresent) { + // Metadata is unavailable; conservatively process all columns. + } else if (writtenColumns.get() == AllColumns.INSTANCE) { + // Exact field IDs are unavailable; conservatively process all columns. + } else { + val fieldIds = writtenColumns.get().asInstanceOf[KnownWrittenColumns].fieldIds() + // Process the exact set of written Paimon field IDs. + } + } + .start() +``` + +A present `KnownWrittenColumns` contains the complete, immutable set of written field IDs in ascending order. The set may be empty; that is a known empty set, not unknown metadata. A present `AllColumns.INSTANCE` means that exact file or schema metadata could not be resolved, so every column must be treated as written. + +An empty `Optional` means that metadata is unavailable, for example because collection was not enabled, the micro-batch is empty, the `Dataset` is not the raw batch from a query with exactly one distinct Paimon streaming source, or its lineage is incomplete or ambiguous. An empty `Optional` does not mean that no columns were written; callers must fall back to processing all columns. + Paimon Structured Streaming supports read row in the form of changelog (add rowkind column in row to represent its change type) in two ways: diff --git a/docs/generated/spark_connector_configuration.html b/docs/generated/spark_connector_configuration.html index cd95fd5fd41a..57708e4db113 100644 --- a/docs/generated/spark_connector_configuration.html +++ b/docs/generated/spark_connector_configuration.html @@ -44,6 +44,12 @@ Boolean Whether to read row in the form of changelog (add rowkind column in row to represent its change type). + +

read.stream.batch-written-columns.enabled
+ false + Boolean + Whether to expose the written field ids of an admitted micro-batch through PaimonSparkMicroBatchMetadata. +
read.stream.maxBytesPerTrigger
(none) diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java b/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java index 5ada67ab0390..aaf6c97bfffb 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/DataEvolutionUtils.java @@ -20,6 +20,7 @@ import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.SpecialFields; import org.apache.paimon.table.source.AllColumns; import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.KnownWrittenColumns; @@ -49,9 +50,7 @@ public class DataEvolutionUtils { public static WrittenColumns collectWrittenColumns( Collection splits, Function schemaLoader) { Set fieldIds = new TreeSet<>(); - Map schemaCache = new HashMap<>(); - Function cachedSchemaLoader = - schemaId -> schemaCache.computeIfAbsent(schemaId, schemaLoader); + Map> fieldIdByNameCache = new HashMap<>(); Map>, Set> fieldIdsCache = new HashMap<>(); for (DataSplit split : splits) { for (DataFileMeta file : split.dataFiles()) { @@ -59,10 +58,10 @@ public static WrittenColumns collectWrittenColumns( Pair> cacheKey = Pair.of(file.schemaId(), file.writeCols()); Set fileFieldIds = fieldIdsCache.get(cacheKey); if (fileFieldIds == null) { - fileFieldIds = computeFileFieldIds(cachedSchemaLoader, file); + fileFieldIds = computeFileFieldIds(schemaLoader, fieldIdByNameCache, file); fieldIdsCache.put(cacheKey, fileFieldIds); + fieldIds.addAll(fileFieldIds); } - fieldIds.addAll(fileFieldIds); } catch (RuntimeException e) { return AllColumns.INSTANCE; } @@ -71,17 +70,44 @@ public static WrittenColumns collectWrittenColumns( return new KnownWrittenColumns(fieldIds); } - /** Resolve a data file's physical columns through the schema the file was written with. */ - public static Set computeFileFieldIds( - Function schemaLoader, DataFileMeta file) { - TableSchema fileSchema = schemaLoader.apply(file.schemaId()); - if (fileSchema == null) { - throw new IllegalArgumentException("Cannot find schema " + file.schemaId()); + private static Set computeFileFieldIds( + Function schemaLoader, + Map> fieldIdByNameCache, + DataFileMeta file) { + Map fieldIdByName = + fieldIdByNameCache.computeIfAbsent( + file.schemaId(), + schemaId -> { + TableSchema fileSchema = schemaLoader.apply(schemaId); + if (fileSchema == null) { + throw new IllegalArgumentException( + "Cannot find schema " + schemaId); + } + + Map fieldIds = new HashMap<>(); + for (DataField field : fileSchema.fields()) { + fieldIds.put(field.name(), field.id()); + } + return fieldIds; + }); + + List writeCols = file.writeCols(); + if (writeCols == null) { + return new TreeSet<>(fieldIdByName.values()); } Set fieldIds = new TreeSet<>(); - for (DataField field : fileSchema.project(file.writeCols()).fields()) { - fieldIds.add(field.id()); + for (String writeCol : writeCols) { + Integer fieldId = fieldIdByName.get(writeCol); + if (fieldId == null) { + checkArgument( + SpecialFields.isSystemField(writeCol), + "Cannot find write column '%s' in schema %s.", + writeCol, + file.schemaId()); + } else { + fieldIds.add(fieldId); + } } return fieldIds; } diff --git a/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java b/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java index 9d2b2fb86f60..d9b556163f28 100644 --- a/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java @@ -45,6 +45,8 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; /** Test for {@link DataEvolutionUtils}. */ @@ -173,6 +175,37 @@ public void testCollectWrittenColumnsFallsBackWhenSchemaResolutionFails() { assertThat(result).isSameAs(AllColumns.INSTANCE); } + @Test + public void testCollectWrittenColumnsFallsBackWhenWriteColumnIsMissing() { + TableSchema schema = tableSchema(1L, new DataField(1, "a", DataTypes.INT())); + DataFileMeta file = dataFile(1L, Collections.singletonList("missing")); + + WrittenColumns result = + DataEvolutionUtils.collectWrittenColumns( + Collections.singletonList(dataSplit(file)), ignored -> schema); + + assertThat(result).isSameAs(AllColumns.INSTANCE); + } + + @Test + public void testCollectWrittenColumnsIgnoresSystemFields() { + TableSchema schema = tableSchema(1L, new DataField(1, "a", DataTypes.INT())); + DataFileMeta file = + dataFile( + 1L, + Arrays.asList( + SpecialFields.ROW_ID.name(), + "a", + SpecialFields.SEQUENCE_NUMBER.name())); + + WrittenColumns result = + DataEvolutionUtils.collectWrittenColumns( + Collections.singletonList(dataSplit(file)), ignored -> schema); + + assertThat(result).isInstanceOf(KnownWrittenColumns.class); + assertThat(((KnownWrittenColumns) result).fieldIds()).containsExactly(1); + } + @Test public void testCollectWrittenColumnsCachesFileSchemaProjection() { TableSchema schema = @@ -199,10 +232,11 @@ public void testCollectWrittenColumnsCachesFileSchemaProjection() { @Test public void testCollectWrittenColumnsCachesSchemaAcrossProjections() { TableSchema schema = - tableSchema( - 1L, - new DataField(1, "a", DataTypes.INT()), - new DataField(2, "b", DataTypes.STRING())); + spy( + tableSchema( + 1L, + new DataField(1, "a", DataTypes.INT()), + new DataField(2, "b", DataTypes.STRING()))); DataFileMeta first = dataFile(1L, Collections.singletonList("a")); DataFileMeta second = dataFile(1L, Collections.singletonList("b")); AtomicInteger schemaLoads = new AtomicInteger(); @@ -217,6 +251,7 @@ public void testCollectWrittenColumnsCachesSchemaAcrossProjections() { assertThat(((KnownWrittenColumns) result).fieldIds()).containsExactly(1, 2); assertThat(schemaLoads).hasValue(1); + verify(schema).fields(); } @Test diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala index b03c77ad471e..7575a74d9b86 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala @@ -29,7 +29,6 @@ import org.apache.spark.sql.paimon.shims.SparkShimLoader import java.util.{IdentityHashMap, Map => JMap, Optional, UUID} -import scala.collection.mutable import scala.util.control.NonFatal /** Driver-side access to metadata planned for a Paimon streaming micro-batch. */ @@ -62,40 +61,78 @@ object PaimonSparkMicroBatchMetadata { } val visited = new IdentityHashMap[RDD[_], java.lang.Boolean]() - val metadata = mutable.ArrayBuffer.empty[PaimonMicroBatchMetadata] - var incompletePaimonSource = false + var only: PaimonMicroBatchMetadata = null + + def inspectOccurrence(dataSourceRDD: DataSourceRDD): Boolean = { + var occurrenceOnly: PaimonMicroBatchMetadata = null + var inputCount = 0 + var valid = true + val partitions = dataSourceRDD.partitions + var partitionIndex = 0 + + while (valid && partitionIndex < partitions.length) { + val inputs = + SparkShimLoader.shim.dataSourceInputPartitions(partitions(partitionIndex)).iterator + while (valid && inputs.hasNext) { + inputs.next() match { + case input: PaimonMicroBatchInputPartition => + val current = input.metadata + if (current eq null) { + valid = false + } else if (occurrenceOnly eq null) { + occurrenceOnly = current + inputCount += 1 + } else if ((occurrenceOnly eq current) || occurrenceOnly == current) { + inputCount += 1 + } else { + valid = false + } + case _: PaimonInputPartition => valid = false + case _ => + } + } + partitionIndex += 1 + } + + if (!valid || ((occurrenceOnly ne null) && inputCount != occurrenceOnly.splitCount)) { + false + } else if (occurrenceOnly eq null) { + true + } else if (only eq null) { + only = occurrenceOnly + true + } else { + (only eq occurrenceOnly) || only == occurrenceOnly + } + } - def visit(rdd: RDD[_]): Unit = { - if (!visited.containsKey(rdd)) { + def visit(rdd: RDD[_]): Boolean = { + if (visited.containsKey(rdd)) { + true + } else { visited.put(rdd, java.lang.Boolean.TRUE) - rdd match { - case dataSourceRDD: DataSourceRDD => - dataSourceRDD.partitions.foreach { - partition => - SparkShimLoader.shim.dataSourceInputPartitions(partition).foreach { - case input: PaimonMicroBatchInputPartition => metadata += input.metadata - case _: PaimonInputPartition => incompletePaimonSource = true - case _ => - } - } - case _ => + val valid = + rdd match { + case dataSourceRDD: DataSourceRDD => inspectOccurrence(dataSourceRDD) + case _ => true + } + if (!valid) { + false + } else { + val dependencies = rdd.dependencies.iterator + var complete = true + while (complete && dependencies.hasNext) { + complete = visit(dependencies.next().rdd) + } + complete } - rdd.dependencies.foreach(dependency => visit(dependency.rdd)) } } - visit(batch.queryExecution.toRdd) - - val distinct = metadata.distinct - if (incompletePaimonSource || distinct.size != 1) { + if (!visit(batch.queryExecution.toRdd) || (only eq null)) { Optional.empty() } else { - val only = distinct.head - if (metadata.size != only.splitCount) { - Optional.empty() - } else { - Optional.of(only.writtenColumns) - } + Optional.of(only.writtenColumns) } } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala index 7d5fe86b1fd5..a49feee617e9 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.connector.read.{InputPartition, PartitionReaderFacto import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset, ReadLimit, SupportsTriggerAvailableNow} import java.lang.{Long => JLong} +import java.util.concurrent.ConcurrentHashMap import java.util.function.Function import scala.collection.JavaConverters._ @@ -107,6 +108,19 @@ class PaimonMicroBatchStream( private lazy val batchWrittenColumnsEnabled: Boolean = options.get(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED) + private[spark] lazy val schemaLoader: Function[JLong, TableSchema] = { + val schemaManager = table.schemaManager() + val schemaCache = new ConcurrentHashMap[JLong, TableSchema]() + val uncachedSchemaLoader = new Function[JLong, TableSchema] { + override def apply(schemaId: JLong): TableSchema = + schemaManager.schema(schemaId.longValue()) + } + new Function[JLong, TableSchema] { + override def apply(schemaId: JLong): TableSchema = + schemaCache.computeIfAbsent(schemaId, uncachedSchemaLoader) + } + } + override def getDefaultReadLimit: ReadLimit = defaultReadLimit override def prepareForTriggerAvailableNow(): Unit = { @@ -167,13 +181,9 @@ class PaimonMicroBatchStream( admittedSplits: Array[IndexedDataSplit]): PlannedMicroBatch = { val writtenColumns = try { - val schemaManager = table.schemaManager() DataEvolutionUtils.collectWrittenColumns( admittedSplits.map(_.entry).toSeq.asJava, - new Function[JLong, TableSchema] { - override def apply(schemaId: JLong): TableSchema = - schemaManager.schema(schemaId.longValue()) - } + schemaLoader ) } catch { case NonFatal(e) => diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala index 90235e9c8d6b..62d9764e9a79 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala @@ -18,14 +18,19 @@ package org.apache.paimon.spark -import org.apache.paimon.spark.sources.PaimonSourceOffset +import org.apache.paimon.schema.{SchemaManager, TableSchema} +import org.apache.paimon.spark.sources.{PaimonMicroBatchStream, PaimonSourceOffset} +import org.apache.paimon.table.DataTable import org.apache.paimon.table.source.{KnownWrittenColumns, WrittenColumns} import org.apache.paimon.utils.InstantiationUtil import org.apache.spark.sql.{Dataset, Row} import org.apache.spark.sql.streaming.{StreamingQueryException, StreamTest, Trigger} import org.junit.jupiter.api.Assertions +import org.mockito.Mockito.{mock, times, verify, when} +import java.lang.{Long => JLong} +import java.util.Collections import java.util.concurrent.TimeUnit import scala.collection.JavaConverters._ @@ -68,6 +73,26 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { assert(restored.metadata == null) } + test("Paimon Source: cache schemas for the stream lifetime") { + val table = mock(classOf[DataTable]) + val schemaManager = mock(classOf[SchemaManager]) + val initialSchema = mock(classOf[TableSchema]) + val evolvedSchema = mock(classOf[TableSchema]) + when(table.options()).thenReturn(Collections.emptyMap[String, String]()) + when(table.schemaManager()).thenReturn(schemaManager) + when(schemaManager.schema(1L)).thenReturn(initialSchema) + when(schemaManager.schema(2L)).thenReturn(evolvedSchema) + + val stream = new PaimonMicroBatchStream(table, null, "checkpoint") + + assert(stream.schemaLoader.apply(JLong.valueOf(1L)) eq initialSchema) + assert(stream.schemaLoader.apply(JLong.valueOf(1L)) eq initialSchema) + assert(stream.schemaLoader.apply(JLong.valueOf(2L)) eq evolvedSchema) + assert(stream.schemaLoader.apply(JLong.valueOf(2L)) eq evolvedSchema) + verify(schemaManager, times(1)).schema(1L) + verify(schemaManager, times(1)).schema(2L) + } + test("Paimon Source: expose written columns to raw foreachBatch") { withTempDir { checkpointDir => @@ -118,6 +143,45 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { } } + test("Paimon Source: expose written columns for a self-union") { + withTempDir { + checkpointDir => + val TableSnapshotState(_, location, _, _, _) = + prepareTableAndGetLocation(1, hasPk = true) + val expectedFieldIds = + loadTable("T").schema().fields().asScala.map(field => Integer.valueOf(field.id())).sorted + @volatile var writtenColumns: WrittenColumns = null + + val source = spark.readStream + .format("paimon") + .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) + .load(location) + val query = source + .union(source) + .writeStream + .option("checkpointLocation", checkpointDir.getCanonicalPath) + .foreachBatch { + (batch: Dataset[Row], _: Long) => + val metadata = PaimonSparkMicroBatchMetadata.writtenColumns(batch) + if (metadata.isPresent) { + writtenColumns = metadata.get() + } + batch.count() + () + } + .start() + + try { + query.processAllAvailable() + assert(writtenColumns.isInstanceOf[KnownWrittenColumns]) + assert( + writtenColumns.asInstanceOf[KnownWrittenColumns].fieldIds() == expectedFieldIds.asJava) + } finally { + query.stop() + } + } + } + test("Paimon Source: written columns metadata is disabled by default") { withTempDir { checkpointDir => From efa8dcaf531eee007d761ecbd8ab17aff532ef65 Mon Sep 17 00:00:00 2001 From: LsomeYeah Date: Mon, 10 Aug 2026 18:57:10 +0800 Subject: [PATCH 4/4] [spark] Plan micro-batch metadata by default --- docs/docs/spark/structured-streaming.md | 5 ++--- .../spark_connector_configuration.html | 6 ------ .../paimon/spark/SparkConnectorOptions.java | 8 -------- .../spark/PaimonSparkMicroBatchMetadata.scala | 5 ++--- .../sources/PaimonMicroBatchStream.scala | 19 +++++-------------- .../paimon/spark/PaimonSourceTest.scala | 9 ++------- 6 files changed, 11 insertions(+), 41 deletions(-) diff --git a/docs/docs/spark/structured-streaming.md b/docs/docs/spark/structured-streaming.md index bc601ed2ada6..54b10274b4d8 100644 --- a/docs/docs/spark/structured-streaming.md +++ b/docs/docs/spark/structured-streaming.md @@ -200,7 +200,7 @@ val query = spark.readStream ### Written Columns of a Micro-Batch -`foreachBatch` consumers can inspect which Paimon field IDs were written by the data files admitted to the current micro-batch. This experimental metadata collection is disabled by default. Enable `read.stream.batch-written-columns.enabled` on the Paimon streaming source, then call `PaimonSparkMicroBatchMetadata.writtenColumns` with the raw `Dataset` passed to `foreachBatch`. +`foreachBatch` consumers can inspect which Paimon field IDs were written by the data files admitted to the current micro-batch. Paimon plans this experimental metadata automatically. Call `PaimonSparkMicroBatchMetadata.writtenColumns` with the raw `Dataset` passed to `foreachBatch`. ```scala import org.apache.paimon.spark.PaimonSparkMicroBatchMetadata @@ -209,7 +209,6 @@ import org.apache.spark.sql.{Dataset, Row} val query = spark.readStream .format("paimon") - .option("read.stream.batch-written-columns.enabled", "true") .table("table_name") .writeStream .option("checkpointLocation", "/path/to/checkpoint") @@ -229,7 +228,7 @@ val query = spark.readStream A present `KnownWrittenColumns` contains the complete, immutable set of written field IDs in ascending order. The set may be empty; that is a known empty set, not unknown metadata. A present `AllColumns.INSTANCE` means that exact file or schema metadata could not be resolved, so every column must be treated as written. -An empty `Optional` means that metadata is unavailable, for example because collection was not enabled, the micro-batch is empty, the `Dataset` is not the raw batch from a query with exactly one distinct Paimon streaming source, or its lineage is incomplete or ambiguous. An empty `Optional` does not mean that no columns were written; callers must fall back to processing all columns. +An empty `Optional` means that metadata is unavailable, for example because the micro-batch is empty, the `Dataset` is not the raw batch from a query with exactly one distinct Paimon streaming source, or its lineage is incomplete or ambiguous. An empty `Optional` does not mean that no columns were written; callers must fall back to processing all columns. Paimon Structured Streaming supports read row in the form of changelog (add rowkind column in row to represent its change type) in two ways: diff --git a/docs/generated/spark_connector_configuration.html b/docs/generated/spark_connector_configuration.html index 57708e4db113..cd95fd5fd41a 100644 --- a/docs/generated/spark_connector_configuration.html +++ b/docs/generated/spark_connector_configuration.html @@ -44,12 +44,6 @@ Boolean Whether to read row in the form of changelog (add rowkind column in row to represent its change type). - -
read.stream.batch-written-columns.enabled
- false - Boolean - Whether to expose the written field ids of an admitted micro-batch through PaimonSparkMicroBatchMetadata. -
read.stream.maxBytesPerTrigger
(none) diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java index b60f04c6f29a..2f315b8df0f5 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java @@ -131,14 +131,6 @@ public class SparkConnectorOptions { .withDescription( "The maximum delay between two adjacent batches, which used to create MinRowsReadLimit with read.stream.minRowsPerTrigger together."); - public static final ConfigOption BATCH_WRITTEN_COLUMNS_ENABLED = - key("read.stream.batch-written-columns.enabled") - .booleanType() - .defaultValue(false) - .withDescription( - "Whether to expose the written field ids of an admitted micro-batch " - + "through PaimonSparkMicroBatchMetadata."); - public static final ConfigOption READ_CHANGELOG = key("read.changelog") .booleanType() diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala index 7575a74d9b86..219bef268b5a 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala @@ -42,9 +42,8 @@ object PaimonSparkMicroBatchMetadata { /** * Returns written columns for a raw foreachBatch Dataset with exactly one Paimon streaming * source. This method only inspects driver-side RDD planning metadata and does not run a Spark - * job. The result is empty when metadata collection was not enabled, the Dataset is not backed by - * a Paimon source, the lineage is incomplete, or multiple Paimon sources make the result - * ambiguous. + * job. The result is empty when the Dataset is not backed by a Paimon source, the lineage is + * incomplete, or multiple Paimon sources make the result ambiguous. */ def writtenColumns(batch: Dataset[_]): Optional[WrittenColumns] = { try { diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala index a49feee617e9..cc9301b8517b 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonMicroBatchStream.scala @@ -21,7 +21,7 @@ package org.apache.paimon.spark.sources import org.apache.paimon.CoreOptions import org.apache.paimon.options.Options import org.apache.paimon.schema.TableSchema -import org.apache.paimon.spark.{PaimonImplicits, PaimonInputPartition, PaimonMicroBatchInputPartition, PaimonMicroBatchMetadata, PaimonPartitionReaderFactory, SparkConnectorOptions} +import org.apache.paimon.spark.{PaimonImplicits, PaimonMicroBatchInputPartition, PaimonMicroBatchMetadata, PaimonPartitionReaderFactory, SparkConnectorOptions} import org.apache.paimon.table.DataTable import org.apache.paimon.table.source.{AllColumns, ReadBuilder} import org.apache.paimon.utils.DataEvolutionUtils @@ -105,9 +105,6 @@ class PaimonMicroBatchStream( private lazy val blobAsDescriptor: Boolean = options.get(CoreOptions.BLOB_AS_DESCRIPTOR) - private lazy val batchWrittenColumnsEnabled: Boolean = - options.get(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED) - private[spark] lazy val schemaLoader: Function[JLong, TableSchema] = { val schemaManager = table.schemaManager() val schemaCache = new ConcurrentHashMap[JLong, TableSchema]() @@ -163,16 +160,10 @@ class PaimonMicroBatchStream( val endOffset = PaimonSourceOffset(end) val admittedSplits = getBatch(startOffset, Some(endOffset), None) - if (!batchWrittenColumnsEnabled) { - admittedSplits - .map(ids => PaimonInputPartition(ids.entry)) - .toArray[InputPartition] - } else { - val plannedBatch = createPlannedMicroBatch(startOffset, endOffset, admittedSplits) - plannedBatch.admittedSplits - .map(ids => PaimonMicroBatchInputPartition(Seq(ids.entry), plannedBatch.metadata)) - .toArray[InputPartition] - } + val plannedBatch = createPlannedMicroBatch(startOffset, endOffset, admittedSplits) + plannedBatch.admittedSplits + .map(ids => PaimonMicroBatchInputPartition(Seq(ids.entry), plannedBatch.metadata)) + .toArray[InputPartition] } private def createPlannedMicroBatch( diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala index 62d9764e9a79..74edd6725345 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSourceTest.scala @@ -105,7 +105,6 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { val query = spark.readStream .format("paimon") - .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) .load(location) .select("a") .writeStream @@ -154,7 +153,6 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { val source = spark.readStream .format("paimon") - .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) .load(location) val query = source .union(source) @@ -182,7 +180,7 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { } } - test("Paimon Source: written columns metadata is disabled by default") { + test("Paimon Source: written columns metadata is available by default") { withTempDir { checkpointDir => val TableSnapshotState(_, location, snapshotData, _, _) = @@ -205,7 +203,7 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { try { query.processAllAvailable() - assert(!metadataAvailable) + assert(metadataAvailable) assert(rowCount == snapshotData.size) } finally { query.stop() @@ -223,10 +221,8 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { spark.sql("INSERT INTO written_columns_source_2 VALUES (2)") val source1 = spark.readStream - .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) .table("written_columns_source_1") val source2 = spark.readStream - .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) .table("written_columns_source_2") @volatile var nonEmptyBatchMetadataPresent = Seq.empty[Boolean] @@ -280,7 +276,6 @@ class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { @volatile var nonEmptyBatchColumns = Seq.empty[WrittenColumns] val query = spark.readStream - .option(SparkConnectorOptions.BATCH_WRITTEN_COLUMNS_ENABLED.key(), true) .option(SparkConnectorOptions.MAX_FILES_PER_TRIGGER.key(), 1) .option("scan.mode", "latest") .table("`T$row_tracking`")