Skip to content

GH-1116: [Java] Fix compressed buffer prefix write and ZSTD dstCapacity #1119

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lidavidm merged 3 commits into from
Jun 30, 2026
Merged

GH-1116: [Java] Fix compressed buffer prefix write and ZSTD dstCapacity #1119

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4c754b1
fix
LuciferYang Apr 18, 2026
d7b6a8f
Address review: drop concurrency framing from compress() comment
LuciferYang Jun 29, 2026
b8e6b46
Address review: drop ZSTD round-trip tests that don't reflect GH-1116
LuciferYang Jun 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBu
long bytesWritten =
Zstd.compressUnsafe(
compressedBuffer.memoryAddress() + CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH,
dstSize,
maxSize,
/*src*/ uncompressedBuffer.memoryAddress(),
/* srcSize= */ uncompressedBuffer.writerIndex(),
/* level= */ this.compressionLevel);
Expand Down
252 changes: 252 additions & 0 deletions compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,9 @@
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.TimeStampMilliVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
Expand All @@ -53,12 +55,15 @@
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
import org.apache.arrow.vector.ipc.message.IpcOption;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
Expand Down Expand Up @@ -347,6 +352,253 @@ void testUnloadCompressed(CompressionUtil.CodecType codec) {
});
}

/**
* Test multi-batch streaming with ZSTD compression, wide schema, VectorSchemaRoot reuse, and
* all-null columns. This reproduces the scenario from GH-1116 where the 8-byte
* uncompressed-length prefix of a compressed buffer could be incorrectly written as 0.
*/
@Test
void testMultiBatchZstdStreamWithWideSchemaAndAllNulls() throws Exception {
final int fieldCount = 100;
final int batchCount = 10;
final int rowsPerBatch = 500;

// Build a wide schema: mix of int, timestamp, and varchar fields
List<Field> fields = new ArrayList<>();
for (int i = 0; i < fieldCount; i++) {
switch (i % 3) {
case 0:
fields.add(Field.nullable("int_" + i, new ArrowType.Int(32, true)));
break;
case 1:
fields.add(
Field.nullable("ts_" + i, new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)));
break;
case 2:
fields.add(Field.nullable("str_" + i, ArrowType.Utf8.INSTANCE));
break;
default:
break;
}
}
Schema schema = new Schema(fields);

ByteArrayOutputStream out = new ByteArrayOutputStream();
try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
ArrowStreamWriter writer =
new ArrowStreamWriter(
root,
new DictionaryProvider.MapDictionaryProvider(),
Channels.newChannel(out),
IpcOption.DEFAULT,
CommonsCompressionFactory.INSTANCE,
CompressionUtil.CodecType.ZSTD)) {
writer.start();

for (int batch = 0; batch < batchCount; batch++) {
// Clear and reallocate — mimics the reporter's reuse pattern
root.clear();
for (FieldVector vector : root.getFieldVectors()) {
vector.allocateNew();
}
root.setRowCount(rowsPerBatch);

for (int col = 0; col < fieldCount; col++) {
FieldVector vector = root.getVector(col);
// Make some batches have all-null columns for certain fields
boolean allNull = (batch % 3 == 0) && (col % 3 == 1); // timestamps in every 3rd batch
switch (col % 3) {
case 0:
{
IntVector iv = (IntVector) vector;
for (int row = 0; row < rowsPerBatch; row++) {
if (allNull || row % 7 == 0) {
iv.setNull(row);
} else {
iv.setSafe(row, batch * rowsPerBatch + row);
}
}
break;
}
case 1:
{
TimeStampMilliVector tv = (TimeStampMilliVector) vector;
for (int row = 0; row < rowsPerBatch; row++) {
if (allNull || row % 5 == 0) {
tv.setNull(row);
} else {
tv.setSafe(row, 1_700_000_000_000L + (long) batch * rowsPerBatch + row);
}
}
break;
}
case 2:
{
VarCharVector sv = (VarCharVector) vector;
for (int row = 0; row < rowsPerBatch; row++) {
if (allNull || row % 9 == 0) {
sv.setNull(row);
} else {
sv.setSafe(row, ("val_" + batch + "_" + row).getBytes(StandardCharsets.UTF_8));
}
}
break;
}
default:
break;
}
vector.setValueCount(rowsPerBatch);
}

writer.writeBatch();
}
writer.end();
}

// Read back and verify all batches round-trip correctly
try (ArrowStreamReader reader =
new ArrowStreamReader(
new ByteArrayReadableSeekableByteChannel(out.toByteArray()),
allocator,
CommonsCompressionFactory.INSTANCE)) {
int batchesRead = 0;
while (reader.loadNextBatch()) {
VectorSchemaRoot readRoot = reader.getVectorSchemaRoot();
assertEquals(rowsPerBatch, readRoot.getRowCount());
assertEquals(fieldCount, readRoot.getFieldVectors().size());

// Verify data values, null patterns, and all-null columns
for (int col = 0; col < fieldCount; col++) {
FieldVector vector = readRoot.getVector(col);
boolean allNull =
(batchesRead % 3 == 0) && (col % 3 == 1); // timestamps in every 3rd batch
if (allNull) {
// The key scenario: all-null columns must survive compression round-trip
assertEquals(
rowsPerBatch,
vector.getNullCount(),
"All-null column col=" + col + " batch=" + batchesRead);
}
for (int row = 0; row < rowsPerBatch; row++) {
switch (col % 3) {
case 0:
{
IntVector iv = (IntVector) vector;
if (allNull || row % 7 == 0) {
assertTrue(
iv.isNull(row),
"Expected null at col=" + col + " row=" + row + " batch=" + batchesRead);
} else {
assertEquals(
batchesRead * rowsPerBatch + row,
iv.get(row),
"Value mismatch at col=" + col + " row=" + row + " batch=" + batchesRead);
}
break;
}
case 1:
{
TimeStampMilliVector tv = (TimeStampMilliVector) vector;
if (allNull || row % 5 == 0) {
assertTrue(
tv.isNull(row),
"Expected null at col=" + col + " row=" + row + " batch=" + batchesRead);
} else {
assertEquals(
1_700_000_000_000L + (long) batchesRead * rowsPerBatch + row,
tv.get(row),
"Value mismatch at col=" + col + " row=" + row + " batch=" + batchesRead);
}
break;
}
case 2:
{
VarCharVector sv = (VarCharVector) vector;
if (allNull || row % 9 == 0) {
assertTrue(
sv.isNull(row),
"Expected null at col=" + col + " row=" + row + " batch=" + batchesRead);
} else {
assertArrayEquals(
("val_" + batchesRead + "_" + row).getBytes(StandardCharsets.UTF_8),
sv.get(row),
"Value mismatch at col=" + col + " row=" + row + " batch=" + batchesRead);
}
break;
}
default:
break;
}
}
}
batchesRead++;
}
assertEquals(batchCount, batchesRead);
}
}

/**
* Test that an all-null fixed-width vector compresses and decompresses correctly. The data buffer
* for such a vector contains all zeros but has a non-zero writerIndex (valueCount * typeWidth).
* The compressed buffer's uncompressed-length prefix must reflect this non-zero size.
*/
@Test
void testAllNullFixedWidthVectorZstdRoundTrip() throws Exception {
final int rowCount = 3469; // same count as the reported issue
final CompressionCodec codec = new ZstdCompressionCodec();

try (TimeStampMilliVector origVec =
new TimeStampMilliVector(
"ts",
FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)),
allocator)) {
origVec.allocateNew(rowCount);
// Set all values to null
for (int i = 0; i < rowCount; i++) {
origVec.setNull(i);
}
origVec.setValueCount(rowCount);

assertEquals(rowCount, origVec.getNullCount());

// Compress and decompress each buffer
List<ArrowBuf> origBuffers = origVec.getFieldBuffers();
assertEquals(2, origBuffers.size());

// The data buffer (index 1) should have non-zero writerIndex even though all values are null
ArrowBuf dataBuffer = origBuffers.get(1);
long expectedDataSize = (long) rowCount * 8; // TimestampMilli = 8 bytes per value
assertEquals(expectedDataSize, dataBuffer.writerIndex());

// Retain buffers before compressing since compress() closes the input buffer.
// This mirrors what VectorUnloader.appendNodes() does.
for (ArrowBuf buf : origBuffers) {
buf.getReferenceManager().retain();
}
List<ArrowBuf> compressedBuffers = compressBuffers(codec, origBuffers);
List<ArrowBuf> decompressedBuffers = deCompressBuffers(codec, compressedBuffers);

assertEquals(2, decompressedBuffers.size());

// The decompressed data buffer should have the same writerIndex as the original
assertEquals(expectedDataSize, decompressedBuffers.get(1).writerIndex());

// Load into a new vector and verify
try (TimeStampMilliVector newVec =
new TimeStampMilliVector(
"ts_new",
FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)),
allocator)) {
newVec.loadFieldBuffers(new ArrowFieldNode(rowCount, rowCount), decompressedBuffers);
assertEquals(rowCount, newVec.getValueCount());
for (int i = 0; i < rowCount; i++) {
assertTrue(newVec.isNull(i));
}
}
AutoCloseables.close(decompressedBuffers);
}
}

void withRoot(
CompressionUtil.CodecType codec,
BiConsumer<CompressionCodec.Factory, VectorSchemaRoot> testBody) {
Expand Down
10 changes: 8 additions & 2 deletions vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java
View file
Open in desktop

@lidavidm lidavidm Jun 29, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While I think this change is fine, nothing should be modifying the buffer during this method. I don't believe any vectors/etc. are safe to use concurrently, so the comment is misleading about the guarantees we can provide.

@LuciferYang LuciferYang Jun 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point — the original framing implied a concurrency guarantee we don't actually provide, and there's no in-tree doCompress implementation that mutates the source buffer's writerIndex. Tightened the comment to state only the engineering invariant (three in-method consumers must observe the same value) and dropped the "shared reference" / doCompress speculation:

// GH-1116: capture writerIndex() once so the empty-buffer check, size
// comparison, and uncompressed-length prefix all see the same value.
long uncompressedLength = uncompressedBuffer.writerIndex();

Pushed in d7b6a8f.

Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,14 @@ public abstract class AbstractCompressionCodec implements CompressionCodec {

@Override
public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) {
if (uncompressedBuffer.writerIndex() == 0L) {
// Capture the uncompressed length once upfront to avoid any inconsistency from
// re-reading writerIndex() at different points. Since the uncompressedBuffer may be
// a shared reference to a vector's internal buffer, reading writerIndex() only once
// ensures the same value is used for the empty-buffer check, compression, size
// comparison, and the 8-byte uncompressed-length prefix.
long uncompressedLength = uncompressedBuffer.writerIndex();

if (uncompressedLength == 0L) {
// shortcut for empty buffer
ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
compressedBuffer.setLong(0, 0);
Expand All @@ -41,7 +48,6 @@ public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer)
ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
long compressedLength =
compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
long uncompressedLength = uncompressedBuffer.writerIndex();

if (compressedLength > uncompressedLength) {
// compressed buffer is larger, send the raw buffer
Expand Down
Loading