peterxcli commented on code in PR #5051:
URL: https://github.com/apache/datafusion-comet/pull/5051#discussion_r3694503535
##########
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:
-
https://github.com/apache/spark/blob/78147abefb29096930346c3f9abaf92bdd5de4bb/sql/core/src/main/scala/org/apache/spark/sql/execution/r/ArrowRRunner.scala#L95-L117
-
https://github.com/apache/spark/blob/78147abefb29096930346c3f9abaf92bdd5de4bb/common/utils/src/main/scala/org/apache/spark/util/SparkErrorUtils.scala#L73-L104
looked through spark's arrow writer usage, FYI.
##########
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:
> reading a metric off `df.queryExecution.executedPlan` after
`checkSparkAnswer(df)` gives zero, because checkSparkAnswer executes its own
copies. My first version passed vacuously with 0 on both sides of the
comparison; it needs the df forced explicitly. The real numbers are 100 rows
decoded with pruning on versus 1000 with it off.
Remove the `println` and add TODO or new issue link would be easier to
followup
--
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]
