sunchao commented on code in PR #7123:
URL: https://github.com/apache/datafusion-comet/pull/7123#discussion_r4324696880
##########
spark/src/test/scala/org/apache/comet/CometBroadcastBuildSideSuite.scala:
##########
@@ -249,4 +262,709 @@ class CometBroadcastBuildSideSuite extends CometTestBase
with CometIcebergTestBa
}
}
}
+
+ 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.convert.broadcastBuildSide.enabled (default on) and
+ // its Text opt-in (on for this suite) the text leaf is bridged to Arrow
so the broadcast and
+ // the join run natively.
+
assert(CometConf.COMET_CONVERT_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_CONVERT_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("a Text build side keeps Spark's bytes unless its opt-in is on
(#7118)") {
+ // The bridge is on by default, but a Text build side also needs
+ // spark.comet.convert.broadcastBuildSide.text.enabled. A Text file can
hold bytes that are not
+ // valid UTF-8, here a Latin-1 "é" (0xE9) and "ï" (0xEF), and converting
them to Arrow replaces
+ // each ill-formed sequence with U+FFFD. Without the opt-in the join stays
on Spark and keeps
+ // the bytes, as in 1.1.0, and the scan records why it was not converted.
+
assert(CometConf.COMET_CONVERT_BROADCAST_BUILD_SIDE_TEXT_ENABLED.defaultValue.contains(false))
+
assert(CometConf.COMET_CONVERT_FROM_TEXT_ENABLED.defaultValue.contains(false))
+ withTempDir { dir =>
+ val factPath = s"${dir.getAbsolutePath}/fact.parquet"
+ val dimFile = new File(dir, "dim.txt")
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 20)
+ .selectExpr("CAST(id % 2 + 1 AS STRING) AS k", "id AS v")
+ .write
+ .parquet(factPath)
+ }
+ Files.write(
+ dimFile.toPath,
+ "1 caf".getBytes(UTF_8) ++ Array(0xe9.toByte) ++ "\n2
na".getBytes(UTF_8) ++
+ Array(0xef.toByte) ++ "ve\n".getBytes(UTF_8))
+ val query =
+ """SELECT /*+ BROADCAST(d) */ f.v, hex(d.value) AS h, length(d.value)
AS n
+ |FROM fact_7118 f JOIN dim_7118 d ON f.k = substring(d.value, 1,
1)""".stripMargin
+ def hexAndLength(df: DataFrame): Set[Row] = df.select("h",
"n").collect().toSet
+ def bridged(df: DataFrame): Boolean =
+ collect(df.queryExecution.executedPlan) { case s:
CometSparkToColumnarExec => s }.nonEmpty
+ for {
+ aqe <- Seq("false", "true")
+ // V1 FileSourceScanExec, then V2 BatchScanExec (TextScan).
+ v1List <- Seq("parquet,text", "parquet")
+ } {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe,
+ SQLConf.USE_V1_SOURCE_LIST.key -> v1List) {
+ spark.read.parquet(factPath).createOrReplaceTempView("fact_7118")
+
spark.read.text(dimFile.getAbsolutePath).createOrReplaceTempView("dim_7118")
+ val ctx = s"aqe=$aqe v1=$v1List"
+
+
withSQLConf(CometConf.COMET_CONVERT_BROADCAST_BUILD_SIDE_TEXT_ENABLED.key ->
"false") {
+ checkSparkAnswerAndFallbackReason(
+ sql(query),
+ "Comet converts a Text scan on a broadcast build side to Arrow
only when")
+ val df = sql(query)
+ assert(hexAndLength(df) == Set(Row("3120636166E9", 6),
Row("32206E61EF7665", 5)), ctx)
+ assert(!bridged(df), s"$ctx: the Text build side must stay on
Spark")
+ }
+
+ // With the opt-in the bridge converts the Text build side, and each
ill-formed byte comes
+ // back as U+FFFD (EF BF BD), which changes `hex` and `length`. The
compatibility guide
+ // quotes these values.
+ val df = sql(query)
+ assert(
+ hexAndLength(df) == Set(Row("3120636166EFBFBD", 6),
Row("32206E61EFBFBD7665", 7)),
+ ctx)
+ assert(bridged(df), s"$ctx: expected the Text build side to be
bridged")
+
+ // spark.comet.convert.text.enabled converts every Text scan, the
build side included.
+ withSQLConf(
+ CometConf.COMET_CONVERT_BROADCAST_BUILD_SIDE_TEXT_ENABLED.key ->
"false",
+ CometConf.COMET_CONVERT_FROM_TEXT_ENABLED.key -> "true") {
+ val df = sql(query)
+ assert(hexAndLength(df).size == 2, ctx)
+ assert(bridged(df), s"$ctx: expected
spark.comet.convert.text.enabled to convert")
+ }
+ }
+ }
+ }
+ }
+
+ 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_CONVERT_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))
+ // `collect` descends AQE query stages, so the bridge on the deep
text scan is reachable
+ // under AQE too (it lives in the build-side shuffle stage). Assert
it in both modes - and
+ // because checkSparkAnswerAndOperator only inspects the initial
plan under AQE, this is
+ // the only thing that checks the final AQE plan.
+ assert(
+ collect(cometPlan) { case s: CometSparkToColumnarExec => s
}.nonEmpty,
+ s"aqe=$aqe: 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],
Review Comment:
[P2] Remove the stale `QueryStageExec` import from `CometExecSuite.scala:42`
when moving this assertion. This relocation removes its final use in that
suite, but leaves the import behind. Running the standard Spark 3.5
strict-warning build now fails during test compilation with `Unused import`,
and all six Java-lint jobs fail Scalafix on the same import. Removing it from
the original suite restores the build checks while retaining this assertion.
Evidence: Exact-head CI reproduced the failure with `./mvnw -B test-compile
-Pspark-3.5,delta -Pstrict-warnings -DskipTests`: `CometExecSuite.scala:42:
Unused import`, followed by `one error found` and `BUILD FAILURE`. Strict
compilation log:
https://github.com/apache/datafusion-comet/actions/runs/39414444629/job/118601923006.
The Spark 4.1 lint log explicitly proposes removing only `QueryStageExec` from
that import:
https://github.com/apache/datafusion-comet/actions/runs/39414444629/job/118601923427.
All six lint-job logs show the same fix.
--
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]
