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

["uros-b (via GitHub)" <gi...@apache.org>]

uros-b commented on code in PR #57395:
URL: https://github.com/apache/spark/pull/57395#discussion_r3715226967


##########
sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala:
##########
@@ -197,8 +197,12 @@ class DefaultCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
       conf: SQLConf): RDD[ColumnarBatch] = {
     val offHeapColumnVectorEnabled = conf.offHeapColumnVectorEnabled
     val outputSchema = DataTypeUtils.fromAttributes(selectedAttributes)
+    // Use AttributeSeq's cached exprId -> ordinal map so each selected attribute is a
+    // constant-time lookup, instead of rebuilding the exprId list and linear-scanning it
+    // per selected attribute.
+    val cacheAttributeSeq = AttributeSeq(cacheAttributes)
     val columnIndices =
-      selectedAttributes.map(a => cacheAttributes.map(o => o.exprId).indexOf(a.exprId)).toArray
+      selectedAttributes.map(a => cacheAttributeSeq.indexOf(a.exprId)).toArray

Review Comment:
   This same computation now appears four times: here, at InMemoryRelation.scala:240-243, and at ArrowCachedBatchSerializer.scala:150-152 and :180-183. Since the point of the PR is to retire the hand-rolled scan, it would be better to leave exactly one implementation behind, so the next serializer that needs column indices calls it instead of copying the pattern a fifth time. Somewhere shared, e.g. the CachedBatchSerializer trait or the columnar package object:
   
   ```
   private[columnar] def cachedColumnIndices(
       cacheAttributes: Seq[Attribute],
       selectedAttributes: Seq[Attribute]): Array[Int] = {
     val cacheAttributeSeq = AttributeSeq(cacheAttributes)
     selectedAttributes.map(a => cacheAttributeSeq.indexOf(a.exprId)).toArray
   }
   ```
   
   The row path in DefaultCachedBatchSerializer needs the data types alongside the ordinals, but it can call the helper and zip with selectedAttributes.map(_.dataType) rather than keep its own copy. This also collapses the four duplicated comments into one.



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org