andygrove commented on code in PR #5368: URL: https://github.com/apache/datafusion-comet/pull/5368#discussion_r3883029263
########## spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala: ########## @@ -0,0 +1,329 @@ +/* + * 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.spark.sql.execution.python + +import java.io.{ByteArrayInputStream, ByteArrayOutputStream, IOException} +import java.nio.ByteBuffer +import java.nio.channels.{Channels, WritableByteChannel} + +import scala.jdk.CollectionConverters._ + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.arrow.memory.{BufferAllocator, RootAllocator} +import org.apache.arrow.vector.{FieldVector, IntVector, NullVector, VarCharVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} +import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter, WriteChannel} +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} +import org.apache.spark.sql.execution.python.CometArrowPythonRunnerBase.serializeBatch + +class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers { + + private def withWriter( + childFields: Seq[Field], + allocator: BufferAllocator, + channel: WritableByteChannel)(f: WritableByteChannel => Unit): Unit = { + val structField = new Field( + "struct", + new FieldType(false, ArrowType.Struct.INSTANCE, null), + childFields.asJava) + val root = VectorSchemaRoot.create(new Schema(Seq(structField).asJava), allocator) + val writer = new ArrowStreamWriter(root, null, channel) + try { + writer.start() + f(channel) + writer.end() + } finally { + writer.close() + root.close() + } + } + + private def withReader(bytes: Array[Byte])(f: ArrowStreamReader => Unit): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val reader = new ArrowStreamReader(new ByteArrayInputStream(bytes), allocator) + try { + f(reader) + } finally { + reader.close() + allocator.close() + } + } + + test("direct batches retain borrowed buffers without copying them into the writer allocator") { Review Comment: All five tests here build their source vectors with a plain `RootAllocator`. In production the source is always FFI-imported, where the buffers are owned by `ReferenceCountedArrowArray` and the validity bitmap is synthesized by `BitVectorHelper.loadValidityBuffer` at import time. That is precisely the reference manager that the refcount and `getAllocatedMemory` assertions are trying to pin down. Would you add one case that exports a vector with `Data.exportVector` and re-imports it into a second allocator before calling `serializeBatch`? I checked this locally against real Comet scan batches and the counts do balance, so this is a coverage question rather than a bug. ########## spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala: ########## @@ -112,10 +112,9 @@ case class EliminateRedundantTransitions(session: SparkSession) // 4.1+ matches the renamed `MapInArrowExec`. // // Falls back to vanilla Spark when `spark.sql.execution.arrow.useLargeVarTypes` is enabled: - // CometArrowPythonRunnerBase.copyVector does raw `setBytes` on each Arrow buffer, but Comet's - // source string/binary vectors always use 4-byte offsets while the destination root is - // allocated with 8-byte offsets when this conf is on. The buffer counts match but the - // offset width does not, so a direct memcpy would corrupt the offsets. + // Comet's source string/binary vectors use 4-byte offsets, while Spark expects 8-byte + // offsets when this conf is on. Direct IPC serialization cannot change the source vectors' + // physical layout, so forwarding them under Spark's widened schema would corrupt the stream. Review Comment: The comment says forwarding Comet's 4-byte-offset vectors under Spark's widened schema "would corrupt the stream". That was true of `copyVector`, but after this change the advertised schema comes from `sourceVectors.map(_.getField)`, so the header and the buffers agree and the stream stays internally consistent. I think the fallback should still stay, but for different reasons: the worker would see `string` where `useLargeVarTypes=true` promises `large_string`, and the worker's `large_*` output would then come back through `ArrowStreamReader` into `CometVector.getVector`. Could you update the comment here and the matching bullet in `pyarrow-udfs.md` to say that instead? ########## spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala: ########## @@ -167,49 +168,37 @@ private[python] trait CometArrowPythonRunnerBase val cometBatch = currentGroup.next() val startData = dataOut.size() + val sourceVectors = (0 until cometBatch.numCols()).map { i => + cometBatch + .column(i) + .asInstanceOf[CometDecodedVector] + .getValueVector + .asInstanceOf[FieldVector] + } if (arrowWriter == null) { - // Build the destination struct root once, sized to the first batch's child fields. + // Build the schema-only struct root once from the first batch's child fields. // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct. // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real // column names from the input schema (the worker reads columns by name, and shaded - // Arrow rejects a null field name). The field types and child structure are kept as-is - // so copyVector still walks the source and destination trees in lockstep. Keeping the - // type as-is also means a TimestampType reaches the worker with Comet's UTC time zone + // Arrow rejects a null field name). Keep the field types and child structure as-is so + // the advertised schema matches the source buffers. Keeping the type as-is also means + // a TimestampType reaches the worker with Comet's UTC time zone // rather than the session zone vanilla Spark would label it with; this is a documented // limitation (see pyarrow-udfs.md), not a value difference, since the stored instant is // identical. val childNames = inputStructType.fieldNames - val childFields = (0 until cometBatch.numCols()).map { i => - val vecField = - cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField - renamed(vecField, childNames(i), forceNullable = true) + val childFields = sourceVectors.zipWithIndex.map { case (vector, i) => + renamed(vector.getField, childNames(i), forceNullable = true) } startWriter(childFields, dataOut) Review Comment: `copyVector` used to fail loudly via its `require(srcBufs.size == dstBufs.size)` if a later batch's layout stopped matching the destination tree. Now that the schema header is written once from the first batch and each subsequent batch's buffers go straight to the wire, a layout change would instead produce a well-formed record batch that the worker silently misreads. I could not find a path where Comet changes layout mid-partition, so this may well be unreachable, but would you consider capturing the child fields in `startWriter` and asserting `sourceVectors.map(_.getField)` still matches them per batch? It keeps the fail-loud behaviour for about three lines. -- 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]
