Skip to content

test: count the partitions and batches the native cache scan reads #6836

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
andygrove merged 2 commits into from
Oct 11, 2026
Merged

test: count the partitions and batches the native cache scan reads #6836

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7f1eeab
test: count the partitions and batches the native cache scan reads
andygrove Oct 10, 2026
459f951
test: return the read counts through a local, since withSQLConf retur…
andygrove Oct 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 17 additions & 55 deletions dev/diffs/3.5.9.diff
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 @@
Expand All @@ -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
Expand Down
72 changes: 17 additions & 55 deletions dev/diffs/4.0.4.diff
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 @@
Expand All @@ -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
Expand Down
72 changes: 17 additions & 55 deletions dev/diffs/4.1.3.diff
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 @@
Expand All @@ -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
Expand Down
72 changes: 17 additions & 55 deletions dev/diffs/4.2.0.diff
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 @@
Expand All @@ -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
Expand Down
Loading
Loading