feat: run broadcast joins natively when an unsupported scan is on the build side

.github/workflows/pr_build_linux.yml

@@ -470,6 +470,7 @@ jobs:
               org.apache.comet.CometVariantProjectionSuite
               org.apache.spark.sql.CometVariantShreddingSuite
               org.apache.comet.CometIcebergNativeSuite
+              org.apache.comet.CometBroadcastBuildSideSuite
               org.apache.comet.CometIcebergEncryptionSuite
               org.apache.comet.CometIcebergRewriteActionSuite
               org.apache.comet.CometIcebergWriteActionSuite

.github/workflows/pr_build_macos.yml

@@ -172,6 +172,7 @@ jobs:
               org.apache.comet.CometVariantProjectionSuite
               org.apache.spark.sql.CometVariantShreddingSuite
               org.apache.comet.CometIcebergNativeSuite
+              org.apache.comet.CometBroadcastBuildSideSuite
               org.apache.comet.CometIcebergEncryptionSuite
               org.apache.comet.CometIcebergRewriteActionSuite
               org.apache.comet.CometIcebergWriteActionSuite

docs/source/user-guide/latest/datasources.md

@@ -74,6 +74,16 @@ format other than Parquet, JSON and CSV, set `spark.comet.sparkToColumnar.enable
 operator in `spark.comet.sparkToColumnar.supportedOperatorList` by its Spark class name without the
 `Exec` suffix, such as `BatchScan` or `FileSourceScan`.
 
+One case is converted automatically: when the build (broadcast) side of a broadcast join is a file
+or V2 scan Comet cannot read natively - for example a small Text lookup table - and every scan on
+the probe side is already native, Comet converts that build-side scan to Arrow so the whole join
+runs natively over the large probe. This is controlled by
+`spark.comet.convert.broadcastBuildSide.enabled`, which is **on by default**, unlike the switches
+above. The build side is usually small, but an explicit `BROADCAST` hint or a selective filter above
+the scan can make the copied volume larger, and with dynamic partition pruning the dimension is read
+a second time for the pruning subquery. Set `spark.comet.convert.broadcastBuildSide.enabled=false` to
+keep such build sides on Spark.
+
 ### Spark-to-Comet conversion types
 
 Spark-to-Comet conversion supports `ARRAY<STRING>` and `MAP<STRING,STRING>` with binary

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -1219,6 +1219,20 @@ object CometConf extends ShimCometConf {
       .toSequence
       .createWithDefault(Nil)
 
+  val COMET_CONVERT_BROADCAST_BUILD_SIDE_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.convert.broadcastBuildSide.enabled")
+      .category(CATEGORY_EXEC)
+      .doc("When enabled, Comet converts a file or V2 scan that it cannot scan natively (for " +
+        "example a Text source) to Arrow when that scan feeds a broadcast join's build side and " +
+        "every scan on the probe side is already native, so the broadcast join runs natively " +
+        "over the large probe. The build side is usually small (auto-broadcasts are capped by " +
+        "the broadcast threshold), so the row to Arrow copy is usually cheap - though an " +
+        "explicit BROADCAST hint or a selective filter above the scan can make it larger. With " +
+        "dynamic partition pruning the dimension is also read a second time for the pruning " +
+        "subquery. Set this to false to keep such build sides on Spark.")
+      .booleanConf
+      .createWithDefault(true)
+
   val COMET_CASE_CONVERSION_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.caseConversion.enabled")
       .category(CATEGORY_EXEC)

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -27,6 +27,7 @@ import scala.jdk.CollectionConverters._
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.{Expression, SortOrder}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge}
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 import org.apache.spark.sql.catalyst.util.sideBySide
@@ -190,6 +191,17 @@ object CometExecRule {
   private val INPUT_FILE_BLOCK_FALLBACK_REASON: String =
     "Spark to Arrow conversion is not compatible with input_file_name, " +
       "input_file_block_start, or input_file_block_length"
+
+  /**
+   * Tag set on a scan leaf that sits on a broadcast join's build side. When
+   * `spark.comet.convert.broadcastBuildSide.enabled` is on, such a scan is bridged to Arrow with
+   * `CometSparkToColumnarExec` even when the general sparkToColumnar path is off, so the build
+   * branch and the broadcast join can run natively. The build side is usually small, so the row
+   * to Arrow copy is usually cheap - see `tagBroadcastBuildSideLeaves` for the caveats. See issue
+   * #6008.
+   */
+  val BROADCAST_BUILD_SIDE_TAG: org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit] =
+    org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.broadcastBuildSideLeaf")
 }
 
 /**
@@ -611,7 +623,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
             newPlan
               .getTagValue(CometExplainInfo.FALLBACK_REASONS)
               .foreach(reasons => withFallbackReasons(plan, reasons))
-            // return the original plan
+            // The join stays on Spark. Any #6008 build-side bridge we inserted is left under a
+            // plain Spark BroadcastExchangeExec here; clearWastedBuildSideBridges undoes it after
+            // the transform. See issue #6008.
             plan
           }
         } else {
@@ -973,7 +987,19 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
       // during the bottom-up conversion. Tags persist through AQE stage creation.
       tagUnsafePartialAggregates(planWithJoinRewritten)
 
-      var newPlan = revertUnsafePartialAggregates(transform(planWithJoinRewritten))
+      // Tag leaves on broadcast join build sides so shouldApplySparkToColumnar can bridge an
+      // unsupported build-side scan (e.g. Text) to Arrow and let the join run natively (#6008).
+      tagBroadcastBuildSideLeaves(planWithJoinRewritten)
+
+      var transformed = transform(planWithJoinRewritten)
+      // A build-side bridge is only worth keeping when its broadcast join went native. Where the
+      // join stayed on Spark, clearWastedBuildSideBridges clears the tag from those scans and we
+      // transform once more so the build side rebuilds exactly as it would without the feature.
+      // See issue #6008.
+      if (clearWastedBuildSideBridges(transformed)) {
+        transformed = transform(planWithJoinRewritten)
+      }
+      var newPlan = revertUnsafePartialAggregates(transformed)
 
       // if the plan cannot be run fully natively then explain why (when appropriate
       // config is enabled)
@@ -1428,15 +1454,17 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
             case _: CSVFileFormat => CometConf.COMET_CONVERT_FROM_CSV_ENABLED.get(conf)
             case _: JsonFileFormat => CometConf.COMET_CONVERT_FROM_JSON_ENABLED.get(conf)
             case _: ParquetFileFormat => CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.get(conf)
-            case _ => isSparkToArrowEnabled(conf, op)
+            // Unsupported format (e.g. Text): bridge it if the general opt-in is on or it feeds a
+            // broadcast build side. The per-format arms above keep honoring their own opt-outs.
+            case _ => shouldBridgeUnsupportedScan(conf, op)
           }
         // Convert Spark DS v2 scan to Arrow format
         case scan: BatchScanExec =>
           scan.scan match {
             case _: CSVScan => CometConf.COMET_CONVERT_FROM_CSV_ENABLED.get(conf)
             case _: JsonScan => CometConf.COMET_CONVERT_FROM_JSON_ENABLED.get(conf)
             case _: ParquetScan => CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.get(conf)
-            case _ => isSparkToArrowEnabled(conf, op)
+            case _ => shouldBridgeUnsupportedScan(conf, op)
           }
         case _: RangeExec =>
           isConversionEnabled(conf, op, CometConf.COMET_CONVERT_FROM_RANGE_ENABLED)
@@ -1480,6 +1508,199 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
     }
   }
 
+  /**
+   * An unsupported file-source scan is bridged to Arrow when the general sparkToColumnar opt-in
+   * covers it, or when it feeds a broadcast join's build side (see issue #6008). This is only
+   * reached from the unsupported-format arms, so the per-format `COMET_CONVERT_FROM_*` opt-outs
+   * for CSV/JSON/Parquet still take precedence.
+   */
+  private def shouldBridgeUnsupportedScan(conf: SQLConf, op: SparkPlan): Boolean =
+    isSparkToArrowEnabled(conf, op) || isBroadcastBuildSideLeaf(op)
+
+  /**
+   * True when the build-side auto-bridge should run: its own config is on and Comet broadcast
+   * exchange conversion is enabled. Without the latter the `BroadcastExchangeExec` cannot become
+   * a `CometBroadcastExchangeExec`, so bridging the build side would only add a row to Arrow copy
+   * for no native join. This does not also check the join serdes
+   * (COMET_EXEC_BROADCAST_HASH_JOIN_ENABLED / COMET_EXEC_BROADCAST_NESTED_LOOP_JOIN_ENABLED): if
+   * the exchange converts but the join does not, `clearWastedBuildSideBridges` clears the tag
+   * after the transform and a second transform rebuilds the Spark build side, so a leftover copy
+   * never reaches runtime. See issue #6008.
+   */
+  private def broadcastBuildSideBridgeEnabled: Boolean =
+    CometConf.COMET_CONVERT_BROADCAST_BUILD_SIDE_ENABLED.get(conf) &&
+      CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED.get(conf)
+
+  /**
+   * True when `op` is a scan on a broadcast join's build side and build-side bridging is enabled.
+   * `tagBroadcastBuildSideLeaves` only sets the tag when bridging is enabled, but the tag lives
+   * on the physical node and can survive onto a reused instance, so re-check here: a stale tag
+   * must not bridge a scan once the feature is turned off. A tagged scan is bridged to Arrow even
+   * when the general sparkToColumnar path is off, letting the build branch and the join run
+   * natively. See issue #6008.
+   */
+  private def isBroadcastBuildSideLeaf(op: SparkPlan): Boolean =
+    broadcastBuildSideBridgeEnabled &&
+      op.getTagValue(CometExecRule.BROADCAST_BUILD_SIDE_TAG).isDefined
+
+  /**
+   * Tag the file-source scan leaves on a broadcast join's build side so
+   * shouldApplySparkToColumnar can bridge an unsupported build-side scan (e.g. a Text scan) to
+   * Arrow, letting the whole join run natively. The build side is usually small (auto-broadcasts
+   * are capped by the broadcast threshold), so the row to Arrow copy is usually cheap - but that
+   * is a heuristic, not a bound: an explicit BROADCAST hint can force a larger build side, and a
+   * selective filter above the scan means the bridge copies the pre-filter scan output. See issue
+   * #6008.
+   *
+   * The bridge is applied ONLY when the probe side is already natively scannable, i.e. the join
+   * can actually become a fully native CometBroadcastHashJoinExec. Bridging a build side under a
+   * join that stays on Spark is pure overhead and, worse, it rewrites the build broadcast's
+   * subtree to Comet and breaks Spark's DPP broadcast reuse (the DPP subquery keeps an unbridged
+   * copy of the same scan, so the two exchanges no longer share a canonical form). CometScanRule
+   * has already run by this point, so a native probe scan is a Comet scan and an unsupported
+   * probe scan is still a plain Spark scan - see `hasOnlyNativeScans`.
+   *
+   * Only `FileSourceScanExec` / `BatchScanExec` are tagged - the scan types the per-format arms
+   * of shouldApplySparkToColumnar handle. A natively scannable file (e.g. Parquet) is normally
+   * already a `CometScanExec` here; if native scan conversion did not take, its
+   * `FileSourceScanExec` is still tagged, but the tag is harmless because the per-format arm
+   * decides that scan before isBroadcastBuildSideLeaf is ever consulted. DPP subquery broadcasts
+   * live in expressions, not `children`, so `foreach` never reaches them - that scope is
+   * intentional.
+   */
+  private def tagBroadcastBuildSideLeaves(plan: SparkPlan): Unit = {
+    if (!broadcastBuildSideBridgeEnabled) {
+      return
+    }
+    plan.foreach {
+      case j: BroadcastHashJoinExec =>
+        tagBuildIfProbeNative(
+          j.buildSide,
+          j.left,
+          j.right,
+          CometConf.COMET_EXEC_BROADCAST_HASH_JOIN_ENABLED)
+      case j: BroadcastNestedLoopJoinExec =>
+        tagBuildIfProbeNative(
+          j.buildSide,
+          j.left,
+          j.right,
+          CometConf.COMET_EXEC_BROADCAST_NESTED_LOOP_JOIN_ENABLED)
+      case _ =>
+    }
+  }
+
+  /**
+   * Tag the build side's file-source scans, but only when the join looks like it can become a
+   * native Comet broadcast join: the probe side is natively scannable and the join's own serde is
+   * enabled. This is a best-effort gate that avoids inserting a bridge we would only tear down
+   * again - it cannot see every reason a join stays on Spark (an unsupported operator above a
+   * native probe scan, or a join type the serde declines), so `clearWastedBuildSideBridges` is
+   * the backstop that undoes any bridge left under a Spark `BroadcastExchangeExec` after the
+   * transform. Skipping here still matters because a bridge under a Spark broadcast breaks
+   * Spark's DPP broadcast reuse until it is torn down (the DPP subquery keeps an unbridged copy).
+   * We skip when:
+   *   - the join's own serde is disabled (`joinEnabled`), so it will stay on Spark; or
+   *   - the build-side `BroadcastExchangeExec` carries `SKIP_COMET_BROADCAST_TAG`, which
+   *     `CometSpark34AqeDppFallbackRule` sets on Spark 3.4 to keep the broadcast Spark-native so
+   *     `PlanAdaptiveDynamicPruningFilters` can match it (bridging it would turn DPP off). See
+   *     issue #6008.
+   */
+  private def tagBuildIfProbeNative(
+      buildSide: BuildSide,
+      left: SparkPlan,
+      right: SparkPlan,
+      joinEnabled: ConfigEntry[Boolean]): Unit = {
+    if (!joinEnabled.get(conf)) {
+      return
+    }
+    val (buildPlan, probePlan) = buildSide match {
+      case BuildLeft => (left, right)
+      case BuildRight => (right, left)
+    }
+    val skipForDpp = buildPlan match {
+      case b: BroadcastExchangeExec =>
+        b.getTagValue(CometExecRule.SKIP_COMET_BROADCAST_TAG).isDefined
+      case _ => false
+    }
+    if (!skipForDpp && hasOnlyNativeScans(probePlan)) {
+      tagBuildSideScans(buildPlan)
+    }
+  }
+
+  /**
+   * Tag the broadcast dimension's own file-source scans, descending through its filters, projects
+   * and aggregations (incl. the aggregate's shuffle) but stopping at a nested join. A nested
+   * join's inputs are not the bounded broadcast dimension - its streamed side can be far larger
+   * than the broadcast output - so bridging them would copy a large input to Arrow for no reason.
+   * That nested join is tagged by its own pre-pass visit based on its own probe. See issue #6008.
+   */
+  private def tagBuildSideScans(plan: SparkPlan): Unit = plan match {
+    case scan @ (_: FileSourceScanExec | _: BatchScanExec) =>
+      scan.setTagValue(CometExecRule.BROADCAST_BUILD_SIDE_TAG, ())
+    case _: BroadcastHashJoinExec | _: BroadcastNestedLoopJoinExec | _: ShuffledHashJoinExec |
+        _: SortMergeJoinExec =>
+    case other => other.children.foreach(tagBuildSideScans)
+  }
+
+  /**
+   * True when `plan` has at least one scan and none of its scan leaves is an unsupported plain
+   * Spark scan - i.e. every scan is already a Comet scan (CometScanRule ran before this rule).
+   * Used to decide whether a broadcast join's probe side can feed a native join. A probe that
+   * itself contains an unsupported scan (e.g. a DSv2 `InMemoryTableWithV2Filter` fact) keeps the
+   * join on Spark, so bridging its build side would only break DPP reuse.
+   */
+  private def hasOnlyNativeScans(plan: SparkPlan): Boolean = {
+    var sawNativeScan = false
+    var sawUnsupportedScan = false
+    plan.foreach {
+      case _: CometScanExec | _: CometNativeScanExec | _: CometBatchScanExec =>
+        sawNativeScan = true
+      case _: FileSourceScanExec | _: BatchScanExec =>
+        sawUnsupportedScan = true
+      case _ =>
+    }
+    sawNativeScan && !sawUnsupportedScan
+  }
+
+  /**
+   * After the main transform, a #6008 build-side bridge is left under a plain Spark
+   * `BroadcastExchangeExec` only when its broadcast join did not go native - an unsupported join
+   * type (e.g. a full outer `BroadcastNestedLoopJoinExec`), a build side the serde declined, or a
+   * probe that is not fully native (e.g. an unsupported operator above a native scan). Clear
+   * `BROADCAST_BUILD_SIDE_TAG` from every scan under such a broadcast so a second transform
+   * rebuilds that build side exactly as it would without the feature, and so a later AQE stage
+   * pass - which prepares the build stage on its own, without the join above it - does not bridge
+   * the scan again. Returns true if any tag was cleared.
+   *
+   * Two guards keep the walk honest. It stops at a `CometBroadcastExchangeExec` - a nested
+   * broadcast join that went native keeps its bridge. And it only clears a tag (and so asks for a
+   * second transform) when the tagged scan sits under a `CometSparkToColumnarExec`, i.e. it was
+   * actually bridged. A CSV/JSON/Parquet scan whose per-format arm left it on Spark still carries
+   * the tag but was never bridged, so touching it would force a pointless second transform. The
+   * descent enters `CometScanWrapper`'s wrapped `originalPlan` by hand because it is a leaf, so a
+   * plain `foreach` would never reach the bridged scan inside it. See issue #6008.
+   */
+  private def clearWastedBuildSideBridges(plan: SparkPlan): Boolean = {
+    var cleared = false
+    def clearScanTags(node: SparkPlan, underBridge: Boolean): Unit = node match {
+      // A nested broadcast join that went native keeps its bridge.
+      case _: CometBroadcastExchangeExec =>
+      case w: CometScanWrapper => clearScanTags(w.originalPlan, underBridge)
+      case bridge: CometSparkToColumnarExec => clearScanTags(bridge.child, underBridge = true)
+      case other =>
+        if (underBridge && other.getTagValue(CometExecRule.BROADCAST_BUILD_SIDE_TAG).isDefined) {
+          other.unsetTagValue(CometExecRule.BROADCAST_BUILD_SIDE_TAG)
+          cleared = true
+        }
+        other.children.foreach(clearScanTags(_, underBridge))
+    }
+    plan.foreach {
+      case b: BroadcastExchangeExec => clearScanTags(b, underBridge = false)
+      case _ =>
+    }
+    cleared
+  }
+
   private def isSparkToArrowEnabled(conf: SQLConf, op: SparkPlan) = {
     COMET_SPARK_TO_ARROW_ENABLED.get(conf) && {
       val simpleClassName = Utils.getSimpleName(op.getClass)