This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 8a6de5b30114 [SPARK-58232][SQL] Simplify cached-batch column-index
resolution with a map lookup
8a6de5b30114 is described below
commit 8a6de5b3011444c18a568fc1efedb518f3975b0d
Author: Ganesha S <[email protected]>
AuthorDate: Tue Aug 11 11:48:12 2026 +0800
[SPARK-58232][SQL] Simplify cached-batch column-index resolution with a map
lookup
### What changes were proposed in this pull request?
When reading from an in-memory (cached) relation,
`DefaultCachedBatchSerializer`
(`sql/core/.../execution/columnar/InMemoryRelation.scala`) maps each
selected output
attribute to its ordinal in the cached schema. Both conversion methods
computed this
with an O(n*m) algorithm that, for each of the `m` selected attributes,
rebuilt the full
list of `n` cached-schema `ExprId`s and linear-scanned it with `indexOf`:
```scala
val columnIndices =
selectedAttributes.map(a => cacheAttributes.map(o =>
o.exprId).indexOf(a.exprId)).toArray
```
This PR builds a single ExprId -> ordinal map once (O(n)) and looks up each
selected
attribute in O(1), for overall O(n+m) and no per-attribute list allocation:
```scala
val cacheAttributeOrdinals =
cacheAttributes.iterator.map(_.exprId).zipWithIndex.toMap
val columnIndices =
selectedAttributes.map(a => cacheAttributeOrdinals.getOrElse(a.exprId,
-1)).toArray
```
The same change is applied to both:
- convertCachedBatchToColumnarBatch (vectorized path)
- convertCachedBatchToInternalRow (row path)
`getOrElse(exprId, -1)` preserves the prior indexOf semantics of returning
-1 when
an attribute is absent.
### Why are the changes needed?
The old implementation is O(n*m) in time and transient allocation: for each
of the `m`
selected attributes it rebuilds the full `n`-element list of cached
`ExprId`s and
linear-scans it. The map-based version is O(n+m). A microbenchmark
isolating the
`columnIndices` computation (all columns selected in reverse order so
`indexOf` cannot
short-circuit; 20k calls per case; asserts the old and new implementations
produce
identical index arrays before timing) shows the expected
quadratic-vs-linear divergence:
```
OpenJDK 64-Bit Server VM 17.0.15+6-Ubuntu-0ubuntu120.04 on Linux
5.4.0-1160-aws-fips
Intel(R) Xeon(R) Platinum 8375C CPU 2.90GHz
columnIndices for 50 columns (20000 calls): Best Time(ms) Avg
Time(ms) Relative
--------------------------------------------------------------------------------------
old: rebuild list + indexOf per attr (O(n*m)) 223
228 1.0X
new: exprId->ordinal map (O(n+m)) 83
87 2.7X
columnIndices for 200 columns (20000 calls): Best Time(ms) Avg
Time(ms) Relative
--------------------------------------------------------------------------------------
old: rebuild list + indexOf per attr (O(n*m)) 2828
2850 1.0X
new: exprId->ordinal map (O(n+m)) 339
344 8.3X
columnIndices for 500 columns (20000 calls): Best Time(ms) Avg
Time(ms) Relative
--------------------------------------------------------------------------------------
old: rebuild list + indexOf per attr (O(n*m)) 18881
18883 1.0X
new: exprId->ordinal map (O(n+m)) 832
844 22.7X
columnIndices for 1000 columns (20000 calls): Best Time(ms) Avg
Time(ms) Relative
--------------------------------------------------------------------------------------
old: rebuild list + indexOf per attr (O(n*m)) 75310
76013 1.0X
new: exprId->ordinal map (O(n+m)) 2146
2148 35.1X
```
To be clear about scope: `columnIndices` is computed once per scan
execution on the
driver , so in absolute terms this saves microseconds per query for typical
tables and up to ~3.7 ms per query only for a pathologically wide
1000-column cached relation. This is primarily a code-quality cleanup,
removing a quadratic algorithm and per-attribute list reallocation in favor
of the
obvious map lookup, with the benchmark included to confirm the algorithmic
improvement
rather than to claim a meaningful end-to-end speedup..
### Does this PR introduce _any_ user-facing change?
No. This is a behavior-preserving internal optimization; output schema,
ordering, and
data types are unchanged.
### How was this patch tested?
- `build/sbt sql/compile` passes.
- Existing `CachedBatchSerializerSuite` and `InMemoryColumnarQuerySuite`
(36 tests,
including column pruning and reordering cases that directly exercise this
mapping)
pass.
- A focused micro-benchmark was written to isolate the column-index
computation, using
real `AttributeReference`/`ExprId` objects with all columns selected in
reverse order
(so `indexOf` cannot short-circuit at position 0), 20k calls per case. It
asserted the
old and new implementations produce identical index arrays before timing,
confirming
behavior is preserved, and measured 2.7x (50 columns) to 35.1x (1000
columns) speedups
(see the table above). The benchmark was a temporary verification aid and
is not
included in this PR.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #57395 from ganeshashree/SPARK-58232.
Lead-authored-by: Ganesha S <[email protected]>
Co-authored-by: Ganesha S <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit 00034886bf5c4ff175b5e6829f249b436320eaec)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../columnar/ArrowCachedBatchSerializer.scala | 11 +++---
.../execution/columnar/CachedColumnIndices.scala | 40 ++++++++++++++++++++++
.../sql/execution/columnar/InMemoryRelation.scala | 14 +++-----
3 files changed, 49 insertions(+), 16 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala
index f90f1a9dcd22..7f8ddb177788 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala
@@ -144,8 +144,7 @@ class ArrowCachedBatchSerializer extends
SimpleMetricsCachedBatchSerializer {
conf: SQLConf): RDD[ColumnarBatch] = {
val cacheSchema = DataTypeUtils.fromAttributes(cacheAttributes)
val selectedSchema = DataTypeUtils.fromAttributes(selectedAttributes)
- val columnIndices =
- selectedAttributes.map(a => cacheAttributes.map(o =>
o.exprId).indexOf(a.exprId)).toArray
+ val columnIndices = CachedColumnIndices(cacheAttributes,
selectedAttributes)
// Capture config on driver
val timeZoneId = conf.sessionLocalTimeZone
val prefetchEnabled = conf.arrowCachePrefetchEnabled
@@ -183,11 +182,6 @@ class ArrowCachedBatchSerializer extends
SimpleMetricsCachedBatchSerializer {
val selectedSchema = DataTypeUtils.fromAttributes(selectedAttributes)
val timeZoneId = conf.sessionLocalTimeZone
- // Calculate column indices for projection
- val selectedIndices = selectedAttributes.map { attr =>
- cacheAttributes.indexWhere(_.exprId == attr.exprId)
- }.toArray
-
// Check if all selected types can use the fast path.
// Types not handled by ArrowColumnReader must use the fallback path.
val needsFallback = selectedSchema.fields.exists { f =>
@@ -227,6 +221,9 @@ class ArrowCachedBatchSerializer extends
SimpleMetricsCachedBatchSerializer {
}
}
} else {
+ // Only the fast path consumes the column indices; the fallback branch
above delegates to
+ // convertCachedBatchToColumnarBatch, which resolves them itself.
+ val selectedIndices = CachedColumnIndices(cacheAttributes,
selectedAttributes)
val prefetchEnabled = conf.arrowCachePrefetchEnabled
input.mapPartitionsInternal { batchIterator =>
new ArrowCachedBatchToInternalRowIterator(
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/CachedColumnIndices.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/CachedColumnIndices.scala
new file mode 100644
index 000000000000..14fd150dfa7f
--- /dev/null
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/CachedColumnIndices.scala
@@ -0,0 +1,40 @@
+/*
+ * 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.execution.columnar
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSeq}
+
+/**
+ * Resolves where the columns a scan requests sit within a cached batch's
schema. Every
+ * `CachedBatchSerializer` needs this mapping to project the requested columns
out of a cached
+ * batch, so it lives here rather than in each serializer.
+ */
+private[columnar] object CachedColumnIndices {
+
+ /**
+ * Returns, for each attribute in `selectedAttributes`, its ordinal in
`cacheAttributes`, or -1
+ * if the attribute is not part of the cached schema.
+ */
+ def apply(cacheAttributes: Seq[Attribute], selectedAttributes:
Seq[Attribute]): Array[Int] = {
+ // The explicit AttributeSeq wrapper is required. On a bare Seq[Attribute]
the inherited
+ // Seq.indexOf[B >: Attribute](elem: B) wins overload resolution with B
inferred as Any, so
+ // the implicit conversion never fires and every column silently resolves
to -1.
+ val cacheAttributeSeq = AttributeSeq(cacheAttributes)
+ selectedAttributes.map(a => cacheAttributeSeq.indexOf(a.exprId)).toArray
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala
index 44197efd0e1b..e6bf65ec89e6 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala
@@ -205,8 +205,7 @@ class DefaultCachedBatchSerializer extends
SimpleMetricsCachedBatchSerializer {
conf: SQLConf): RDD[ColumnarBatch] = {
val offHeapColumnVectorEnabled = conf.offHeapColumnVectorEnabled
val outputSchema = DataTypeUtils.fromAttributes(selectedAttributes)
- val columnIndices =
- selectedAttributes.map(a => cacheAttributes.map(o =>
o.exprId).indexOf(a.exprId)).toArray
+ val columnIndices = CachedColumnIndices(cacheAttributes,
selectedAttributes)
def createAndDecompressColumn(cb: CachedBatch): ColumnarBatch = {
val cachedColumnarBatch = cb.asInstanceOf[DefaultCachedBatch]
@@ -239,21 +238,18 @@ class DefaultCachedBatchSerializer extends
SimpleMetricsCachedBatchSerializer {
selectedAttributes: Seq[Attribute],
conf: SQLConf): RDD[InternalRow] = {
// Find the ordinals and data types of the requested columns.
- val (requestedColumnIndices, requestedColumnDataTypes) =
- selectedAttributes.map { a =>
- cacheAttributes.map(_.exprId).indexOf(a.exprId) -> a.dataType
- }.unzip
+ val requestedColumnIndices = CachedColumnIndices(cacheAttributes,
selectedAttributes)
- val columnTypes = requestedColumnDataTypes.map {
+ val columnTypes = selectedAttributes.map(_.dataType match {
case udt: UserDefinedType[_] => udt.sqlType
case other => other
- }.toArray
+ }).toArray
input.mapPartitionsInternal { cachedBatchIterator =>
val columnarIterator =
GenerateColumnAccessor.generate(columnTypes.toImmutableArraySeq)
columnarIterator.initialize(cachedBatchIterator.asInstanceOf[Iterator[DefaultCachedBatch]],
columnTypes,
- requestedColumnIndices.toArray)
+ requestedColumnIndices)
columnarIterator
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
