dwsmith1983 commented on code in PR #5880:
URL: https://github.com/apache/datafusion-comet/pull/5880#discussion_r3996929100
##########
spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala:
##########
@@ -1012,6 +1090,288 @@ class CometTaskMetricsSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("native scan left unconsumed by a limit still reports task input
metrics") {
+ withTempPath { dir =>
+ spark
+ .createDataFrame((0 until 20000).map(i => (i, s"big_$i")))
+ .repartition(4)
+ .write
+ .parquet(dir.getAbsolutePath)
+
spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("limit_big")
+ spark
+ .createDataFrame((0 until 100).map(i => (i * 7, s"local_$i")))
+ .createOrReplaceTempView("limit_local")
+
+ // The broadcast side is a JVM input to the join block, so native
execution polls and only
+ // publishes scan metrics on the update interval. The limit stops
pulling before the scan
+ // is exhausted, leaving the final publish to the iterator's
completion-time close.
+ val query = "SELECT /*+ BROADCAST(limit_local) */ * FROM limit_big JOIN
limit_local " +
+ "ON limit_big._1 = limit_local._1 LIMIT 3"
+ val localBroadcast = Seq(
+ CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true",
+ CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key ->
"LocalTableScan")
+
+ Seq("-1",
CometConf.COMET_METRICS_UPDATE_INTERVAL.defaultValueString).foreach { interval
=>
+ val confs = localBroadcast :+
(CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> interval)
+ val (cometBytes, cometRecords, cometPlan) =
+ collectInputMetrics(query, (CometConf.COMET_ENABLED.key -> "true")
+: confs: _*)
+
+ val join = find(cometPlan)(_.isInstanceOf[CometBroadcastHashJoinExec])
+ assert(
+ join.isDefined,
+ s"Expected CometBroadcastHashJoinExec in
plan:\n${cometPlan.treeString}")
+ val streamedNativeScan = join.get.children.exists { child =>
+ find(child)(_.isInstanceOf[CometBroadcastExchangeExec]).isEmpty &&
+ find(child)(_.isInstanceOf[CometNativeScanExec]).isDefined
+ }
+ assert(
+ streamedNativeScan,
+ s"Expected a native scan on the streamed side of the
join:\n${cometPlan.treeString}")
+
+ assert(cometBytes > 0, s"bytesRead should be > 0 at interval
$interval, got $cometBytes")
+ assert(
+ cometRecords >= 3 && cometRecords <= 20000,
+ s"recordsRead should cover at least the limit at interval $interval,
got $cometRecords")
+ }
+ }
+ }
+
+ test("task input metrics keep bytes read by a fallback Spark scan in the
same task") {
+ withTempPath { parquetDir =>
+ withTempPath { jsonDir =>
+ spark
+ .createDataFrame((0 until 5000).map(i => (i, s"parquet_$i")))
+ .repartition(1)
+ .write
+ .parquet(parquetDir.getAbsolutePath)
+ spark
+ .createDataFrame((5000 until 10000).map(i => (i, s"json_$i")))
+ .repartition(1)
+ .write
+ .json(jsonDir.getAbsolutePath)
+
spark.read.parquet(parquetDir.getAbsolutePath).createOrReplaceTempView("mixed_parquet")
+
spark.read.json(jsonDir.getAbsolutePath).createOrReplaceTempView("mixed_json")
+
+ // Coalescing the union to one partition computes the native scan and
the fallback JSON
+ // scan inside the same task, so Spark's own input metrics and Comet's
must add up.
+ val query = "SELECT /*+ COALESCE(1) */ * FROM (SELECT _1 FROM
mixed_parquet " +
+ "UNION ALL SELECT CAST(_1 AS INT) FROM mixed_json)"
+ val convertJson = CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key ->
"true"
+
+ val (sparkBytes, sparkRecords, _) =
+ collectInputMetrics(query, CometConf.COMET_ENABLED.key -> "false",
convertJson)
+ val (cometBytes, cometRecords, cometPlan) =
+ collectInputMetrics(query, CometConf.COMET_ENABLED.key -> "true",
convertJson)
+
+ assert(
+ find(cometPlan)(_.isInstanceOf[CometNativeScanExec]).isDefined,
+ s"Expected CometNativeScanExec in plan:\n${cometPlan.treeString}")
+ assert(
+ find(cometPlan)(_.isInstanceOf[CometSparkToColumnarExec]).isDefined,
+ s"Expected CometSparkToColumnarExec in
plan:\n${cometPlan.treeString}")
+
+ assert(sparkRecords > 0, s"Spark recordsRead should be > 0, got
$sparkRecords")
+ assert(
+ cometRecords == sparkRecords,
+ s"recordsRead mismatch: comet=$cometRecords, spark=$sparkRecords")
+ assert(sparkBytes > 0, s"Spark bytesRead should be > 0, got
$sparkBytes")
+ assertCometBytesReadInRange(cometBytes, sparkBytes)
Review Comment:
> Could this measure the two sides separately first and then assert the
coalesced run is at least their sum?
Done for the fallback, coalesced-union and cached-input tests. Each side is
measured on its own with Comet enabled and the coalesced run must be at least
the sum. With inc reverted to set and the records assertions muted, the bytes
assertions fail on their own: the fallback test reports comet=20033 against
native=20033 plus fallback=146148, the cached test comet=45277 against
parquet=45277 plus cached=82280, and the two-scan union drops to a 0.47 ratio.
--
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]
