feat(spark): SupportsReportPartitioning via optional reportPartitioni…

…ng hook Add a `reportPartitioning(opts)` default method on `FfiProviderFactory` that returns a `ReportedPartitioning` (Transform[] keys, derived numPartitions) or null for unknown. `DatafusionScan` now implements `SupportsReportPartitioning` and emits `KeyGroupedPartitioning(keys, partitions.length)` when the bridge opts in, `UnknownPartitioning` otherwise — letting Spark's optimizer skip shuffles ahead of compatible joins/aggregations. Move the `listPartitions` + `reportPartitioning` call sites into `DatafusionScanBuilder.build()` so the Scan has both facts cached up front; `DatafusionBatch.planInputPartitions` now reuses the cached `PartitionInfo[]` instead of re-invoking the factory. `ReportedPartitioning` ships with `identity(cols…)` and `bucket(n, cols…)` convenience builders for the common cases. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
daa3ba541a79e1915b5e4b82e2b11c6e7672baae
commit daa3ba541a79e1915b5e4b82e2b11c6e7672baae
diff --git a/examples/SPARK_INTEGRATION.md b/examples/SPARK_INTEGRATION.md
@@ -120,6 +120,16 @@ public final class MyBridgeProviderFactory implements FfiProviderFactory {
         // The driver-side schema probe passes an empty array; honour it.
         return MyBridgeNative.createFfiProvider(optionsProtoBytes, partitionBytes);
     }
+
+    @Override
+    public ReportedPartitioning reportPartitioning(byte[] optionsProtoBytes) {
+        // Optional. Return non-null only when each PartitionInfo's rows all share the same
+        // key tuple under the declared transforms — Spark elides shuffles ahead of joins/aggs
+        // grouped on those keys. Return null when the layout is unknown or row-key mapping
+        // would be lossy.
+        return ReportedPartitioning.identity("device_id");
+        // or: ReportedPartitioning.bucket(numBuckets, "user_id");
+    }
 }
 ```
 
@@ -168,7 +178,9 @@ df = (spark.read.format("my_format")
 | Phase                       | Where     | Path |
 | --------------------------- | --------- | ---- |
 | `inferSchema`               | Driver    | `factory.encodeOptions` → `factory.createProvider(opts, EMPTY)` → widen → `registerFfiTable` → `ctx.tableSchema` |
-| `planInputPartitions`       | Driver    | `factory.listPartitions(optionsBytes)` → one task per `PartitionInfo`; each task gets that entry's `partitionBytes` + `preferredLocations` |
+| `ScanBuilder.build`         | Driver    | `factory.listPartitions(optionsBytes)` (cached on Scan) + `factory.reportPartitioning(optionsBytes)` (cached on Scan) |
+| `outputPartitioning`        | Driver    | `KeyGroupedPartitioning(reported.keys, partitions.length)` when bridge declared one; `UnknownPartitioning(partitions.length)` otherwise. Spark may elide shuffles when keys line up with downstream join/agg grouping. |
+| `planInputPartitions`       | Driver    | Reuses the cached `PartitionInfo[]`; one task per entry with that entry's `partitionBytes` + `preferredLocations` |
 | Predicate translation       | Driver    | `SparkPredicateTranslator.translate(Predicate)` → `LogicalExprNode` proto bytes (each pushed predicate is independent) |
 | Per-task scan               | Executor  | Same factory → `createProvider(opts, partitionBytes)` → widen → `registerFfiTable` → `ctx.sql("SELECT proj FROM t")` → fold `DataFrame.filterFromProto(bytes)` over pushed predicates → `executeStream` |
 

diff --git a/spark/src/main/java/io/datafusion/spark/FfiProviderFactory.java b/spark/src/main/java/io/datafusion/spark/FfiProviderFactory.java
@@ -81,4 +81,23 @@ public interface FfiProviderFactory {
    *     DatafusionSource.inferSchema}.
    */
   long createProvider(byte[] optionsProtoBytes, byte[] partitionBytes);
+
+  /**
+   * Declare how rows are partitioned across the {@link PartitionInfo} entries returned by {@link
+   * #listPartitions(byte[])}. Driver-side only.
+   *
+   * <p>When non-null, the connector surfaces a {@code KeyGroupedPartitioning(keys,
+   * listPartitions(...).length)} to Spark via {@code SupportsReportPartitioning} so the optimizer
+   * can elide shuffles ahead of joins/aggregations on the declared keys.
+   *
+   * <p>Default returns {@code null} — no partitioning guarantees, Spark plans as if the scan's
+   * output ordering and grouping are unknown.
+   *
+   * <p>If a bridge implements this, it must hold the {@link ReportedPartitioning} contract: every
+   * row in a given partition evaluates to the same tuple of key values under the declared
+   * transforms.
+   */
+  default ReportedPartitioning reportPartitioning(byte[] optionsProtoBytes) {
+    return null;
+  }
 }
diff --git a/spark/src/main/java/io/datafusion/spark/ReportedPartitioning.java b/spark/src/main/java/io/datafusion/spark/ReportedPartitioning.java
@@ -0,0 +1,87 @@
+/*
+ * 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 io.datafusion.spark;
+
+import java.util.Arrays;
+
+import org.apache.spark.sql.connector.expressions.Expressions;
+import org.apache.spark.sql.connector.expressions.Transform;
+
+/**
+ * Driver-side declaration of how a bridge's data is partitioned on the key columns. When supplied
+ * via {@link FfiProviderFactory#reportPartitioning(byte[])}, the connector surfaces a {@link
+ * org.apache.spark.sql.connector.read.partitioning.KeyGroupedPartitioning} from {@link
+ * org.apache.spark.sql.connector.read.SupportsReportPartitioning#outputPartitioning()} — Spark's
+ * optimizer can then skip the shuffle ahead of joins/aggregations whose grouping keys line up
+ * with these transforms.
+ *
+ * <p>Contract: for any partition reported by {@link FfiProviderFactory#listPartitions(byte[])},
+ * every row produced by that partition must evaluate to the same tuple of key values under these
+ * transforms. Different partitions <i>may</i> share key values (Spark will fuse them); they
+ * <b>must not</b> straddle key values.
+ *
+ * <p>The partition count Spark sees is {@code listPartitions(...).length}; it is not carried here
+ * to keep a single source of truth.
+ */
+public final class ReportedPartitioning {
+
+  private final Transform[] keys;
+
+  public ReportedPartitioning(Transform[] keys) {
+    if (keys == null || keys.length == 0) {
+      throw new IllegalArgumentException(
+          "ReportedPartitioning: keys must contain at least one transform");
+    }
+    this.keys = keys;
+  }
+
+  public Transform[] keys() {
+    return keys;
+  }
+
+  /**
+   * Convenience: declare identity partitioning on one or more columns (a row in partition P has
+   * the same {@code (col1, col2, …)} values as every other row in P).
+   */
+  public static ReportedPartitioning identity(String... columns) {
+    if (columns == null || columns.length == 0) {
+      throw new IllegalArgumentException(
+          "ReportedPartitioning.identity: at least one column required");
+    }
+    Transform[] ts = Arrays.stream(columns).map(Expressions::identity).toArray(Transform[]::new);
+    return new ReportedPartitioning(ts);
+  }
+
+  /**
+   * Convenience: declare hash-bucket partitioning. Mirrors Spark's {@code bucket(N, cols…)}
+   * transform — each row is assigned to bucket {@code hash(cols) mod numBuckets}.
+   */
+  public static ReportedPartitioning bucket(int numBuckets, String... columns) {
+    if (numBuckets <= 0) {
+      throw new IllegalArgumentException(
+          "ReportedPartitioning.bucket: numBuckets must be > 0, got " + numBuckets);
+    }
+    if (columns == null || columns.length == 0) {
+      throw new IllegalArgumentException(
+          "ReportedPartitioning.bucket: at least one column required");
+    }
+    return new ReportedPartitioning(new Transform[] {Expressions.bucket(numBuckets, columns)});
+  }
+}
diff --git a/spark/src/main/scala/io/datafusion/spark/DatafusionBatch.scala b/spark/src/main/scala/io/datafusion/spark/DatafusionBatch.scala
@@ -23,27 +23,18 @@ import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionRead
 
 /**
  * Spark `Batch` for a DataFusion-backed scan. Owns:
- *   - partition planning (driver-side: `factory.listPartitions` enumerates `PartitionInfo`
- *     entries → one task per entry; each task receives that partition's `partitionBytes` and
- *     `preferredLocations`)
+ *   - partition planning (driver-side: reuses the `PartitionInfo[]` already resolved by
+ *     [[DatafusionScanBuilder]] — one task per entry; each task receives that entry's
+ *     `partitionBytes` + `preferredLocations`)
  *   - per-task reader factory ([[DatafusionPartitionReaderFactory]])
  */
 class DatafusionBatch(val scan: DatafusionScan) extends Batch {
 
   override def planInputPartitions(): Array[InputPartition] = {
-    val factory = instantiateFactory(scan.factoryFqcn)
-    val partitions: Array[PartitionInfo] = factory.listPartitions(scan.optionsProtoBytes)
-
-    if (partitions == null || partitions.isEmpty) {
-      throw new IllegalStateException(
-        s"FfiProviderFactory '${scan.factoryFqcn}' returned no partitions to scan"
-      )
-    }
-
     val projection = scan.prunedSchema.fieldNames
     val filterBytes: Array[Array[Byte]] = scan.pushedPredicateBytes
 
-    partitions.iterator.map { p =>
+    scan.partitions.iterator.map { p =>
       DatafusionInputPartition(
         factoryFqcn = scan.factoryFqcn,
         optionsProtoBytes = scan.optionsProtoBytes,
@@ -58,9 +49,4 @@ class DatafusionBatch(val scan: DatafusionScan) extends Batch {
 
   override def createReaderFactory(): PartitionReaderFactory =
     new DatafusionPartitionReaderFactory(scan.prunedSchema)
-
-  private def instantiateFactory(fqcn: String): FfiProviderFactory = {
-    val cls = Class.forName(fqcn)
-    cls.getDeclaredConstructor().newInstance().asInstanceOf[FfiProviderFactory]
-  }
 }
diff --git a/spark/src/main/scala/io/datafusion/spark/DatafusionScan.scala b/spark/src/main/scala/io/datafusion/spark/DatafusionScan.scala
@@ -20,28 +20,47 @@
 package io.datafusion.spark
 
 import org.apache.spark.sql.connector.expressions.filter.Predicate
-import org.apache.spark.sql.connector.read.{Batch, Scan}
+import org.apache.spark.sql.connector.read.{Batch, Scan, SupportsReportPartitioning}
+import org.apache.spark.sql.connector.read.partitioning.{
+  KeyGroupedPartitioning,
+  Partitioning,
+  UnknownPartitioning
+}
 import org.apache.spark.sql.types.StructType
 
 /**
  * Read plan for a DataFusion-backed scan. Holds pruning state, the pushed predicates (for
  * `description()` / `explain(True)`), and the corresponding `LogicalExprNode` proto byte arrays
  * the executor applies via `DataFrame.filterFromProto`.
+ *
+ * Also carries the driver-resolved `PartitionInfo[]` (so [[DatafusionBatch]] doesn't re-call
+ * `listPartitions`) and the optional bridge-declared [[ReportedPartitioning]]; when present, the
+ * scan surfaces a `KeyGroupedPartitioning` via `SupportsReportPartitioning` so Spark's optimizer
+ * can skip shuffles ahead of compatible joins/aggregations. When absent, an
+ * `UnknownPartitioning(partitions.length)` is reported (still correct, just no shuffle elision).
  */
 class DatafusionScan(
     val factoryFqcn: String,
     val optionsProtoBytes: Array[Byte],
     val fullSchema: StructType,
     val prunedSchema: StructType,
     val pushedPredicates: Array[Predicate],
-    val pushedPredicateBytes: Array[Array[Byte]]
-) extends Scan {
+    val pushedPredicateBytes: Array[Array[Byte]],
+    val partitions: Array[PartitionInfo],
+    val reportedPartitioning: ReportedPartitioning
+) extends Scan
+    with SupportsReportPartitioning {
 
   override def readSchema(): StructType = prunedSchema
 
   override def description(): String =
     s"DatafusionScan(factory=$factoryFqcn, projection=${prunedSchema.fieldNames.mkString(",")}," +
-      s" pushedPredicates=${pushedPredicates.length})"
+      s" pushedPredicates=${pushedPredicates.length}, partitions=${partitions.length}," +
+      s" reportedPartitioning=${if (reportedPartitioning == null) "unknown" else "key-grouped"})"
 
   override def toBatch: Batch = new DatafusionBatch(this)
+
+  override def outputPartitioning(): Partitioning =
+    if (reportedPartitioning == null) new UnknownPartitioning(partitions.length)
+    else new KeyGroupedPartitioning(reportedPartitioning.keys().toArray, partitions.length)
 }
diff --git a/spark/src/main/scala/io/datafusion/spark/DatafusionScanBuilder.scala b/spark/src/main/scala/io/datafusion/spark/DatafusionScanBuilder.scala
@@ -30,6 +30,12 @@ import org.apache.spark.sql.types.StructType
  * Pushdown discipline: over-claiming Exact = wrong results, under-claiming = full scans. The
  * translator (see [[SparkPredicateTranslator]]) only emits proto for predicates it can encode
  * losslessly — anything else returns `None` and lands in residuals.
+ *
+ * `build()` is also where we resolve driver-side facts that the optimizer needs *before* it
+ * starts asking the [[DatafusionScan]] about its output partitioning: the partition list
+ * (`listPartitions`) and the bridge's optional [[ReportedPartitioning]]. Resolving both here once
+ * and threading them onto the Scan keeps `DatafusionBatch.planInputPartitions` shuffle-free and
+ * lets `outputPartitioning()` answer without an extra factory call per query.
  */
 class DatafusionScanBuilder(
     factoryFqcn: String,
@@ -71,6 +77,29 @@ class DatafusionScanBuilder(
     pruned = requiredSchema
   }
 
-  override def build(): Scan =
-    new DatafusionScan(factoryFqcn, optionsProtoBytes, fullSchema, pruned, pushed, pushedBytes)
+  override def build(): Scan = {
+    val factory = instantiateFactory(factoryFqcn)
+    val partitions: Array[PartitionInfo] = factory.listPartitions(optionsProtoBytes)
+    if (partitions == null || partitions.isEmpty) {
+      throw new IllegalStateException(
+        s"FfiProviderFactory '$factoryFqcn' returned no partitions to scan"
+      )
+    }
+    val reported: ReportedPartitioning = factory.reportPartitioning(optionsProtoBytes)
+    new DatafusionScan(
+      factoryFqcn,
+      optionsProtoBytes,
+      fullSchema,
+      pruned,
+      pushed,
+      pushedBytes,
+      partitions,
+      reported
+    )
+  }
+
+  private def instantiateFactory(fqcn: String): FfiProviderFactory = {
+    val cls = Class.forName(fqcn)
+    cls.getDeclaredConstructor().newInstance().asInstanceOf[FfiProviderFactory]
+  }
 }