sunchao commented on code in PR #5051:
URL: https://github.com/apache/datafusion-comet/pull/5051#discussion_r3867601976
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -268,6 +268,37 @@ object Utils extends CometTypeShim with Logging {
}
}
+ /**
+ * Serializes each column of `batch` into its own compressed Arrow IPC
stream, in column order.
+ *
+ * [[serializeBatches]] writes one stream covering every column, so a reader
has to inflate all
+ * of them before it can project. Comet's in-memory cache stores columns
separately instead, so
+ * a scan decodes only the ones it selected. Each stream is self-contained,
including its schema
+ * and any dictionaries the column needs.
+ *
+ * The row count is not recoverable from the result when `batch` has no
columns, so callers keep
+ * it alongside. As with [[serializeBatches]], the batch's vectors are
cleared once written.
+ */
+ def serializeBatchColumns(batch: ColumnarBatch): Array[ChunkedByteBuffer] = {
+ val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch)
+ val provider = batchProviderOpt.getOrElse(new CDataDictionaryProvider)
+ val codec = CompressionCodec.createCodec(SparkEnv.get.conf)
+
+ fieldVectors.map { fieldVector =>
+ val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024,
ByteBuffer.allocate)
+ val out = new DataOutputStream(codec.compressedOutputStream(cbbos))
+
+ val root = new VectorSchemaRoot(Seq(fieldVector).asJava)
+ val writer = new ArrowStreamWriter(root, provider,
Channels.newChannel(out))
+ writer.start()
+ writer.writeBatch()
Review Comment:
**[P2] Use each column's dictionary provider when serializing it**
The new decoder opens one Arrow reader per column, so its dictionary-backed
output columns have independent providers. `getBatchFieldVectors` returns only
the first dictionary column's provider, which is then passed to every writer
here. Re-encoding a decoded cache batch cannot resolve the later columns'
dictionary IDs.
I reproduced this through normal Spark operations: with Comet caching and
`spark.comet.shuffle.mode=jvm`, cache a repartitioned DataFrame with two
low-cardinality string columns. Then disable `spark.comet.exec.enabled` and
`spark.comet.exec.inMemoryCache.enabled` and run
`first.union(first).cache().count()`. Spark's columnar Union passes the decoded
cached batches back to this serializer, and the second cache fails with
`IllegalArgumentException: Could not find dictionary with ID 1`. The original
cache reads correctly and its columns have separate provider namespaces. Please
obtain the provider associated with each column rather than reusing the first
one.
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala:
##########
@@ -0,0 +1,169 @@
+/*
+ * 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.{CachedBatch, CachedBatchSerializer}
+import org.apache.spark.sql.execution.LeafExecNode
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.types._
+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,
+ cachedBuffers: RDD[CachedBatch],
+ 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"))
+
+ // For an empty-projection scan (`SELECT count(*)`) this is empty while
`scanOutput` holds one
+ // placeholder column, so the emitted batches are wider than the declared
output. That is safe
+ // because the only consumer of an empty-output scan is a count-style
aggregate, which reads the
+ // row count rather than any column; see `scanOutputFor` for why the scan
cannot simply be empty.
+ override def output: Seq[Attribute] = originalPlan.output
+
+ // 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")
+
+ 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 = scanOutputFor(op).flatMap(attr =>
serializeDataType(attr.dataType))
+
+ val scanBuilder = OperatorOuterClass.Scan
+ .newBuilder()
+ .setSource(op.getClass.getSimpleName)
+ .addAllFields(scanTypes.asJava)
+
+ Some(builder.setScan(scanBuilder).build())
+ }
+
+ // Reuse Spark's InMemoryRelation metadata so cache materialization,
pruning, and storage
+ // behavior remain controlled by Spark's cache manager.
+ override def createExec(nativeOp: Operator, op: InMemoryTableScanExec):
CometNativeExec = {
+ val relation = op.relation
+
+ CometScanWrapper(
+ nativeOp,
+ CometInMemoryTableScanExec(
+ op,
+ relation.cacheBuilder.serializer,
+ relation.cacheBuilder.cachedColumnBuffers,
+ relation.output,
+ scanOutputFor(op)))
+ }
+
+ /**
+ * Columns the cache scan asks the serializer to decode.
+ *
+ * An empty-output scan (`SELECT count(*)`) still needs a non-empty schema
for native planning,
+ * and the batches the node emits have to match that schema. Falling back to
the whole cache
+ * schema satisfies both, but the serializer decodes exactly what it is
asked for, so the
+ * cheapest query in the workload would decode every cached column. One
column is enough: the
+ * aggregate above an empty-output scan reads the row count and never a
value, so pick the
+ * cheapest to decode rather than all of them.
+ *
+ * `convert` and `createExec` must choose identically, or the native scan's
declared schema and
+ * the batches fed to it disagree.
+ */
+ private def scanOutputFor(op: InMemoryTableScanExec): Seq[Attribute] = {
+ if (op.output.nonEmpty) {
+ op.output
+ } else if (op.relation.output.isEmpty) {
+ Nil
+ } else {
+ Seq(op.relation.output.minBy(a => decodeCostRank(a.dataType)))
Review Comment:
**[P1] Preserve zero-column output for joins**
An empty-output cache scan can feed a join, not only a count-style
aggregate. Adding this column changes the native join's column positions while
`output` still declares the original empty schema. With native caching enabled
and AQE disabled, I reproduced:
```scala
val left = spark.range(10L, 13L).cache()
left.collect()
left.createOrReplaceTempView("cached_left")
spark.sql("""
SELECT /*+ BROADCAST(r) */ sum(r.id)
FROM cached_left l JOIN range(2) r ON true
""").collect()
```
The result is **66**, versus the correct **3** when native cache scanning is
disabled. The native join receives the hidden left column before `r.id`, so the
aggregate's bound ordinal 0 reads the wrong values. Selecting `r.id` directly
also fails with `Output column count mismatch: expected 1, got 2`. Please
preserve the zero-column scan schema/batches instead of introducing an
unreported column. A genuinely zero-column cached relation already returns the
correct count on this head.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala:
##########
@@ -0,0 +1,512 @@
+/*
+ * 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.execution.arrow
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.TaskContext
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression,
GenericInternalRow, IsNotNull, IsNull, UnsafeProjection}
+import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch,
SimpleMetricsCachedBatchSerializer}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.storage.StorageLevel
+import org.apache.spark.unsafe.types.{ByteArray, UTF8String}
+import org.apache.spark.util.io.ChunkedByteBuffer
+
+import org.apache.comet.CometArrowAllocator
+
+/**
+ * Cached batch format used when Comet writes Spark in-memory cache data.
+ *
+ * `columns` holds one compressed Arrow stream per cached column, in
cache-schema order, produced
+ * by `Utils.serializeBatchColumns`. Storing columns separately is what lets a
scan decode only
+ * the ones it projected; a single stream covering the whole batch would have
to be inflated in
+ * full before any projection could be applied. The cache manager still owns
storage and eviction;
+ * this class only changes the cached payload.
+ */
+private case class CometCachedBatch(
+ override val numRows: Int,
+ override val sizeInBytes: Long,
+ override val stats: InternalRow,
+ columns: Array[ChunkedByteBuffer])
+ extends SimpleMetricsCachedBatch
+
+/**
+ * Cache serializer that stores Comet-compatible Arrow batches in Spark's
in-memory cache.
+ *
+ * The cached payload format is decided by the schema alone. A relation whose
schema Comet's Arrow
+ * writer supports is stored as `CometCachedBatch`, and every other relation
is delegated in full
+ * to Spark's `DefaultCachedBatchSerializer`. The format deliberately does not
depend on any
+ * runtime config: `spark.sql.cache.serializer` is a static conf, so
installing this serializer is
+ * already a per-application decision, and a relation whose format could flip
mid-session cannot
+ * be read back reliably. `spark.comet.exec.inMemoryCache.enabled` still
governs whether a scan
+ * over the cache runs natively, and its value at startup is what makes
`CometDriverPlugin`
+ * install this serializer in the first place.
+ *
+ * Reads of `CometCachedBatch` keep working when the native scan is disabled,
because Spark then
+ * reads the same cached data through the SparkToColumnar fallback path.
+ */
+class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
+
+ import ArrowCachedBatchSerializer.supportsSchema
+
+ private val fallback = new DefaultCachedBatchSerializer()
+
+ // Bounds and null counts per column, gathered before the batch is
serialized: serializing
+ // clears the batch's vectors, and the per-column byte sizes that complete
the statistics row
+ // are only known afterwards. See statsRow.
+ private def gatherColumnStats(
+ batch: ColumnarBatch,
+ attrs: Seq[Attribute]): (Array[Any], Array[Any], Array[Int]) = {
+ val numCols = attrs.length
+ val lower = new Array[Any](numCols)
+ val upper = new Array[Any](numCols)
+ val nulls = Array.fill[Int](numCols)(0)
+ val numRows = batch.numRows()
+
+ var c = 0
+ while (c < numCols) {
+ val dt = attrs(c).dataType
+ val col = batch.column(c)
+ var r = 0
+ while (r < numRows) {
+ if (col.isNullAt(r)) {
+ nulls(c) += 1
+ } else if (tracksBounds(dt)) {
+ val value = readValue(col, dt, r)
+ if (lower(c) == null || compare(dt, value, lower(c)) < 0) {
+ lower(c) = value
+ }
+ if (upper(c) == null || compare(dt, value, upper(c)) > 0) {
+ upper(c) = value
+ }
+ }
+ r += 1
+ }
+ c += 1
+ }
+
+ (lower, upper, nulls)
+ }
+
+ // Build the statistics row expected by SimpleMetricsCachedBatchSerializer.
+ // For each cached column Spark expects five values in this order:
+ // lower bound, upper bound, null count, row count, and size in bytes.
+ private def statsRow(
+ lower: Array[Any],
+ upper: Array[Any],
+ nulls: Array[Int],
+ numRows: Int,
+ columnSizes: Array[Long]): InternalRow = {
+ val numCols = lower.length
+ val values = new Array[Any](numCols * 5)
+ var c = 0
+ while (c < numCols) {
+ val base = c * 5
+ values(base) = lower(c)
+ values(base + 1) = upper(c)
+ values(base + 2) = nulls(c)
+ values(base + 3) = numRows
+ // Each column is its own compressed stream, so its size is known
exactly. Cache pruning
+ // uses bounds/null-count/row-count rather than this field, but Spark
reserves it and
+ // reports it, so record the real value.
+ values(base + 4) = columnSizes(c)
+ c += 1
+ }
+
+ new GenericInternalRow(values)
+ }
+
+ // Spark can prune cache batches only for types whose bounds can be compared.
+ // Other types still report null count and row count but leave bounds as
null.
+ private def tracksBounds(dt: DataType): Boolean = dt match {
+ case BooleanType | ByteType | ShortType | IntegerType | LongType |
FloatType | DoubleType |
+ _: DecimalType | StringType | DateType | TimestampType |
TimestampNTZType =>
+ true
+ case _ => false
+ }
+
+ // Read a non-null value from a ColumnVector using Spark's internal value
type
+ // for the corresponding DataType.
+ private def readValue(col: ColumnVector, dt: DataType, rowId: Int): Any = dt
match {
+ case BooleanType => col.getBoolean(rowId)
+ case ByteType => col.getByte(rowId)
+ case ShortType => col.getShort(rowId)
+ case IntegerType | DateType => col.getInt(rowId)
+ case LongType | TimestampType | TimestampNTZType => col.getLong(rowId)
+ case FloatType => col.getFloat(rowId)
+ case DoubleType => col.getDouble(rowId)
+ case d: DecimalType => col.getDecimal(rowId, d.precision, d.scale)
+ case StringType => col.getUTF8String(rowId).copy()
+ case _ => null
+ }
+
+ // Compare values using the same physical representation used in the stats
row.
+ private def compare(dt: DataType, left: Any, right: Any): Int = dt match {
+ case BooleanType =>
+ java.lang.Boolean.compare(left.asInstanceOf[Boolean],
right.asInstanceOf[Boolean])
+ case ByteType =>
+ java.lang.Byte.compare(left.asInstanceOf[Byte], right.asInstanceOf[Byte])
+ case ShortType =>
+ java.lang.Short.compare(left.asInstanceOf[Short],
right.asInstanceOf[Short])
+ case IntegerType | DateType =>
+ java.lang.Integer.compare(left.asInstanceOf[Int],
right.asInstanceOf[Int])
+ case LongType | TimestampType | TimestampNTZType =>
+ java.lang.Long.compare(left.asInstanceOf[Long], right.asInstanceOf[Long])
+ case FloatType =>
+ java.lang.Float.compare(left.asInstanceOf[Float],
right.asInstanceOf[Float])
+ case DoubleType =>
+ java.lang.Double.compare(left.asInstanceOf[Double],
right.asInstanceOf[Double])
+ case _: DecimalType =>
+ left.asInstanceOf[Decimal].compare(right.asInstanceOf[Decimal])
+ case StringType =>
+ ByteArray.compareBinary(
+ left.asInstanceOf[UTF8String].getBytes,
+ right.asInstanceOf[UTF8String].getBytes)
+ case other =>
+ throw new IllegalStateException(s"compare called for unsupported type
$other")
+ }
+
+ // Compute Spark-compatible cache stats before serializing each batch to
Arrow.
+ // The stats are stored beside the Arrow bytes so Spark's cache filter can
prune
+ // CometCachedBatch without decoding the batch first.
+ //
+ // A columnar input batch is not guaranteed to be Arrow-backed; see
supportsColumnarInput for
+ // why. Batches that are not get copied into Arrow first, since
Utils.serializeBatches only
+ // writes CometVector columns.
+ private def encodeBatches(
+ batches: Iterator[ColumnarBatch],
+ attrs: Seq[Attribute]): Iterator[CachedBatch] = {
+ val arrowSchema =
+ Utils.toArrowSchema(Utils.fromAttributes(attrs),
CometArrowStream.NATIVE_TIMEZONE)
+
+ batches.map { batch =>
+ // Bounds and null counts are read from the input batch, which
serializing then clears, so
+ // they have to be gathered first. The row is only assembled once the
per-column sizes are
+ // known.
+ val (lower, upper, nulls) = gatherColumnStats(batch, attrs)
+ val numRows = batch.numRows()
+
+ val columns = if (Utils.isArrowBacked(batch)) {
+ Utils.serializeBatchColumns(batch)
+ } else {
+ val arrowBatch =
+ CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema,
CometArrowAllocator)
+ try Utils.serializeBatchColumns(arrowBatch)
+ finally arrowBatch.close()
Review Comment:
**[P2] Handle large-offset Arrow vectors before taking the direct write
path**
`Utils.isArrowBacked` accepts any `CometVector`, including a
`CometPlainVector` wrapping `LargeVarCharVector` or `LargeVarBinaryVector`.
However, `serializeBatchColumns` calls `Utils.getFieldVector`, which rejects
both representations. Thus `supportsColumnarInput` accepts the Spark
`StringType`/`BinaryType` schema, but materializing the cache fails.
At the serializer boundary, I reproduced both large-vector cases with just
`hello` and `world`: Spark's default cache round-trips them, while this
serializer throws `Unsupported Arrow Vector for serialize`. Ordinary
`VarCharVector` succeeds on both paths. Accelerated `mapInArrow` is a relevant
producer because its existing runner preserves returned Arrow vectors,
including `pa.large_string()`/`pa.large_binary()`. That Python integration path
was source-traced, not executed end to end; the serializer failure and
default-cache comparison were run. Please normalize these physical
representations through the conversion path or support them in the direct
writer.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala:
##########
@@ -0,0 +1,512 @@
+/*
+ * 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.execution.arrow
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.TaskContext
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression,
GenericInternalRow, IsNotNull, IsNull, UnsafeProjection}
+import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch,
SimpleMetricsCachedBatchSerializer}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.storage.StorageLevel
+import org.apache.spark.unsafe.types.{ByteArray, UTF8String}
+import org.apache.spark.util.io.ChunkedByteBuffer
+
+import org.apache.comet.CometArrowAllocator
+
+/**
+ * Cached batch format used when Comet writes Spark in-memory cache data.
+ *
+ * `columns` holds one compressed Arrow stream per cached column, in
cache-schema order, produced
+ * by `Utils.serializeBatchColumns`. Storing columns separately is what lets a
scan decode only
+ * the ones it projected; a single stream covering the whole batch would have
to be inflated in
+ * full before any projection could be applied. The cache manager still owns
storage and eviction;
+ * this class only changes the cached payload.
+ */
+private case class CometCachedBatch(
+ override val numRows: Int,
+ override val sizeInBytes: Long,
+ override val stats: InternalRow,
+ columns: Array[ChunkedByteBuffer])
+ extends SimpleMetricsCachedBatch
+
+/**
+ * Cache serializer that stores Comet-compatible Arrow batches in Spark's
in-memory cache.
+ *
+ * The cached payload format is decided by the schema alone. A relation whose
schema Comet's Arrow
+ * writer supports is stored as `CometCachedBatch`, and every other relation
is delegated in full
+ * to Spark's `DefaultCachedBatchSerializer`. The format deliberately does not
depend on any
+ * runtime config: `spark.sql.cache.serializer` is a static conf, so
installing this serializer is
+ * already a per-application decision, and a relation whose format could flip
mid-session cannot
+ * be read back reliably. `spark.comet.exec.inMemoryCache.enabled` still
governs whether a scan
+ * over the cache runs natively, and its value at startup is what makes
`CometDriverPlugin`
+ * install this serializer in the first place.
+ *
+ * Reads of `CometCachedBatch` keep working when the native scan is disabled,
because Spark then
+ * reads the same cached data through the SparkToColumnar fallback path.
+ */
+class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
+
+ import ArrowCachedBatchSerializer.supportsSchema
+
+ private val fallback = new DefaultCachedBatchSerializer()
+
+ // Bounds and null counts per column, gathered before the batch is
serialized: serializing
+ // clears the batch's vectors, and the per-column byte sizes that complete
the statistics row
+ // are only known afterwards. See statsRow.
+ private def gatherColumnStats(
+ batch: ColumnarBatch,
+ attrs: Seq[Attribute]): (Array[Any], Array[Any], Array[Int]) = {
+ val numCols = attrs.length
+ val lower = new Array[Any](numCols)
+ val upper = new Array[Any](numCols)
+ val nulls = Array.fill[Int](numCols)(0)
+ val numRows = batch.numRows()
+
+ var c = 0
+ while (c < numCols) {
+ val dt = attrs(c).dataType
+ val col = batch.column(c)
+ var r = 0
+ while (r < numRows) {
+ if (col.isNullAt(r)) {
+ nulls(c) += 1
+ } else if (tracksBounds(dt)) {
+ val value = readValue(col, dt, r)
+ if (lower(c) == null || compare(dt, value, lower(c)) < 0) {
+ lower(c) = value
+ }
+ if (upper(c) == null || compare(dt, value, upper(c)) > 0) {
+ upper(c) = value
+ }
+ }
+ r += 1
+ }
+ c += 1
+ }
+
+ (lower, upper, nulls)
+ }
+
+ // Build the statistics row expected by SimpleMetricsCachedBatchSerializer.
+ // For each cached column Spark expects five values in this order:
+ // lower bound, upper bound, null count, row count, and size in bytes.
+ private def statsRow(
+ lower: Array[Any],
+ upper: Array[Any],
+ nulls: Array[Int],
+ numRows: Int,
+ columnSizes: Array[Long]): InternalRow = {
+ val numCols = lower.length
+ val values = new Array[Any](numCols * 5)
+ var c = 0
+ while (c < numCols) {
+ val base = c * 5
+ values(base) = lower(c)
+ values(base + 1) = upper(c)
+ values(base + 2) = nulls(c)
+ values(base + 3) = numRows
+ // Each column is its own compressed stream, so its size is known
exactly. Cache pruning
+ // uses bounds/null-count/row-count rather than this field, but Spark
reserves it and
+ // reports it, so record the real value.
+ values(base + 4) = columnSizes(c)
+ c += 1
+ }
+
+ new GenericInternalRow(values)
+ }
+
+ // Spark can prune cache batches only for types whose bounds can be compared.
+ // Other types still report null count and row count but leave bounds as
null.
+ private def tracksBounds(dt: DataType): Boolean = dt match {
+ case BooleanType | ByteType | ShortType | IntegerType | LongType |
FloatType | DoubleType |
+ _: DecimalType | StringType | DateType | TimestampType |
TimestampNTZType =>
+ true
+ case _ => false
+ }
+
+ // Read a non-null value from a ColumnVector using Spark's internal value
type
+ // for the corresponding DataType.
+ private def readValue(col: ColumnVector, dt: DataType, rowId: Int): Any = dt
match {
+ case BooleanType => col.getBoolean(rowId)
+ case ByteType => col.getByte(rowId)
+ case ShortType => col.getShort(rowId)
+ case IntegerType | DateType => col.getInt(rowId)
+ case LongType | TimestampType | TimestampNTZType => col.getLong(rowId)
+ case FloatType => col.getFloat(rowId)
+ case DoubleType => col.getDouble(rowId)
+ case d: DecimalType => col.getDecimal(rowId, d.precision, d.scale)
+ case StringType => col.getUTF8String(rowId).copy()
+ case _ => null
+ }
+
+ // Compare values using the same physical representation used in the stats
row.
+ private def compare(dt: DataType, left: Any, right: Any): Int = dt match {
+ case BooleanType =>
+ java.lang.Boolean.compare(left.asInstanceOf[Boolean],
right.asInstanceOf[Boolean])
+ case ByteType =>
+ java.lang.Byte.compare(left.asInstanceOf[Byte], right.asInstanceOf[Byte])
+ case ShortType =>
+ java.lang.Short.compare(left.asInstanceOf[Short],
right.asInstanceOf[Short])
+ case IntegerType | DateType =>
+ java.lang.Integer.compare(left.asInstanceOf[Int],
right.asInstanceOf[Int])
+ case LongType | TimestampType | TimestampNTZType =>
+ java.lang.Long.compare(left.asInstanceOf[Long], right.asInstanceOf[Long])
+ case FloatType =>
+ java.lang.Float.compare(left.asInstanceOf[Float],
right.asInstanceOf[Float])
+ case DoubleType =>
+ java.lang.Double.compare(left.asInstanceOf[Double],
right.asInstanceOf[Double])
+ case _: DecimalType =>
+ left.asInstanceOf[Decimal].compare(right.asInstanceOf[Decimal])
+ case StringType =>
+ ByteArray.compareBinary(
+ left.asInstanceOf[UTF8String].getBytes,
+ right.asInstanceOf[UTF8String].getBytes)
+ case other =>
+ throw new IllegalStateException(s"compare called for unsupported type
$other")
+ }
+
+ // Compute Spark-compatible cache stats before serializing each batch to
Arrow.
+ // The stats are stored beside the Arrow bytes so Spark's cache filter can
prune
+ // CometCachedBatch without decoding the batch first.
+ //
+ // A columnar input batch is not guaranteed to be Arrow-backed; see
supportsColumnarInput for
+ // why. Batches that are not get copied into Arrow first, since
Utils.serializeBatches only
+ // writes CometVector columns.
+ private def encodeBatches(
+ batches: Iterator[ColumnarBatch],
+ attrs: Seq[Attribute]): Iterator[CachedBatch] = {
+ val arrowSchema =
+ Utils.toArrowSchema(Utils.fromAttributes(attrs),
CometArrowStream.NATIVE_TIMEZONE)
+
+ batches.map { batch =>
+ // Bounds and null counts are read from the input batch, which
serializing then clears, so
+ // they have to be gathered first. The row is only assembled once the
per-column sizes are
+ // known.
+ val (lower, upper, nulls) = gatherColumnStats(batch, attrs)
+ val numRows = batch.numRows()
+
+ val columns = if (Utils.isArrowBacked(batch)) {
+ Utils.serializeBatchColumns(batch)
+ } else {
+ val arrowBatch =
+ CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema,
CometArrowAllocator)
+ try Utils.serializeBatchColumns(arrowBatch)
+ finally arrowBatch.close()
+ }
+
+ val columnSizes = columns.map(_.size)
+ CometCachedBatch(
+ numRows = numRows,
+ sizeInBytes = columnSizes.sum,
+ stats = statsRow(lower, upper, nulls, numRows, columnSizes),
+ columns = columns)
+ }
+ }
+
+ // Resolve requested columns by exprId, not by name, because aliases may
reuse names.
+ //
+ // An empty selection stays empty rather than expanding to every column.
Spark asks for no
+ // columns when the query only needs the row count (SELECT count(*)), and
since projection now
+ // decides what gets decoded, expanding it would turn the cheapest possible
read into the most
+ // expensive one.
+ private def selectedIndices(
+ cacheAttributes: Seq[Attribute],
+ selectedAttributes: Seq[Attribute]): Array[Int] = {
+ val byExprId = cacheAttributes.zipWithIndex.map { case (attr, idx) =>
+ attr.exprId -> idx
+ }.toMap
+
+ selectedAttributes.map { attr =>
+ byExprId.getOrElse(
+ attr.exprId,
+ throw new IllegalStateException(
+ s"Could not resolve selected attribute ${attr.name} from cache
attributes"))
+ }.toArray
+ }
+
+ // Spark's SimpleMetricsCachedBatchSerializer prunes a batch when the
generated partition filter
+ // does not evaluate to true against the stats row. Bounds are only computed
for the types
+ // tracksBounds accepts, and for every other column the lower and upper
bounds stay null, which
+ // makes a comparison against them evaluate to null and therefore prune the
batch. That would
+ // silently drop rows, so predicates over columns without bounds are not
pushed down at all.
+ // Null counts and row counts are recorded for every column, so IsNull and
IsNotNull stay safe.
+ override def buildFilter(
+ predicates: Seq[Expression],
+ cachedAttributes: Seq[Attribute]): (Int, Iterator[CachedBatch]) =>
Iterator[CachedBatch] = {
+ val prunable = cachedAttributes.collect {
+ case a if tracksBounds(a.dataType) => a.exprId
+ }.toSet
+
+ val prunablePredicates = predicates.filter {
+ case _: IsNull | _: IsNotNull => true
+ case p => p.references.forall(a => prunable.contains(a.exprId))
+ }
+
+ super.buildFilter(prunablePredicates, cachedAttributes)
+ }
+
+ // Comet's Arrow writer only handles the types listed in supportsSchema.
Reporting false here
+ // sends the relation down the row path, where it is delegated to Spark's
default serializer,
+ // instead of failing at cache materialization inside Utils.serializeBatches.
+ //
+ // This answer is schema-only, because attributes are all Spark gives us; it
says nothing about
+ // the vectors. Returning true also makes InMemoryRelation strip the
ColumnarToRow above the
+ // cached plan, so convertColumnarBatchToCachedBatch then receives whatever
that plan produces:
+ // a Comet scan's CometVectors, but equally Spark's vectorized Parquet/ORC
reader or a
+ // connector's own vectors. encodeBatches converts the non-Arrow ones; that
conversion is load
+ // bearing, not defensive.
+ override def supportsColumnarInput(schema: Seq[Attribute]): Boolean =
supportsSchema(schema)
+
+ // A relation Comet stores is always readable as columnar Arrow. Anything
else holds
+ // DefaultCachedBatch, so defer to Spark, which only claims columnar output
for the primitive
+ // types its ColumnAccessor.decompress path can actually decode.
+ override def supportsColumnarOutput(schema: StructType): Boolean = {
+ if (schema.fields.forall(f =>
ArrowCachedBatchSerializer.supportsType(f.dataType))) {
+ true
+ } else {
+ fallback.supportsColumnarOutput(schema)
+ }
+ }
+
+ // Columnar Comet output is stored as compressed Arrow stream bytes. Spark
only calls this when
+ // supportsColumnarInput returned true, so the schema is known to be
Comet-writable here.
+ override def convertColumnarBatchToCachedBatch(
+ input: RDD[ColumnarBatch],
+ schema: Seq[Attribute],
+ storageLevel: StorageLevel,
+ conf: SQLConf): RDD[CachedBatch] = {
+
+ input.mapPartitions { batches =>
+ encodeBatches(batches, schema)
+ }
+ }
+
+ override def convertCachedBatchToColumnarBatch(
+ input: RDD[CachedBatch],
+ cacheAttributes: Seq[Attribute],
+ selectedAttributes: Seq[Attribute],
+ conf: SQLConf): RDD[ColumnarBatch] = {
+ if (!supportsSchema(cacheAttributes)) {
+ return fallback.convertCachedBatchToColumnarBatch(
+ input,
+ cacheAttributes,
+ selectedAttributes,
+ conf)
+ }
+
+ val indices = selectedIndices(cacheAttributes, selectedAttributes)
+
+ input.mapPartitions { it =>
+ // A ColumnReaders closes its readers (releasing the vectors they are
holding) only when the
+ // batch it produced has been consumed. A consumer that stops early --
LIMIT, take(), or a
+ // cancelled task -- leaves the readers for the batch in flight open, so
close them on task
+ // completion. Spark's own ArrowCachedBatchSerializer registers a
listener for the same
+ // reason.
+ //
+ // flatMap consumes each inner iterator fully before building the next,
so at most one batch
+ // is open at a time and tracking the current one is enough. close() is
idempotent, so
+ // closing one that already released itself is a no-op.
+ @volatile var current: ColumnReaders = null
+ Option(TaskContext.get()).foreach { tc =>
+ tc.addTaskCompletionListener[Unit] { _ =>
+ val readers = current
+ current = null
+ if (readers != null) {
+ readers.close()
+ }
+ }
+ }
+
+ it.flatMap {
+ case cb: CometCachedBatch =>
+ if (indices.isEmpty) {
+ // Nothing to decode: the row count is the whole answer, and it is
already here.
+ Iterator.single(new ColumnarBatch(Array.empty[ColumnVector],
cb.numRows))
+ } else {
+ val readers = new ColumnReaders(indices.map(i => cb.columns(i)),
cb.numRows)
+ current = readers
+ readers.batches
+ }
+
+ case other =>
+ throw new IllegalStateException(
+ s"Unsupported cached batch type ${other.getClass.getName}")
+ }
+ }
+ }
+
+ // Decodes one selected column stream apiece and stitches the results back
into a single batch.
+ //
+ // Each stream is self-contained, so the columns a scan did not select are
never inflated. The
+ // decoded vectors stay owned by their readers: closing them releases the
batch, which is why
+ // this yields a single-element iterator that closes on exhaustion, matching
what
+ // ArrowReaderIterator did when the payload was one stream.
+ private class ColumnReaders(buffers: Array[ChunkedByteBuffer], numRows: Int)
{
+ private val readers: Array[Iterator[ColumnarBatch]] =
+ buffers.map(Utils.decodeBatches(_, "CometCache"))
Review Comment:
**[P2] Close previously opened readers when initialization fails**
`decodeBatches` opens and decodes a column eagerly. If opening a later
column throws, the readers already created by this `map` are lost. The
task-completion listener cannot release them because the new `ColumnReaders`
instance is assigned to `current` only after this constructor returns.
I reproduced this with valid cached data containing two integer columns and
an allocator limit that permits the first column but rejects the second
allocation. After task completion, the first column's 512 bytes remain
allocated; three failed tasks leave **512, 1024, and 1536 bytes** respectively,
starting from zero. This turns allocation failures into persistent off-heap
leaks and worsens subsequent memory pressure. Please close every successfully
opened reader if constructing the remaining readers fails.
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala:
##########
@@ -0,0 +1,169 @@
+/*
+ * 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.{CachedBatch, CachedBatchSerializer}
+import org.apache.spark.sql.execution.LeafExecNode
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.types._
+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,
+ cachedBuffers: RDD[CachedBatch],
+ 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"))
+
+ // For an empty-projection scan (`SELECT count(*)`) this is empty while
`scanOutput` holds one
+ // placeholder column, so the emitted batches are wider than the declared
output. That is safe
+ // because the only consumer of an empty-output scan is a count-style
aggregate, which reads the
+ // row count rather than any column; see `scanOutputFor` for why the scan
cannot simply be empty.
+ override def output: Seq[Attribute] = originalPlan.output
+
+ // 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")
+
+ 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 = scanOutputFor(op).flatMap(attr =>
serializeDataType(attr.dataType))
+
+ val scanBuilder = OperatorOuterClass.Scan
+ .newBuilder()
+ .setSource(op.getClass.getSimpleName)
+ .addAllFields(scanTypes.asJava)
+
+ Some(builder.setScan(scanBuilder).build())
+ }
+
+ // Reuse Spark's InMemoryRelation metadata so cache materialization,
pruning, and storage
+ // behavior remain controlled by Spark's cache manager.
+ override def createExec(nativeOp: Operator, op: InMemoryTableScanExec):
CometNativeExec = {
+ val relation = op.relation
+
+ CometScanWrapper(
+ nativeOp,
+ CometInMemoryTableScanExec(
+ op,
+ relation.cacheBuilder.serializer,
+ relation.cacheBuilder.cachedColumnBuffers,
Review Comment:
**[P2] Defer cache RDD construction until execution**
Accessing `cachedColumnBuffers` here is not a passive metadata lookup:
Spark's cache builder constructs it by calling
`cachedPlan.execute`/`executeColumnar`. If that plan is adaptive, this can
execute shuffle stages while Comet is still planning the outer query.
With AQE and native caching enabled, I reproduced:
```scala
val cached = spark.range(100).repartition(2).cache()
cached.createOrReplaceTempView("cached_adaptive")
spark.sql("SELECT * FROM cached_adaptive").explain()
```
`explain()` starts **one Spark job** and changes the cached adaptive plan to
`isFinalPlan=true`; with native cache scanning disabled and the same cache
serializer, it starts **zero jobs** and the cached plan stays unexecuted.
Planning/EXPLAIN can therefore perform expensive computation or fail on
execution errors. Please retain the relation/cache builder and obtain its RDD
inside `doExecuteColumnar()`.
--
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]
