sunchao commented on code in PR #5051: URL: https://github.com/apache/datafusion-comet/pull/5051#discussion_r3877987279
########## spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala: ########## @@ -0,0 +1,158 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet + +import scala.collection.JavaConverters._ + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.columnar.CachedBatchSerializer +import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan} +import org.apache.spark.sql.execution.columnar.{CachedRDDBuilder, InMemoryTableScanExec} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.CometConf +import org.apache.comet.serde.CometOperatorSerde +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.QueryPlanSerde.serializeDataType + +/** + * Reads Spark cached table data when the cache was written by Comet's cache serializer. + * + * Spark stores cached data through `CachedBatchSerializer`. This node keeps the scan inside Comet + * by asking the serializer to decode cached batches directly into `ColumnarBatch` output, + * avoiding the extra Spark columnar-to-Comet columnar conversion used by the default path. + * + * `relationOutput` is the full schema stored in the cache. `scanOutput` is the subset requested + * by this scan after pruning. + */ +case class CometInMemoryTableScanExec( + originalPlan: InMemoryTableScanExec, + serializer: CachedBatchSerializer, + cacheBuilder: CachedRDDBuilder, + relationOutput: Seq[Attribute], + scanOutput: Seq[Attribute]) + extends CometExec + with LeafExecNode { + + override lazy val metrics: Map[String, SQLMetric] = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + + // `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 + // the declared output, and a consumer that reads by ordinal rather than by row count -- a join, + // for instance -- then reads the wrong column. + override def output: Seq[Attribute] = originalPlan.output + + // `originalPlan` is a plan-typed field rather than a child, so QueryPlan's canonicalization + // walks straight past it: its attributes and predicates keep the expression IDs of whichever + // occurrence of the cached relation produced them. Two scans of one cache then compare unequal, + // and since sameResult is what exchange and broadcast reuse are keyed on, a UNION of two + // identical aggregates over a cached table runs two shuffles where Spark's own cache scan runs + // one and reuses it. + // + // Defer to `InMemoryTableScanExec`, which normalizes its own attributes and predicates against + // the relation's output. Dropping the field instead would also make the scans compare equal, + // but it would equate scans carrying different pruning predicates along with them. + override protected def doCanonicalize(): SparkPlan = + super + .doCanonicalize() + .asInstanceOf[CometInMemoryTableScanExec] + .copy(originalPlan = originalPlan.canonicalized.asInstanceOf[InMemoryTableScanExec]) + + // Use the serializer's vector types because the cached batch layout is owned by the serializer. + override def vectorTypes: Option[Seq[String]] = + serializer.vectorTypes(scanOutput, conf) + + // Apply Spark's cache batch filter before decoding. Spark's InMemoryTableScanExec does this in + // filteredCachedBatches(), but that method is private. Reusing the serializer's buildFilter here + // keeps Comet on the same stats-based pruning path instead of decoding every cached batch. + // + // Gated on conf.inMemoryPartitionPruning the same way Spark's filteredCachedBatches is, so + // spark.sql.inMemoryColumnarStorage.partitionPruning=false disables pruning here too. Pruning is + // normally a win, but the config exists to be able to turn it off -- for debugging a suspected + // 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] = { + val numOutputRows = longMetric("numOutputRows") + + // Resolved here rather than at planning time. CachedRDDBuilder.cachedColumnBuffers is not a + // metadata lookup: it builds the RDD by calling execute/executeColumnar on the cached plan, + // so touching it while Comet is still planning the outer query runs jobs during planning -- + // visibly, an EXPLAIN of a query over an adaptively-cached relation would launch a job and + // finalize that plan. + val cachedBuffers = cacheBuilder.cachedColumnBuffers + + val filteredBuffers = + if (originalPlan.predicates.nonEmpty && conf.inMemoryPartitionPruning) { + val filter = serializer.buildFilter(originalPlan.predicates, relationOutput) + cachedBuffers.mapPartitionsWithIndex(filter) + } else { + cachedBuffers + } + + serializer + .convertCachedBatchToColumnarBatch(filteredBuffers, relationOutput, scanOutput, conf) + .map { cb => + numOutputRows += cb.numRows() + cb + } + } +} + +object CometInMemoryTableScanExec extends CometOperatorSerde[InMemoryTableScanExec] { + + override def enabledConfig: Option[org.apache.comet.ConfigEntry[Boolean]] = + Some(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED) + + override def convert( + op: InMemoryTableScanExec, + builder: OperatorOuterClass.Operator.Builder, + childOp: Operator*): Option[Operator] = { + + val scanTypes = op.output.flatMap(attr => serializeDataType(attr.dataType)) Review Comment: **[P2] Widen nested nullability in the cache scan schema** This serializes containsNull=false/valueContainsNull=false into the native scan schema, reintroducing the nested-type failures already fixed for CometLocalTableScanExec. On Spark 4.1.3, caching Seq(Seq(1, 2, 3), Seq(4, 5)) and evaluating slice(x, 2, 2) fails because spark_array_slice returns List(non-null Int32) while its declared result is List(Int32). map_entries over a cached Map[Int, Int] similarly panics on the value field's nullability. Both work uncached and with native cache scanning disabled over the same payload. An isolated change to serializeDataType(attr.dataType.asNullable) fixes both. Please preserve that normalization at this scan boundary and add cache regressions. ########## spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala: ########## @@ -398,33 +447,109 @@ object Utils extends CometTypeShim with Logging { } } + /** + * Whether every column in `batch` is an Arrow-backed `CometVector`, so [[getBatchFieldVectors]] + * can hand out its vectors directly. Callers that may receive batches from a plan they did not + * build (e.g. Comet's cache serializer, which Spark hands the cached plan's columnar output) + * use this to convert foreign vectors to Arrow instead of tripping the exception below. + * + * Stricter than what [[getBatchFieldVectors]] accepts: a `ConstantColumnVector` is rejected + * here even though that method materializes one, so such a batch takes the conversion path + * rather than being materialized column by column. + */ + def isArrowBacked(batch: ColumnarBatch): Boolean = + (0 until batch.numCols()).forall { i => + batch.column(i) match { + // Not every CometVector can be handed to getFieldVector: a CometPlainVector can wrap a + // LargeVarCharVector or LargeVarBinaryVector (an accelerated mapInArrow returning + // pa.large_string(), for instance), which it rejects. Answering true for those would + // send a batch down the direct write path that then fails, so check the vector itself + // and let the caller convert instead. + case v: CometVector => isSupportedFieldVector(v.getValueVector) + case _ => false + } + } + def getBatchFieldVectors( batch: ColumnarBatch): (Seq[FieldVector], Option[DictionaryProvider]) = { - var provider: Option[DictionaryProvider] = None + val columns = getBatchFieldVectorsWithProviders(batch) + (columns.map(_._1), combineDictionaryProviders(columns)) + } + + /** + * The dictionaries every dictionary-encoded column of `columns` refers to, as one provider. + * + * Columns of a batch need not share a provider. Comet's cache decodes each column from its own + * Arrow stream, so a dictionary-backed column arrives carrying the provider its reader built, + * and a batch that reaches [[serializeBatches]] -- a native broadcast of a cache scan, say -- + * can hold several. Writing the whole batch emits one schema covering every column and resolves + * each column's dictionary ID against the single provider the writer was given, so handing it + * any one column's provider fails with "Could not find dictionary with ID n" for the others. + */ + private def combineDictionaryProviders( + columns: Seq[(FieldVector, Option[DictionaryProvider])]): Option[DictionaryProvider] = { + val dictionaries = scala.collection.mutable.LinkedHashMap.empty[Long, Dictionary] + + columns.foreach { case (vector, providerOpt) => + val encoding = vector.getField.getDictionary + if (encoding != null) { + val id = encoding.getId + val dictionary = providerOpt.map(_.lookup(id)).orNull + if (dictionary == null) { + throw new SparkException( + s"Column ${vector.getField.getName} is dictionary encoded with ID $id, but no " + + "dictionary with that ID was provided") + } + dictionaries.get(id) match { + // Every provider seen here descends from one upstream reader, which numbers the + // dictionaries it hands out, so two columns sharing an ID share the dictionary itself. + // A genuine clash would need renumbering, which means rewriting each vector's field, + // so refuse rather than silently decode one column against another's dictionary. + case Some(existing) if existing.getVector ne dictionary.getVector => + throw new SparkException( + s"Columns of the same batch carry different dictionaries under ID $id") Review Comment: **[P2] Handle equivalent dictionaries from repeated cached columns** Selecting one cached dictionary column twice opens two readers for the same IPC stream. Their dictionary IDs and values match, but their vector objects differ, so this identity check rejects valid data. With a dictionary-encoded cache produced through JVM shuffle and AQE disabled, spark.range(1).join(broadcast(cached.select(cached("s1"), cached("s1"))), lit(true), "inner").collect() fails with 'Columns of the same batch carry different dictionaries under ID 0'. Both plans use CometBroadcastExchange; disabling only native cache scanning returns all 2,000 rows correctly. Please support equivalent dictionaries or normalize IDs instead of requiring object identity. ########## spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala: ########## @@ -0,0 +1,192 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.spark.SparkConf +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +object CometInMemoryCacheBenchmark extends CometBenchmarkBase { + private val numRows = 5 * 1000 * 1000 + private val cacheTable = "comet_cache_bench" + private val sourceTable = "comet_cache_bench_src" + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName("CometInMemoryCacheBenchmark") + .set("spark.master", "local[1]") + .setIfMissing("spark.driver.memory", "3g") + .setIfMissing("spark.executor.memory", "3g") + .set("spark.plugins", "org.apache.spark.CometPlugin") + .set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set( + "spark.sql.cache.serializer", + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + + val sparkSession = SparkSession + .builder() + .config(conf) + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + + sparkSession.conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") + sparkSession.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") + sparkSession.conf.set(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, "true") + sparkSession.conf.set(SQLConf.ANSI_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") + sparkSession + } + + override def runCometBenchmark(args: Array[String]): Unit = { + withTempTable(sourceTable, cacheTable) { + spark + .range(0, numRows, 1, 16) + .selectExpr( + "id", + "id % 1000 AS k", + "id + 1 AS v", + "concat('str_a_', cast(id % 100000 as string)) AS s1", + "concat('str_b_', cast(id % 7919 as string)) AS s2", + "concat('str_c_', cast(id as string)) AS s3") + .createOrReplaceTempView(sourceTable) + + runCacheBenchmark( + "in-memory cache repeated scan", + s"SELECT sum(id), sum(k), sum(v) FROM $cacheTable") + + runCacheBenchmark( + "in-memory cache selective filter", + s""" + |SELECT sum(id), sum(k), sum(v) + |FROM $cacheTable + |WHERE id >= 4500000 AND id < 4750000 + """.stripMargin) + + // A CometCachedBatch stores each column as its own stream, so a scan decodes only what it + // projected and cost tracks the width of the projection. These three cases span that range + // over one cached relation: no columns, one column, and all six. + runCacheBenchmark( + "in-memory cache row count only (0 of 6 columns)", + s"SELECT count(*) FROM $cacheTable") + + runCacheBenchmark( + "in-memory cache narrow projection (1 of 6 columns)", + s"SELECT count(k) FROM $cacheTable") + + runCacheBenchmark( + "in-memory cache full projection (6 of 6 columns)", + s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $cacheTable") Review Comment: **[P3] Make the full projection benchmark consume all six columns** The Range-derived id, v, and s3 columns are non-nullable, so Catalyst rewrites count(id), count(v), and count(s3) to count(1). Both the native and fallback benchmark plans therefore scan only k, s1, and s2, despite this case being labeled 6 of 6 columns. This also omits the highest-cardinality string column s3. Use expressions that consume all six columns and assert the actual cache-scan projection, otherwise the benchmark does not measure the full-width per-column decoding cost it claims. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala: ########## @@ -0,0 +1,158 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet + +import scala.collection.JavaConverters._ + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.columnar.CachedBatchSerializer +import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan} +import org.apache.spark.sql.execution.columnar.{CachedRDDBuilder, InMemoryTableScanExec} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.CometConf +import org.apache.comet.serde.CometOperatorSerde +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.QueryPlanSerde.serializeDataType + +/** + * Reads Spark cached table data when the cache was written by Comet's cache serializer. + * + * Spark stores cached data through `CachedBatchSerializer`. This node keeps the scan inside Comet + * by asking the serializer to decode cached batches directly into `ColumnarBatch` output, + * avoiding the extra Spark columnar-to-Comet columnar conversion used by the default path. + * + * `relationOutput` is the full schema stored in the cache. `scanOutput` is the subset requested + * by this scan after pruning. + */ +case class CometInMemoryTableScanExec( + originalPlan: InMemoryTableScanExec, + serializer: CachedBatchSerializer, + cacheBuilder: CachedRDDBuilder, + relationOutput: Seq[Attribute], + scanOutput: Seq[Attribute]) + extends CometExec + with LeafExecNode { Review Comment: **[P2] Preserve AQE cache materialization and reoptimization** Replacing Spark's InMemoryTableScanLike with an ordinary leaf during AQE preparation prevents TableCacheQueryStageExec creation. On Spark 4.1.3, first reading an adaptive cached join through a grouped aggregate executes an extra outer shuffle with native caching enabled: native disabled materializes the cache and removes that shuffle, while enabled retains it. A join against a cold cache already partitioned by its join key similarly runs two outer shuffles instead of one. Warm-cache queries avoid the difference. Results are correct, but the cold path unnecessarily reshuffles data. Please preserve materialization and reoptimization at the cache boundary. This is an executed failure of the scenarios requested in #5245, not just a coverage concern. ########## spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala: ########## @@ -0,0 +1,191 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.exec + +import org.apache.spark.SparkConf +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.storage.StorageLevel + +import org.apache.comet.{CometConf, CometKryoRegistrator} + +/** + * Covers Comet's cached batch format under `spark.kryo.registrationRequired=true`. + * + * Kryo then rejects any class it has not been told about, and Spark serializes a `CachedBatch` + * whenever a cached block leaves the heap: the disk half of the default `MEMORY_AND_DISK`, the + * `_SER` levels, replication, and cross-executor fetches. So this is not a `DISK_ONLY`-only + * concern -- a plain `df.cache()` that spills is enough to reach it. Spark registers its own + * `ArrowCachedBatch` in `KryoSerializer.loadableSparkClasses`; Comet cannot add to that list, so + * [[CometKryoRegistrator]] has to be set explicitly, and this suite is what proves it is + * sufficient. + * + * This needs its own suite because `spark.serializer` and `spark.kryo.registrator` are read when + * `SparkEnv` builds the serializer, so they cannot be changed per test. + */ +class CometInMemoryCacheKryoSuite extends CometTestBase { Review Comment: **[P2] Add the new Kryo suite to both CI test matrices** This new suite is absent from the suite lists in both .github/workflows/pr_build_linux.yml and .github/workflows/pr_build_macos.yml. python3 dev/ci/check-suites.py exits 255 with 'Suite not found in workflow ... CometInMemoryCacheKryoSuite'; the current GitHub Preflight job fails for exactly this reason and the platform/Spark test jobs are skipped. Add the suite to both matrices so CI can validate the PR and keep exercising the new Kryo regressions. -- 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]
