[SPARK-58232][SQL] Simplify cached-batch column-index resolution with a map lookup

[ganeshashree]

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 ExprIds and linear-scanned it with indexOf:

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:

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 ExprIds 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)

[uros-b]

On another note, PTAL at ArrowCachedBatchSerializer.scala:148 - the sibling serializer's convertCachedBatchToColumnarBatch still carries the identical O(n·m) pattern (selectedAttributes.map(a => cacheAttributes.map(o => o.exprId).indexOf(a.exprId)).toArray) that this PR removes from DefaultCachedBatchSerializer.

Since the stated goal is removing the quadratic column-index computation, leaving it in the Arrow fast path (the serializer used when Arrow caching is enabled, SPARK-57268, and at least as relevant for wide tables) makes the cleanup incomplete. The fix is a direct copy of the same map lookup.

[ganeshashree]

On another note, PTAL at ArrowCachedBatchSerializer.scala:148 - the sibling serializer's convertCachedBatchToColumnarBatch still carries the identical O(n·m) pattern (selectedAttributes.map(a => cacheAttributes.map(o => o.exprId).indexOf(a.exprId)).toArray) that this PR removes from DefaultCachedBatchSerializer.

Since the stated goal is removing the quadratic column-index computation, leaving it in the Arrow fast path (the serializer used when Arrow caching is enabled, SPARK-57268, and at least as relevant for wide tables) makes the cleanup incomplete. The fix is a direct copy of the same map lookup.

ArrowCachedBatchSerializer.convertCachedBatchToColumnarBatch now uses AttributeSeq.indexOf as well, so the quadratic pattern is gone from the Arrow fast path too. While there, I noticed the sibling convertCachedBatchToInternalRow in the same serializer had the same defect in a slightly different form: cacheAttributes.indexWhere(_.exprId == attr.exprId), so I converted that one too. Both Arrow paths now go through the shared AttributeSeq map lookup, matching DefaultCachedBatchSerializer.
Verified with ArrowCachedBatchSerializerSuite (covering both the columnar and row paths, including prefetch), all green.