HyukjinKwon commented on code in PR #59933: URL: https://github.com/apache/spark/pull/59933#discussion_r4419261349
########## 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: Done in 0962134: the user-facing section now describes the copied values that skip the checks of a conversion (a decimal that does not fit the precision is written instead of failing with `NUMERIC_VALUE_OUT_OF_RANGE`) and the smaller batches, and the testing section lists all the new tests of `ColumnarDataSourceSuite`. ########## 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: Done in 0962134: "a copied batch ends at the byte limit as a batch of rows does" now also writes a batch whose names are `large_utf8`, so converted, and compares its batches and rows with those of the row path at each limit. Below 3 times the size of 2 rows, its windows start again at the last batch, with `WindowColumn.first` at that row. ########## 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: Done in 0962134: the test intercepts `SparkArithmeticException` and checks `NUMERIC_VALUE_OUT_OF_RANGE.WITHOUT_SUGGESTION`. ########## 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: Done in 0962134: `batchLength` also returns whether the batch holds `maxBytesPerBatch` bytes, and a window only starts again at a batch that ends with it without holding them. The new test "a window does not convert the rows of a batch that ends at the byte limit with it" counts the conversions of a column (10 for 10 rows; 16 without the change). ########## 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: Done in 0962134: the comment now mentions the value buffers of booleans, and the offset of the row `start` for strings and binaries, as suggested. -- 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]
