comphead commented on code in PR #7347:
URL: https://github.com/apache/datafusion-comet/pull/7347#discussion_r4486202206
##########
spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala:
##########
@@ -112,6 +112,76 @@ class CometInMemoryCacheSuite extends CometTestBase {
private def withConversions(pairs: (String, String)*)(f: => Unit): Unit =
withSQLConf(pairs ++ sparkToArrowConversionConfs(enabled = true): _*)(f)
+ test("cache scan reports the concrete decoded vector types") {
+ withConversions(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val cached = spark
+ .sql(
+ "SELECT 1 AS i, 'text' AS s, array(1, 2) AS a, " +
+ "named_struct('v', 3) AS st, map('k', 4) AS m, true AS b, " +
+ "CAST(1 AS TINYINT) AS ti, CAST(2 AS SMALLINT) AS sh, 3L AS l, " +
+ "CAST(1 AS FLOAT) AS f, CAST(2 AS DOUBLE) AS d, " +
+ "CAST(12.3 AS DECIMAL(10, 1)) AS dec, DATE '2024-01-02' AS day, " +
+ "TIMESTAMP '2024-01-02 03:04:05' AS ts, " +
+ "CAST('2024-01-02 03:04:05' AS TIMESTAMP_NTZ) AS ntz, " +
+ "CAST('bytes' AS BINARY) AS bin, NULL AS n, 42 AS unused")
+ .cache()
+ try {
+ cached.count()
+ // Reorder and prune the columns, so types follow scan output, not
cache-schema order.
+ val query = cached.select(
+ "m",
+ "i",
+ "st",
+ "a",
+ "s",
+ "b",
+ "ti",
+ "sh",
+ "l",
+ "f",
+ "d",
+ "dec",
+ "day",
+ "ts",
+ "ntz",
+ "bin",
+ "n")
+ query.collect()
+ val scans = collect(query.queryExecution.executedPlan) {
+ case scan: CometInMemoryTableScanExec => scan
+ }
+ assert(scans.size == 1, query.queryExecution.executedPlan)
+ val scan = scans.head
+ val actual = scan
+ .executeColumnar()
+ .mapPartitions { batches =>
+ batches.map { batch =>
+ (0 until batch.numCols()).map(i =>
batch.column(i).getClass.getName).toArray
+ }
+ }
+ .collect()
+ assert(actual.nonEmpty)
+ val reported = scan.vectorTypes.getOrElse(fail("cache scan reported no
vector types"))
+ actual.foreach(types => assert(types.toSeq == reported))
+ } finally {
Review Comment:
This compares the batch classes with `vectorTypes`, but nothing shows the
codegen cast runs: the result of `query.collect()` is not checked against an
expected answer, and there is no assertion that a `CometColumnarToRowExec`
inside a `WholeStageCodegenExec` reads this scan. The source is also one
constant row. Could the test assert that plan shape and use a multi-batch
cache, for example from Parquet with repeated strings?
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala:
##########
@@ -543,6 +543,29 @@ class ArrowCachedBatchSerializer extends
SimpleMetricsCachedBatchSerializer {
}
}
+ /**
+ * Report vector types for a cache scan whose full stored schema is known.
Spark's serializer
+ * `vectorTypes` hook receives only the projected attributes, which is
insufficient here: an
+ * unsupported unselected column makes the entire cached relation use
Spark's fallback format.
+ */
+ def vectorTypesForScan(
+ cacheAttributes: Seq[Attribute],
+ selectedAttributes: Seq[Attribute],
+ conf: SQLConf): Option[Seq[String]] = {
+ if (!supportsSchema(cacheAttributes)) {
+ fallback.vectorTypes(selectedAttributes, conf)
+ } else {
+ Some(selectedAttributes.map { attr =>
+ attr.dataType match {
+ case _: StructType => classOf[CometStructVector].getName
+ case _: MapType => classOf[CometMapVector].getName
+ case _: ArrayType => classOf[CometListVector].getName
+ case _ => classOf[CometPlainVector].getName
+ }
+ })
+ }
Review Comment:
I do not think the fallback branch or the `cacheAttributes` parameter can be
reached from a planned scan. `CometExecRule.readsCometCacheFormat` builds a
`CometInMemoryTableScanExec` only when the serializer is an
`ArrowCachedBatchSerializer` and `supportsSchema(scan.relation.output)` holds,
and `createExec` passes `relation.output` as `relationOutput`. So the
unsupported unselected column case never produces this scan, and the `case
other` arm in `CometInMemoryTableScanExec.vectorTypes` is unreachable for the
same reason. Could the scan compute the names from `scanOutput` itself and
leave the serializer unchanged? That would also drop the test around
`CometInMemoryCacheSuite.scala` line 173 that compares `vectorTypesForScan`
with the `fallback.vectorTypes` it delegates to.
--
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]
