andygrove commented on code in PR #5051:
URL: https://github.com/apache/datafusion-comet/pull/5051#discussion_r3709229295
##########
spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala:
##########
@@ -412,6 +413,114 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
+ 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
+ // spark.sql.inMemoryColumnarStorage.partitionPruning, so Comet must too.
+ //
+ // Pruning is transparent in the results, so it is observed through the
scan's numOutputRows:
+ // that counts the rows in the batches actually decoded, so pruning fewer
batches means fewer
+ // rows. With pruning off, every cached row must be decoded.
+ def scanRowsFor(pruning: Boolean): (Long, Long) = {
+ var result: (Long, Long) = (0L, 0L)
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true",
+ "spark.comet.sparkToColumnar.enabled" -> "true",
+ "spark.sql.inMemoryColumnarStorage.batchSize" -> "100",
+ SQLConf.IN_MEMORY_PARTITION_PRUNING.key -> pruning.toString) {
+
+ spark.catalog.clearCache()
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id as key", "id % 7 as value")
+ .createOrReplaceTempView("prune_conf_cache")
+ spark.catalog.cacheTable("prune_conf_cache")
+ val totalRows = spark.table("prune_conf_cache").count()
+
+ val df =
+ spark.sql("SELECT key, value FROM prune_conf_cache WHERE key >= 900
AND key < 905")
+ checkSparkAnswer(df)
+
+ val scans = df.queryExecution.executedPlan.collect {
+ case s: org.apache.spark.sql.comet.CometInMemoryTableScanExec => s
+ }
+ assert(scans.length == 1, s"expected one CometInMemoryTableScan, got
${scans.length}")
+ // scalastyle:off println
+ println(
+ "DIAG rows=" + df.collect().length + " metrics=" + scans.head.metrics
+ .map { case (k, v) => k + "=" + v.value }
+ .mkString(","))
+ println("DIAG plan=" + df.queryExecution.executedPlan.getClass.getName)
+ // scalastyle:on println
Review Comment:
Removed in ebefc1cf4. One thing worth keeping visible: the `df.collect()`
buried in that `println` was load-bearing, not diagnostic. `checkSparkAnswer`
takes its argument by name and executes its own copies of the query, so without
an explicit force this `df`'s plan instance has never run and every metric on
it reads zero, which makes the pruned-vs-unpruned comparison pass vacuously as
0 == 0. So the `collect()` is now an explicit call with that reason in a
comment above it, rather than a side effect of a debug line.
Same commit fixes an unrelated failure I hit while re-running the suite:
"cache a non-Arrow-backed Spark columnar plan with complex types" built its
binary column with `cast(id as binary)`, which ANSI mode rejects
(`DATATYPE_MISMATCH.CAST_WITH_CONF_SUGGESTION`). It passed when I wrote it on
spark-3.5 and would have failed the Spark 4.x jobs. Now cast via string. All 21
tests pass on the default 4.1 profile.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala:
##########
@@ -64,15 +66,23 @@ object CometArrowConverters extends Logging {
override def next(): ColumnarBatch = {
val root = VectorSchemaRoot.create(arrowSchema, allocator)
- val writer = ArrowWriter.create(root)
- var rowCount = 0L
- while (rowIter.hasNext &&
- (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch)) {
- writer.write(rowIter.next())
- rowCount += 1
+ // Same ownership rule as columnarBatchToArrowBatch: the caller only
owns the batch that
+ // rootAsBatch returns, so a throw from writing a row has to release
the root here.
+ try {
+ val writer = ArrowWriter.create(root)
+ var rowCount = 0L
+ while (rowIter.hasNext &&
+ (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch)) {
+ writer.write(rowIter.next())
+ rowCount += 1
+ }
+ writer.finish()
+ NativeUtil.rootAsBatch(root)
+ } catch {
+ case NonFatal(e) =>
+ root.close()
+ throw e
Review Comment:
Useful reference. `tryWithSafeFinally` itself does not fit here, since
ownership of the root transfers to the returned batch on success and it must
deliberately stay open, so the release is failure-only rather than a `finally`.
But the part of it I was missing does apply: my guard closed the root with a
bare `root.close()`, and an Arrow root can throw from `close`
(`IllegalStateException` for outstanding child allocations). That exception
would have replaced the failure that actually caused the abort, which is the
less useful of the two.
4b9a5167a adopts the suppressed-exception semantics: the close error is
attached to the original with `addSuppressed` and the original is rethrown. The
guard is now one helper shared by `rowToArrowBatchIter` and
`columnarBatchToArrowBatch` instead of two copies.
--
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]
