viirya commented on code in PR #59581: URL: https://github.com/apache/spark/pull/59581#discussion_r4362032988
########## sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedColumnVector.java: ########## @@ -0,0 +1,37 @@ +/* + * 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.vectorized; + +/** + * A marker for a {@link org.apache.spark.sql.vectorized.ColumnVector} that its producer may + * release or overwrite after returning it in a batch: for example, a reader that refills the same + * vectors for every batch, or frees the vectors of a batch once the next one is read or once it is + * closed. A consumer may only read such a vector until it asks the producer for the next batch, + * and must copy what it needs past that point. + * + * <p>There is no other agreement between the producer and the consumer of a batch, so the type of + * its vectors is the protocol, as for {@code ColumnVector.closeIfFreeable}: the producer owns the + * lifetime of a vector carrying this marker, and its {@code closeIfFreeable} is a no-op. A vector + * without it may be held by the consumer past the next batch without a copy, and the consumer Review Comment: Agreed that the contract and the Arrow cache should not disagree. Following your comment on L36, the `ReusedColumnVector` interface moved to the follow-up in 09bae80321d, so this PR no longer states that contract. #59582, where the Arrow columnar path starts classifying its input, also changes `ArrowCachedBatchToColumnarBatchIterator`. It no longer closes the previous root in `next()`. Each root is closed once all columns of its batch are released through `closeIfFreeable`, and the rest are closed when the task completes. That keeps cache + UDF on the Arrow columnar path, rather than marking the cache vectors and sending it to the slower row queue path. #59582 has a test for the 2,000,000-row case you describe. Happy to discuss marking them instead there. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala: ########## @@ -113,10 +113,11 @@ case class ArrowEvalPythonExec( // Columnar path: delegate to doExecuteColumnar, flatten to Review Comment: Thanks, fixed in 09bae80321d together with your comment on L120: `doExecute` is now only the row path, and whether this node takes columnar input is fixed by the transitions the planner inserted. I added a test that builds the physical plan with columnar input enabled and executes it for the first time after disabling it. Without the change it fails with `ClassCastException: class org.apache.spark.sql.vectorized.ColumnarBatch cannot be cast to class org.apache.spark.sql.catalyst.InternalRow`. One detail: collecting the same Dataset a second time does not fail, because `SparkPlan.execute()` caches its RDD (`executeRDD`), so the second action reuses the RDD built on the columnar path. A plan has to be executed for the first time after the change, as in the test. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala: ########## @@ -242,9 +249,10 @@ abstract class BaseSliceArrowOutputProcessor( batch } + // The slices of a batch are closed once the next one is produced, so they are reused vectors. Review Comment: Thanks, I updated the comment in 09bae80321d: the vectors are either the reader's root, into which the next batch is loaded, or a slice of it, which is closed once the next one is produced. The closing of the reader's root on a root-to-slice transition is fixed in #59616 (SPARK-60367), together with the duplicated rows from your comment on L249. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala: ########## @@ -242,9 +249,10 @@ abstract class BaseSliceArrowOutputProcessor( batch Review Comment: Thanks, this is fixed in #59616 (SPARK-60367), together with the closing of the reader's root from your comment on L252. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowOutputProcessorSuite.scala: ########## @@ -0,0 +1,103 @@ +/* + * 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} + +import scala.collection.mutable.ArrayBuffer + +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.execution.arrow.ArrowWriter +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.execution.vectorized.ReusedColumnVector +import org.apache.spark.sql.types.{IntegerType, StringType, StructType} +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.unsafe.types.UTF8String + +class ArrowOutputProcessorSuite extends SparkFunSuite { + + private val schema = new StructType() + .add("s", new StructType().add("a", IntegerType).add("b", StringType)) + private val numBatches = 3 + private val rowsPerBatch = 3 + + /** Writes `numBatches` batches of a struct column as an Arrow stream, as a Python worker does. */ + private def arrowStream(): Array[Byte] = { + val allocator = + ArrowUtils.rootAllocator.newChildAllocator("arrow-output-writer", 0, Long.MaxValue) + val root = VectorSchemaRoot.create( + ArrowUtils.toArrowSchema(schema, null, errorOnDuplicatedFieldNames = true, + largeVarTypes = false), + allocator) + try { + val out = new ByteArrayOutputStream() + val writer = new ArrowStreamWriter(root, null, out) + val arrowWriter = ArrowWriter.create(root) + writer.start() + (0 until numBatches * rowsPerBatch).grouped(rowsPerBatch).foreach { values => + values.foreach { v => + arrowWriter.write(InternalRow(InternalRow(v, UTF8String.fromString(s"v$v")))) + } + arrowWriter.finish() + writer.writeBatch() + arrowWriter.reset() + } + writer.end() + out.toByteArray + } finally { + root.close() + allocator.close() + } + } + + test("Struct output columns stay valid after a consumer releases each batch") { + // The output processor loads every batch into the same vectors. A consumer releases each + // batch with closeIfFreeable once it is done with it, which must leave the vectors, including + // the child columns of a struct, to the processor. + val allocator = + ArrowUtils.rootAllocator.newChildAllocator("arrow-output-reader", 0, Long.MaxValue) + val reader = new ArrowStreamReader(new ByteArrayInputStream(arrowStream()), allocator) + val processor = new ArrowOutputProcessorImpl(reader, metrics()) Review Comment: Added in 09bae80321d: the test now also runs `SliceRecordsArrowOutputProcessorImpl` with `rowsPerBatch` and `SliceBytesArrowOutputProcessorImpl` with `Int.MaxValue`. With plain `ArrowColumnVector`s in `BaseSliceArrowOutputProcessor.getVectors`, both fail with `should have as many children as in the schema: found 0 expected 2`. After your comment on L120, a `ColumnarToRowExec` sits above this node, so whole-stage codegen does make a difference now: the UDF test runs with it on and off, and under a limit, where the generated code closes the current batch. -- 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]
