dongjoon-hyun commented on code in PR #59933: URL: https://github.com/apache/spark/pull/59933#discussion_r4419050074
########## docs/sql-native-data-sources.md: ########## @@ -37,7 +37,16 @@ package, and loads it when a query uses the data source. Spark adapts it to Data is exchanged only as columnar batches, through the [Arrow C data interface](https://arrow.apache.org/docs/format/CDataInterface.html), so Spark reads -the data that a library returns without copying it. +the data that a library returns without copying it. When a write gets the batches of a scan as +they are, as in a plain copy from one native data source to another, Spark writes them without +converting them to rows: it copies the buffers of their columns, since a library may keep the data +that it writes until it closes its data writer, and converts value by value the columns of other +Arrow types than those of the schema, such as `utf8_view` strings. The copied values are written as +the scan returned them, without the checks of a conversion, such as the precision of a decimal. Review Comment: **1. [Low] The PR description still says that the written data is the same** With the change for [5](https://github.com/apache/spark/pull/59933#discussion_r4417897756) of my second review, this sentence says that the copied values are written without the checks of a conversion. That is a user-facing change. Before this PR, `ColumnarToRowExec` read a decimal column with `ArrowColumnVector.getDecimal(rowId, precision, scale)`, and `Decimal.set` fails with `NUMERIC_VALUE_OUT_OF_RANGE.WITHOUT_SUGGESTION` when a value does not fit the precision of the schema. So a copy from a native data source that returned such a value failed, and now it writes the value. The new test "release the converted columns when a later column fails to convert" relies on this check on the conversion path. The section "Does this PR introduce _any_ user-facing change?" of the PR description still says "No. The written data is the same." Could it describe this change instead? Also, "How was this patch tested?" lists two of the six new tests of `ColumnarDataSourceSuite`. The other four are "write the pending rows of converted batches before a copied batch", "a copied batch ends at the byte limit as a batch of rows does", "an empty columnar batch does not end the pending batch of rows" and "release the converted columns when a later column fails to convert". ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala: ########## @@ -128,6 +142,354 @@ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { } } + test("write the columnar batches of a query") { + // Runs a write, and returns the plan of the query that it writes. + def writtenPlan(write: => Unit): SparkPlan = { + QueryTest.withQueryExecutionsCaptured(spark)(write).flatMap { qe => + qe.executedPlan.collectFirst { case writeExec: V2TableWriteExec => writeExec.query } + }.head + } + // Writes the ids [0, 10) and their names, read in Arrow batches of `arrowBatchSize` rows, or in + // a batch of on-heap vectors for each of the two partitions of the scan. + def copy(table: String, arrowBatchSize: Int = 0, keep: Boolean = false): SparkPlan = { + writtenPlan { + spark.read.format(format).option("end", "10").option("arrowBatchSize", arrowBatchSize) + .load() + .write.format(format).option("table", table).option("keep", keep).mode("append").save() + } + } + val expected = (0L until 10L).map(id => (id, s"name$id")) + + withSQLConf( + SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "3", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_FORCE_APPLY.key -> "false") { + // The data writers get the batches of the scan, without a transition to rows. On-heap + // vectors are converted to Arrow value by value, into batches that end as those of rows do. + assert(copy("columnar_write").isInstanceOf[BatchScanExec]) + assert(TestColumnarSink.rows("columnar_write") == expected) + assert(TestColumnarSink.batchSizes("columnar_write") == Seq(3, 2, 3, 2)) + + // Arrow batches are copied buffer by buffer, each in batches of at most 3 rows. + assert(copy("arrow_write", arrowBatchSize = 2).isInstanceOf[BatchScanExec]) + assert(TestColumnarSink.rows("arrow_write") == expected) + assert(TestColumnarSink.batchSizes("arrow_write") == Seq(2, 2, 1, 2, 2, 1)) + copy("large_arrow_write", arrowBatchSize = 5) + assert(TestColumnarSink.rows("large_arrow_write") == expected) + assert(TestColumnarSink.batchSizes("large_arrow_write") == Seq(3, 2, 3, 2)) + + // A data writer may keep the data of the batches that it gets, though the scan overwrites + // its Arrow vectors with each batch. + copy("kept_arrow_write", arrowBatchSize = 2, keep = true) + assert(TestColumnarSink.rows("kept_arrow_write") == expected) + + // A query that does not produce columnar batches is written as rows. + val plan = writtenPlan { + spark.read.format(format).option("end", "10").load() + .selectExpr("id", "upper(name) AS name") + .write.format(format).option("table", "row_write").mode("append").save() + } + assert(plan.exists(_.isInstanceOf[ColumnarToRowExec]), plan) + + // With adaptive query execution, a query with a shuffle is an adaptive plan, which is not + // columnar, so its batches are converted to rows. + val shuffled = writtenPlan { + spark.read.format(format).option("end", "10").option("arrowBatchSize", 2).load() + .repartition(2) + .write.format(format).option("table", "shuffled_write").mode("append").save() + } + assert(shuffled.isInstanceOf[AdaptiveSparkPlanExec], shuffled) + assert(find(shuffled)(_.isInstanceOf[ColumnarToRowExec]).isDefined, shuffled) + assert(TestColumnarSink.rows("shuffled_write").sorted == expected) + + // A query that supports rows too, such as the scan of a cached DataFrame, is written as rows. + val cached = spark.read.format(format).option("end", "10").load().select("id").cache() + try { + cached.count() + val writes = QueryTest.withQueryExecutionsCaptured(spark) { + cached.write.format(format).option("table", "cached_write").mode("append").save() + }.flatMap(_.executedPlan.collectFirst { case writeExec: V2TableWriteExec => writeExec }) + val write = writes.head + assert(write.supportsColumnarInput, write) + assert(write.query.isInstanceOf[InMemoryTableScanExec], write.query) + assert(write.query.supportsColumnar && write.query.supportsRowBased, write.query) + assert(!write.writesColumnar, write) + assert(TestColumnarSink.rows("cached_write").sorted == expected.map { case (id, _) => + (id, "") + }) + } finally { + cached.unpersist() + } + assert(TestColumnarSink.rows("row_write") == expected.map { case (id, name) => + (id, name.toUpperCase(Locale.ROOT)) + }) + } + } + + test("copy the Arrow columns of every type, or convert the columns of a batch") { + val nullable = Seq( + "id % 2 = 0", "CAST(id AS TINYINT)", "CAST(id AS SMALLINT)", "CAST(id AS INT)", "id", + "CAST(id AS FLOAT)", "CAST(id AS DOUBLE)", "CAST(id AS DECIMAL(10, 2))", + "CAST(id AS DECIMAL(38, 10))", "CAST(id AS STRING)", "CAST(CAST(id AS STRING) AS BINARY)", + "DATE_ADD(DATE'2020-01-01', CAST(id AS INT))", "TIMESTAMP_SECONDS(id)", + "CAST(TIMESTAMP_SECONDS(id) AS TIMESTAMP_NTZ)", "MAKE_YM_INTERVAL(CAST(id AS INT), 1)", + "MAKE_DT_INTERVAL(CAST(id AS INT))", "MAKE_INTERVAL(0, 1, 0, CAST(id AS INT))", + "PARSE_JSON(CAST(id AS STRING))", "ARRAY(id, NULL)", "ARRAY(CAST(id AS STRING))", + "MAP(CAST(id AS STRING), id)", "MAP(id, ARRAY(CAST(id AS STRING)))", + "NAMED_STRUCT('a', id, 'b', CAST(id AS STRING))", "ARRAY(NAMED_STRUCT('x', id))") + // c0 is not nullable, c1 has the null type, and the other columns have nulls. + val df = spark.range(7).selectExpr( + Seq("id AS c0", "NULL AS c1") ++ nullable.zipWithIndex.map { case (expr, i) => + s"IF(id % 3 = 1, NULL, $expr) AS c${i + 2}" + }: _*) + val schema = df.schema + val rows = df.queryExecution.toRdd.map(_.copy()).collect().toSeq + val allocator = ArrowUtils.rootAllocator.newChildAllocator("test", 0, Long.MaxValue) + try { + // Writes a batch with the session time zone America/Los_Angeles, and returns the sizes of + // the batches that the data writer got, the time zone of their timestamp column c14, which + // a copy keeps, the Arrow types of their string column c11, and their rows. The data writer + // keeps the data of the batches, and reads it once `overwrite` overwrote the batch. + def writeAndRead( + batch: ColumnarBatch, + overwrite: () => Unit, + schema: StructType = schema, + maxRecordsPerBatch: Int = 10, + maxBytesPerBatch: Long = Int.MaxValue) + : (Seq[Int], Seq[String], Seq[ArrowType], Seq[InternalRow]) = { + val sizes = ArrayBuffer.empty[Int] + val timeZones = ArrayBuffer.empty[String] + val stringTypes = ArrayBuffer.empty[ArrowType] + val kept = ArrayBuffer.empty[TestArrowExport] + def field(batch: ColumnarBatch, name: String) = batch.column(schema.fieldIndex(name)) + .asInstanceOf[ArrowColumnVector].getValueVector.getField + val dataWriter = new DataWriter[ColumnarBatch] { + override def write(batch: ColumnarBatch): Unit = { + sizes += batch.numRows() + timeZones += field(batch, "c14").getType.asInstanceOf[ArrowType.Timestamp].getTimezone + stringTypes += field(batch, "c11").getType + kept += TestArrowExport(batch) + } + override def commit(): WriterCommitMessage = null + override def abort(): Unit = {} + override def close(): Unit = kept.foreach(_.close()) + } + val writer = new ArrowBatchingDataWriter(dataWriter, ArrowBatchingOptions( + schema, "America/Los_Angeles", maxRecordsPerBatch, maxBytesPerBatch)) + try { + writer.writeBatch(batch) + writer.commit() + overwrite() + val toUnsafe = UnsafeProjection.create(schema) + val written = kept.toSeq.flatMap { exported => + exported.read(_.rowIterator().asScala.map(toUnsafe(_).copy()).toSeq) + } + kept.clear() + (sizes.toSeq, timeZones.toSeq, stringTypes.toSeq, written) + } finally { + writer.close() + } + } + + // Runs `f` with a batch of Arrow vectors with the time zone UTC, and a function that + // overwrites the vectors with the rows in reverse order. + def withArrowBatch(largeVarTypes: Boolean)(f: (ColumnarBatch, () => Unit) => Unit): Unit = { + val root = VectorSchemaRoot.create( + ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = largeVarTypes), + allocator) + try { + val arrowWriter = ArrowWriter.create(root) + rows.foreach(arrowWriter.write) + arrowWriter.finish() + val batch = new ColumnarBatch( + root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray, + rows.length) + f(batch, () => { + arrowWriter.reset() + rows.reverse.foreach(arrowWriter.write) + arrowWriter.finish() + }) + } finally { + root.close() + } + } + + val expected = rows.map(UnsafeProjection.create(schema)(_).copy()) + val utf8 = ArrowType.Utf8.INSTANCE + Seq(false, true).foreach { largeVarTypes => + withArrowBatch(largeVarTypes) { (batch, _) => + // The size of a range of rows, which the batches of a copy end at, is the size of a slice + // of these rows. + (0 until batch.numCols()).foreach { i => + val vector = batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector + Seq((0, 7), (1, 4), (3, 7), (6, 7), (2, 2)).foreach { case (start, end) => + val pair = vector.getTransferPair(vector.getAllocator) + pair.splitAndTransfer(start, end - start) + val sliceSize = + try pair.getTo.getBufferSizeFor(end - start) finally pair.getTo.close() + assert(ArrowBatchingDataWriter.rangeSize(vector, start, end) == sliceSize, + s"${vector.getField} [$start, $end)") + } + } + } + } + withArrowBatch(largeVarTypes = false) { (batch, overwrite) => + // The columns are Arrow vectors with the types of the schema, so they are copied. + assert(writeAndRead(batch, overwrite) == ((Seq(7), Seq("UTC"), Seq(utf8), expected))) + } + withArrowBatch(largeVarTypes = false) { (batch, overwrite) => + // In slices of at most 3 rows, which start in the middle of the bytes of validity. + assert(writeAndRead(batch, overwrite, maxRecordsPerBatch = 3) == + ((Seq(3, 3, 1), Seq.fill(3)("UTC"), Seq.fill(3)(utf8), expected))) + } + withArrowBatch(largeVarTypes = false) { (batch, overwrite) => + // A copied batch also ends after the first row with which it holds the maximum size. + assert(writeAndRead(batch, overwrite, maxBytesPerBatch = 1) == + ((Seq.fill(7)(1), Seq.fill(7)("UTC"), Seq.fill(7)(utf8), expected))) + } + withArrowBatch(largeVarTypes = true) { (batch, overwrite) => + // Each column of another Arrow type than that of the schema, here a string or a binary, is + // converted on its own, and the other columns are copied. + assert(writeAndRead(batch, overwrite, maxRecordsPerBatch = 3) == + ((Seq(3, 3, 1), Seq.fill(3)("UTC"), Seq.fill(3)(utf8), expected))) + } + withArrowBatch(largeVarTypes = true) { (batch, overwrite) => + // The size of the converted columns counts too. + assert(writeAndRead(batch, overwrite, maxBytesPerBatch = 1) == + ((Seq.fill(7)(1), Seq.fill(7)("UTC"), Seq.fill(7)(utf8), expected))) + } + + // On-heap vectors are converted too. They do not support the null type. + val onHeapSchema = StructType(schema.filterNot(_.name == "c1")) + val onHeapRows = rows.map(row => InternalRow.fromSeq(row.toSeq(schema).patch(1, Nil, 1))) + val vectors: Seq[WritableColumnVector] = + OnHeapColumnVector.allocateColumns(onHeapRows.length, onHeapSchema).toSeq + try { + val converter = new RowToColumnConverter(onHeapSchema) + onHeapRows.foreach(converter.convert(_, vectors.toArray)) + val batch = new ColumnarBatch(vectors.map(v => v: ColumnVector).toArray, onHeapRows.length) + assert(writeAndRead(batch, () => (), onHeapSchema, maxRecordsPerBatch = 3) == + ((Seq(3, 3, 1), Seq.fill(3)("America/Los_Angeles"), Seq.fill(3)(utf8), + onHeapRows.map(UnsafeProjection.create(onHeapSchema)(_).copy())))) + } finally { + vectors.foreach(_.close()) + } + } finally { + allocator.close() + } + } + + test("write the pending rows of converted batches before a copied batch") { + val schema = StructType.fromDDL("id BIGINT, name STRING") + val writer = new ArrowBatchingDataWriter( + new TestDataWriter(fail = false), ArrowBatchingOptions(schema, "UTC", 10, Int.MaxValue)) + val onHeap = new TestBatchReader(0, 3, schema) + val arrow = new TestArrowBatchReader(3, 7, schema, 4) + try { + // The rows of the on-heap batch stay in the batch of rows until the Arrow batch is copied. + assert(onHeap.next()) + writer.writeBatch(onHeap.get()) + assert(arrow.next()) + writer.writeBatch(arrow.get()) + val message = writer.commit().asInstanceOf[TestCommitMessage] + assert(message.rows == (0L until 7L).map(id => (id, s"name$id"))) + assert(message.batchSizes == Seq(3, 4)) + } finally { + writer.close() + onHeap.close() + arrow.close() + } + } + + test("a copied batch ends at the byte limit as a batch of rows does") { + val schema = StructType.fromDDL("id BIGINT, name STRING") + def batchSizes(batch: ColumnarBatch, maxBytesPerBatch: Long): Seq[Int] = { + val writer = new ArrowBatchingDataWriter(new TestDataWriter(fail = false), + ArrowBatchingOptions(schema, "UTC", 5, maxBytesPerBatch)) + try { + writer.writeBatch(batch) + writer.commit().asInstanceOf[TestCommitMessage].batchSizes + } finally { + writer.close() + } + } + val onHeap = new TestBatchReader(0, 20, schema) + val arrow = new TestArrowBatchReader(0, 20, schema, 20) + try { + assert(onHeap.next() && arrow.next()) + val vectors = (0 until 2).map { + arrow.get().column(_).asInstanceOf[ArrowColumnVector].getValueVector + } + // The size of the first 2 rows. The names of the ids from 10 are 1 byte longer. + val twoRows = vectors.map(ArrowBatchingDataWriter.rangeSize(_, 0, 2)).sum + Seq(twoRows - 1, twoRows, twoRows + 1, 3 * twoRows).foreach { limit => + val rowSizes = batchSizes(onHeap.get(), limit) + assert(rowSizes.sum == 20) + assert(batchSizes(arrow.get(), limit) == rowSizes, s"limit $limit") + } + // The limit ends the batches every 2 rows, also across the windows of 5 rows. + assert(batchSizes(arrow.get(), twoRows) == Seq.fill(10)(2)) + } finally { + onHeap.close() + arrow.close() + } + } + + test("an empty columnar batch does not end the pending batch of rows") { + val schema = StructType.fromDDL("id BIGINT, name STRING") + val writer = new ArrowBatchingDataWriter( + new TestDataWriter(fail = false), ArrowBatchingOptions(schema, "UTC", 10, Int.MaxValue)) + val first = new TestBatchReader(0, 2, schema) + val empty = new TestArrowBatchReader(0, 0, schema, 1) + val second = new TestBatchReader(2, 4, schema) + try { + assert(first.next() && second.next()) + writer.writeBatch(first.get()) + writer.writeBatch(empty.get()) + writer.writeBatch(second.get()) + assert(writer.commit().asInstanceOf[TestCommitMessage].batchSizes == Seq(4)) + } finally { + writer.close() + first.close() + empty.close() + second.close() + } + } + + test("release the converted columns when a later column fails to convert") { + val schema = StructType.fromDDL("id BIGINT, a DECIMAL(10, 2), b DECIMAL(10, 2)") + val allocator = ArrowUtils.rootAllocator.newChildAllocator("test", 0, Long.MaxValue) + val root = VectorSchemaRoot.create(new Schema(Seq( + Field.notNullable("id", new ArrowType.Int(64, true)), + Field.nullable("a", new ArrowType.Decimal(38, 2, 128)), + Field.nullable("b", new ArrowType.Decimal(38, 2, 128))).asJava), allocator) + try { + root.allocateNew() + root.getVector("id").asInstanceOf[BigIntVector].setSafe(0, 0) + // A decimal of another precision than that of the schema is converted, and the value of b + // does not fit the precision of the schema. + root.getVector("a").asInstanceOf[DecimalVector].setSafe(0, new java.math.BigDecimal("1.00")) + root.getVector("b").asInstanceOf[DecimalVector] + .setSafe(0, new java.math.BigDecimal("1000000000000.00")) + root.setRowCount(1) + val batch = new ColumnarBatch( + root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray, 1) + val allocated = ArrowUtils.rootAllocator.getAllocatedMemory + val writer = new ArrowBatchingDataWriter( + new TestDataWriter(fail = false), ArrowBatchingOptions(schema, "UTC", 10, Int.MaxValue)) + try { + intercept[Exception](writer.writeBatch(batch)) Review Comment: **3. [Nit] The test accepts any exception** `intercept[Exception]` also passes when `writeBatch` fails for another reason, for example before `a` is converted. The check of the allocator at L486 then holds without any converted vector to release. Checking the error of `b` ties the failure to the case that the test is about: ```suggestion val e = intercept[SparkArithmeticException](writer.writeBatch(batch)) assert(e.getCondition == "NUMERIC_VALUE_OUT_OF_RANGE.WITHOUT_SUGGESTION") ``` This needs `org.apache.spark.SparkArithmeticException` in the imports. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala: ########## @@ -128,6 +142,354 @@ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { } } + test("write the columnar batches of a query") { + // Runs a write, and returns the plan of the query that it writes. + def writtenPlan(write: => Unit): SparkPlan = { + QueryTest.withQueryExecutionsCaptured(spark)(write).flatMap { qe => + qe.executedPlan.collectFirst { case writeExec: V2TableWriteExec => writeExec.query } + }.head + } + // Writes the ids [0, 10) and their names, read in Arrow batches of `arrowBatchSize` rows, or in + // a batch of on-heap vectors for each of the two partitions of the scan. + def copy(table: String, arrowBatchSize: Int = 0, keep: Boolean = false): SparkPlan = { + writtenPlan { + spark.read.format(format).option("end", "10").option("arrowBatchSize", arrowBatchSize) + .load() + .write.format(format).option("table", table).option("keep", keep).mode("append").save() + } + } + val expected = (0L until 10L).map(id => (id, s"name$id")) + + withSQLConf( + SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "3", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_FORCE_APPLY.key -> "false") { + // The data writers get the batches of the scan, without a transition to rows. On-heap + // vectors are converted to Arrow value by value, into batches that end as those of rows do. + assert(copy("columnar_write").isInstanceOf[BatchScanExec]) + assert(TestColumnarSink.rows("columnar_write") == expected) + assert(TestColumnarSink.batchSizes("columnar_write") == Seq(3, 2, 3, 2)) + + // Arrow batches are copied buffer by buffer, each in batches of at most 3 rows. + assert(copy("arrow_write", arrowBatchSize = 2).isInstanceOf[BatchScanExec]) + assert(TestColumnarSink.rows("arrow_write") == expected) + assert(TestColumnarSink.batchSizes("arrow_write") == Seq(2, 2, 1, 2, 2, 1)) + copy("large_arrow_write", arrowBatchSize = 5) + assert(TestColumnarSink.rows("large_arrow_write") == expected) + assert(TestColumnarSink.batchSizes("large_arrow_write") == Seq(3, 2, 3, 2)) + + // A data writer may keep the data of the batches that it gets, though the scan overwrites + // its Arrow vectors with each batch. + copy("kept_arrow_write", arrowBatchSize = 2, keep = true) + assert(TestColumnarSink.rows("kept_arrow_write") == expected) + + // A query that does not produce columnar batches is written as rows. + val plan = writtenPlan { + spark.read.format(format).option("end", "10").load() + .selectExpr("id", "upper(name) AS name") + .write.format(format).option("table", "row_write").mode("append").save() + } + assert(plan.exists(_.isInstanceOf[ColumnarToRowExec]), plan) + + // With adaptive query execution, a query with a shuffle is an adaptive plan, which is not + // columnar, so its batches are converted to rows. + val shuffled = writtenPlan { + spark.read.format(format).option("end", "10").option("arrowBatchSize", 2).load() + .repartition(2) + .write.format(format).option("table", "shuffled_write").mode("append").save() + } + assert(shuffled.isInstanceOf[AdaptiveSparkPlanExec], shuffled) + assert(find(shuffled)(_.isInstanceOf[ColumnarToRowExec]).isDefined, shuffled) + assert(TestColumnarSink.rows("shuffled_write").sorted == expected) + + // A query that supports rows too, such as the scan of a cached DataFrame, is written as rows. + val cached = spark.read.format(format).option("end", "10").load().select("id").cache() + try { + cached.count() + val writes = QueryTest.withQueryExecutionsCaptured(spark) { + cached.write.format(format).option("table", "cached_write").mode("append").save() + }.flatMap(_.executedPlan.collectFirst { case writeExec: V2TableWriteExec => writeExec }) + val write = writes.head + assert(write.supportsColumnarInput, write) + assert(write.query.isInstanceOf[InMemoryTableScanExec], write.query) + assert(write.query.supportsColumnar && write.query.supportsRowBased, write.query) + assert(!write.writesColumnar, write) + assert(TestColumnarSink.rows("cached_write").sorted == expected.map { case (id, _) => + (id, "") + }) + } finally { + cached.unpersist() + } + assert(TestColumnarSink.rows("row_write") == expected.map { case (id, name) => + (id, name.toUpperCase(Locale.ROOT)) + }) + } + } + + test("copy the Arrow columns of every type, or convert the columns of a batch") { + val nullable = Seq( + "id % 2 = 0", "CAST(id AS TINYINT)", "CAST(id AS SMALLINT)", "CAST(id AS INT)", "id", + "CAST(id AS FLOAT)", "CAST(id AS DOUBLE)", "CAST(id AS DECIMAL(10, 2))", + "CAST(id AS DECIMAL(38, 10))", "CAST(id AS STRING)", "CAST(CAST(id AS STRING) AS BINARY)", + "DATE_ADD(DATE'2020-01-01', CAST(id AS INT))", "TIMESTAMP_SECONDS(id)", + "CAST(TIMESTAMP_SECONDS(id) AS TIMESTAMP_NTZ)", "MAKE_YM_INTERVAL(CAST(id AS INT), 1)", + "MAKE_DT_INTERVAL(CAST(id AS INT))", "MAKE_INTERVAL(0, 1, 0, CAST(id AS INT))", + "PARSE_JSON(CAST(id AS STRING))", "ARRAY(id, NULL)", "ARRAY(CAST(id AS STRING))", + "MAP(CAST(id AS STRING), id)", "MAP(id, ARRAY(CAST(id AS STRING)))", + "NAMED_STRUCT('a', id, 'b', CAST(id AS STRING))", "ARRAY(NAMED_STRUCT('x', id))") + // c0 is not nullable, c1 has the null type, and the other columns have nulls. + val df = spark.range(7).selectExpr( + Seq("id AS c0", "NULL AS c1") ++ nullable.zipWithIndex.map { case (expr, i) => + s"IF(id % 3 = 1, NULL, $expr) AS c${i + 2}" + }: _*) + val schema = df.schema + val rows = df.queryExecution.toRdd.map(_.copy()).collect().toSeq + val allocator = ArrowUtils.rootAllocator.newChildAllocator("test", 0, Long.MaxValue) + try { + // Writes a batch with the session time zone America/Los_Angeles, and returns the sizes of + // the batches that the data writer got, the time zone of their timestamp column c14, which + // a copy keeps, the Arrow types of their string column c11, and their rows. The data writer + // keeps the data of the batches, and reads it once `overwrite` overwrote the batch. + def writeAndRead( + batch: ColumnarBatch, + overwrite: () => Unit, + schema: StructType = schema, + maxRecordsPerBatch: Int = 10, + maxBytesPerBatch: Long = Int.MaxValue) + : (Seq[Int], Seq[String], Seq[ArrowType], Seq[InternalRow]) = { + val sizes = ArrayBuffer.empty[Int] + val timeZones = ArrayBuffer.empty[String] + val stringTypes = ArrayBuffer.empty[ArrowType] + val kept = ArrayBuffer.empty[TestArrowExport] + def field(batch: ColumnarBatch, name: String) = batch.column(schema.fieldIndex(name)) + .asInstanceOf[ArrowColumnVector].getValueVector.getField + val dataWriter = new DataWriter[ColumnarBatch] { + override def write(batch: ColumnarBatch): Unit = { + sizes += batch.numRows() + timeZones += field(batch, "c14").getType.asInstanceOf[ArrowType.Timestamp].getTimezone + stringTypes += field(batch, "c11").getType + kept += TestArrowExport(batch) + } + override def commit(): WriterCommitMessage = null + override def abort(): Unit = {} + override def close(): Unit = kept.foreach(_.close()) + } + val writer = new ArrowBatchingDataWriter(dataWriter, ArrowBatchingOptions( + schema, "America/Los_Angeles", maxRecordsPerBatch, maxBytesPerBatch)) + try { + writer.writeBatch(batch) + writer.commit() + overwrite() + val toUnsafe = UnsafeProjection.create(schema) + val written = kept.toSeq.flatMap { exported => + exported.read(_.rowIterator().asScala.map(toUnsafe(_).copy()).toSeq) + } + kept.clear() + (sizes.toSeq, timeZones.toSeq, stringTypes.toSeq, written) + } finally { + writer.close() + } + } + + // Runs `f` with a batch of Arrow vectors with the time zone UTC, and a function that + // overwrites the vectors with the rows in reverse order. + def withArrowBatch(largeVarTypes: Boolean)(f: (ColumnarBatch, () => Unit) => Unit): Unit = { + val root = VectorSchemaRoot.create( + ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = largeVarTypes), + allocator) + try { + val arrowWriter = ArrowWriter.create(root) + rows.foreach(arrowWriter.write) + arrowWriter.finish() + val batch = new ColumnarBatch( + root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray, + rows.length) + f(batch, () => { + arrowWriter.reset() + rows.reverse.foreach(arrowWriter.write) + arrowWriter.finish() + }) + } finally { + root.close() + } + } + + val expected = rows.map(UnsafeProjection.create(schema)(_).copy()) + val utf8 = ArrowType.Utf8.INSTANCE + Seq(false, true).foreach { largeVarTypes => + withArrowBatch(largeVarTypes) { (batch, _) => + // The size of a range of rows, which the batches of a copy end at, is the size of a slice + // of these rows. + (0 until batch.numCols()).foreach { i => + val vector = batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector + Seq((0, 7), (1, 4), (3, 7), (6, 7), (2, 2)).foreach { case (start, end) => + val pair = vector.getTransferPair(vector.getAllocator) + pair.splitAndTransfer(start, end - start) + val sliceSize = + try pair.getTo.getBufferSizeFor(end - start) finally pair.getTo.close() + assert(ArrowBatchingDataWriter.rangeSize(vector, start, end) == sliceSize, + s"${vector.getField} [$start, $end)") + } + } + } + } + withArrowBatch(largeVarTypes = false) { (batch, overwrite) => + // The columns are Arrow vectors with the types of the schema, so they are copied. + assert(writeAndRead(batch, overwrite) == ((Seq(7), Seq("UTC"), Seq(utf8), expected))) + } + withArrowBatch(largeVarTypes = false) { (batch, overwrite) => + // In slices of at most 3 rows, which start in the middle of the bytes of validity. + assert(writeAndRead(batch, overwrite, maxRecordsPerBatch = 3) == + ((Seq(3, 3, 1), Seq.fill(3)("UTC"), Seq.fill(3)(utf8), expected))) + } + withArrowBatch(largeVarTypes = false) { (batch, overwrite) => + // A copied batch also ends after the first row with which it holds the maximum size. + assert(writeAndRead(batch, overwrite, maxBytesPerBatch = 1) == + ((Seq.fill(7)(1), Seq.fill(7)("UTC"), Seq.fill(7)(utf8), expected))) + } + withArrowBatch(largeVarTypes = true) { (batch, overwrite) => + // Each column of another Arrow type than that of the schema, here a string or a binary, is + // converted on its own, and the other columns are copied. + assert(writeAndRead(batch, overwrite, maxRecordsPerBatch = 3) == + ((Seq(3, 3, 1), Seq.fill(3)("UTC"), Seq.fill(3)(utf8), expected))) + } + withArrowBatch(largeVarTypes = true) { (batch, overwrite) => + // The size of the converted columns counts too. + assert(writeAndRead(batch, overwrite, maxBytesPerBatch = 1) == + ((Seq.fill(7)(1), Seq.fill(7)("UTC"), Seq.fill(7)(utf8), expected))) + } + + // On-heap vectors are converted too. They do not support the null type. + val onHeapSchema = StructType(schema.filterNot(_.name == "c1")) + val onHeapRows = rows.map(row => InternalRow.fromSeq(row.toSeq(schema).patch(1, Nil, 1))) + val vectors: Seq[WritableColumnVector] = + OnHeapColumnVector.allocateColumns(onHeapRows.length, onHeapSchema).toSeq + try { + val converter = new RowToColumnConverter(onHeapSchema) + onHeapRows.foreach(converter.convert(_, vectors.toArray)) + val batch = new ColumnarBatch(vectors.map(v => v: ColumnVector).toArray, onHeapRows.length) + assert(writeAndRead(batch, () => (), onHeapSchema, maxRecordsPerBatch = 3) == + ((Seq(3, 3, 1), Seq.fill(3)("America/Los_Angeles"), Seq.fill(3)(utf8), + onHeapRows.map(UnsafeProjection.create(onHeapSchema)(_).copy())))) + } finally { + vectors.foreach(_.close()) + } + } finally { + allocator.close() + } + } + + test("write the pending rows of converted batches before a copied batch") { + val schema = StructType.fromDDL("id BIGINT, name STRING") + val writer = new ArrowBatchingDataWriter( + new TestDataWriter(fail = false), ArrowBatchingOptions(schema, "UTC", 10, Int.MaxValue)) + val onHeap = new TestBatchReader(0, 3, schema) + val arrow = new TestArrowBatchReader(3, 7, schema, 4) + try { + // The rows of the on-heap batch stay in the batch of rows until the Arrow batch is copied. + assert(onHeap.next()) + writer.writeBatch(onHeap.get()) + assert(arrow.next()) + writer.writeBatch(arrow.get()) + val message = writer.commit().asInstanceOf[TestCommitMessage] + assert(message.rows == (0L until 7L).map(id => (id, s"name$id"))) + assert(message.batchSizes == Seq(3, 4)) + } finally { + writer.close() + onHeap.close() + arrow.close() + } + } + + test("a copied batch ends at the byte limit as a batch of rows does") { Review Comment: **2. [Nit] No test restarts a window that has a converted column** Thanks for this test, the change for [6](https://github.com/apache/spark/pull/59933#discussion_r4417897763) of my second review. It covers the restart of a window in `ArrowBatchingDataWriter.writeBatch`, but both of its columns are copied, so the restarted window converts nothing. In "copy the Arrow columns of every type, or convert the columns of a batch", the cases with `largeVarTypes = true` convert the string and binary columns, but either without a byte limit, where every batch starts a window, or in a single window of 10 rows over 7 rows, where no window is restarted. So no test runs a window that starts in the middle of the previous one and converts its rows again, with `WindowColumn.first` at that row. Could this test also compare a batch with a converted column, such as a `large_utf8` name column, with the row path? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala: ########## @@ -214,15 +235,216 @@ class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options: Arro } } - private def flush(): Unit = { - if (root != null) { - arrowWriter.finish() - val columns = root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector) + override def writeBatch(batch: ColumnarBatch): Unit = { + val numBatchRows = batch.numRows() + val copied = copiedVectors(batch) + if (copied.forall(_.isEmpty)) { + // Without columns to copy, the values are converted as rows are. + var i = 0 + while (i < numBatchRows) { + write(batch.getRow(i)) + i += 1 + } + } else if (numBatchRows > 0) { + // A batch without rows leaves the pending batch of rows as it is. + flush() + // A window of up to `maxRecordsPerBatch` rows of the batch, whose columns that are not copied + // are converted at once. The next window starts where the last batch written ended, so that + // the batches end as those of rows do. + var windowStart = 0 + while (windowStart < numBatchRows) { + val first = windowStart + val windowEnd = math.min(numBatchRows, first + options.maxRecordsPerBatch) + val converted = ArrayBuffer.empty[FieldVector] + windowStart = Utils.tryWithSafeFinally { + val columns = copied.indices.map { i => + copied(i) match { + case Some(vector) => WindowColumn(vector, 0, copied = true) + case None => + val vector = convertColumn(batch, i, first, windowEnd) + converted += vector + WindowColumn(vector, first, copied = false) + } + } + var start = first + var next = -1 + while (next < 0) { + val length = batchLength(columns, start, windowEnd) + if (start > first && start + length == windowEnd && windowEnd < numBatchRows) { Review Comment: **4. [Nit] A batch that the byte limit ends at the end of a window also restarts the window** Thanks for the change for [4](https://github.com/apache/spark/pull/59933#discussion_r4417897746) of my second review. The reply there says that a batch that the byte limit did not end starts the next window, but this condition does not check why the batch ended. So a batch that the byte limit ends exactly at `windowEnd` restarts the window too. For example, with `maxRecordsPerBatch` of 6 and a byte limit of the size of 3 rows whose values have the same size: - After [0, 3), `batchLength(3, 6)` returns 3, since the rows [3, 6) reach the limit with their last row, so the window restarts at row 3. - The window [3, 9) converts the rows [3, 6) again, writes them, and then restarts at row 6 for the same reason, and so on. The written batches are the same, but almost every row of a converted column is converted twice. If `batchLength` also returned whether the byte limit ended the batch, a window would only restart when it did not. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala: ########## @@ -214,15 +235,216 @@ class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options: Arro } } - private def flush(): Unit = { - if (root != null) { - arrowWriter.finish() - val columns = root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector) + override def writeBatch(batch: ColumnarBatch): Unit = { + val numBatchRows = batch.numRows() + val copied = copiedVectors(batch) + if (copied.forall(_.isEmpty)) { + // Without columns to copy, the values are converted as rows are. + var i = 0 + while (i < numBatchRows) { + write(batch.getRow(i)) + i += 1 + } + } else if (numBatchRows > 0) { + // A batch without rows leaves the pending batch of rows as it is. + flush() + // A window of up to `maxRecordsPerBatch` rows of the batch, whose columns that are not copied + // are converted at once. The next window starts where the last batch written ended, so that + // the batches end as those of rows do. + var windowStart = 0 + while (windowStart < numBatchRows) { + val first = windowStart + val windowEnd = math.min(numBatchRows, first + options.maxRecordsPerBatch) + val converted = ArrayBuffer.empty[FieldVector] + windowStart = Utils.tryWithSafeFinally { + val columns = copied.indices.map { i => + copied(i) match { + case Some(vector) => WindowColumn(vector, 0, copied = true) + case None => + val vector = convertColumn(batch, i, first, windowEnd) + converted += vector + WindowColumn(vector, first, copied = false) + } + } + var start = first + var next = -1 + while (next < 0) { + val length = batchLength(columns, start, windowEnd) + if (start > first && start + length == windowEnd && windowEnd < numBatchRows) { + // The batch may hold rows after the window, so it starts the next window, whose + // conversion converts its rows again. + next = start + } else { + writeRows(columns, start, length) + start += length + if (start == windowEnd) next = start + } + } + next + } { + converted.foreach(_.close()) + } + } + } + } + + /** + * Returns, for each column of a batch, its Arrow vector if it can be copied buffer by buffer: an + * Arrow vector with as many values as the batch has rows, whose field is physically congruent + * with the one of `arrowSchema`, though it may differ in the names, nullability and metadata of + * its fields, and in the time zones of timestamps, which the copy keeps. The other columns are + * converted value by value. + */ + private def copiedVectors(batch: ColumnarBatch): Seq[Option[FieldVector]] = { + val fields = arrowSchema.getFields.asScala + if (batch.numCols() != fields.length) { + Seq.fill(fields.length)(None) + } else { + fields.indices.map { i => + batch.column(i) match { + case column: ArrowColumnVector => column.getValueVector match { + case vector: FieldVector if vector.getValueCount == batch.numRows() && + ArrowUtils.isCompatibleWithDeclaredField(vector.getField, fields(i)) => + Some(vector) + case _ => None + } + case _ => None + } + } + } + } + + /** Converts the values of the rows [start, end) of a column into a vector of `arrowSchema`. */ + private def convertColumn( + batch: ColumnarBatch, + ordinal: Int, + start: Int, + end: Int): FieldVector = { + val vector = arrowSchema.getFields.get(ordinal).createVector(allocator) + try { + vector.setInitialCapacity(end - start) + vector.allocateNew() + val fieldWriter = ArrowWriter.createFieldWriter(vector) + var i = start + while (i < end) { + fieldWriter.write(batch.getRow(i), ordinal) + i += 1 + } + fieldWriter.finish() + vector + } catch { + case e: Throwable => + vector.close() + throw e + } + } + + /** + * Returns the number of rows of the batch that starts at the row `start` of the columns: up to + * `end`, or as a batch of rows ends, up to the first row with which it holds `maxBytesPerBatch` + * bytes, by the measure of `getBufferSizeFor`, which `rangeSize` computes without slicing. + */ + private def batchLength(columns: Seq[WindowColumn], start: Int, end: Int): Int = { + // The size is not computed when the limit is the largest possible one, as for rows. + if (options.maxBytesPerBatch >= Int.MaxValue) { + end - start + } else { + def sizeInBytes(rows: Int): Long = columns.map { column => + ArrowBatchingDataWriter.rangeSize( + column.vector, start - column.first, start - column.first + rows) + }.sum + if (sizeInBytes(end - start) < options.maxBytesPerBatch) { + end - start + } else { + // The size grows with the number of rows. + var lo = 1 + var hi = end - start + while (lo < hi) { + val mid = (lo + hi) >>> 1 + if (sizeInBytes(mid) >= options.maxBytesPerBatch) hi = mid else lo = mid + 1 + } + lo + } + } + } + + /** + * Writes the rows [start, start + length) of the columns. The copied columns are sliced, which + * shares most buffers of the batch, though it allocates new offsets for lists, and for strings + * and binaries whose offsets do not start at 0, and new validity buffers unless `start` is a + * multiple of 8. Then each buffer of the slice is copied at its size. Arrow Java reads the + * offsets of a vector of variable width, or of a list, as if they started at 0, which those of a + * slice do. A converted column, whose vector nothing else writes to, is only sliced. Review Comment: **5. [Nit] The comment does not mention the value buffers of booleans** Thanks for the change for [10](https://github.com/apache/spark/pull/59933#discussion_r4417897801) of my second review. In Arrow 19, `BitVector.splitAndTransferTo` passes both the validity buffer and the value buffer to `splitAndTransferBuffer`, so a slice of a boolean column also allocates a new value buffer unless `start` is a multiple of 8, as with `maxRecordsPerBatch = 3`. Also, "strings and binaries whose offsets do not start at 0" can be read as the first offset of the vector, while Arrow checks the offset of the row `start`: a slice that starts after a non-empty value gets new offsets even when the offsets of the vector start at 0. Maybe: ```suggestion * shares most buffers of the batch, though it allocates new offsets for lists, and for strings * and binaries unless the offset of the row `start` is 0, and new validity buffers, and value * buffers for booleans, unless `start` is a multiple of 8. Then each buffer of the slice is * copied at its size. Arrow Java reads the offsets of a vector of variable width, or of a list, * as if they started at 0, which those of a slice do. A converted column, whose vector nothing * else writes to, is only sliced. ``` -- 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]
