parthchandra commented on code in PR #6786:
URL: https://github.com/apache/datafusion-comet/pull/6786#discussion_r4235358846
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1518,6 +1535,126 @@ 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, the broadcast arm returns
the original plan and
+ * EliminateRedundantTransitions removes the now-unused bridge, so a
leftover copy never reaches
Review Comment:
Agreed. When the join doesn't convert, the rule now reverts the build side
to its original Spark plan (revertBridgedBroadcasts). That handles the
FULL/LEFT OUTER and broadcastHashJoin.enabled=false cases, and the 3.4 case
too, so the dimension isn't broadcast twice.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1518,6 +1535,126 @@ 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, the broadcast arm returns
the original plan and
+ * EliminateRedundantTransitions removes the now-unused bridge, so a
leftover copy never reaches
+ * runtime. See issue #6008.
+ */
+ private def broadcastBuildSideBridgeEnabled: Boolean =
+ CometConf.COMET_SPARK_TO_ARROW_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)
+ case j: BroadcastNestedLoopJoinExec =>
tagBuildIfProbeNative(j.buildSide, j.left, j.right)
+ case _ =>
+ }
+ }
+
+ /** Tag the build side's file-source scans, but only if the probe side is
natively scannable. */
+ private def tagBuildIfProbeNative(
+ buildSide: BuildSide,
+ left: SparkPlan,
+ right: SparkPlan): Unit = {
+ val (buildPlan, probePlan) = buildSide match {
+ case BuildLeft => (left, right)
+ case BuildRight => (right, left)
+ }
+ if (hasOnlyNativeScans(probePlan)) {
Review Comment:
Good catch. tagBuildIfProbeNative now skips a build side whose
BroadcastExchangeExec has SKIP_COMET_BROADCAST_TAG, so on 3.4 the text scan
isn't bridged and DPP stays on.
##########
spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala:
##########
@@ -351,6 +360,26 @@ case object CometPlanAdaptiveDynamicPruningFilters
}
}
+ /**
+ * True when `plan` (a DPP subquery's build, taken from an
AdaptiveSparkPlanExec.executedPlan)
+ * produces native Comet columnar output, so it is safe as a
CometBroadcastExchangeExec child.
+ * Unwraps AQE stage wrappers and any columnar->row transition (mirrors the
stripping in
+ * CometExecRule.rewriteInSubqueryPlan), then checks for a CometNativeExec.
See issue #6008.
+ */
+ private def isNativeBuildSide(plan: SparkPlan): Boolean = {
+ val stripped = stripAQEPlan(plan) match {
+ case c2r: CometNativeColumnarToRowExec => c2r.child
+ case c2r: CometColumnarToRowExec => c2r.child
+ case WholeStageCodegenExec(c2r: CometColumnarToRowExec) =>
+ c2r.child match {
+ case InputAdapter(child) => child
+ case other => other
+ }
+ case other => other
+ }
+ stripped.isInstanceOf[CometNativeExec]
Review Comment:
Fixed — changed it to stripped.isInstanceOf[CometPlan] &&
stripped.supportsColumnar, so
CometUnion/CometCoalesce/CometTakeOrderedAndProject are accepted and the
UNION-ALL dimension keeps its reuse. Thanks for filing #6815 for the non-AQE
twin.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -301,6 +301,366 @@ class CometExecSuite extends CometTestBase {
}
}
+ test("broadcast build side with unsupported text source goes native
(#6008)") {
Review Comment:
Yes please, I'll take the probe. I can add an ORC build-side test (nulls
across types, struct, ARRAY<STRING>, MAP, timestamp/ntz, LA timezone) and an
Iceberg-declined one from it.
##########
spark/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -1166,6 +1166,20 @@ object CometConf extends ShimCometConf {
.toSequence
.createWithDefault(Nil)
+ val COMET_SPARK_TO_ARROW_BROADCAST_BUILD_SIDE_ENABLED: ConfigEntry[Boolean] =
+ conf("spark.comet.sparkToColumnar.broadcastBuildSide.enabled")
Review Comment:
Renamed to `spark.comet.convert.broadcastBuildSide.enabled`. Reworded the
doc too: it only covers file/V2 scans, only when every probe-side scan is
native, and notes the dimension is read twice under DPP.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1518,6 +1535,126 @@ 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, the broadcast arm returns
the original plan and
+ * EliminateRedundantTransitions removes the now-unused bridge, so a
leftover copy never reaches
+ * runtime. See issue #6008.
+ */
+ private def broadcastBuildSideBridgeEnabled: Boolean =
+ CometConf.COMET_SPARK_TO_ARROW_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)
+ case j: BroadcastNestedLoopJoinExec =>
tagBuildIfProbeNative(j.buildSide, j.left, j.right)
+ case _ =>
+ }
+ }
+
+ /** Tag the build side's file-source scans, but only if the probe side is
natively scannable. */
+ private def tagBuildIfProbeNative(
+ buildSide: BuildSide,
+ left: SparkPlan,
+ right: SparkPlan): Unit = {
+ val (buildPlan, probePlan) = buildSide match {
+ case BuildLeft => (left, right)
+ case BuildRight => (right, left)
+ }
+ if (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 = {
Review Comment:
I wanted to keep this minimal so I would say this is out of scope for this
PR, opened https://github.com/apache/datafusion-comet/issues/6832 to track it.
It's a missed optimization (only the inner join goes native), not wrong results.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -301,6 +301,366 @@ class CometExecSuite extends CometTestBase {
}
}
+ test("broadcast build side with unsupported text source goes native
(#6008)") {
+ // A tiny lookup table read from a Text file (which Comet cannot scan
natively) sits on the
+ // build side of a broadcast join over a large native probe. Without the
auto-bridge the whole
+ // join stays on Spark. With
spark.comet.sparkToColumnar.broadcastBuildSide.enabled (default on)
+ // the text leaf is bridged to Arrow so the broadcast and the join run
natively.
+ assert(
+
CometConf.COMET_SPARK_TO_ARROW_BROADCAST_BUILD_SIDE_ENABLED.defaultValue.contains(true))
+ withTempDir { dir =>
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ spark
+ .range(0, 5)
+ .selectExpr("CAST(id AS STRING) AS value")
+ .write
+ .text(allowPath)
+ }
+
+ // Cover both the AQE and non-AQE planning paths - the tag pre-pass runs
per stage under AQE.
+ for (aqe <- Seq(false, true)) {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ spark.read.parquet(probePath).createOrReplaceTempView("probe_6008")
+ spark.read.text(allowPath).createOrReplaceTempView("allow_6008")
+ val query =
+ "SELECT /*+ BROADCAST(a) */ p.v FROM probe_6008 p JOIN allow_6008
a ON p.key = a.value"
+
+ // Default: the text build side is bridged to Arrow and the join
goes native. The plan is
+ // fully native (the bridge counts as a native leaf), so use
checkSparkAnswerAndOperator -
+ // checkSparkAnswer would fail under COMET_STRICT_TESTING on a
fully-native plan.
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query))
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s
}.nonEmpty,
+ s"aqe=$aqe: expected a CometSparkToColumnarExec on the text build
side:\n" +
+ cometPlan.treeString)
+ assert(
+ collect(cometPlan) { case b: CometBroadcastExchangeExec => b
}.nonEmpty,
+ s"aqe=$aqe: expected a
CometBroadcastExchangeExec:\n${cometPlan.treeString}")
+ assert(
+ collect(cometPlan) { case j: CometBroadcastHashJoinExec => j
}.nonEmpty,
+ s"aqe=$aqe: expected a
CometBroadcastHashJoinExec:\n${cometPlan.treeString}")
+
+ // Off: the unsupported text leaf blocks native execution of the
broadcast and the join.
+ withSQLConf(
+ CometConf.COMET_SPARK_TO_ARROW_BROADCAST_BUILD_SIDE_ENABLED.key ->
"false") {
+ val (_, cometPlanOff) = checkSparkAnswer(sql(query))
+ assert(
+ collect(cometPlanOff) { case s: CometSparkToColumnarExec => s
}.isEmpty,
+ s"aqe=$aqe: expected no CometSparkToColumnarExec when
disabled:\n" +
+ cometPlanOff.treeString)
+ assert(
+ collect(cometPlanOff) { case b: CometBroadcastExchangeExec => b
}.isEmpty,
+ s"aqe=$aqe: expected no CometBroadcastExchangeExec when
disabled:\n" +
+ cometPlanOff.treeString)
+ assert(
+ collect(cometPlanOff) { case j: CometBroadcastHashJoinExec => j
}.isEmpty,
+ s"aqe=$aqe: expected no CometBroadcastHashJoinExec when
disabled:\n" +
+ cometPlanOff.treeString)
+ }
+ }
+ }
+ }
+ }
+
+ test("broadcast build side bridge is gated on broadcast-exchange conversion
(#6008)") {
+ // The bridge is pointless when Comet broadcast exchange conversion is
off: the broadcast (and
+ // so the join) cannot go native. With
COMET_EXEC_BROADCAST_EXCHANGE_ENABLED=false the text
+ // build side must not be bridged even though broadcastBuildSide is on.
+ withTempDir { dir =>
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ spark.range(0, 5).selectExpr("CAST(id AS STRING) AS
value").write.text(allowPath)
+ }
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_SPARK_TO_ARROW_BROADCAST_BUILD_SIDE_ENABLED.key ->
"true",
+ CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED.key -> "false") {
+
spark.read.parquet(probePath).createOrReplaceTempView("probe_gate_6008")
+ spark.read.text(allowPath).createOrReplaceTempView("allow_gate_6008")
+ val (_, cometPlan) = checkSparkAnswer(sql("""
+ |SELECT /*+ BROADCAST(a) */ p.v
+ |FROM probe_gate_6008 p JOIN allow_gate_6008 a ON p.key = a.value
+ |""".stripMargin))
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s }.isEmpty,
+ "with broadcast-exchange conversion off the build side must not be
bridged:\n" +
+ cometPlan.treeString)
+ }
+ }
+ }
+
+ test("broadcast build side scan below an aggregate is bridged, shuffle stage
is not (#6008)") {
+ // A file scan can sit BELOW build-side operators. Broadcasting an
aggregated Text table puts a
+ // Text FileSourceScanExec under a partial aggregate + shuffle on the
build side. The tag
+ // pre-pass walks the whole broadcast subtree, so the deep text scan must
be bridged to Arrow,
+ // yet no query stage (e.g. the build-side shuffle under AQE) may be
wrapped in the bridge.
+ withTempDir { dir =>
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.range(0, 20).selectExpr("CAST(id % 5 AS STRING) AS
value").write.text(allowPath)
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id % 5 AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ }
+ spark.read.text(allowPath).createOrReplaceTempView("allow_agg_6008")
+ spark.read.parquet(probePath).createOrReplaceTempView("probe_agg_6008")
+ val query = """
+ |SELECT /*+ BROADCAST(d) */ p.v, d.n
+ |FROM probe_agg_6008 p
+ |JOIN (SELECT value AS k, COUNT(*) AS n FROM allow_agg_6008 GROUP BY
value) d
+ | ON p.key = d.k
+ |""".stripMargin
+ for (aqe <- Seq(false, true)) {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> (10 * 1024 *
1024).toString) {
+ // The plan is fully native (the bridged text scan under the
aggregate is a native leaf),
+ // so checkSparkAnswerAndOperator also pins that the text scan did
get bridged - an
+ // un-bridged Spark scan would read as non-native and fail here. It
is also strict-testing
+ // safe, unlike checkSparkAnswer on a fully-native plan.
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query))
+ // With AQE off the whole build subtree is visible in one tree, so
the bridge on the deep
+ // text scan is collectable. (Under AQE it lives inside a
materialized shuffle input stage
+ // and may not be reachable from the final plan, so only assert it
when AQE is off.)
+ if (!aqe) {
Review Comment:
Done — dropped it; the bridge is asserted in both AQE modes now, so the
final AQE plan is covered.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -301,6 +301,366 @@ class CometExecSuite extends CometTestBase {
}
}
+ test("broadcast build side with unsupported text source goes native
(#6008)") {
+ // A tiny lookup table read from a Text file (which Comet cannot scan
natively) sits on the
+ // build side of a broadcast join over a large native probe. Without the
auto-bridge the whole
+ // join stays on Spark. With
spark.comet.sparkToColumnar.broadcastBuildSide.enabled (default on)
+ // the text leaf is bridged to Arrow so the broadcast and the join run
natively.
+ assert(
+
CometConf.COMET_SPARK_TO_ARROW_BROADCAST_BUILD_SIDE_ENABLED.defaultValue.contains(true))
+ withTempDir { dir =>
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ spark
+ .range(0, 5)
+ .selectExpr("CAST(id AS STRING) AS value")
+ .write
+ .text(allowPath)
+ }
+
+ // Cover both the AQE and non-AQE planning paths - the tag pre-pass runs
per stage under AQE.
+ for (aqe <- Seq(false, true)) {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ spark.read.parquet(probePath).createOrReplaceTempView("probe_6008")
+ spark.read.text(allowPath).createOrReplaceTempView("allow_6008")
+ val query =
+ "SELECT /*+ BROADCAST(a) */ p.v FROM probe_6008 p JOIN allow_6008
a ON p.key = a.value"
+
+ // Default: the text build side is bridged to Arrow and the join
goes native. The plan is
+ // fully native (the bridge counts as a native leaf), so use
checkSparkAnswerAndOperator -
+ // checkSparkAnswer would fail under COMET_STRICT_TESTING on a
fully-native plan.
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query))
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s
}.nonEmpty,
+ s"aqe=$aqe: expected a CometSparkToColumnarExec on the text build
side:\n" +
+ cometPlan.treeString)
+ assert(
+ collect(cometPlan) { case b: CometBroadcastExchangeExec => b
}.nonEmpty,
+ s"aqe=$aqe: expected a
CometBroadcastExchangeExec:\n${cometPlan.treeString}")
+ assert(
+ collect(cometPlan) { case j: CometBroadcastHashJoinExec => j
}.nonEmpty,
+ s"aqe=$aqe: expected a
CometBroadcastHashJoinExec:\n${cometPlan.treeString}")
+
+ // Off: the unsupported text leaf blocks native execution of the
broadcast and the join.
+ withSQLConf(
+ CometConf.COMET_SPARK_TO_ARROW_BROADCAST_BUILD_SIDE_ENABLED.key ->
"false") {
+ val (_, cometPlanOff) = checkSparkAnswer(sql(query))
+ assert(
+ collect(cometPlanOff) { case s: CometSparkToColumnarExec => s
}.isEmpty,
+ s"aqe=$aqe: expected no CometSparkToColumnarExec when
disabled:\n" +
+ cometPlanOff.treeString)
+ assert(
+ collect(cometPlanOff) { case b: CometBroadcastExchangeExec => b
}.isEmpty,
+ s"aqe=$aqe: expected no CometBroadcastExchangeExec when
disabled:\n" +
+ cometPlanOff.treeString)
+ assert(
+ collect(cometPlanOff) { case j: CometBroadcastHashJoinExec => j
}.isEmpty,
+ s"aqe=$aqe: expected no CometBroadcastHashJoinExec when
disabled:\n" +
+ cometPlanOff.treeString)
+ }
+ }
+ }
+ }
+ }
+
+ test("broadcast build side bridge is gated on broadcast-exchange conversion
(#6008)") {
+ // The bridge is pointless when Comet broadcast exchange conversion is
off: the broadcast (and
+ // so the join) cannot go native. With
COMET_EXEC_BROADCAST_EXCHANGE_ENABLED=false the text
+ // build side must not be bridged even though broadcastBuildSide is on.
+ withTempDir { dir =>
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ spark.range(0, 5).selectExpr("CAST(id AS STRING) AS
value").write.text(allowPath)
+ }
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_SPARK_TO_ARROW_BROADCAST_BUILD_SIDE_ENABLED.key ->
"true",
+ CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED.key -> "false") {
+
spark.read.parquet(probePath).createOrReplaceTempView("probe_gate_6008")
+ spark.read.text(allowPath).createOrReplaceTempView("allow_gate_6008")
+ val (_, cometPlan) = checkSparkAnswer(sql("""
+ |SELECT /*+ BROADCAST(a) */ p.v
+ |FROM probe_gate_6008 p JOIN allow_gate_6008 a ON p.key = a.value
+ |""".stripMargin))
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s }.isEmpty,
+ "with broadcast-exchange conversion off the build side must not be
bridged:\n" +
+ cometPlan.treeString)
+ }
+ }
+ }
+
+ test("broadcast build side scan below an aggregate is bridged, shuffle stage
is not (#6008)") {
+ // A file scan can sit BELOW build-side operators. Broadcasting an
aggregated Text table puts a
+ // Text FileSourceScanExec under a partial aggregate + shuffle on the
build side. The tag
+ // pre-pass walks the whole broadcast subtree, so the deep text scan must
be bridged to Arrow,
+ // yet no query stage (e.g. the build-side shuffle under AQE) may be
wrapped in the bridge.
+ withTempDir { dir =>
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.range(0, 20).selectExpr("CAST(id % 5 AS STRING) AS
value").write.text(allowPath)
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id % 5 AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ }
+ spark.read.text(allowPath).createOrReplaceTempView("allow_agg_6008")
+ spark.read.parquet(probePath).createOrReplaceTempView("probe_agg_6008")
+ val query = """
+ |SELECT /*+ BROADCAST(d) */ p.v, d.n
+ |FROM probe_agg_6008 p
+ |JOIN (SELECT value AS k, COUNT(*) AS n FROM allow_agg_6008 GROUP BY
value) d
+ | ON p.key = d.k
+ |""".stripMargin
+ for (aqe <- Seq(false, true)) {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> (10 * 1024 *
1024).toString) {
+ // The plan is fully native (the bridged text scan under the
aggregate is a native leaf),
+ // so checkSparkAnswerAndOperator also pins that the text scan did
get bridged - an
+ // un-bridged Spark scan would read as non-native and fail here. It
is also strict-testing
+ // safe, unlike checkSparkAnswer on a fully-native plan.
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query))
+ // With AQE off the whole build subtree is visible in one tree, so
the bridge on the deep
+ // text scan is collectable. (Under AQE it lives inside a
materialized shuffle input stage
+ // and may not be reachable from the final plan, so only assert it
when AQE is off.)
+ if (!aqe) {
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s
}.nonEmpty,
+ s"expected the build-side text scan to be
bridged:\n${cometPlan.treeString}")
+ }
+ // A query stage (the build-side shuffle) must never be wrapped in
the bridge.
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s.child
}.foreach { child =>
+ assert(
+ !child.isInstanceOf[QueryStageExec],
+ s"aqe=$aqe: a QueryStageExec must not be bridged to
Arrow:\n${cometPlan.treeString}")
+ }
+ }
+ }
+ }
+ }
+
+ test("per-format opt-out still wins on a broadcast build side (#6008)") {
+ // A CSV build side stays a Spark FileSourceScanExec and is tagged, but
the CSV arm of
+ // shouldApplySparkToColumnar honors COMET_CONVERT_FROM_CSV_ENABLED and
never calls the
+ // broadcast-build-side bypass. So with convert-from-csv off, the CSV
build side must NOT be
+ // bridged and the join must stay on Spark, even with broadcastBuildSide
enabled.
+ withTempDir { dir =>
+ val allowPath = s"${dir.getAbsolutePath}/allow.csv"
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.range(0, 5).selectExpr("CAST(id AS STRING) AS
value").write.csv(allowPath)
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ }
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_CONVERT_FROM_CSV_ENABLED.key -> "false",
+ CometConf.COMET_SPARK_TO_ARROW_BROADCAST_BUILD_SIDE_ENABLED.key ->
"true") {
+ spark.read.csv(allowPath).createOrReplaceTempView("allow_csv_6008")
+ spark.read.parquet(probePath).createOrReplaceTempView("probe_csv_6008")
+ val (_, cometPlan) = checkSparkAnswer(sql("""
+ |SELECT /*+ BROADCAST(a) */ p.v
+ |FROM probe_csv_6008 p JOIN allow_csv_6008 a ON p.key = a._c0
+ |""".stripMargin))
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s }.isEmpty,
+ "a CSV build side with convert-from-csv disabled must not be
bridged:\n" +
+ cometPlan.treeString)
+ }
+ }
+ }
+
+ test("broadcast build side with an unsupported DSv2 scan goes native
(#6008)") {
+ // Exercises the BatchScanExec arm: an unsupported-format scan (Text read
as DSv2) on the build
+ // side is bridged so the join goes native. Keep parquet on the V1 list so
the probe stays a
+ // native CometScan; drop text from the list so it is read as a DSv2
BatchScan/TextScan.
+ withTempDir { dir =>
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.range(0, 5).selectExpr("CAST(id AS STRING) AS
value").write.text(allowPath)
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ }
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.USE_V1_SOURCE_LIST.key -> "parquet") {
+ spark.read.text(allowPath).createOrReplaceTempView("allow_v2_6008")
+ spark.read.parquet(probePath).createOrReplaceTempView("probe_v2_6008")
+ // Fully native (bridged DSv2 text build side + native parquet probe +
native join), so use
+ // checkSparkAnswerAndOperator - strict-testing safe, unlike
checkSparkAnswer here.
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql("""
+ |SELECT /*+ BROADCAST(a) */ p.v
+ |FROM probe_v2_6008 p JOIN allow_v2_6008 a ON p.key = a.value
+ |""".stripMargin))
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s
}.nonEmpty,
+ s"expected the DSv2 text build side to be
bridged:\n${cometPlan.treeString}")
+ assert(
+ collect(cometPlan) { case j: CometBroadcastHashJoinExec => j
}.nonEmpty,
+ s"expected the join to go native over the DSv2 build
side:\n${cometPlan.treeString}")
+ }
+ }
+ }
+
+ test("broadcast join with an unsupported probe stays on Spark (#6008)") {
+ // Both sides are unsupported Text sources. The probe is not natively
scannable, so the join
+ // cannot go native - and because the probe is not native, the build side
must NOT be bridged
+ // either (bridging it would be useless and would break Spark's
broadcast/DPP reuse). This pins
+ // the probe-native gate that fixed the DynamicPartitionPruningV2
regression.
+ withTempDir { dir =>
+ val probePath = s"${dir.getAbsolutePath}/probe.txt"
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.range(0, 100).selectExpr("CAST(id AS STRING) AS
value").write.text(probePath)
+ spark.range(0, 5).selectExpr("CAST(id AS STRING) AS
value").write.text(allowPath)
+ }
+ for (aqe <- Seq(false, true)) {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ spark.read.text(probePath).createOrReplaceTempView("probe_txt_6008")
+ spark.read.text(allowPath).createOrReplaceTempView("allow_txt_6008")
+ val (_, cometPlan) = checkSparkAnswer(sql("""
+ |SELECT /*+ BROADCAST(a) */ p.value
+ |FROM probe_txt_6008 p JOIN allow_txt_6008 a ON p.value = a.value
+ |""".stripMargin))
+ // The non-native probe must leave the build side un-bridged (the
probe-native gate).
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s
}.isEmpty,
+ s"aqe=$aqe: a non-native probe must leave the build side
un-bridged:\n" +
+ cometPlan.treeString)
+ assert(
+ collect(cometPlan) { case j: CometBroadcastHashJoinExec => j
}.isEmpty,
+ s"aqe=$aqe: the join must stay on Spark when the probe is not
native:\n" +
+ cometPlan.treeString)
+ assert(
+ collect(cometPlan) { case b: CometBroadcastExchangeExec => b
}.isEmpty,
+ s"aqe=$aqe: the broadcast must stay on Spark when the join does
not convert:\n" +
+ cometPlan.treeString)
+ }
+ }
+ }
+ }
+
+ test("broadcast build side on the left (BuildLeft) with a text source goes
native (#6008)") {
+ // All other tests broadcast the right table (BuildRight). Put the small
text table first and
+ // hint it so it becomes the left build side, exercising the BuildLeft arm
of tagBuildIfProbeNative.
+ withTempDir { dir =>
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id AS STRING) AS key", "id AS v")
+ .write
+ .parquet(probePath)
+ spark.range(0, 5).selectExpr("CAST(id AS STRING) AS
value").write.text(allowPath)
+ }
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ spark.read.text(allowPath).createOrReplaceTempView("allow_bl_6008")
+ spark.read.parquet(probePath).createOrReplaceTempView("probe_bl_6008")
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql("""
+ |SELECT /*+ BROADCAST(a) */ p.v
+ |FROM allow_bl_6008 a JOIN probe_bl_6008 p ON a.value = p.key
+ |""".stripMargin))
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s
}.nonEmpty,
+ s"expected the text build side to be
bridged:\n${cometPlan.treeString}")
+ assert(
+ collect(cometPlan) { case j: CometBroadcastHashJoinExec => j
}.nonEmpty,
+ s"expected a CometBroadcastHashJoinExec:\n${cometPlan.treeString}")
+ }
+ }
+ }
+
+ test("broadcast nested loop join build side with a text source goes native
(#6008)") {
+ // Exercises the BroadcastNestedLoopJoinExec arm of
tagBroadcastBuildSideLeaves: an inequality
+ // join forces a BNLJ; the text build side is bridged so the BNLJ runs
natively.
+ withTempDir { dir =>
+ val probePath = s"${dir.getAbsolutePath}/probe.parquet"
+ val allowPath = s"${dir.getAbsolutePath}/allow.txt"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.range(0, 100).selectExpr("id AS v", "CAST(id AS INT) AS
k").write.parquet(probePath)
+ spark.range(0, 5).selectExpr("CAST(id AS STRING) AS
value").write.text(allowPath)
+ }
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ spark.read.text(allowPath).createOrReplaceTempView("allow_bnlj_6008")
+
spark.read.parquet(probePath).createOrReplaceTempView("probe_bnlj_6008")
+ val (_, cometPlan) = checkSparkAnswer(sql("""
+ |SELECT /*+ BROADCAST(a) */ p.v
+ |FROM probe_bnlj_6008 p JOIN allow_bnlj_6008 a ON p.k >
CAST(a.value AS INT)
+ |""".stripMargin))
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s
}.nonEmpty,
+ s"expected the text build side to be
bridged:\n${cometPlan.treeString}")
+ assert(
+ collect(cometPlan) { case j: CometBroadcastNestedLoopJoinExec => j
}.nonEmpty,
+ s"expected a
CometBroadcastNestedLoopJoinExec:\n${cometPlan.treeString}")
+ }
+ }
+ }
+
+ test("AQE DPP with a text dimension on the broadcast build side runs
natively (#6008)") {
+ // A Text dimension on a broadcast build side joined to a PARTITIONED
Parquet fact triggers DPP.
+ // The build-side bridge makes the join a CometBroadcastHashJoinExec; the
DPP subquery's own
+ // build (the text scan) is not bridged, so
CometPlanAdaptiveDynamicPruningFilters must NOT wrap
+ // it in a CometBroadcastExchangeExec (that crashed with "Comet execution
only takes Arrow
+ // Arrays"). It falls back to a Spark DPP broadcast while the join stays
native. See issue #6008.
+ assume(isSpark35Plus, "Native AQE DPP requires Spark 3.5+")
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
+ withTempDir { dir =>
+ val factPath = s"${dir.getAbsolutePath}/fact.parquet"
+ val dimPath = s"${dir.getAbsolutePath}/dim.txt"
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 100)
+ .selectExpr("CAST(id % 10 AS STRING) AS fk", "id AS v")
+ .write
+ .partitionBy("fk")
+ .parquet(factPath)
+ spark.range(0, 10).selectExpr("CAST(id AS STRING) AS
value").write.text(dimPath)
+ }
+ spark.read.parquet(factPath).createOrReplaceTempView("dpp_fact_txt")
+ spark.read.text(dimPath).createOrReplaceTempView("dpp_dim_txt")
+ val df = sql("""
+ |SELECT /*+ BROADCAST(d) */ f.v
+ |FROM dpp_fact_txt f JOIN dpp_dim_txt d ON f.fk = d.value
+ |WHERE d.value < '5'
Review Comment:
Right. Filter is now CAST(d.value AS INT) % 2 = 0 so Spark can't push it to
the fact, and the test asserts the fact scan still has a live DPP
SubqueryBroadcast.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
