Copilot commented on code in PR #13280:
URL: https://github.com/apache/gluten/pull/13280#discussion_r4333756015
##########
backends-velox/src/main/scala/org/apache/spark/sql/execution/utils/ExecUtil.scala:
##########
@@ -142,18 +144,20 @@ object ExecUtil {
.filter(cb => cb.numRows != 0 && cb.numCols != 0)
.map {
cb =>
- val pidVec = ArrowWritableColumnVector
- .allocateColumns(cb.numRows, new StructType().add("pid",
IntegerType))
- .head
+ val pidRoot =
+ ArrowColumnVectors.allocateRoot(new StructType().add("pid",
IntegerType))
+ val pidVec = pidRoot.getVector(0).asInstanceOf[IntVector]
+ pidVec.allocateNew(cb.numRows)
convertColumnarToRow(cb).zipWithIndex.foreach {
case (row, i) =>
val pid =
rangePartitioner.get.getPartition(partitionKeyExtractor(row))
- pidVec.putInt(i, pid)
+ pidVec.set(i, pid)
}
+ pidRoot.setRowCount(cb.numRows)
val pidBatch = VeloxColumnarBatches.toVeloxBatch(
ColumnarBatches.offload(
ArrowBufferAllocators.contextInstance(),
- new ColumnarBatch(Array[ColumnVector](pidVec),
cb.numRows)))
+ ArrowColumnVectors.toBatch(pidRoot)))
Review Comment:
Similar to `ColumnarRangeExec`, this fills an `IntVector` but doesn’t set
`pidVec.setValueCount(cb.numRows)` (only `pidRoot.setRowCount`). If downstream
export/offload uses the vectors’ `valueCount`, this can produce incorrect/empty
output. Set the vector `valueCount` explicitly or centralize the fix in
`ArrowColumnVectors.toBatch`.
##########
gluten-arrow/src/main/java/org/apache/gluten/columnarbatch/ColumnarBatches.java:
##########
@@ -227,6 +266,12 @@ public static ColumnarBatch load(BufferAllocator
allocator, ColumnarBatch input)
}
public static ColumnarBatch offload(BufferAllocator allocator, ColumnarBatch
input) {
+ if (isForeignArrowBatch(input)) {
+ // A batch of Spark ArrowColumnVectors (e.g. the output of a Spark Arrow
operator) is owned
+ // by its producer and only valid until the next one is requested: take
ownership of its
+ // vectors in a Gluten batch and offload that one, leaving the input to
its producer.
+ return offload(allocator, adopt(allocator, input));
+ }
Review Comment:
The comment says the input is left to its producer, but `adopt(...)` may do
a zero-copy `transfer()` when allocator roots match, which can *mutate/empty*
the input vectors even though the input batch isn’t closed here. Consider
clarifying the comment to reflect that adoption may move buffers (leaving input
vectors emptied) vs copy, or change adoption behavior to always copy for
foreign batches if preserving producer-owned vectors is required (with the
obvious performance tradeoff).
##########
gluten-arrow/src/main/scala/org/apache/gluten/vectorized/ArrowColumnVectors.scala:
##########
@@ -0,0 +1,96 @@
+/*
+ * 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.gluten.vectorized
+
+import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.utils.{SparkArrowUtil, SparkSchemaUtil}
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.{FieldVector, ValueVector, VectorSchemaRoot}
+import org.apache.arrow.vector.util.VectorAppender
+
+import scala.collection.JavaConverters._
+
+/** Utilities to build Gluten's Java Arrow batches of
[[RefCountedArrowColumnVector]]s. */
+object ArrowColumnVectors {
+
+ /** Allocates an empty [[VectorSchemaRoot]] for `schema`. */
+ def allocateRoot(
+ schema: StructType,
+ allocator: BufferAllocator = ArrowBufferAllocators.contextInstance()):
VectorSchemaRoot = {
+ val arrowSchema = SparkArrowUtil.toArrowSchema(schema,
SparkSchemaUtil.getLocalTimezoneID)
+ VectorSchemaRoot.create(arrowSchema, allocator)
+ }
+
+ /** Wraps `vectors`, whose ownership moves to the returned column vectors. */
+ def wrap(vectors: java.util.List[FieldVector]):
Array[RefCountedArrowColumnVector] = {
+ vectors.asScala.map(new RefCountedArrowColumnVector(_)).toArray
+ }
+
+ /** Wraps the vectors of `root` into a batch, which owns them. */
+ def toBatch(root: VectorSchemaRoot): ColumnarBatch = {
+ new ColumnarBatch(
+ wrap(root.getFieldVectors).map(_.asInstanceOf[ColumnVector]),
+ root.getRowCount)
+ }
Review Comment:
`toBatch` relies on `root.getRowCount`, but it does not ensure each
underlying `FieldVector` has `valueCount` set. Several new call sites only call
`root.setRowCount(...)` after filling vectors, which may not propagate
`valueCount` to the vectors; export/offload paths typically use
`vector.getValueCount`. Consider setting `setValueCount(root.getRowCount)` on
all `root.getFieldVectors` inside `toBatch` (or require callers to set vector
value counts explicitly and update those call sites).
##########
backends-velox/src/main/scala/org/apache/gluten/execution/ColumnarRangeExec.scala:
##########
@@ -98,17 +100,17 @@ case class ColumnarRangeExec(range: LogicalRange) extends
ColumnarRangeBaseExec
batchSize
)
- val vectors =
ArrowWritableColumnVector.allocateColumns(numRows, schema)
-
+ val root = ArrowColumnVectors.allocateRoot(schema)
+ val vector = root.getVector(0).asInstanceOf[BigIntVector]
+ vector.allocateNew(numRows)
for (i <- 0 until numRows) {
val value = current + i * step
- vectors(0).putLong(i, getSafeMargin(value))
+ vector.set(i, getSafeMargin(value))
}
- vectors.foreach(_.setValueCount(numRows))
+ root.setRowCount(numRows)
current += numRows * step
- val batch = new
ColumnarBatch(vectors.asInstanceOf[Array[ColumnVector]], numRows)
- batch
+ ArrowColumnVectors.toBatch(root)
Review Comment:
This code sets `root.setRowCount(numRows)` but never sets
`vector.setValueCount(numRows)`. If `VectorSchemaRoot.setRowCount` doesn’t
propagate value counts to vectors (common in Arrow usage patterns),
exported/offloaded batches can appear empty (`valueCount = 0`). Fix by setting
`vector.setValueCount(numRows)` (or by making `ArrowColumnVectors.toBatch` set
value counts for all vectors, and ensuring that behavior is covered by tests).
--
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]
