[spark] Expose written columns for streaming micro-batches

docs/docs/spark/structured-streaming.md

@@ -198,6 +198,35 @@ 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. Call `PaimonSparkMicroBatchMetadata.writtenColumnIds` with the raw `Dataset` passed to `foreachBatch`. Paimon resolves the file metadata lazily when this method is called.
+
+```scala
+import org.apache.paimon.spark.PaimonSparkMicroBatchMetadata
+import org.apache.spark.sql.{Dataset, Row}
+
+val query = spark.readStream
+  .format("paimon")
+  .table("table_name")
+  .writeStream
+  .option("checkpointLocation", "/path/to/checkpoint")
+  .foreachBatch { (batch: Dataset[Row], _: Long) =>
+    val writtenColumnIds = PaimonSparkMicroBatchMetadata.writtenColumnIds(batch)
+    if (!writtenColumnIds.isPresent) {
+      // Metadata is unavailable; conservatively process all columns.
+    } else {
+      val fieldIds = writtenColumnIds.get()
+      // Process the exact set of written Paimon field IDs.
+    }
+  }
+  .start()
+```
+
+A present `Optional` contains the complete, immutable list of written field IDs in ascending order. The list may be empty; that is a known empty set, not unknown metadata.
+
+An empty `Optional` means that metadata is unavailable, for example because a file or schema cannot be resolved, 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:
 

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.DataSplit;
 import org.apache.paimon.types.DataField;
 
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.TreeSet;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -38,21 +46,76 @@
 /** Util class for data evolution. */
 public class DataEvolutionUtils {
 
+    /**
+     * Collect exact written field ids; an empty list is exact and an empty optional is unresolved.
+     */
+    public static Optional<List<Integer>> collectWrittenColumnIds(
+            Collection<DataSplit> splits, Function<Long, TableSchema> schemaLoader) {
+        Set<Integer> fieldIds = new TreeSet<>();
+        Map<Long, List<DataField>> schemaFieldsCache = new HashMap<>();
+        Map<Pair<Long, List<String>>, Set<Integer>> fieldIdsCache = new HashMap<>();
+        try {
+            for (DataSplit split : splits) {
+                for (DataFileMeta file : split.dataFiles()) {
+                    Pair<Long, List<String>> cacheKey = Pair.of(file.schemaId(), file.writeCols());
+                    Set<Integer> fileFieldIds = fieldIdsCache.get(cacheKey);
+                    if (fileFieldIds == null) {
+                        List<DataField> schemaFields =
+                                schemaFieldsCache.computeIfAbsent(
+                                        file.schemaId(),
+                                        schemaId -> {
+                                            TableSchema schema = schemaLoader.apply(schemaId);
+                                            checkArgument(
+                                                    schema != null,
+                                                    "Cannot find schema %s.",
+                                                    schemaId);
+                                            return schema.fields();
+                                        });
+                        fileFieldIds = resolveFileFieldIds(schemaFields, file, true);
+                        fieldIdsCache.put(cacheKey, fileFieldIds);
+                    }
+                    fieldIds.addAll(fileFieldIds);
+                }
+            }
+        } catch (RuntimeException e) {
+            return Optional.empty();
+        }
+        return Optional.of(Collections.unmodifiableList(new ArrayList<>(fieldIds)));
+    }
+
     /**
      * Table field ids physically present in a file, resolved through the schema used to write it.
      */
     public static Set<Integer> fileFieldIds(
             Function<Long, TableSchema> scanTableSchema, DataFileMeta file) {
-        TableSchema schema = scanTableSchema.apply(file.schemaId());
+        return resolveFileFieldIds(scanTableSchema.apply(file.schemaId()).fields(), file, false);
+    }
+
+    private static Set<Integer> resolveFileFieldIds(
+            List<DataField> schemaFields, DataFileMeta file, boolean strict) {
         List<String> writeCols = file.writeCols();
         Set<String> writeColNames = writeCols == null ? null : new HashSet<>(writeCols);
+        Set<String> unresolved =
+                strict && writeColNames != null ? new HashSet<>(writeColNames) : null;
         Set<Integer> ids = new HashSet<>();
-        for (DataField field : schema.fields()) {
+        for (DataField field : schemaFields) {
             // writeCols may also contain physical row-tracking fields outside the table schema.
             if (writeColNames == null || writeColNames.contains(field.name())) {
                 ids.add(field.id());
+                if (unresolved != null) {
+                    unresolved.remove(field.name());
+                }
             }
         }
+
+        if (unresolved != null) {
+            unresolved.removeIf(SpecialFields::isSystemField);
+            checkArgument(
+                    unresolved.isEmpty(),
+                    "Cannot find write columns %s in schema %s.",
+                    unresolved,
+                    file.schemaId());
+        }
         return ids;
     }
 

paimon-core/src/test/java/org/apache/paimon/utils/DataEvolutionUtilsTest.java

@@ -18,11 +18,15 @@
 
 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.DataSplit;
 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 +35,17 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+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 {
@@ -107,6 +118,128 @@ public void testFileFieldIdsHandlesFullEmptyAndUnrelatedWrites() {
                                         1,
                                         Collections.singletonList("other"))))
                 .containsExactly(2);
+        assertThat(
+                        DataEvolutionUtils.fileFieldIds(
+                                ignored -> schema,
+                                dataFile(
+                                        "unknown.parquet",
+                                        1,
+                                        Collections.singletonList("unknown"))))
+                .isEmpty();
+    }
+
+    @Test
+    public void testCollectWrittenColumnIdsAcrossSchemas() {
+        Map<Long, TableSchema> 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"));
+
+        assertThat(collectWrittenColumnIds(schemas::get, oldSchemaFile, newSchemaFile))
+                .hasValue(Arrays.asList(1, 2, 3));
+    }
+
+    @Test
+    public void testCollectWrittenColumnIdsFallsBackWhenResolutionFails() {
+        DataFileMeta unknownSchemaFile = dataFile(99L, Collections.singletonList("a"));
+        assertThat(collectWrittenColumnIds(ignored -> null, unknownSchemaFile))
+                .as("unknown schema")
+                .isEmpty();
+
+        DataFileMeta unresolvedSchemaFile = dataFile(1L, Collections.singletonList("missing"));
+        assertThat(
+                        collectWrittenColumnIds(
+                                ignored -> {
+                                    throw new IllegalArgumentException("schema cannot be resolved");
+                                },
+                                unresolvedSchemaFile))
+                .as("schema loader failure")
+                .isEmpty();
+
+        TableSchema schema = tableSchema(1L, new DataField(1, "a", DataTypes.INT()));
+        DataFileMeta unknownColumnFile = dataFile(1L, Collections.singletonList("missing"));
+        assertThat(collectWrittenColumnIds(ignored -> schema, unknownColumnFile))
+                .as("unknown non-system write column")
+                .isEmpty();
+    }
+
+    @Test
+    public void testCollectWrittenColumnIdsIgnoresSystemFields() {
+        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()));
+
+        assertThat(collectWrittenColumnIds(ignored -> schema, file))
+                .hasValue(Collections.singletonList(1));
+
+        assertThat(
+                        collectWrittenColumnIds(
+                                ignored -> schema,
+                                dataFile(
+                                        1L,
+                                        Arrays.asList(
+                                                SpecialFields.ROW_ID.name(),
+                                                SpecialFields.SEQUENCE_NUMBER.name()))))
+                .hasValue(Collections.emptyList());
+    }
+
+    @Test
+    public void testCollectWrittenColumnIdsCachesSchemaAcrossProjections() {
+        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"));
+        DataFileMeta repeated = dataFile(1L, Collections.singletonList("a"));
+        AtomicInteger schemaLoads = new AtomicInteger();
+
+        Optional<List<Integer>> result =
+                collectWrittenColumnIds(
+                        ignored -> {
+                            schemaLoads.incrementAndGet();
+                            return schema;
+                        },
+                        first,
+                        second,
+                        repeated);
+
+        assertThat(result.get()).containsExactly(1, 2);
+        assertThat(schemaLoads).hasValue(1);
+        verify(schema).fields();
+        verify(repeated).writeCols();
+    }
+
+    @Test
+    public void testCollectWrittenColumnIdsExpandsLegacyFileSchema() {
+        TableSchema schema =
+                tableSchema(
+                        1L,
+                        new DataField(1, "a", DataTypes.INT()),
+                        new DataField(2, "b", DataTypes.STRING()));
+        DataFileMeta legacyFile = dataFile(1L, null);
+
+        assertThat(collectWrittenColumnIds(ignored -> schema, legacyFile))
+                .hasValue(Arrays.asList(1, 2));
     }
 
     @Test
@@ -171,4 +304,38 @@ private static DataFileMeta dataFile(
                 0L,
                 writeCols);
     }
+
+    private static DataFileMeta dataFile(long schemaId, java.util.List<String> 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 Optional<List<Integer>> collectWrittenColumnIds(
+            Function<Long, TableSchema> schemaLoader, DataFileMeta... files) {
+        return DataEvolutionUtils.collectWrittenColumnIds(
+                Collections.singletonList(dataSplit(files)), schemaLoader);
+    }
+
+    private static TableSchema tableSchema(long id, DataField... fields) {
+        return TableSchema.create(
+                id,
+                new Schema(
+                        Arrays.asList(fields),
+                        Collections.emptyList(),
+                        Collections.emptyList(),
+                        Collections.emptyMap(),
+                        null));
+    }
 }

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala

@@ -24,6 +24,10 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
 import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, SupportsReportPartitioning}
 
+import java.util.{List => JList, Objects, Optional}
+
+import scala.util.control.NonFatal
+
 trait PaimonInputPartition extends InputPartition {
   def splits: Seq[Split]
 
@@ -36,6 +40,48 @@ trait PaimonInputPartition extends InputPartition {
 }
 
 case class SimplePaimonInputPartition(splits: Seq[Split]) extends PaimonInputPartition
+
+final private[spark] class PaimonMicroBatchMetadata private[spark] (
+    val sourceId: String,
+    val startOffset: String,
+    val endOffset: String,
+    val splitCount: Int,
+    @transient private var writtenColumnIdsThunk: () => Optional[JList[Integer]])
+  extends Serializable {
+
+  @transient private lazy val cachedWrittenColumnIds: Optional[JList[Integer]] = {
+    try {
+      val thunk = writtenColumnIdsThunk
+      writtenColumnIdsThunk = null
+      val supplied = if (thunk == null) null else thunk()
+      if (supplied == null) Optional.empty() else supplied
+    } catch {
+      case NonFatal(_) => Optional.empty()
+      case _: LinkageError => Optional.empty()
+    }
+  }
+
+  def writtenColumnIds: Optional[JList[Integer]] = cachedWrittenColumnIds
+
+  override def equals(other: Any): Boolean =
+    other match {
+      case that: PaimonMicroBatchMetadata =>
+        sourceId == that.sourceId &&
+        startOffset == that.startOffset &&
+        endOffset == that.endOffset &&
+        splitCount == that.splitCount
+      case _ => false
+    }
+
+  override def hashCode(): Int =
+    Objects.hash(sourceId, startOffset, endOffset, Integer.valueOf(splitCount))
+}
+
+private[spark] case class PaimonMicroBatchInputPartition(
+    splits: Seq[Split],
+    @transient metadata: PaimonMicroBatchMetadata)
+  extends PaimonInputPartition
+
 object PaimonInputPartition {
   def apply(split: Split): PaimonInputPartition = {
     SimplePaimonInputPartition(Seq(split))