feat: add experimental native support for in-memory cache, disabled by default #5051
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,143 @@ | ||
| /* | ||
| * 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.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 the | ||
| // full cache schema, 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; `convert` and `createExec` deliberately fall back to | ||
| // the cache schema in that case because the native plan still needs a non-empty scan schema. | ||
| 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] = { | ||
|
Member
There was a problem hiding this comment. pruning configuration ignored: still prunes when not sure enable
Member
Author
There was a problem hiding this comment. Real bug, fixed in 1367665 — thanks. Spark's To your "not sure enable partitionPruning is always better than disable": pruning is normally the win, but that is not really the point — the config exists so a user can turn it off, typically to rule out a stats bug when results look wrong. Ignoring it means the one knob they are reaching for does nothing, and Comet diverges from Spark precisely when someone is debugging. So honoring it is right regardless of which default is faster. Covered by a new test that observes the scan's |
||
| 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] = { | ||
|
|
||
| // Empty-output scans still need a schema for native planning, so fall back to the cache schema. | ||
| val actualOutput = | ||
| if (op.output.nonEmpty) op.output | ||
| else op.relation.output | ||
|
|
||
| val scanTypes = actualOutput.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 | ||
|
|
||
| val actualOutput = | ||
| if (op.output.nonEmpty) op.output | ||
| else relation.output | ||
|
|
||
| CometScanWrapper( | ||
| nativeOp, | ||
| CometInMemoryTableScanExec( | ||
| op, | ||
| relation.cacheBuilder.serializer, | ||
| relation.cacheBuilder.cachedColumnBuffers, | ||
| relation.output, | ||
| actualOutput)) | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Add AQE test:
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not done, and I would rather flag that than half-do it. Porting the three AQE scenarios you linked (SPARK-42101 cold/warm materialization, the
TableCacheQueryStageExecjoin, SPARK-37742) is a meaningful chunk of work and touches AQE plan-shape assertions rather than the cache format this PR changes, so I have left it out of this round rather than rushing it alongside the four correctness fixes.You said your comments could be follow-ups, so unless you would rather block on it, my suggestion is a tracking issue for AQE coverage of
CometInMemoryTableScanExec. I have not filed one yet — say the word and I will, or feel free to open it yourself if you have a preference for how it should be scoped.