Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,8 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) {
return targetVector; // nothing to append, return
}

validateDelta(deltaVector);

int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount();

int targetDataSize =
Expand All @@ -134,7 +136,7 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) {
.getOffsetBuffer()
.getInt((long) deltaVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH)
- deltaDataStart;
int newValueCapacity = targetDataSize + deltaDataSize;
int newValueCapacity = Math.addExact(targetDataSize, deltaDataSize);

// make sure there is enough capacity
while (targetVector.getValueCapacity() < newValueCount) {
Expand All @@ -153,10 +155,9 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) {
targetVector.getValidityBuffer());

// append data buffer
MemoryUtil.copyMemory(
deltaVector.getDataBuffer().memoryAddress() + deltaDataStart,
targetVector.getDataBuffer().memoryAddress() + targetDataSize,
deltaDataSize);
targetVector

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.

Nice improvement, this turns a raw out-of-bounds native copy into a checked one. One concern is ordering. By the time this setBytes runs, the target's data buffer has been reallocated for deltaDataSize and concatBits has already written the delta's validity bits into the target.

If the check throws, the target is left with extra validity bits while valueCount, offsets and lastSet are unchanged.

Could we validate the delta (last offset against its data buffer capacity) before any realloc or write, so a rejected delta leaves the target untouched?

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.

Agreed. Added a validateDelta pass at the top of both visits that runs before any realloc or concatBits, so a rejected delta now leaves the target as it was. The test checks that too (see below).

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.

Also, only the data copy is checked now. The validity copy via concatBits, the offset buffer MemoryUtil.copyMemory, and the offset-rebasing loop still read the delta without bounds checks. A delta whose offset buffer is shorter than (valueCount + 1) * 4 bytes, or whose validity buffer is shorter than valueCount bits, would still read out-of-bounds native memory.

Is it worth validating those buffer sizes up front too?

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.

Yes, worth it. validateDelta now checks the validity buffer against getValidityBufferSizeFromCount(valueCount) and the offset buffer against (valueCount + 1) * OFFSET_WIDTH before anything reads them, so concatBits, the offset copyMemory and the rebasing loop only run on buffers that cover the value count.

.getDataBuffer()
.setBytes(targetDataSize, deltaVector.getDataBuffer(), deltaDataStart, deltaDataSize);

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.

deltaDataSize is derived from the delta's offsets without validation. A last offset smaller than the first gives a negative size, and a very large one can overflow the int when added to targetDataSize.
Intermediate offsets are not checked against the data size either, so later reads of the appender vector could go out of bounds even when the final offset passes.

It would be good to guard 0 <= start <= end <= dataBuffer.capacity() and use Math.addExact or similar.

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.

Done. It walks the offsets once and requires 0 <= first, non-decreasing offsets throughout, and last <= data buffer capacity, so intermediate offsets are covered too. targetDataSize + deltaDataSize uses Math.addExact now.


// copy offset buffer
MemoryUtil.copyMemory(
Expand Down Expand Up @@ -194,6 +195,8 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) {
return targetVector; // nothing to append, return
}

validateDelta(deltaVector);

int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount();

long targetDataSize =
Expand All @@ -210,7 +213,7 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) {
.getLong(
(long) deltaVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH)
- deltaDataStart;
long newValueCapacity = targetDataSize + deltaDataSize;
long newValueCapacity = Math.addExact(targetDataSize, deltaDataSize);

// make sure there is enough capacity
while (targetVector.getValueCapacity() < newValueCount) {
Expand All @@ -229,10 +232,9 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) {
targetVector.getValidityBuffer());

// append data buffer
MemoryUtil.copyMemory(
deltaVector.getDataBuffer().memoryAddress() + deltaDataStart,
targetVector.getDataBuffer().memoryAddress() + targetDataSize,
deltaDataSize);
targetVector

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.

The same comments applies here 😄

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.

Same changes on the large path, with long offsets and Math.addExact(long, long).

.getDataBuffer()
.setBytes(targetDataSize, deltaVector.getDataBuffer(), deltaDataStart, deltaDataSize);

// copy offset buffer
MemoryUtil.copyMemory(
Expand Down Expand Up @@ -261,6 +263,70 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) {
return targetVector;
}

/**
* Checks that the delta's validity and offset buffers cover its value count, and that its offsets
* are non-negative, non-decreasing and within its data buffer. This runs before the target is
* touched, so a malformed delta leaves the target unchanged.
*/
private static void validateDelta(BaseVariableWidthVector deltaVector) {
int valueCount = deltaVector.getValueCount();
checkCapacity(
deltaVector.getValidityBuffer(),
BitVectorHelper.getValidityBufferSizeFromCount(valueCount),
"validity");
ArrowBuf offsets = deltaVector.getOffsetBuffer();
checkCapacity(offsets, (valueCount + 1L) * BaseVariableWidthVector.OFFSET_WIDTH, "offset");
int prev = offsets.getInt(0);
if (prev < 0) {
throw new IndexOutOfBoundsException("Delta vector has negative start offset " + prev);
}
for (int i = 1; i <= valueCount; i++) {
int cur = offsets.getInt((long) i * BaseVariableWidthVector.OFFSET_WIDTH);
if (cur < prev) {
throw new IndexOutOfBoundsException(
"Delta vector offsets decrease at index " + i + ": " + prev + " > " + cur);
}
prev = cur;
}
checkCapacity(deltaVector.getDataBuffer(), prev, "data");
}

/** See {@link #validateDelta(BaseVariableWidthVector)}. */
private static void validateDelta(BaseLargeVariableWidthVector deltaVector) {
int valueCount = deltaVector.getValueCount();
checkCapacity(
deltaVector.getValidityBuffer(),
BitVectorHelper.getValidityBufferSizeFromCount(valueCount),
"validity");
ArrowBuf offsets = deltaVector.getOffsetBuffer();
checkCapacity(offsets, (valueCount + 1L) * BaseLargeVariableWidthVector.OFFSET_WIDTH, "offset");
long prev = offsets.getLong(0);
if (prev < 0) {
throw new IndexOutOfBoundsException("Delta vector has negative start offset " + prev);
}
for (int i = 1; i <= valueCount; i++) {
long cur = offsets.getLong((long) i * BaseLargeVariableWidthVector.OFFSET_WIDTH);
if (cur < prev) {
throw new IndexOutOfBoundsException(
"Delta vector offsets decrease at index " + i + ": " + prev + " > " + cur);
}
prev = cur;
}
checkCapacity(deltaVector.getDataBuffer(), prev, "data");
}

private static void checkCapacity(ArrowBuf buffer, long required, String name) {
if (buffer.capacity() < required) {
throw new IndexOutOfBoundsException(
"Delta vector "
+ name
+ " buffer has capacity "
+ buffer.capacity()
+ ", needs "
+ required);
}
}

@Override
public ValueVector visit(BaseVariableWidthViewVector deltaVector, Void value) {
Preconditions.checkArgument(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
import java.util.Map;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.arrow.flatbuf.FieldNode;
import org.apache.arrow.flatbuf.Message;
import org.apache.arrow.flatbuf.RecordBatch;
Expand All @@ -53,12 +54,15 @@
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.util.Collections2;
import org.apache.arrow.vector.BitVectorHelper;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.LargeVarCharVector;
import org.apache.arrow.vector.NullVector;
import org.apache.arrow.vector.TestUtils;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VariableWidthFieldVector;
import org.apache.arrow.vector.VectorLoader;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.VectorUnloader;
Expand Down Expand Up @@ -89,6 +93,10 @@
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;
import org.junit.jupiter.params.provider.ValueSource;

public class TestArrowReaderWriter {

Expand Down Expand Up @@ -649,6 +657,129 @@ public void testDeltaDictionary() throws Exception {
AutoCloseables.close(closeableList);
}

@ParameterizedTest
@ValueSource(booleans = {false, true})
public void testDeltaDictionaryWithNonZeroStartOffset(boolean large) throws Exception {
// a valid delta whose offsets start at 3, so its data begins after a 3-byte prefix
byte[] stream = writeDeltaDictionaryStream(large, new long[] {3, 6, 9}, "xxxbazqux");
try (ArrowStreamReader reader =
new ArrowStreamReader(new ByteArrayReadableSeekableByteChannel(stream), allocator)) {
assertFalse(reader.loadNextBatch());
FieldVector dictionaryVector = reader.getDictionaryVectors().get(1L).getVector();
assertEquals(4, dictionaryVector.getValueCount());
assertEquals("foo", dictionaryVector.getObject(0).toString());
assertEquals("bar", dictionaryVector.getObject(1).toString());
assertEquals("baz", dictionaryVector.getObject(2).toString());
assertEquals("qux", dictionaryVector.getObject(3).toString());
}
}

@ParameterizedTest
@MethodSource("malformedDeltaOffsets")
public void testDeltaDictionaryWithMalformedOffsets(boolean large, long[] deltaOffsets)
throws Exception {
byte[] stream = writeDeltaDictionaryStream(large, deltaOffsets, "bazzquxx");
try (ArrowStreamReader reader =
new ArrowStreamReader(new ByteArrayReadableSeekableByteChannel(stream), allocator)) {
assertThrows(IndexOutOfBoundsException.class, reader::loadNextBatch);
// the rejected delta must leave the original dictionary untouched
FieldVector dictionaryVector = reader.getDictionaryVectors().get(1L).getVector();
assertEquals(2, dictionaryVector.getValueCount());
assertEquals("foo", dictionaryVector.getObject(0).toString());
assertEquals("bar", dictionaryVector.getObject(1).toString());
assertEquals(0, BitVectorHelper.get(dictionaryVector.getValidityBuffer(), 2));
assertEquals(0, BitVectorHelper.get(dictionaryVector.getValidityBuffer(), 3));
}
}

static Stream<Arguments> malformedDeltaOffsets() {
List<long[]> offsets =
Arrays.asList(
new long[] {0, 4, 4096}, // last offset past the 8-byte data buffer
new long[] {0, 6, 4}, // decreasing intermediate offset
new long[] {4, 0, 0}, // last offset before the first
new long[] {-4, 0, 4}); // negative start offset
return Stream.of(false, true)
.flatMap(large -> offsets.stream().map(o -> Arguments.of(large, o)));
}

/**
* Writes a stream holding a {"foo", "bar"} dictionary followed by a delta dictionary batch built
* directly from the given offsets and data bytes, bypassing any vector-level validation.
*/
private byte[] writeDeltaDictionaryStream(boolean large, long[] deltaOffsets, String deltaData)
throws IOException {
List<AutoCloseable> closeables = new ArrayList<>();
try {
FieldVector dictVector =
large
? newVector(LargeVarCharVector.class, "dict", MinorType.LARGEVARCHAR, allocator)
: newVarCharVector("dict", allocator);
closeables.add(dictVector);
dictVector.allocateNew();
((VariableWidthFieldVector) dictVector).setSafe(0, "foo".getBytes(StandardCharsets.UTF_8));
((VariableWidthFieldVector) dictVector).setSafe(1, "bar".getBytes(StandardCharsets.UTF_8));
dictVector.setValueCount(2);
Dictionary dictionary = new Dictionary(dictVector, new DictionaryEncoding(1L, false, null));
DictionaryProvider.MapDictionaryProvider provider =
new DictionaryProvider.MapDictionaryProvider();
provider.put(dictionary);
Field encodedField =
new Field(
"f",
new FieldType(true, new ArrowType.Int(32, true), dictionary.getEncoding()),
null);
Schema schema =
new Schema(
Collections.singletonList(
DictionaryUtility.toMessageFormat(encodedField, provider, new HashSet<>())));

ByteArrayOutputStream outStream = new ByteArrayOutputStream();
WriteChannel out = new WriteChannel(newChannel(outStream));
MessageSerializer.serialize(out, schema);
serializeDictionaryBatch(out, dictionary, false, closeables);

int offsetWidth = large ? 8 : 4;
ArrowBuf validity = allocator.buffer(0);
closeables.add(validity);
ArrowBuf offsets = allocator.buffer((long) deltaOffsets.length * offsetWidth);
closeables.add(offsets);
for (int i = 0; i < deltaOffsets.length; i++) {
if (large) {
offsets.setLong((long) i * offsetWidth, deltaOffsets[i]);
} else {
offsets.setInt((long) i * offsetWidth, (int) deltaOffsets[i]);
}
}
offsets.writerIndex((long) deltaOffsets.length * offsetWidth);
byte[] dataBytes = deltaData.getBytes(StandardCharsets.UTF_8);
ArrowBuf data = allocator.buffer(dataBytes.length);
closeables.add(data);
data.setBytes(0, dataBytes);
data.writerIndex(dataBytes.length);

int deltaValueCount = deltaOffsets.length - 1;
ArrowRecordBatch deltaBatch =
new ArrowRecordBatch(
deltaValueCount,
Collections.singletonList(new ArrowFieldNode(deltaValueCount, 0)),
Arrays.asList(validity, offsets, data));
ArrowDictionaryBatch deltaDictionaryBatch = new ArrowDictionaryBatch(1L, deltaBatch, true);
closeables.add(deltaDictionaryBatch);
MessageSerializer.serialize(out, deltaDictionaryBatch);

// write eos
out.writeIntLittleEndian(0);
return outStream.toByteArray();
} finally {
try {
AutoCloseables.close(closeables);
} catch (Exception e) {
throw new IOException(e);
}
}
}

// Tests that the ArrowStreamWriter re-emits dictionaries when they change
@Test
public void testWriteReadStreamWithDictionaryReplacement() throws Exception {
Expand Down
Loading