Repository navigation
perf(spark): append rows straight into the Arrow batch - #7
Conversation
VortexDataWriter allocated every batch at Arrow's default capacity of a few thousand values per vector, and the native writer keeps each batch until a row block of 8192 rows is complete, so small vortex.write.batch.size values retained hundreds of kilobytes of direct memory per batch and exhausted the executor's direct memory. Set the initial capacity of each vector to the batch size. Signed-off-by: gitRasheed <[email protected]>
VortexDataWriter copied every InternalRow into a list and converted the list to Arrow vectors once a batch was full. The copy is only there because Spark reuses the row object, and it costs an UnsafeRow plus its bytes per row (about 1 KiB per row for 100 columns), while the batch conversion dispatched on the Spark type once per value, looked vectors and struct children up per value, and copied each string twice. Append each row into the batch vectors as it arrives, through one ArrowFieldWriter per field in the shape of Spark's own ArrowWriter: typed once, with a per-writer position so list elements and map entries land in the right slot, struct and map children resolved once, and strings copied straight from the row into the Arrow buffer. The schema is exported once per writer; the native side only reads it. Vectors are still allocated fresh per batch, since the native writer imports the exported buffers zero-copy and releases them from its own thread. Because the buffers are now allocated at the first row of a batch rather than when it is full, the vectors start at 64 slots instead of the batch size: Arrow doubles on demand and keeps the grown capacity for the next allocateNew(), so a writer that fills its batches reaches the batch size after its first one, while a partition writer that sees a few rows holds 131 KiB rather than 3.3 MiB for 100 columns (200 open writers: 25 MiB instead of 630 MiB). A write() that fails part-way through a row poisons the writer, since the batch columns are misaligned from then on, so later write() and commit() calls fail until it is aborted. The exported schema is released before it is closed, which frees the child structs the export allocated (the native side only borrows the schema and never releases them). Writer-only, 2M rows, release native build: nested -31% to -38% wall, strings -35% to -39%, 100 numeric columns -24%, 5 longs within noise; allocation per row 116 -> 8 B (5 longs), 1072 -> 122 B (100 columns), 495 -> 9 B (strings), 761 -> 298 B (nested). Signed-off-by: gitRasheed <[email protected]>
ChristopherWirt
left a comment
ChristopherWirt
left a comment
There was a problem hiding this comment.
One regression to fix before merge (on commit()). The rest are non-blocking.
I checked the append logic against Arrow 19 and the native writeBatch, and ran a round-trip of 10k rows at batch sizes 1/2048/65536 with nulls at every nesting level. It reads back identical, and the Spark 3.5 and 4.1 suites pass.
| */ | ||
| @Override | ||
| public WriterCommitMessage commit() throws IOException { | ||
| checkNotFailed(); |
There was a problem hiding this comment.
Blocking. This throws before any cleanup, so a failed writer's native writer, allocator, vectors and exported schema all leak. Develop's commit() always cleaned up.
Abort doesn't rescue it: PartitionedVortexDataWriter.commit() sets closed = true even when a child throws, so its abort() returns early and never reaches the failed child.
Could a failed writer release its resources first, e.g. call abort() here, then throw?
| // For unsupported types, set null | ||
| throw new IllegalArgumentException("Unsupported data type: " + dataType); | ||
| for (ArrowFieldWriter writer : fieldWriters) { | ||
| writer.reset(); |
There was a problem hiding this comment.
Non-blocking: Arrow never shrinks the capacity it remembers. One batch with a 50 MB string makes every later batch allocate 50 MB at its first row. Develop behaves the same, but this PR's memory argument relies on that capacity. Worth re-applying setInitialCapacity here, or at least noting it in a comment?
| // Arrow's default capacity is a few thousand values per vector regardless of how many | ||
| // rows a batch holds, and the native writer keeps every batch it receives until a row | ||
| // block is complete, so the batch bounds the capacity. Start well below it, though: | ||
| // the buffers are allocated at the first row of every batch, and a partition writer | ||
| // that sees a handful of rows must not hold a whole batch of direct memory for its | ||
| // lifetime. Arrow doubles on demand and remembers the grown capacity for the next | ||
| // allocateNew(), so a writer that fills its batches reaches the batch size once. | ||
| int initialCapacity = Math.min(batchSize, INITIAL_VECTOR_CAPACITY); |
There was a problem hiding this comment.
Nit: this mostly repeats the commit message. Something like // Small start so writers that see few rows hold little memory; Arrow grows on demand. would do.
|
|
||
| @Override | ||
| void setValue(SpecializedGetters row, int ordinal) { | ||
| decimals.setSafe(count, row.getDecimal(ordinal, precision, scale).toJavaBigDecimal()); |
There was a problem hiding this comment.
Non-blocking: this allocates a BigDecimal per value. For precision <= 18, row.getDecimal(...).toUnscaledLong() avoids the allocation. Only wider decimals need BigDecimal.
|
|
||
| @Override | ||
| void setValue(SpecializedGetters row, int ordinal) { | ||
| binaries.setSafe(count, row.getBinary(ordinal)); |
There was a problem hiding this comment.
Non-blocking: getBinary allocates a byte[] and setSafe copies it again. The StringWriter approach (setValueLengthSafe, then copy straight from the row) would avoid the per-value garbage here too.
…ls as longs A commit after a failed write threw before any cleanup, and PartitionedVortexDataWriter does not abort a child whose commit threw, so the native writer, allocator, vectors and exported schema leaked. Abort first, then throw. Decimals of up to 18 digits are written from their unscaled long instead of a BigDecimal per value. The capacity comment is shorter and notes that Arrow keeps the largest capacity a batch needed. Signed-off-by: gitRasheed <[email protected]>
|
Fixed in 910abcc.
|
Summary
Fixes #3
VortexDataWritercopied every row into a list, because Spark reuses the row object, and converted the list to Arrow when the batch was full. Writing each row straight into the Arrow vectors removes the copy and the per-value type dispatch. Writer alone, 2M rows, release native build:Changes
ArrowFieldWriterappends one field of a row to its Arrow vector, in the shape of Spark'sArrowWriter: typed once, struct and map children resolved once, strings copied once from the row into the Arrow buffer.write()that fails part-way through a row marks the writer failed, so laterwrite()andcommit()calls throw instead of writing misaligned columns.VortexDataWriterTestcovers memory growing with rows rather than batches, and rows the caller reuses afterwrite(): on-heap, off-heap, and aGenericInternalRowwith a binary column whosebyte[]is overwritten. Develop read the overwritten bytes back in the last case.Half or more of the writer thread's time is now the native encoder, which runs on the Spark task thread. Write worker threads are a separate change.
VortexOutputWriterin vortex-data#9658 still copies rows the same way, so this change applies there too.Checked locally:
spotlessApplythenspotlessCheck,./gradlew --parallel javadoc checkwith Docker running (213 vortex-spark tests pass on Spark 3.5.9 and 4.1.2, includingVortexDataSourceS3MockTest), the same tests on Spark 4.2.0 with-PsparkTestVersion=4.2.0, and a write of 400k rows into 200 partition directories (1600 open writers, every column type, with nulls) that reads back identical to the same data written as Parquet.I used AI tools to help write and review this change.