You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/11 12:56:52 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6836-9a7b1c1b99ad9a1ddd8657b659723f6fdacf35dd created (now 8081c7c0fb)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6836-9a7b1c1b99ad9a1ddd8657b659723f6fdacf35dd
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 8081c7c0fb test: count the partitions and batches the native cache scan reads (#6836)
This branch includes the following new commits:
new 8081c7c0fb test: count the partitions and batches the native cache scan reads (#6836)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: test: count the partitions and batches the native cache scan reads (#6836)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6836-9a7b1c1b99ad9a1ddd8657b659723f6fdacf35dd
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 8081c7c0fb13f60962ca9e069db75abd3e0ac26b
Author: Andy Grove <ag...@apache.org>
AuthorDate: Sun Oct 11 12:56:29 2026 +0000
test: count the partitions and batches the native cache scan reads (#6836)
* test: count the partitions and batches the native cache scan reads
Spark's PartitionBatchPruningSuite checks cache pruning through the readPartitions
and readBatches accumulators of InMemoryTableScanExec, which Spark's scan fills
under spark.sql.inMemoryTableScanStatistics.enable. CometInMemoryTableScanExec
had none, so the Spark SQL diffs made the suite assert only the answer when a
Comet scan was present. That verifies nothing about pruning: the answer is the
same whether a batch was pruned or decoded and filtered. With the cache on by
default since #5634, Spark's forty-odd pruning cases ran against Comet's format
without checking that any batch was skipped.
Give the Comet scan the same two accumulators, counted after the statistics
filter as Spark counts them, and restore the suite's original assertions in the
3.5.9, 4.0.4, 4.1.3 and 4.2.0 diffs with a case for the Comet scan. The 3.4.3
diff never touched the suite because the cache is opt-in there. The new
CometInMemoryCacheSuite test checks the counts with pruning on and off, and
with AQE on and off, since on Spark 3.5+ the scan then sits in a table cache
stage and the instance that ran has to be the one the final plan holds.
* test: return the read counts through a local, since withSQLConf returns Unit on Spark 3.x
---
dev/diffs/3.5.9.diff | 72 +++++-----------------
dev/diffs/4.0.4.diff | 72 +++++-----------------
dev/diffs/4.1.3.diff | 72 +++++-----------------
dev/diffs/4.2.0.diff | 72 +++++-----------------
.../sql/comet/CometInMemoryTableScanExec.scala | 33 +++++++++-
.../comet/exec/CometInMemoryCacheSuite.scala | 51 +++++++++++++++
6 files changed, 151 insertions(+), 221 deletions(-)
diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff
index 4ee26ebe93..7549996767 100644
--- a/dev/diffs/3.5.9.diff
+++ b/dev/diffs/3.5.9.diff
@@ -2461,7 +2461,7 @@ index 2c73622739a..5d0efeb263a 100644
val d1 = spark.range(1).withColumn("key", expr("id % 100"))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
-index 885286843a1..f0f805a6cd9 100644
+index 885286843a1..6e9d6690374 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
@@ -17,6 +17,7 @@
@@ -2472,60 +2472,22 @@ index 885286843a1..f0f805a6cd9 100644
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
-@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
- val result = df.collect().map(_(0)).toArray
- assert(result.length === 1)
-
-- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
-- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
-- }.head
-- assert(readPartitions === 5)
-- assert(readBatches === 10)
-+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in }
-+ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count.
-+ if (scans.isEmpty) {
-+ assert(collect(df.queryExecution.executedPlan) {
-+ case c: CometInMemoryTableScanExec => c
-+ }.nonEmpty)
-+ } else {
-+ assert(scans.head.readPartitions.value === 5)
-+ assert(scans.head.readBatches.value === 10)
-+ }
- }
-
- def checkBatchPruning(
-@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
- df.collect().map(_(0)).toArray
- }
-
-- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
-- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
-- }.head
--
-- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution")
-- assert(
-- readPartitions === expectedReadPartitions,
-- s"Wrong number of read partitions: $queryExecution")
-+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in }
-+ // Comet's cache scan prunes on the same statistics, but has none of these test-only
-+ // accumulators to read, so only the answer above is checked for it.
-+ if (scans.isEmpty) {
-+ assert(collect(df.queryExecution.executedPlan) {
-+ case c: CometInMemoryTableScanExec => c
-+ }.nonEmpty)
-+ } else {
-+ val readPartitions = scans.head.readPartitions.value
-+ val readBatches = scans.head.readBatches.value
-+ assert(
-+ readBatches === expectedReadBatches,
-+ s"Wrong number of read batches: $queryExecution")
-+ assert(
-+ readPartitions === expectedReadPartitions,
-+ s"Wrong number of read partitions: $queryExecution")
-+ }
- }
- }
- }
+@@ -183,6 +184,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
+
+ val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
+ case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
+ }.head
+ assert(readPartitions === 5)
+ assert(readBatches === 10)
+@@ -204,6 +206,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
+
+ val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
+ case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
+ }.head
+
+ assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
index fd52d038ca6..154c800be67 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff
index 895f6bd3db..2133589aa9 100644
--- a/dev/diffs/4.0.4.diff
+++ b/dev/diffs/4.0.4.diff
@@ -3228,7 +3228,7 @@ index 2c73622739a..5d0efeb263a 100644
val d1 = spark.range(1).withColumn("key", expr("id % 100"))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
-index 88ff51d0ff4..7d00fca3b0a 100644
+index 88ff51d0ff4..3763fad99c8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
@@ -17,6 +17,7 @@
@@ -3239,60 +3239,22 @@ index 88ff51d0ff4..7d00fca3b0a 100644
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
-@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
- val result = df.collect().map(_(0)).toArray
- assert(result.length === 1)
-
-- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
-- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
-- }.head
-- assert(readPartitions === 5)
-- assert(readBatches === 10)
-+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in }
-+ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count.
-+ if (scans.isEmpty) {
-+ assert(collect(df.queryExecution.executedPlan) {
-+ case c: CometInMemoryTableScanExec => c
-+ }.nonEmpty)
-+ } else {
-+ assert(scans.head.readPartitions.value === 5)
-+ assert(scans.head.readBatches.value === 10)
-+ }
- }
-
- def checkBatchPruning(
-@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
- df.collect().map(_(0)).toArray
- }
-
-- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
-- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
-- }.head
--
-- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution")
-- assert(
-- readPartitions === expectedReadPartitions,
-- s"Wrong number of read partitions: $queryExecution")
-+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in }
-+ // Comet's cache scan prunes on the same statistics, but has none of these test-only
-+ // accumulators to read, so only the answer above is checked for it.
-+ if (scans.isEmpty) {
-+ assert(collect(df.queryExecution.executedPlan) {
-+ case c: CometInMemoryTableScanExec => c
-+ }.nonEmpty)
-+ } else {
-+ val readPartitions = scans.head.readPartitions.value
-+ val readBatches = scans.head.readBatches.value
-+ assert(
-+ readBatches === expectedReadBatches,
-+ s"Wrong number of read batches: $queryExecution")
-+ assert(
-+ readPartitions === expectedReadPartitions,
-+ s"Wrong number of read partitions: $queryExecution")
-+ }
- }
- }
- }
+@@ -183,6 +184,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
+
+ val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
+ case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
+ }.head
+ assert(readPartitions === 5)
+ assert(readBatches === 10)
+@@ -204,6 +206,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
+
+ val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
+ case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
+ }.head
+
+ assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
index 269990d7d14..140ee4112b1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff
index 0b814aa1b1..929fbfd09d 100644
--- a/dev/diffs/4.1.3.diff
+++ b/dev/diffs/4.1.3.diff
@@ -3447,7 +3447,7 @@ index 2c73622739a..5d0efeb263a 100644
val d1 = spark.range(1).withColumn("key", expr("id % 100"))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
-index 88ff51d0ff4..7d00fca3b0a 100644
+index 88ff51d0ff4..3763fad99c8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
@@ -17,6 +17,7 @@
@@ -3458,60 +3458,22 @@ index 88ff51d0ff4..7d00fca3b0a 100644
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
-@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
- val result = df.collect().map(_(0)).toArray
- assert(result.length === 1)
-
-- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
-- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
-- }.head
-- assert(readPartitions === 5)
-- assert(readBatches === 10)
-+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in }
-+ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count.
-+ if (scans.isEmpty) {
-+ assert(collect(df.queryExecution.executedPlan) {
-+ case c: CometInMemoryTableScanExec => c
-+ }.nonEmpty)
-+ } else {
-+ assert(scans.head.readPartitions.value === 5)
-+ assert(scans.head.readBatches.value === 10)
-+ }
- }
-
- def checkBatchPruning(
-@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
- df.collect().map(_(0)).toArray
- }
-
-- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
-- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
-- }.head
--
-- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution")
-- assert(
-- readPartitions === expectedReadPartitions,
-- s"Wrong number of read partitions: $queryExecution")
-+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in }
-+ // Comet's cache scan prunes on the same statistics, but has none of these test-only
-+ // accumulators to read, so only the answer above is checked for it.
-+ if (scans.isEmpty) {
-+ assert(collect(df.queryExecution.executedPlan) {
-+ case c: CometInMemoryTableScanExec => c
-+ }.nonEmpty)
-+ } else {
-+ val readPartitions = scans.head.readPartitions.value
-+ val readBatches = scans.head.readBatches.value
-+ assert(
-+ readBatches === expectedReadBatches,
-+ s"Wrong number of read batches: $queryExecution")
-+ assert(
-+ readPartitions === expectedReadPartitions,
-+ s"Wrong number of read partitions: $queryExecution")
-+ }
- }
- }
- }
+@@ -183,6 +184,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
+
+ val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
+ case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
+ }.head
+ assert(readPartitions === 5)
+ assert(readBatches === 10)
+@@ -204,6 +206,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
+
+ val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
+ case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
+ }.head
+
+ assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
index 269990d7d14..140ee4112b1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff
index f58f210fbe..549799d941 100644
--- a/dev/diffs/4.2.0.diff
+++ b/dev/diffs/4.2.0.diff
@@ -3518,7 +3518,7 @@ index 2c73622739a..5d0efeb263a 100644
val d1 = spark.range(1).withColumn("key", expr("id % 100"))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
-index 88ff51d0ff4..7d00fca3b0a 100644
+index 88ff51d0ff4..3763fad99c8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
@@ -17,6 +17,7 @@
@@ -3529,60 +3529,22 @@ index 88ff51d0ff4..7d00fca3b0a 100644
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
-@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
- val result = df.collect().map(_(0)).toArray
- assert(result.length === 1)
-
-- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
-- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
-- }.head
-- assert(readPartitions === 5)
-- assert(readBatches === 10)
-+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in }
-+ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count.
-+ if (scans.isEmpty) {
-+ assert(collect(df.queryExecution.executedPlan) {
-+ case c: CometInMemoryTableScanExec => c
-+ }.nonEmpty)
-+ } else {
-+ assert(scans.head.readPartitions.value === 5)
-+ assert(scans.head.readBatches.value === 10)
-+ }
- }
-
- def checkBatchPruning(
-@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
- df.collect().map(_(0)).toArray
- }
-
-- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
-- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
-- }.head
--
-- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution")
-- assert(
-- readPartitions === expectedReadPartitions,
-- s"Wrong number of read partitions: $queryExecution")
-+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in }
-+ // Comet's cache scan prunes on the same statistics, but has none of these test-only
-+ // accumulators to read, so only the answer above is checked for it.
-+ if (scans.isEmpty) {
-+ assert(collect(df.queryExecution.executedPlan) {
-+ case c: CometInMemoryTableScanExec => c
-+ }.nonEmpty)
-+ } else {
-+ val readPartitions = scans.head.readPartitions.value
-+ val readBatches = scans.head.readBatches.value
-+ assert(
-+ readBatches === expectedReadBatches,
-+ s"Wrong number of read batches: $queryExecution")
-+ assert(
-+ readPartitions === expectedReadPartitions,
-+ s"Wrong number of read partitions: $queryExecution")
-+ }
- }
- }
- }
+@@ -183,6 +184,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
+
+ val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
+ case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
+ }.head
+ assert(readPartitions === 5)
+ assert(readBatches === 10)
+@@ -204,6 +206,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl
+
+ val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) {
+ case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value)
+ }.head
+
+ assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
index fd8d1308e99..0e1f80045a3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala
index 207513f0a7..97f041033d 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.columnar.{CachedRDDBuilder, InMemoryRelation, InMemoryTableScanExec}
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.util.LongAccumulator
import org.apache.comet.CometConf
import org.apache.comet.serde.CometOperatorSerde
@@ -68,6 +69,16 @@ case class CometInMemoryTableScanExec(
override lazy val metrics: Map[String, SQLMetric] = Map(
"numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"))
+ // The partitions and batches a scan read after pruning, as Spark's InMemoryTableScanExec counts
+ // them under spark.sql.inMemoryTableScanStatistics.enable. Spark's PartitionBatchPruningSuite
+ // reads these to check that pruning skipped what it should have, and only they can show that:
+ // the answer is the same whether a batch was pruned or decoded and filtered.
+ lazy val enableAccumulatorsForTest: Boolean = conf.inMemoryTableScanStatisticsEnabled
+
+ lazy val readPartitions: LongAccumulator = sparkContext.longAccumulator
+
+ lazy val readBatches: LongAccumulator = sparkContext.longAccumulator
+
// `scanOutput` always equals this, including when it is empty. An empty-output scan
// (`SELECT count(*)`) emits genuinely zero-column batches carrying only a row count: widening it
// to a placeholder column, or to the whole cache schema, makes the emitted batches disagree with
@@ -125,6 +136,11 @@ case class CometInMemoryTableScanExec(
// stats bug, for instance -- and silently ignoring it would make Comet diverge from Spark on a
// knob a user reaching for it is specifically trying to control.
override def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ if (enableAccumulatorsForTest) {
+ readPartitions.setValue(0)
+ readBatches.setValue(0)
+ }
+
val numOutputRows = longMetric("numOutputRows")
// Resolved here rather than at planning time. CachedRDDBuilder.cachedColumnBuffers is not a
@@ -142,8 +158,23 @@ case class CometInMemoryTableScanExec(
cachedBuffers
}
+ // Counted after pruning, as Spark counts them, so a pruned batch is not a read one.
+ val countedBuffers = if (enableAccumulatorsForTest) {
+ filteredBuffers.mapPartitionsInternal { iter =>
+ if (iter.hasNext) {
+ readPartitions.add(1)
+ }
+ iter.map { batch =>
+ readBatches.add(1)
+ batch
+ }
+ }
+ } else {
+ filteredBuffers
+ }
+
serializer
- .convertCachedBatchToColumnarBatch(filteredBuffers, relationOutput, scanOutput, conf)
+ .convertCachedBatchToColumnarBatch(countedBuffers, relationOutput, scanOutput, conf)
.map { cb =>
numOutputRows += cb.numRows().toLong
cb
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
index aba97d8021..be0eb109db 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -1214,6 +1214,57 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
+ test("Comet in-memory cache counts the partitions and batches a scan reads") {
+ // Spark's PartitionBatchPruningSuite reads these accumulators to check that pruning skipped
+ // the batches it should have; the answer alone cannot tell a pruned batch from a decoded and
+ // filtered one.
+ // With AQE the scan sits inside a table cache stage on Spark 3.5+, and the instance that ran
+ // has to be the one the final plan holds.
+ Seq(false, true).foreach { adaptive =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ SQLConf.IN_MEMORY_TABLE_SCAN_STATISTICS_ENABLED.key -> "true",
+ SQLConf.COLUMN_BATCH_SIZE.key -> "50",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") {
+ try {
+ // 10 partitions of 2 batches each.
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id as key", "id % 7 as value")
+ .createOrReplaceTempView("counted_cache")
+ spark.catalog.cacheTable("counted_cache")
+ spark.table("counted_cache").count()
+
+ // withSQLConf returns Unit on Spark 3.x, so the counts leave it through a local.
+ def readCounts(query: String, pruning: Boolean): (Long, Long) = {
+ var counts: (Long, Long) = null
+ withSQLConf(SQLConf.IN_MEMORY_PARTITION_PRUNING.key -> pruning.toString) {
+ val df = spark.sql(query)
+ df.collect()
+ val scans = collect(df.queryExecution.executedPlan) {
+ case s: CometInMemoryTableScanExec => s
+ }
+ assert(scans.size == 1, df.queryExecution.executedPlan)
+ counts = (scans.head.readPartitions.value, scans.head.readBatches.value)
+ }
+ counts
+ }
+
+ val narrow = "SELECT key FROM counted_cache WHERE key >= 900 AND key < 905"
+ assert(readCounts(narrow, pruning = true) == ((1L, 1L)))
+ assert(readCounts(narrow, pruning = false) == ((10L, 20L)))
+ assert(
+ readCounts("SELECT key FROM counted_cache WHERE key < 0", pruning = true) == ((
+ 0L,
+ 0L)))
+ assert(readCounts("SELECT key FROM counted_cache", pruning = true) == ((10L, 20L)))
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+
test("Comet in-memory cache honors inMemoryColumnarStorage.partitionPruning=false") {
// CometInMemoryTableScanExec applies the serializer's stats filter before decoding, the same
// way Spark's InMemoryTableScanExec.filteredCachedBatches does. Spark gates that on
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org