Copilot commented on code in PR #13280:
URL: https://github.com/apache/gluten/pull/13280#discussion_r4329728356
##########
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:
Same issue as in `ColumnarRangeExec`: `pidVec.set(i, pid)` doesn’t update
`valueCount`, and `pidRoot.setRowCount` may not propagate to vectors. Ensure
`pidVec.setValueCount(cb.numRows)` (or centralize this in
`ArrowColumnVectors.toBatch`) to avoid out-of-bounds errors when Spark/Arrow
reads the vector.
##########
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:
`VectorSchemaRoot.setRowCount` does not guarantee the underlying Arrow
vectors’ `valueCount` is updated. Since `vector.set(i, ...)` doesn’t set
`valueCount` either, reads via Spark/Arrow can see `valueCount == 0` and fail
bounds checks. Set `vector.setValueCount(numRows)` (or have
`ArrowColumnVectors.toBatch` set `valueCount` on all vectors based on
`root.getRowCount`) before exposing the batch.
##########
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:
For “foreign” Spark `ArrowColumnVector` batches, `ArrowColumnVectors.adopt`
may choose a zero-copy `transfer()` path when allocator roots match; that will
steal buffers from Spark’s producer-owned vectors and can break producers that
reuse their `VectorSchemaRoot` for subsequent batches. For foreign batches,
adoption should be a deep copy regardless of allocator-root equality; a
concrete approach is adding an `allowTransfer`/`forceCopy` parameter to
`ArrowColumnVectors.adopt` and passing `forceCopy=true` from this call site.
##########
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 `close()` implementation is not atomic with respect to concurrent
callers: two threads can both observe `refCnt.get() != 0`, then both decrement,
potentially driving the count negative and making the lifecycle inconsistent.
If `retain/close` can be invoked from different threads, use an atomic
decrement-and-check pattern that clamps at 0 (e.g., CAS loop or
`getAndDecrement` with validation) and consider rejecting `retain()` after the
count reaches 0.
##########
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)
+ }
Review Comment:
`toBatch` currently trusts that each underlying `FieldVector` already has
the correct `valueCount`. That’s true for batches produced by
`ArrowWriter.finish()`, but it is not true for call sites that fill vectors
with `set(...)` and only call `root.setRowCount(...)` (e.g.
`ColumnarRangeExec`, `ExecUtil`). Consider setting `valueCount` for every
vector to `root.getRowCount` inside `toBatch` to enforce a consistent invariant
for all batch builders.
##########
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)
+ }
Review Comment:
If `rows.foreach(writer.write)` throws (bad input, expression bug, etc.),
`root` (and its allocated buffers) won’t be released, leading to Arrow memory
leaks. Wrap the write/finish section with exception-safe cleanup (e.g.,
close/clear the root on failure) so partially-built batches don’t leak native
memory.
--
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]
