diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 18dc8e4086c..4a640a03b46 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -342,6 +342,7 @@ jobs: org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite org.apache.spark.sql.comet.CometMapInBatchSuite + org.apache.spark.sql.execution.python.CometArrowPythonRunnerSuite org.apache.comet.CometNativeSuite org.apache.comet.CometConfSuite org.apache.comet.CometPublicApiSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index fa82e4ebca7..8ae0e3f135a 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -158,6 +158,7 @@ jobs: org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite org.apache.spark.sql.comet.CometMapInBatchSuite + org.apache.spark.sql.execution.python.CometArrowPythonRunnerSuite org.apache.comet.CometNativeSuite org.apache.comet.CometConfSuite org.apache.comet.CometPublicApiSuite diff --git a/.github/workflows/pyarrow_udf_test.yml b/.github/workflows/pyarrow_udf_test.yml index 1a03962685d..8fff5d414cd 100644 --- a/.github/workflows/pyarrow_udf_test.yml +++ b/.github/workflows/pyarrow_udf_test.yml @@ -39,10 +39,12 @@ on: - "spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala" - "spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala" - "spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala" + - "spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala" - "spark/src/test/resources/pyspark/conftest.py" - "spark/src/test/resources/pyspark/test_pyarrow_udf.py" - "spark/src/test/spark-3.5/org/apache/spark/sql/comet/CometMapInBatchSuite.scala" - "spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala" + - "spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala" - ".github/workflows/pyarrow_udf_test.yml" pull_request: paths: *feature-paths diff --git a/docs/source/user-guide/latest/pyarrow-udfs.md b/docs/source/user-guide/latest/pyarrow-udfs.md index ff493de23c0..14975c42c31 100644 --- a/docs/source/user-guide/latest/pyarrow-udfs.md +++ b/docs/source/user-guide/latest/pyarrow-udfs.md @@ -203,18 +203,14 @@ on the unoptimized path. session time zone such a UDF can diverge from the unoptimized path. Set `spark.comet.exec.pyarrowUDF.enabled=false` for those UDFs. - `spark.sql.execution.arrow.useLargeVarTypes=true` is not supported. With this conf enabled, - Spark widens `StringType` and `BinaryType` to Arrow's 8-byte-offset variants in the - destination IPC root, while Comet's source vectors always use 4-byte offsets. The buffer-copy - path cannot bridge that mismatch, so `EliminateRedundantTransitions` skips the rewrite and - vanilla Spark handles the operation. -- Each batch is copied twice on the JVM side: once from Comet's vectors into Spark's - destination IPC root (per-buffer `setBytes`), and a second time inside the IPC writer when - `VectorUnloader` / `MessageSerializer.serialize` walks the root and writes bytes to the - pipe to the Python worker. The pipe write is structural (Spark's transport to Python is - fork + pipe + Arrow IPC, so the buffer bytes must reach the pipe at least once); dropping - the first copy by serialising directly from Comet's vectors is tracked in - [#4294](https://github.com/apache/datafusion-comet/issues/4294). Even after that, - true zero-copy at the JVM boundary is blocked because Comet's source `FieldVector`s are - imported from native via Arrow C Data Interface (their buffers route `release` through FFI), - while Spark's destination IPC root is a child of `ArrowUtils.rootAllocator`. The two - reference managers cannot share buffers via `TransferPair`. + Spark supplies `large_string` and `large_binary` input columns with 8-byte offsets. Native + Comet vectors use 4-byte offsets, and direct serialization advertises their matching `string` + and `binary` types. This produces a valid IPC stream, but does not preserve the input types + requested by the configuration. `EliminateRedundantTransitions` therefore skips the rewrite + and vanilla Spark handles the operation. Comet can read `large_string` and `large_binary` + columns returned by a Python worker; that output support does not widen the input vectors. +- Comet writes input Arrow IPC record batches directly from its existing vector buffers. The + only additional Arrow buffer is the validity bitmap for the non-null struct that wraps the + input columns. Writing the IPC bytes to the Python worker's pipe still requires one copy; + that copy is inherent to Spark's process-based Python transport. This path does not transfer + buffers between Arrow allocators or change their ownership. diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index ec277cfc7bc..45fc26caee9 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -112,10 +112,10 @@ 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. + // Native Comet string/binary vectors use 4-byte offsets. The IPC schema follows these + // vectors, so the stream is internally consistent, but the worker would receive string / + // binary instead of the large_string / large_binary input types requested by this conf. + // Keep the fallback to preserve Spark's input type contract. // // `EligibleMapInBatch` matches whenever the operator would run natively if the feature were // enabled. When it is disabled (the default) we leave the vanilla Spark operator in place diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala index b22b0a792c3..943287cde1f 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala @@ -21,13 +21,16 @@ package org.apache.spark.sql.execution.python import java.io.{DataInputStream, DataOutputStream} import java.nio.channels.Channels +import java.util.ArrayList import java.util.concurrent.atomic.AtomicBoolean import scala.jdk.CollectionConverters._ -import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot} -import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector} -import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} +import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.vector.complex.StructVector +import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter, WriteChannel} +import org.apache.arrow.vector.ipc.message.{ArrowFieldNode, ArrowRecordBatch, MessageSerializer} import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType} import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths} @@ -36,7 +39,6 @@ import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.spark.unsafe.Platform import org.apache.comet.CometArrowAllocator import org.apache.comet.vector.{CometDecodedVector, CometVector} @@ -53,7 +55,7 @@ import org.apache.comet.vector.{CometDecodedVector, CometVector} * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary: - * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the + * - Input: each Comet `ColumnarBatch` is written directly from its shaded Arrow vectors to the * worker with a shaded `ArrowStreamWriter`. * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators @@ -109,7 +111,7 @@ private[python] trait CometArrowPythonRunnerBase private var currentGroup: Iterator[ColumnarBatch] = _ private var arrowWriter: ArrowStreamWriter = _ private var writeRoot: VectorSchemaRoot = _ - private var structVec: StructVector = _ + private var streamFields: Seq[Field] = _ // The runner's input schema is a single struct column ("struct") whose children are the // user's input columns (see `schema` above). Cast once here rather than at each use site. @@ -132,14 +134,14 @@ private[python] trait CometArrowPythonRunnerBase writeUDF(dataOut) } - /** Build the destination struct root and start the writer from the given child fields. */ + /** Build the schema-only struct root and start the writer from the given child fields. */ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = { val structField = new Field( "struct", new FieldType(false, ArrowType.Struct.INSTANCE, null), childFields.asJava) - structVec = structField.createVector(allocator).asInstanceOf[StructVector] + val structVec = structField.createVector(allocator).asInstanceOf[StructVector] writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava) arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut)) arrowWriter.start() @@ -167,49 +169,44 @@ 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] + } + val batchFields = sourceVectors.map(_.getField) 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) + streamFields = batchFields.zipWithIndex.map { case (field, i) => + renamed(field, childNames(i), forceNullable = true) } - startWriter(childFields, dataOut) + startWriter(streamFields, dataOut) } - var i = 0 - while (i < cometBatch.numCols()) { - val src = cometBatch - .column(i) - .asInstanceOf[CometDecodedVector] - .getValueVector - .asInstanceOf[FieldVector] - val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector] - copyVector(src, dst) - i += 1 - } - val numRows = cometBatch.numRows() - structVec.setValueCount(numRows) - // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly - // allocated and zero-initialised, so without this Python would see an all-null struct. - val validityBytes = (numRows + 7) / 8 - Platform.setMemory( - structVec.getValidityBuffer.memoryAddress(), - 0xff.toByte, - validityBytes) - writeRoot.setRowCount(numRows) - arrowWriter.writeBatch() + // Union branches may differ in names, nullability, or descriptive metadata. Only + // differences that change how the advertised schema interprets the buffers are invalid. + require( + CometArrowPythonRunnerBase.hasCompatibleSchema(streamFields, batchFields), + s"Arrow input schema changed between batches: expected $streamFields, got $batchFields") + + CometArrowPythonRunnerBase.serializeBatch( + new WriteChannel(Channels.newChannel(dataOut)), + sourceVectors, + cometBatch.numRows(), + allocator) pythonMetrics("pythonDataSent") += dataOut.size() - startData true @@ -291,14 +288,17 @@ private[python] trait CometArrowPythonRunnerBase /** * Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child * whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can - * materialize the struct. Keeping the type and structure intact means the destination tree - * still mirrors the Comet source tree for [[copyVector]]. + * materialize the struct. Keeping the type and structure intact means the advertised schema + * still mirrors the Comet source vectors serialized directly into each record batch. */ - private def renamed(field: Field, name: String, forceNullable: Boolean): Field = { - // A Map's descendants must keep their original nullability: Arrow requires the entries struct - // (and its key) to be non-nullable, and `MapVector.createVector` rejects a nullable entries - // struct. Stop forcing nullable once we enter a Map subtree. - val childrenForceNullable = forceNullable && !field.getType.isInstanceOf[ArrowType.Map] + private def renamed( + field: Field, + name: String, + forceNullable: Boolean, + isMapEntry: Boolean = false): Field = { + // Map entries and keys must stay non-nullable, but values and fields nested inside a + // complex key may be nullable. Do not propagate the restriction through the whole subtree. + val isMap = field.getType.isInstanceOf[ArrowType.Map] val children = field.getChildren val newChildren = if (children.isEmpty) children @@ -311,80 +311,91 @@ private[python] trait CometArrowPythonRunnerBase renamed( child, if (child.getName == null) s"_$idx" else child.getName, - childrenForceNullable) + forceNullable = !isMap && !(isMapEntry && idx == 0), + isMapEntry = isMap) }.asJava // Force the field nullable where allowed. Comet's FFI-imported vectors may carry a // non-nullable Arrow `Field` even for columns that contain nulls (Comet uses positional schema // and does not round-trip Spark's nullability), and the worker rejects a null value under a // non-nullable field (`from_pandas(pdf, schema=batch.schema)` raises). Marking the field - // nullable is a safe superset; `copyVector` fills an all-valid validity buffer when the source - // has no nulls. + // nullable is a safe superset; Arrow IPC permits an empty validity buffer when its field node + // reports no null values. val ft = field.getFieldType val nullable = forceNullable || ft.isNullable val newFt = new FieldType(nullable, ft.getType, ft.getDictionary, ft.getMetadata) new Field(name, newFt, newChildren) } +} - /** - * Copy a Comet column into the destination FieldVector. Walks both trees in lockstep: sizes - * each destination node from the source, copies every buffer with `ArrowBuf.setBytes`, then - * sets value counts bottom-up so `setValueCount` does not rewrite the offset bytes we just - * copied. Both source and destination are Comet's (shaded) Arrow vectors, so no shaded / - * unshaded type crosses. - */ - private def copyVector(src: FieldVector, dst: FieldVector): Unit = { - val valueCount = src.getValueCount - - dst match { - case bfwv: BaseFixedWidthVector => - bfwv.allocateNew(valueCount) - case bvwv: BaseVariableWidthVector => - bvwv.allocateNew(src.getDataBuffer.readableBytes, valueCount) - case blvwv: BaseLargeVariableWidthVector => - blvwv.allocateNew(src.getDataBuffer.readableBytes, valueCount) - case _ => - dst.setInitialCapacity(valueCount) - dst.allocateNew() - } - - val srcBufs = src.getFieldBuffers - val dstBufs = dst.getFieldBuffers - require( - srcBufs.size == dstBufs.size, - s"buffer count mismatch for ${dst.getField}: src=${srcBufs.size}, dst=${dstBufs.size}") - srcBufs.asScala.zip(dstBufs.asScala).foreach { case (s, d) => - d.setBytes(0, s, 0, s.readableBytes) - } - - val srcChildren = src.getChildrenFromFields - val dstChildren = dst.getChildrenFromFields - require( - srcChildren.size == dstChildren.size, - s"child count mismatch for ${dst.getField}: src=${srcChildren.size}, dst=${dstChildren.size}") - srcChildren.asScala.zip(dstChildren.asScala).foreach { case (sc, dc) => - copyVector(sc.asInstanceOf[FieldVector], dc.asInstanceOf[FieldVector]) +private[python] object CometArrowPythonRunnerBase { + + // Extensions can change interpretation even when their underlying storage types match. + private val extensionMetadataKeys = Seq( + ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME, + ArrowType.ExtensionType.EXTENSION_METADATA_KEY_METADATA) + + /** Names, nullability and ordinary field metadata do not change the IPC buffer layout. */ + private[python] def hasCompatibleSchema(expected: Seq[Field], actual: Seq[Field]): Boolean = { + expected.size == actual.size && expected.zip(actual).forall { case (left, right) => + left.getType == right.getType && + left.getDictionary == right.getDictionary && + extensionMetadataKeys.forall(key => + left.getMetadata.get(key) == right.getMetadata.get(key)) && + hasCompatibleSchema(left.getChildren.asScala.toSeq, right.getChildren.asScala.toSeq) } + } - // For vectors that fill offset-buffer "holes" in setValueCount (variable-width and list - // types), set lastSet = vc - 1 first so fillHoles is a no-op and the already-copied offset - // bytes are preserved. - dst match { - case v: BaseVariableWidthVector => v.setLastSet(valueCount - 1) - case v: BaseLargeVariableWidthVector => v.setLastSet(valueCount - 1) - case v: ListVector => v.setLastSet(valueCount - 1) - case v: LargeListVector => v.setLastSet(valueCount - 1) - case _ => - } - dst.setValueCount(valueCount) - - // Every destination field is nullable (see `renamed`), so the worker reads the validity - // buffer. When the source has no nulls its validity buffer may be empty (Comet omits it), - // which would otherwise leave the freshly-allocated destination validity all-zero and make - // the worker see every value as null. Set all-valid in that case. Done after setValueCount, - // which can rewrite validity, mirroring the struct-level all-valid fill in writeNextInput. - if (valueCount > 0 && dst.getField.isNullable && src.getNullCount == 0) { - val validityBytes = (valueCount + 7) / 8 - Platform.setMemory(dst.getValidityBuffer.memoryAddress(), 0xff.toByte, validityBytes) + /** + * Serialize source vectors directly beneath the non-null struct advertised in the IPC stream. + * + * VectorUnloader recursively retains the source buffers without moving them between allocators. + * The wrapping record batch takes its own references, so closing both temporary batches + * restores the original reference counts after the synchronous pipe write. The borrowed + * VectorSchemaRoot must never be closed because its vectors are owned by the input + * ColumnarBatch. + */ + private[python] def serializeBatch( + writeChannel: WriteChannel, + sourceVectors: Seq[FieldVector], + numRows: Int, + allocator: BufferAllocator): Unit = { + val sourceRoot = + new VectorSchemaRoot(sourceVectors.map(_.getField).asJava, sourceVectors.asJava, numRows) + val sourceBatch = new VectorUnloader(sourceRoot).getRecordBatch + try { + val validityBytes = (numRows.toLong + 7L) / 8L + val structValidity = allocator.buffer(validityBytes) + try { + if (validityBytes > 0) { + structValidity.setOne(0L, validityBytes) + } + structValidity.writerIndex(validityBytes) + + val nodes = new ArrayList[ArrowFieldNode](sourceBatch.getNodes.size() + 1) + nodes.add(new ArrowFieldNode(numRows, 0)) + nodes.addAll(sourceBatch.getNodes) + + val buffers = new ArrayList[ArrowBuf](sourceBatch.getBuffers.size() + 1) + buffers.add(structValidity) + buffers.addAll(sourceBatch.getBuffers) + + val wrappedBatch = new ArrowRecordBatch( + numRows, + nodes, + buffers, + sourceBatch.getBodyCompression, + sourceBatch.getVariadicBufferCounts, + true) + try { + MessageSerializer.serialize(writeChannel, wrappedBatch) + } finally { + wrappedBatch.close() + } + } finally { + structValidity.close() + } + } finally { + sourceBatch.close() } } } diff --git a/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py b/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py index 19d9dac5d32..7d21c9822b3 100644 --- a/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py +++ b/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py @@ -31,8 +31,8 @@ * vanilla: CometScan -> ColumnarToRow + UnsafeProjection -> ArrowPythonRunner (per-row InternalRow.getXXX() loop inside ArrowWriter.write) * optimized: CometScan -> CometMapInBatchExec -> CometArrowPythonRunner - (per-buffer Unsafe.copyMemory from Comet's vectors into the - runner's persistent VectorSchemaRoot; no row materialization) + (Arrow IPC serialization directly from Comet's source vectors; + no row materialization or intermediate vector-buffer copy) Results are wall-clock seconds, so they include Python interpreter, Arrow IPC, and downstream count() costs. That's intentional: the @@ -47,9 +47,9 @@ Usage: # Build Comet (release for representative numbers): - make release + make release PROFILES='-Pspark-4.0 -Pscala-2.13' - pip install pyspark==3.5.9 pyarrow pandas + pip install pyspark==4.0.4 pyarrow pandas python3 spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py @@ -84,6 +84,10 @@ def _build_spark() -> SparkSession: .config("spark.plugins", "org.apache.spark.CometPlugin") .config("spark.comet.enabled", "true") .config("spark.comet.exec.enabled", "true") + .config( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", + ) .config("spark.memory.offHeap.enabled", "true") .config("spark.memory.offHeap.size", "4g") .config("spark.driver.memory", "4g") @@ -157,6 +161,10 @@ def _time_run(spark: SparkSession, parquet_path: str, accelerate: bool, api: str df = df.mapInArrow(_passthrough_arrow, schema) else: df = df.mapInPandas(_passthrough_pandas, schema) + plan = df._jdf.queryExecution().executedPlan().toString() + if ("CometMapInBatch" in plan) != accelerate: + expected = "CometMapInBatch" if accelerate else "vanilla Python execution" + raise RuntimeError(f"Expected {expected} for {api}, but found:\n{plan}") t0 = time.perf_counter() df.count() return time.perf_counter() - t0 diff --git a/spark/src/test/resources/pyspark/test_pyarrow_udf.py b/spark/src/test/resources/pyspark/test_pyarrow_udf.py index 84f405e8b2c..5c25cb03ba5 100644 --- a/spark/src/test/resources/pyspark/test_pyarrow_udf.py +++ b/spark/src/test/resources/pyspark/test_pyarrow_udf.py @@ -39,6 +39,7 @@ import datetime as dt import os +from collections import Counter from decimal import Decimal import pyarrow as pa @@ -361,11 +362,11 @@ def test_map_in_arrow_decimal_precision_sweep( spark, tmp_path, accelerated, precision, scale ): """ - The Arrow `DecimalVector` that `copyVector` touches is always 16 bytes wide regardless of - precision, so there is no buffer-width boundary on the Arrow path (the 8-byte long-backed form - is Spark's `UnsafeRow` encoding, a layer this Arrow buffer copy never sees). This sweep instead - guards the precision/scale extremes and the 18/19 point where Spark's own decimal handling - changes representation, keeping the round trip value-exact. Scale extremes: 0, half, max. + Arrow's `DecimalVector` is always 16 bytes wide regardless of precision, so there is no + buffer-width boundary on the direct Arrow path (the 8-byte long-backed form is Spark's + `UnsafeRow` encoding, a layer this path never sees). This sweep instead guards the + precision/scale extremes and the 18/19 point where Spark's own decimal handling changes + representation, keeping the round trip value-exact. Scale extremes: 0, half, max. """ schema_in = T.StructType( [ @@ -405,10 +406,9 @@ def test_map_in_arrow_null_density_sweep( spark, tmp_path, accelerated, null_fraction ): """ - Validity-buffer memcpy is where Arrow Java vector copies historically break. Sweep null - density across the corner cases: all-non-null, sparse-null, half-null, sparse-non-null, - all-null. Catches off-by-one in validity packing and edge cases where source/destination - null counts diverge. + Sweep validity-buffer density across the corner cases: all-non-null, sparse-null, half-null, + sparse-non-null, all-null. Catches off-by-one errors in validity packing and mismatches + between field-node null counts and serialized validity buffers. """ schema_in = T.StructType( [ @@ -437,11 +437,10 @@ def passthrough(iterator): def test_map_in_arrow_multi_batch_per_partition(spark, tmp_path, accelerated): """ - Force many small batches in a single partition so the writer runs its per-batch - allocate/copy/write loop hundreds of times against a reused struct container (the leaf - buffers are reallocated each batch today; see #4383). Catches errors that only appear across - the batch boundary: stale value counts, offset/validity sizing on the second and later - batches, and variable-width data-buffer sizing as row content changes batch to batch. + Exercise the stream across many rows with a small Spark Arrow batch limit. The accelerated + path writes complete Comet source batches, so its source-vector turnover is covered separately + by test_map_in_arrow_nested_source_batches below. Catches stale value counts and changing + variable-width content on the Spark fallback path. """ schema_in = T.StructType( [ @@ -472,10 +471,222 @@ def passthrough(iterator): spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", prev_records) +def test_map_in_arrow_nested_source_batches(spark, tmp_path, accelerated): + """Serialize changing nested/null source vectors directly across actual worker batches.""" + item_type = T.StructType( + [ + T.StructField("label", T.StringType()), + T.StructField("score", T.IntegerType()), + ] + ) + payload_type = T.StructType( + [ + T.StructField("items", T.ArrayType(item_type, containsNull=True)), + T.StructField( + "attrs", + T.MapType(T.StringType(), T.IntegerType(), valueContainsNull=True), + ), + ] + ) + input_schema = T.StructType( + [ + T.StructField("id", T.LongType(), nullable=False), + T.StructField("payload", payload_type), + ] + ) + output_schema = T.StructType( + [ + *input_schema.fields, + T.StructField("batch_index", T.IntegerType(), nullable=False), + T.StructField("batch_size", T.IntegerType(), nullable=False), + ] + ) + + rows = [] + for index in range(37): + if index % 11 == 0: + payload = None + else: + if index % 5 == 0: + items = None + else: + items = [ + None + if (index + offset) % 6 == 0 + else { + "label": None + if (index + offset) % 4 == 0 + else f"item_{index}_{'x' * offset}", + "score": None + if (index + offset) % 3 == 0 + else index * 10 + offset, + } + for offset in range(index % 4) + ] + + if index % 7 == 0: + attrs = None + elif index % 6 == 0: + attrs = {} + else: + attrs = { + f"key_{index % 3}": None if index % 4 == 0 else index, + "marker": index * 10, + } + payload = {"items": items, "attrs": attrs} + rows.append((index, payload)) + + src = str(tmp_path / "nested_source_batches.parquet") + spark.createDataFrame(rows, input_schema).coalesce(1).write.parquet(src) + + batch_size = 7 + previous_comet_batch_size = spark.conf.get("spark.comet.batchSize", "8192") + previous_arrow_batch_size = spark.conf.get( + "spark.sql.execution.arrow.maxRecordsPerBatch" + ) + spark.conf.set("spark.comet.batchSize", str(batch_size)) + spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", str(batch_size)) + try: + + def annotate_batches(iterator): + for batch_index, batch in enumerate(iterator): + assert batch.schema.names == ["id", "payload"] + yield pa.RecordBatch.from_arrays( + [ + *batch.columns, + pa.array([batch_index] * batch.num_rows, type=pa.int32()), + pa.array([batch.num_rows] * batch.num_rows, type=pa.int32()), + ], + names=[*batch.schema.names, "batch_index", "batch_size"], + ) + + result_df = spark.read.parquet(src).mapInArrow(annotate_batches, output_schema) + _assert_plan_matches_mode(_executed_plan(result_df), accelerated) + + output = result_df.collect() + assert len(output) == len(rows) + expected_payloads = dict(rows) + observed_batches = {} + for row in output: + payload = row["payload"] + actual_payload = None if payload is None else payload.asDict(recursive=True) + assert actual_payload == expected_payloads[row["id"]] + observed_batches.setdefault(row["batch_index"], []).append(row) + + assert sorted(observed_batches) == list(range(6)) + expected_batch_sizes = [batch_size] * 5 + [2] + observed_batch_sizes = [len(observed_batches[index]) for index in range(6)] + assert observed_batch_sizes == expected_batch_sizes + for batch_rows in observed_batches.values(): + assert {row["batch_size"] for row in batch_rows} == {len(batch_rows)} + finally: + spark.conf.set("spark.comet.batchSize", previous_comet_batch_size) + spark.conf.set( + "spark.sql.execution.arrow.maxRecordsPerBatch", previous_arrow_batch_size + ) + + +@pytest.mark.parametrize( + "wrapper,rename_field,api", + [ + (wrapper, rename_field, api) + for wrapper, rename_field in [ + ("{}", False), + ("array({})", False), + ("map(1,{})", False), + ("map({},1)", False), + ("map(1,array({}))", False), + ("{}", True), + ("array({})", True), + ("map(1,{})", True), + ] + for api in ["mapInArrow", "mapInPandas"] + # Spark's pandas conversion cannot use dict-valued struct keys in a map. + if api == "mapInArrow" or wrapper != "map({},1)" + ], +) +def test_map_in_batch_union_with_compatible_nested_fields( + spark, tmp_path, accelerated, api, wrapper, rename_field +): + """Union batches can vary in names/nullability without changing their buffer layout.""" + rows = [(0, None), (1, 2), (2, 3)] + src = str(tmp_path / "union_fields.parquet") + spark.createDataFrame(rows, "id int, value int").coalesce(1).write.parquet(src) + source = spark.read.parquet(src) + left_value = "named_struct('id', id, 'value', 1)" + right_name = "renamed_value" if rename_field else "value" + right_value = f"named_struct('id', id, '{right_name}', value)" + if wrapper == "map(1,{})": + right_value = f"IF(value IS NULL, NULL, {right_value})" + left = source.selectExpr(f"{wrapper.format(left_value)} AS payload") + right = source.selectExpr(f"{wrapper.format(right_value)} AS payload") + combined = left.union(right).coalesce(1) + expected = Counter(map(repr, combined.collect())) + + def passthrough(iterator): + for batch in iterator: + if api == "mapInArrow": + data_type = batch.schema.field("payload").type + if pa.types.is_map(data_type): + assert not data_type.field(0).nullable # entries + assert not data_type.key_field.nullable + nested = data_type.key_type + if not pa.types.is_struct(nested): + nested = data_type.item_type + if pa.types.is_struct(nested): + assert data_type.item_field.nullable + if pa.types.is_list(nested): + nested = nested.value_type + assert nested.field("value").nullable + yield batch + + result = getattr(combined, api)(passthrough, combined.schema) + plan = _executed_plan(result) + _assert_plan_matches_mode( + plan, + accelerated, + vanilla_node="MapInArrow" if api == "mapInArrow" else "MapInPandas", + ) + if accelerated: + assert "CometUnion" in plan + assert Counter(map(repr, result.collect())) == expected + + +@pytest.mark.parametrize("api", ["mapInArrow", "mapInPandas"]) +def test_map_in_batch_union_with_different_field_metadata( + spark, tmp_path, accelerated, api +): + """Parquet field IDs describe source columns, not the outgoing IPC buffer layout.""" + sources = [] + for field_id in (1, 2): + nested = T.StructType( + [T.StructField("value", T.IntegerType(), True, {"parquet.field.id": field_id})] + ) + schema = T.StructType([T.StructField("payload", nested)]) + path = str(tmp_path / f"field_id_{field_id}.parquet") + spark.createDataFrame([((10,),), ((None,),)], schema).coalesce(1).write.parquet(path) + sources.append(spark.read.parquet(path)) + combined = sources[0].union(sources[1]).coalesce(1) + + def passthrough(iterator): + yield from iterator + + result = getattr(combined, api)(passthrough, combined.schema) + plan = _executed_plan(result) + _assert_plan_matches_mode( + plan, + accelerated, + vanilla_node="MapInArrow" if api == "mapInArrow" else "MapInPandas", + ) + if accelerated: + assert "CometUnion" in plan + assert Counter(row.payload.value for row in result.collect()) == Counter({10: 2, None: 2}) + + def test_map_in_arrow_wide_schema(spark, tmp_path, accelerated): """ - 50-column mixed-type schema. The bulk-copy path walks a flattened addresses[] array indexed - across the whole vector tree; off-by-one in flattening logic surfaces at depth * width. + 50-column mixed-type schema. Direct serialization must preserve every source vector and its + independent null bitmap when building the IPC field-node and buffer lists. """ fields = [T.StructField("id", T.LongType())] for i in range(15): @@ -664,10 +875,9 @@ def _normalize(row): def test_map_in_arrow_numeric_scalars(spark, tmp_path, accelerated): """ - Covers the BaseFixedWidthVector branch in CometColumnarPythonInput.copyVector for - every fixed-width primitive Comet's scan supports beyond the long/double/int already - exercised by other tests: boolean, byte, short, float. Each has a distinct buffer - size, and the validity bit handling is independent per column. + Covers every fixed-width primitive Comet's scan supports beyond the long/double/int already + exercised by other tests: boolean, byte, short, float. Each has a distinct buffer size, and + the validity bit handling is independent per directly serialized source vector. """ schema_in = T.StructType( [ @@ -781,8 +991,8 @@ def test_map_in_arrow_map_type(spark, tmp_path, accelerated): """ MapType is encoded in Arrow as a List> with extra metadata. The buffer layout (offsets + struct child + key/value children) is distinct from a plain - list, and CometMapVector is a separate vector class from CometListVector. Without - this test the recursive copy path through map-typed columns is unexercised. + list, and CometMapVector is a separate vector class from CometListVector. Without this test, + direct recursive serialization of map-typed source vectors is unexercised. """ schema_in = T.StructType( [ @@ -830,11 +1040,10 @@ def _normalize(row): def test_map_in_arrow_deeply_nested(spark, tmp_path, accelerated): """ - Exercises the recursive descent in CometColumnarPythonInput.copyVector at depth > 1, - in every nesting combination: array-of-array, array-of-struct, struct-of-array, - struct-of-struct. Single-level nesting is covered by test_map_in_arrow_array_and_struct; - the bug surface here is that setLastSet / setValueCount must be applied bottom-up - correctly at every level. + Exercises direct source-vector serialization at depth > 1, in every nesting combination: + array-of-array, array-of-struct, struct-of-array, struct-of-struct. Single-level nesting is + covered by test_map_in_arrow_array_and_struct; the bug surface here is that Arrow field + nodes and buffers must remain in the expected depth-first order at every level. """ schema_in = T.StructType( [ @@ -952,9 +1161,9 @@ def _norm_input_config(c): def test_map_in_arrow_falls_back_when_use_large_var_types(spark, tmp_path): """ `spark.sql.execution.arrow.useLargeVarTypes=true` widens StringType / BinaryType to - LargeUtf8 / LargeBinary in the destination IPC root (8-byte offsets). Comet's source - vectors always use 4-byte offsets; CometColumnarPythonInput.copyVector does a raw - setBytes per buffer and would corrupt the offset buffer in this configuration. + LargeUtf8 / LargeBinary in Spark's input IPC schema (8-byte offsets). Native Comet + vectors use 4-byte offsets; direct serialization advertises their matching types, + producing a valid stream but not the large input types requested by the configuration. EliminateRedundantTransitions must skip the rewrite in that case so vanilla Spark handles the operation. This test does not use the `accelerated` fixture: it sets pyarrowUDF.enabled=true AND useLargeVarTypes=true and asserts the plan still falls diff --git a/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala b/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala index 09802dddaa4..8d8617fae33 100644 --- a/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala +++ b/spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala @@ -161,37 +161,4 @@ class CometMapInBatchSuite extends CometTestBase { } } - test("end-to-end: rewrite-on output matches rewrite-off output for primitives + varchar") { - // This test needs PySpark workers; only run if PYSPARK_PYTHON is set in the env. - assume( - sys.env.contains("PYSPARK_PYTHON"), - "set PYSPARK_PYTHON to enable end-to-end pyarrow UDF tests") - - withTempPath { path => - val pathStr = path.getCanonicalPath - spark - .range(0, 1000, 1, 4) - .selectExpr( - "id AS id", - "CAST(id AS DOUBLE) * 1.5 AS dbl", - "CASE WHEN id % 10 = 0 THEN NULL ELSE CONCAT('row_', CAST(id AS STRING)) END AS s") - .write - .mode("overwrite") - .parquet(pathStr) - - // Baseline: rewrite disabled, vanilla MapInArrowExec runs. - val baseline = withSQLConf(CometConf.COMET_PYARROW_UDF_ENABLED.key -> "false") { - spark.read.parquet(pathStr).collect().map(_.toSeq).toSet - } - - // Optimized: rewrite enabled, CometMapInBatchExec + CometArrowPythonRunner runs. - withSQLConf(CometConf.COMET_PYARROW_UDF_ENABLED.key -> "true") { - val df = spark.read.parquet(pathStr) - val result = df.collect().map(_.toSeq).toSet - assert( - result == baseline, - s"optimized output differs from baseline:\noptimized=$result\nbaseline=$baseline") - } - } - } } diff --git a/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala new file mode 100644 index 00000000000..c20deff471e --- /dev/null +++ b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala @@ -0,0 +1,481 @@ +/* + * 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 java.nio.file.{Files, Paths} + +import scala.jdk.CollectionConverters._ + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.arrow.c.{ArrowArray, ArrowSchema, Data} +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, DictionaryEncoding, Field, FieldType, Schema} +import org.apache.spark.sql.execution.python.CometArrowPythonRunnerBase.{hasCompatibleSchema, 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("input schema compatibility preserves physical types and nested layouts") { + val intType = new ArrowType.Int(32, true) + def nested(dataType: ArrowType): Seq[Field] = Seq( + new Field( + "outer", + FieldType.nullable(ArrowType.Struct.INSTANCE), + Seq(new Field("value", FieldType.nullable(dataType), null)).asJava)) + + val renamed = Seq( + new Field( + "renamed", + FieldType.notNullable(ArrowType.Struct.INSTANCE), + Seq(new Field("other", FieldType.notNullable(intType), null)).asJava)) + hasCompatibleSchema(nested(intType), renamed) shouldBe true + hasCompatibleSchema(nested(intType), Seq.empty) shouldBe false + hasCompatibleSchema( + nested(intType), + Seq(new Field("outer", FieldType.nullable(ArrowType.Struct.INSTANCE), null))) shouldBe false + + val incompatibleTypes: Seq[(ArrowType, ArrowType)] = Seq( + (intType, new ArrowType.Int(64, true)), + (intType, new ArrowType.Int(32, false)), + (ArrowType.Utf8.INSTANCE, ArrowType.LargeUtf8.INSTANCE), + (ArrowType.Binary.INSTANCE, ArrowType.LargeBinary.INSTANCE), + (new ArrowType.Decimal(10, 2, 128), new ArrowType.Decimal(10, 3, 128))) + incompatibleTypes.foreach { case (expected, actual) => + hasCompatibleSchema(nested(expected), nested(actual)) shouldBe false + } + } + + test("input schema compatibility preserves extension and dictionary interpretation") { + val intType = new ArrowType.Int(32, true) + def fields( + metadata: Map[String, String] = Map.empty, + dictionary: DictionaryEncoding = null): Seq[Field] = + Seq(new Field("value", new FieldType(true, intType, dictionary, metadata.asJava), null)) + + hasCompatibleSchema( + fields(Map("PARQUET:field_id" -> "1")), + fields(Map("PARQUET:field_id" -> "2"))) shouldBe true + Seq( + ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME, + ArrowType.ExtensionType.EXTENSION_METADATA_KEY_METADATA).foreach { key => + hasCompatibleSchema( + fields(Map(key -> "before")), + fields(Map(key -> "after"))) shouldBe false + } + val dictionary = new DictionaryEncoding(1L, false, intType) + hasCompatibleSchema(fields(dictionary = dictionary), fields()) shouldBe false + hasCompatibleSchema( + fields(dictionary = dictionary), + fields(dictionary = new DictionaryEncoding(2L, false, intType))) shouldBe false + } + + test("direct batches retain borrowed buffers without copying them into the writer allocator") { + val sourceAllocator = new RootAllocator(Long.MaxValue) + val writerAllocator = new RootAllocator(1024) + val vector = new VarCharVector("source_name", sourceAllocator) + val output = new ByteArrayOutputStream() + try { + val payload = Array.fill[Byte](16 * 1024)('x'.toByte) + vector.allocateNew(payload.length.toLong, 2) + vector.setSafe(0, payload) + vector.setNull(1) + vector.setValueCount(2) + + val field = new Field("payload", vector.getField.getFieldType, vector.getField.getChildren) + val buffers = vector.getFieldBuffers.asScala.toSeq + val originalReferenceCounts = buffers.map(_.refCnt()) + val originalLastSet = vector.getLastSet + + withWriter(Seq(field), writerAllocator, Channels.newChannel(output)) { channel => + val originalWriterAllocation = writerAllocator.getAllocatedMemory + serializeBatch(new WriteChannel(channel), Seq(vector), 2, writerAllocator) + + writerAllocator.getAllocatedMemory shouldBe originalWriterAllocation + buffers.map(_.refCnt()) shouldBe originalReferenceCounts + vector.getLastSet shouldBe originalLastSet + vector.getValueCount shouldBe 2 + vector.get(0) shouldBe payload + vector.isNull(1) shouldBe true + } + + withReader(output.toByteArray) { reader => + reader.loadNextBatch() shouldBe true + val struct = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector] + struct.getNullCount shouldBe 0 + val result = struct.getChild("payload").asInstanceOf[VarCharVector] + result.get(0) shouldBe payload + result.isNull(1) shouldBe true + reader.loadNextBatch() shouldBe false + } + } finally { + vector.close() + writerAllocator.close() + sourceAllocator.close() + } + } + + for (failSerialization <- Seq(false, true)) { + test(s"direct FFI batches release temporary references (write failure: $failSerialization)") { + // Arrow's JNI loader extracts its library here; Maven's target/tmp may not exist yet. + Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + val sourceAllocator = new RootAllocator(Long.MaxValue) + val importAllocator = new RootAllocator(Long.MaxValue) + val writerAllocator = new RootAllocator(1024) + val source = new VarCharVector("payload", sourceAllocator) + val array = ArrowArray.allocateNew(sourceAllocator) + val schema = ArrowSchema.allocateNew(sourceAllocator) + var imported: VarCharVector = null + var failWrites = false + val output = new ByteArrayOutputStream() { + override def write(bytes: Array[Byte], offset: Int, length: Int): Unit = { + if (failWrites) { + throw new IOException("injected Arrow IPC write failure") + } + super.write(bytes, offset, length) + } + } + try { + val payload = Array.fill[Byte](16 * 1024)('x'.toByte) + source.allocateNew(payload.length.toLong, 2) + source.setSafe(0, payload) + source.setNull(1) + source.setValueCount(2) + + Data.exportVector(sourceAllocator, source, null, array, schema) + imported = + Data.importVector(importAllocator, array, schema, null).asInstanceOf[VarCharVector] + imported.getDataBuffer.memoryAddress() shouldBe source.getDataBuffer.memoryAddress() + // Only the C Data Interface release callback now keeps the original buffers alive. + source.close() + + val buffers = imported.getFieldBuffers.asScala.toSeq + val originalReferenceCounts = buffers.map(_.refCnt()) + val originalImportAllocation = importAllocator.getAllocatedMemory + val originalSourceAllocation = sourceAllocator.getAllocatedMemory + originalSourceAllocation should be > 0L + + withWriter(Seq(imported.getField), writerAllocator, Channels.newChannel(output)) { + channel => + val originalWriterAllocation = writerAllocator.getAllocatedMemory + failWrites = failSerialization + try { + if (failSerialization) { + val error = intercept[IOException] { + serializeBatch(new WriteChannel(channel), Seq(imported), 2, writerAllocator) + } + error.getMessage shouldBe "injected Arrow IPC write failure" + } else { + serializeBatch(new WriteChannel(channel), Seq(imported), 2, writerAllocator) + } + } finally { + failWrites = false + } + + buffers.map(_.refCnt()) shouldBe originalReferenceCounts + importAllocator.getAllocatedMemory shouldBe originalImportAllocation + sourceAllocator.getAllocatedMemory shouldBe originalSourceAllocation + writerAllocator.getAllocatedMemory shouldBe originalWriterAllocation + imported.getValueCount shouldBe 2 + imported.get(0) shouldBe payload + imported.isNull(1) shouldBe true + } + + if (!failSerialization) { + withReader(output.toByteArray) { reader => + reader.loadNextBatch() shouldBe true + val struct = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector] + val result = struct.getChild("payload").asInstanceOf[VarCharVector] + result.get(0) shouldBe payload + result.isNull(1) shouldBe true + reader.loadNextBatch() shouldBe false + } + } + + imported.close() + imported = null + importAllocator.getAllocatedMemory shouldBe 0L + sourceAllocator.getAllocatedMemory shouldBe 0L + writerAllocator.getAllocatedMemory shouldBe 0L + } finally { + if (imported != null) { + imported.close() + } + schema.close() + array.close() + source.close() + writerAllocator.close() + importAllocator.close() + sourceAllocator.close() + } + } + } + + test("direct batches preserve nested list, struct, map, and null field layouts") { + val sourceAllocator = new RootAllocator(Long.MaxValue) + val writerAllocator = new RootAllocator(Long.MaxValue) + val list = ListVector.empty("items", sourceAllocator) + val struct = StructVector.empty("details", sourceAllocator) + val map = MapVector.empty("mapping", sourceAllocator, false) + val nulls = new NullVector("nulls", 3) + val output = new ByteArrayOutputStream() + try { + val listWriter = list.getWriter + listWriter.setPosition(0) + listWriter.startList() + listWriter.integer().writeInt(11) + listWriter.integer().writeInt(12) + listWriter.endList() + listWriter.setPosition(1) + listWriter.writeNull() + listWriter.setPosition(2) + listWriter.startList() + listWriter.integer().writeInt(13) + listWriter.endList() + listWriter.setValueCount(3) + + val structWriter = struct.getWriter + structWriter.setPosition(0) + structWriter.start() + structWriter.integer("count").writeInt(21) + structWriter.end() + structWriter.setPosition(1) + structWriter.writeNull() + structWriter.setPosition(2) + structWriter.start() + structWriter.integer("count").writeNull() + structWriter.end() + structWriter.setValueCount(3) + + val mapWriter = map.getWriter + mapWriter.setPosition(0) + mapWriter.startMap() + mapWriter.startEntry() + mapWriter.key().integer().writeInt(31) + mapWriter.value().integer().writeInt(32) + mapWriter.endEntry() + mapWriter.endMap() + mapWriter.setPosition(1) + mapWriter.writeNull() + mapWriter.setPosition(2) + mapWriter.startMap() + mapWriter.startEntry() + mapWriter.key().integer().writeInt(33) + mapWriter.value().integer().writeNull() + mapWriter.endEntry() + mapWriter.endMap() + mapWriter.setValueCount(3) + + val vectors = Seq[FieldVector](list, struct, map, nulls) + withWriter(vectors.map(_.getField), writerAllocator, Channels.newChannel(output)) { + channel => + serializeBatch(new WriteChannel(channel), vectors, 3, writerAllocator) + } + + withReader(output.toByteArray) { reader => + reader.loadNextBatch() shouldBe true + val result = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector] + result.getNullCount shouldBe 0 + + val resultList = result.getChild("items").asInstanceOf[ListVector] + resultList.getObject(0).asScala.toSeq shouldBe Seq(11, 12) + resultList.isNull(1) shouldBe true + resultList.getObject(2).asScala.toSeq shouldBe Seq(13) + + val resultStruct = result.getChild("details").asInstanceOf[StructVector] + resultStruct.getChild("count").asInstanceOf[IntVector].get(0) shouldBe 21 + resultStruct.isNull(1) shouldBe true + resultStruct.getChild("count").isNull(2) shouldBe true + + val resultMap = result.getChild("mapping").asInstanceOf[MapVector] + val entries = resultMap.getDataVector.asInstanceOf[StructVector] + entries.getChildByOrdinal(0).getField.getName shouldBe MapVector.KEY_NAME + entries.getChildByOrdinal(1).getField.getName shouldBe MapVector.VALUE_NAME + entries.getChildByOrdinal(0).asInstanceOf[IntVector].get(0) shouldBe 31 + entries.getChildByOrdinal(1).asInstanceOf[IntVector].get(0) shouldBe 32 + resultMap.isNull(1) shouldBe true + entries.getChildByOrdinal(1).isNull(1) shouldBe true + + val resultNulls = result.getChild("nulls").asInstanceOf[NullVector] + resultNulls.getNullCount shouldBe 3 + reader.loadNextBatch() shouldBe false + } + } finally { + nulls.close() + map.close() + struct.close() + list.close() + writerAllocator.close() + sourceAllocator.close() + } + } + + test("direct batches preserve zero-row batches between populated batches") { + val sourceAllocator = new RootAllocator(Long.MaxValue) + val writerAllocator = new RootAllocator(Long.MaxValue) + val first = new IntVector("value", sourceAllocator) + val empty = new IntVector("value", sourceAllocator) + val last = new IntVector("value", sourceAllocator) + val output = new ByteArrayOutputStream() + try { + first.allocateNew(2) + first.setSafe(0, 41) + first.setSafe(1, 42) + first.setValueCount(2) + empty.setValueCount(0) + last.allocateNew(1) + last.setSafe(0, 43) + last.setValueCount(1) + + withWriter(Seq(first.getField), writerAllocator, Channels.newChannel(output)) { channel => + serializeBatch(new WriteChannel(channel), Seq(first), 2, writerAllocator) + serializeBatch(new WriteChannel(channel), Seq(empty), 0, writerAllocator) + serializeBatch(new WriteChannel(channel), Seq(last), 1, writerAllocator) + } + + withReader(output.toByteArray) { reader => + reader.loadNextBatch() shouldBe true + reader.getVectorSchemaRoot.getRowCount shouldBe 2 + reader.loadNextBatch() shouldBe true + reader.getVectorSchemaRoot.getRowCount shouldBe 0 + reader.loadNextBatch() shouldBe true + reader.getVectorSchemaRoot.getRowCount shouldBe 1 + val struct = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector] + struct.getChild("value").asInstanceOf[IntVector].get(0) shouldBe 43 + reader.loadNextBatch() shouldBe false + } + } finally { + last.close() + empty.close() + first.close() + writerAllocator.close() + sourceAllocator.close() + } + } + + test("direct batches represent non-null structs with no child columns") { + val allocator = new RootAllocator(Long.MaxValue) + val output = new ByteArrayOutputStream() + try { + withWriter(Seq.empty, allocator, Channels.newChannel(output)) { channel => + serializeBatch(new WriteChannel(channel), Seq.empty, 3, allocator) + } + + withReader(output.toByteArray) { reader => + reader.loadNextBatch() shouldBe true + val struct = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector] + struct.getValueCount shouldBe 3 + struct.getNullCount shouldBe 0 + struct.getChildrenFromFields.isEmpty shouldBe true + reader.loadNextBatch() shouldBe false + } + } finally { + allocator.close() + } + } + + test("direct batches release temporary references when writing the stream fails") { + val sourceAllocator = new RootAllocator(Long.MaxValue) + val writerAllocator = new RootAllocator(Long.MaxValue) + val source = new IntVector("value", sourceAllocator) + val output = new ByteArrayOutputStream() + var failWrites = false + val channel = new WritableByteChannel { + private var open = true + + override def isOpen: Boolean = open + + override def close(): Unit = open = false + + override def write(buffer: ByteBuffer): Int = { + if (failWrites) { + throw new IOException("injected Arrow IPC write failure") + } + val bytes = new Array[Byte](buffer.remaining()) + buffer.get(bytes) + output.write(bytes) + bytes.length + } + } + try { + source.allocateNew(1) + source.setSafe(0, 51) + source.setValueCount(1) + + withWriter(Seq(source.getField), writerAllocator, channel) { channel => + val originalReferenceCounts = source.getFieldBuffers.asScala.map(_.refCnt()).toSeq + val originalWriterAllocation = writerAllocator.getAllocatedMemory + failWrites = true + try { + val error = intercept[IOException] { + serializeBatch(new WriteChannel(channel), Seq(source), 1, writerAllocator) + } + error.getMessage shouldBe "injected Arrow IPC write failure" + } finally { + failWrites = false + } + source.getFieldBuffers.asScala.map(_.refCnt()).toSeq shouldBe originalReferenceCounts + writerAllocator.getAllocatedMemory shouldBe originalWriterAllocation + source.get(0) shouldBe 51 + } + } finally { + source.close() + writerAllocator.close() + sourceAllocator.close() + } + } +}