Copilot commented on code in PR #13280: URL: https://github.com/apache/gluten/pull/13280#discussion_r4268007491
########## 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 ref-counting logic is not thread-safe: two concurrent `close()` calls can both pass the `refCnt.get() == 0` check and decrement, driving `refCnt` negative (e.g., 1 → 0 → -1) and leaving the object in an invalid state. Similarly, `retain()` can increment a ref count from 0 after underlying resources have been released. Use a CAS loop for both `retain()` and `close()` to prevent refCnt from going below 0 and to forbid retain-after-release (e.g., only increment when current > 0; only decrement when current > 0; call `super.close()` only on the transition 1 → 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) + } + + /** 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: `fromRows` allocates Arrow resources (`VectorSchemaRoot` and likely `ArrowWriter`) but does not ensure they are cleaned up on exceptions during `rows.foreach(writer.write)` or `writer.finish()`. This can leak Arrow buffers in failure paths. Wrap the write in try/finally so that on failure you close the writer (if it’s closeable) and close the root to release vectors/buffers; on success, transfer ownership to the returned batch (and avoid double-closing). ########## backends-velox/src/test/scala/org/apache/gluten/execution/python/ArrowEvalPythonExecSuite.scala: ########## @@ -116,6 +117,40 @@ class ArrowEvalPythonExecSuite extends WholeStageTransformerSuite { checkAnswer(df, expected) } + testWithMinSparkVersion("arrow_udf test: chained UDFs over small batches with nulls", "4.0") { + withSQLConf("spark.gluten.sql.columnar.maxBatchSize" -> "64") { + val base = spark + .range(0, 1000, 1, 2) + .selectExpr( + "id", + "if(id % 11 = 0, null, cast(id % 7 as string)) as s", + "if(id % 5 = 0, null, id * 1.5) as d") + // Velox operators both below and above the UDF, over many small batches. + def query() = base + .filter("id > 10") + .withColumn("p_s", pyarrowTestUDFString(base("s"))) Review Comment: Within `query()`, `withColumn` references `base(\"s\")` rather than the current DataFrame’s column (e.g., using the column named `s` from the evolving plan). While it may work here since `query()` is derived from `base`, it makes the transformation harder to read and more brittle if the query structure changes. Prefer referencing the column from the current DataFrame (e.g., `col(\"s\")` / `$\"s\"` / `expr(\"s\")`) to keep the expression tied to the DataFrame being transformed. -- 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]
