Copilot commented on code in PR #13280:
URL: https://github.com/apache/gluten/pull/13280#discussion_r4330282880
##########
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 path sets `root.setRowCount(numRows)` but never sets
`vector.setValueCount(numRows)`. For Arrow vectors written via `set(i, ...)`,
failing to set the vector value count can result in downstream readers seeing
`valueCount == 0` (empty batch / invalid reads), even if the root row count is
set. Fix by setting the value count on the vector(s) before wrapping (or
alternatively make `ArrowColumnVectors.toBatch` set value counts for all field
vectors to `root.getRowCount`).
##########
gluten-arrow/src/main/java/org/apache/gluten/vectorized/RefCountedArrowColumnVector.java:
##########
@@ -0,0 +1,57 @@
+/*
+ * 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.arrow.vector.ValueVector;
+import org.apache.spark.sql.vectorized.ArrowColumnVector;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Spark's {@link ArrowColumnVector} with a reference count, used by Gluten's
Java Arrow batches
+ * (see {@link org.apache.gluten.columnarbatch.ColumnarBatches}). The count
starts at 1; {@link
+ * #retain()} increments it and {@link #close()} decrements it. The Arrow
vector is released when
+ * the count reaches 0.
+ *
+ * <p>Reading goes through Spark's accessors. To write, fill the Arrow vectors
(e.g. with Spark's
+ * {@code ArrowWriter}) and wrap them, see {@link ArrowColumnVectors}.
+ */
+public final class RefCountedArrowColumnVector extends ArrowColumnVector {
+ private final AtomicLong refCnt = new AtomicLong(1L);
+
+ public RefCountedArrowColumnVector(ValueVector vector) {
+ super(vector);
+ }
+
+ public void retain() {
+ refCnt.getAndIncrement();
+ }
+
+ public long refCnt() {
+ return refCnt.get();
+ }
+
+ @Override
+ public void close() {
+ if (refCnt.get() == 0) {
+ return;
+ }
+ if (refCnt.decrementAndGet() == 0) {
+ super.close();
+ }
+ }
Review Comment:
The refcount transitions are not atomic with respect to the `0` boundary:
`close()` does a non-atomic `get()` check followed by `decrementAndGet()`,
which can race under concurrent `close()` calls and drive the refcount
negative. Also, `retain()` can resurrect the count from 0 after the underlying
Arrow vector has been closed. Consider using a CAS loop to guard transitions at
0 (e.g., only decrement if current > 0, and only increment if current > 0), and
either throw or no-op on `retain()` after the vector is fully released.
##########
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 code writes into `pidVec` but only sets
`pidRoot.setRowCount(...)`. It should also set
`pidVec.setValueCount(cb.numRows)` (or ensure `ArrowColumnVectors.toBatch`
enforces vector value counts). Without this, consumers that rely on the
vector's `valueCount` may observe an empty/invalid column.
##########
gluten-arrow/src/main/java/org/apache/gluten/columnarbatch/ColumnarBatches.java:
##########
@@ -78,13 +80,41 @@ private static BatchType identifyBatchType(ColumnarBatch
batch) {
// it's likely a heavy batch
for (int i = 0; i < batch.numCols(); i++) {
ColumnVector col = batch.column(i);
- if (!(col instanceof ArrowWritableColumnVector)) {
+ if (!(col instanceof RefCountedArrowColumnVector)) {
throw new IllegalStateException("Heavy batch should consist of arrow
vectors");
}
}
return BatchType.HEAVY;
}
+ /**
+ * Whether the batch consists of Spark {@link ArrowColumnVector}s that are
not Gluten's
+ * reference-counted ones.
+ */
+ private static boolean isForeignArrowBatch(ColumnarBatch batch) {
+ if (batch.numCols() == 0) {
+ return false;
+ }
+ boolean allArrow = true;
+ boolean allRefCounted = true;
+ for (int i = 0; i < batch.numCols(); i++) {
+ ColumnVector col = batch.column(i);
+ allArrow &= col instanceof ArrowColumnVector;
+ allRefCounted &= col instanceof RefCountedArrowColumnVector;
+ }
+ return allArrow && !allRefCounted;
+ }
+
+ /** Moves (or copies) the Arrow vectors of `input` into a new batch owned by
Gluten. */
+ private static ColumnarBatch adopt(BufferAllocator allocator, ColumnarBatch
input) {
+ ColumnVector[] columns = new ColumnVector[input.numCols()];
+ for (int i = 0; i < input.numCols(); i++) {
+ ValueVector vector = ((ArrowColumnVector)
input.column(i)).getValueVector();
+ columns[i] = new
RefCountedArrowColumnVector(ArrowColumnVectors.adopt(vector, allocator));
+ }
+ return new ColumnarBatch(columns, input.numRows());
+ }
Review Comment:
When `input` is a \"foreign\" Spark Arrow batch, calling
`ArrowColumnVectors.adopt(vector, allocator)` may `transfer()` buffers if the
allocator roots match. That path *moves* buffers out of the producer-owned
vectors, mutating/invalidating the input batch even though the comment says the
input is left to its producer. For foreign batches, adoption should always be a
deep copy (never a transfer-from-source). A concrete fix is to add a dedicated
`adoptCopy(...)` (or a `forceCopy` flag) and use it from
`ColumnarBatches.adopt(...)` when `isForeignArrowBatch(input)` is true.
##########
gluten-arrow/src/main/scala/org/apache/gluten/vectorized/ArrowColumnVectors.scala:
##########
@@ -0,0 +1,85 @@
+/*
+ * 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)
+ }
+
+ /** Writes `rows` into a new batch with Spark's [[ArrowWriter]]. */
+ def fromRows(
+ schema: StructType,
+ rows: Iterator[InternalRow],
+ allocator: BufferAllocator = ArrowBufferAllocators.contextInstance()):
ColumnarBatch = {
+ val root = allocateRoot(schema, allocator)
+ val writer = ArrowWriter.create(root)
+ rows.foreach(writer.write)
+ writer.finish()
+ toBatch(root)
+ }
+
+ /**
+ * Takes ownership of `vector` in `allocator`: its buffers are moved to a
new vector without
+ * copying when it shares the allocator's root, and copied otherwise (Arrow
buffers can't move
+ * across allocator roots, e.g. from vectors allocated by Spark).
+ */
+ def adopt(vector: ValueVector, allocator: BufferAllocator): FieldVector = {
+ if (vector.getAllocator.getRoot == allocator.getRoot) {
+ val pair = vector.getTransferPair(allocator)
+ pair.transfer()
+ pair.getTo.asInstanceOf[FieldVector]
+ } else {
+ val target = vector.getField.createVector(allocator)
+ target.allocateNew()
+ vector.accept(new VectorAppender(target), null)
+ target
+ }
+ }
Review Comment:
In the cross-root copy path, `target.allocateNew()` allocates using
defaults, which can cause repeated reallocations when appending large vectors.
You can reduce reallocations by sizing the allocation based on the source
(e.g., initial capacity/value count, and for variable-width vectors also total
buffer sizes when available). Also consider explicitly setting
`target.setValueCount(vector.getValueCount)` after appending (even if
`VectorAppender` typically advances counts) to make the resulting vector state
unambiguous for downstream consumers.
--
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]
