sunchao commented on code in PR #5368:
URL: https://github.com/apache/datafusion-comet/pull/5368#discussion_r3883361482
##########
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:
Updated in 1de4491c. The code comment, user guide, and fallback test
docstring now explain that the IPC schema matches the source vectors, so the
stream is valid, but it does not provide the large input types requested by the
configuration. The fallback stays in place to preserve that input contract.
I also clarified that Comet already reads large_string / large_binary output
through Utils and CometPlainVector; that output support does not widen the
input vectors.
##########
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:
Added in 1de4491c. The writer snapshots the first batch's raw Arrow fields
and requires every subsequent batch to match before serializing it. I compare
raw fields with raw fields because the advertised fields have their names and
nullability normalized.
The full Spark 4.0 Python suite passed, including the multiple-batch,
nested-source, and chained-UDF cases.
##########
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:
Added two FFI cases in 1de4491c: successful serialization and an injected
write failure. Each exports with Data.exportVector, imports into a second
allocator, verifies the shared data address, and closes the original vector
before serialization so the imported ownership keeps the buffers alive.
Both cases check balanced buffer reference counts, unchanged allocator
usage, readable imported values, and zero outstanding allocations after the
imported vector is closed. A 16 KiB payload is serialized with a 1 KiB writer
allocator to rule out a payload copy. All seven Arrow tests pass on Spark 4.0.
--
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]
