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/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..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,13 +20,21 @@ 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; +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 +46,72 @@ /** 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> fieldIdByNameCache = new HashMap<>(); + 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, fieldIdByNameCache, file); + fieldIdsCache.put(cacheKey, fileFieldIds); + fieldIds.addAll(fileFieldIds); + } + } catch (RuntimeException e) { + return AllColumns.INSTANCE; + } + } + } + return new KnownWrittenColumns(fieldIds); + } + + 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 (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; + } + /** * 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..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 @@ -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,16 @@ 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.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; /** Test for {@link DataEvolutionUtils}. */ public class DataEvolutionUtilsTest { @@ -109,6 +122,155 @@ 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 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 = + 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 testCollectWrittenColumnsCachesSchemaAcrossProjections() { + TableSchema schema = + 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(); + + 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); + verify(schema).fields(); + } + + @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 +333,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-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/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..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 @@ -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], + @transient 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..7575a74d9b86 --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSparkMicroBatchMetadata.scala @@ -0,0 +1,172 @@ +/* + * 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.spark.sources.PaimonMicroBatchStream +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, Map => JMap, Optional, UUID} + +import scala.util.control.NonFatal + +/** Driver-side access to metadata planned for a Paimon streaming micro-batch. */ +@Experimental +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 + * 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] = { + if (!hasExactlyOnePaimonSource(batch)) { + return Optional.empty() + } + + val visited = new IdentityHashMap[RDD[_], java.lang.Boolean]() + 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[_]): Boolean = { + if (visited.containsKey(rdd)) { + true + } else { + visited.put(rdd, java.lang.Boolean.TRUE) + 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 + } + } + } + + if (!visit(batch.queryExecution.toRdd) || (only eq null)) { + Optional.empty() + } else { + Optional.of(only.writtenColumns) + } + } + + 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-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..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 @@ -20,15 +20,27 @@ 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.concurrent.ConcurrentHashMap +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 +105,22 @@ 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]() + 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 = { @@ -134,9 +162,42 @@ 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 { + DataEvolutionUtils.collectWrittenColumns( + admittedSplits.map(_.entry).toSeq.asJava, + schemaLoader + ) + } 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..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,23 @@ 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.Row +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._ + class PaimonSourceTest extends PaimonSparkTestBase with StreamTest { import testImplicits._ @@ -48,6 +57,266 @@ 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: 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 => + 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: 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 => + 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: 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") { + 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) }