From 2077c8a5abf15c4da0a6c2fb8baa283f614e4cd6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Wed, 14 Oct 2026 23:58:40 +0100 Subject: [PATCH] BDX-1590: Preserve dictionary deltas in Arrow 19 --- .../org/apache/arrow/flight/FlightStream.java | 23 +- .../TestFlightStreamDictionaryDelta.java | 606 ++++++++++++++++++ .../vector/dictionary/DictionaryEncoder.java | 8 +- .../dictionary/StructSubfieldEncoder.java | 3 +- .../arrow/vector/util/VectorAppender.java | 110 +++- .../arrow/vector/TestDictionaryVector.java | 79 ++- .../arrow/vector/util/TestVectorAppender.java | 254 ++++++++ 7 files changed, 1035 insertions(+), 48 deletions(-) create mode 100644 flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightStreamDictionaryDelta.java diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java index 15cfd6ba85..82cfdd2311 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java @@ -46,6 +46,7 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.DictionaryUtility; +import org.apache.arrow.vector.util.VectorBatchAppender; import org.apache.arrow.vector.validate.MetadataV4UnionChecker; /** An adaptor between protobuf streams and flight data streams. */ @@ -291,13 +292,21 @@ public boolean next() { } final FieldVector vector = dictionary.getVector(); - final VectorSchemaRoot dictionaryRoot = - new VectorSchemaRoot( - Collections.singletonList(vector.getField()), - Collections.singletonList(vector), - 0); - final VectorLoader dictionaryLoader = new VectorLoader(dictionaryRoot); - dictionaryLoader.load(arb.getDictionary()); + if (arb.isDelta()) { + try (FieldVector deltaVector = vector.getField().createVector(allocator)) { + final VectorSchemaRoot deltaRoot = VectorSchemaRoot.of(deltaVector); + new VectorLoader(deltaRoot).load(arb.getDictionary()); + VectorBatchAppender.batchAppend(vector, deltaVector); + } + } else { + final VectorSchemaRoot dictionaryRoot = + new VectorSchemaRoot( + Collections.singletonList(vector.getField()), + Collections.singletonList(vector), + 0); + final VectorLoader dictionaryLoader = new VectorLoader(dictionaryRoot); + dictionaryLoader.load(arb.getDictionary()); + } } return next(); } else { diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightStreamDictionaryDelta.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightStreamDictionaryDelta.java new file mode 100644 index 0000000000..bbc7210054 --- /dev/null +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightStreamDictionaryDelta.java @@ -0,0 +1,606 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.flight; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.grpc.MethodDescriptor; +import io.grpc.stub.StreamObserver; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +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.ValueVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.complex.LargeListVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.dictionary.DictionaryEncoder; +import org.apache.arrow.vector.dictionary.DictionaryProvider; +import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; +import org.apache.arrow.vector.ipc.message.ArrowFieldNode; +import org.apache.arrow.vector.ipc.message.IpcOption; +import org.apache.arrow.vector.testing.ValueVectorDataPopulator; +import org.apache.arrow.vector.types.Types; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.DictionaryEncoding; +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.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.ValueSource; + +/** Tests dictionary updates received as framed Flight messages. */ +public class TestFlightStreamDictionaryDelta { + private static final DictionaryEncoding ENCODING = + new DictionaryEncoding(7, false, new ArrowType.Int(32, true)); + private static final Schema WIRE_SCHEMA = schema(ArrowType.Utf8.INSTANCE, ENCODING); + private static final Schema MEMORY_SCHEMA = schema(new ArrowType.Int(32, true), ENCODING); + + @Test + public void nullDictionaryDeltaAndReplacement() throws Exception { + try (TestStream incoming = new TestStream(schema(ArrowType.Null.INSTANCE, ENCODING), false)) { + final DictionaryProvider provider = incoming.stream.getDictionaryProvider(); + assertEquals(MEMORY_SCHEMA, incoming.stream.getSchema()); + incoming.nullDictionary(false, 1); + incoming.batch(0); + assertTrue(incoming.stream.next()); + final List captured = incoming.decode(); + assertEquals(Collections.singletonList(null), captured); + assertEquals(Collections.singletonList(null), values(provider.lookup(7).getVector())); + + incoming.nullDictionary(true, 0); + incoming.nullDictionary(true, 1); + incoming.batch(1); + assertTrue(incoming.stream.next()); + assertEquals(Collections.singletonList(null), incoming.decode()); + assertEquals(Arrays.asList(null, null), values(provider.lookup(7).getVector())); + assertSame(provider, incoming.stream.getDictionaryProvider()); + + incoming.nullDictionary(false, 1); + incoming.batch(0); + assertTrue(incoming.stream.next()); + assertEquals(Collections.singletonList(null), incoming.decode()); + assertEquals(Collections.singletonList(null), values(provider.lookup(7).getVector())); + incoming.batch(1); + assertTrue(incoming.stream.next()); + assertThrows(IllegalArgumentException.class, incoming::decode); + assertEquals(Collections.singletonList(null), captured); + incoming.observer.onCompleted(); + assertFalse(incoming.stream.next()); + } + } + + @Test + public void dictionaryDeltaAndReplacement() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + final DictionaryProvider provider = incoming.stream.getDictionaryProvider(); + assertEquals(MEMORY_SCHEMA, incoming.stream.getSchema()); + assertEquals(Collections.emptyList(), values(provider.lookup(7).getVector())); + incoming.dictionary(7, false, "alpha", "beta"); + incoming.batch(0, 1); + assertTrue(incoming.stream.next()); + final List captured = incoming.decode(); + assertEquals(Arrays.asList("alpha", "beta"), captured); + + incoming.dictionary(7, true); + incoming.dictionary(7, true, "gamma"); + incoming.batch(2, 0, null); + assertTrue(incoming.stream.next()); + assertSame(provider, incoming.stream.getDictionaryProvider()); + assertEquals(Arrays.asList("alpha", "beta", "gamma"), values(provider.lookup(7).getVector())); + assertEquals(Arrays.asList("gamma", "alpha", null), incoming.decode()); + + incoming.dictionary(7, false, "delta", "epsilon"); + incoming.batch(1, 0); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("epsilon", "delta"), incoming.decode()); + assertEquals(Arrays.asList("delta", "epsilon"), values(provider.lookup(7).getVector())); + assertEquals(Arrays.asList("alpha", "beta"), captured); + incoming.observer.onCompleted(); + assertFalse(incoming.stream.next()); + } + } + + @Test + public void multipleDeltasAppendToEmptyDictionary() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, false); + incoming.batch((Integer) null); + assertTrue(incoming.stream.next()); + assertEquals(Collections.singletonList(null), incoming.decode()); + incoming.dictionary(7, true, "alpha", null); + incoming.dictionary(7, true); + incoming.dictionary(7, true, "beta"); + incoming.batch(2, 0, 1, null); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("beta", "alpha", null, null), incoming.decode()); + assertEquals( + Arrays.asList("alpha", null, "beta"), + values(incoming.stream.getDictionaryProvider().lookup(7).getVector())); + incoming.observer.onCompleted(); + assertFalse(incoming.stream.next()); + } + } + + @Test + public void deltaIntoEmptySchemaDictionaryUsesSharedAppender() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, true, "alpha"); + incoming.dictionary(7, true, "beta"); + incoming.batch(0, 1); + incoming.observer.onCompleted(); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("alpha", "beta"), incoming.decode()); + assertFalse(incoming.stream.next()); + } + } + + @ParameterizedTest + @ValueSource(longs = {-1, 2, 4294967296L, Long.MAX_VALUE, Long.MIN_VALUE}) + public void invalidLongDictionaryIndicesReleaseBuffers(long index) throws Exception { + DictionaryEncoding encoding = new DictionaryEncoding(7, false, new ArrowType.Int(64, true)); + try (TestStream incoming = new TestStream(schema(ArrowType.Utf8.INSTANCE, encoding), false)) { + incoming.dictionary(7, false, "alpha", "beta"); + incoming.longBatch(1, 0); + assertTrue(incoming.stream.next()); + List captured = incoming.decode(); + assertEquals(Arrays.asList("beta", "alpha"), captured); + incoming.longBatch(index); + assertTrue(incoming.stream.next()); + assertEquals( + "Provided dictionary does not contain value for index " + index, + assertThrows(IllegalArgumentException.class, incoming::decode).getMessage()); + assertEquals(Arrays.asList("beta", "alpha"), captured); + assertEquals( + Arrays.asList("alpha", "beta"), + values(incoming.stream.getDictionaryProvider().lookup(7).getVector())); + } + } + + @Test + public void malformedDeltaPreservesDictionaryAndReleasesBuffers() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, false, "alpha"); + incoming.batch(0); + assertTrue(incoming.stream.next()); + assertEquals(Collections.singletonList("alpha"), incoming.decode()); + try (IntVector vector = new IntVector("wrong-type", incoming.allocator); + VectorSchemaRoot root = VectorSchemaRoot.of(vector)) { + vector.allocateNew(1); + vector.setSafe(0, 42); + root.setRowCount(1); + try (ArrowDictionaryBatch batch = + new ArrowDictionaryBatch(7, new VectorUnloader(root).getRecordBatch(), true)) { + incoming.push(new ArrowMessage(batch, IpcOption.DEFAULT)); + } + } + incoming.batch(0); + assertThrows(RuntimeException.class, incoming.stream::next); + assertEquals( + Collections.singletonList("alpha"), + values(incoming.stream.getDictionaryProvider().lookup(7).getVector())); + } + } + + @ParameterizedTest + @EnumSource( + value = Types.MinorType.class, + names = {"VARCHAR", "LARGEVARCHAR", "LIST", "LARGELIST"}) + public void slicedDictionaryDeltas(Types.MinorType type) throws Exception { + try (RootAllocator allocator = new RootAllocator(1 << 20)) { + try (FieldVector source = + type.getNewVector( + new Field( + "dictionary", FieldType.nullable(type.getType()), Collections.emptyList()), + allocator, + null); + ArrowBuf validity = allocator.buffer(1)) { + switch (type) { + case VARCHAR: + ValueVectorDataPopulator.setVector( + (VarCharVector) source, "skip", "beta", null, "gamma", "tail"); + break; + case LARGEVARCHAR: + ValueVectorDataPopulator.setVector( + (LargeVarCharVector) source, "skip", "beta", null, "gamma", "tail"); + break; + case LIST: + ValueVectorDataPopulator.setVector( + (ListVector) source, + List.of(-1, -2), + List.of(10, 11), + null, + List.of(), + List.of(12, 13), + List.of(99)); + break; + case LARGELIST: + ValueVectorDataPopulator.setVector( + (LargeListVector) source, + List.of(-1, -2), + List.of(10, 11), + null, + List.of(), + List.of(12, 13), + List.of(99)); + break; + default: + throw new AssertionError(type); + } + boolean list = type == Types.MinorType.LIST || type == Types.MinorType.LARGELIST; + int count = list ? 4 : 3; + int width = + type == Types.MinorType.LARGEVARCHAR || type == Types.MinorType.LARGELIST ? 8 : 4; + validity.setByte(0, list ? 13 : 5); + Schema wireSchema = + new Schema( + Collections.singletonList( + new Field( + "word", + new FieldType(true, type.getType(), ENCODING, null), + source.getField().getChildren()))); + try (FieldVector delta = source.getField().createVector(allocator); + TestStream incoming = new TestStream(wireSchema, false)) { + for (int i = 0; i < source.getChildrenFromFields().size(); i++) { + FieldVector child = source.getChildrenFromFields().get(i); + delta + .getChildrenFromFields() + .get(i) + .loadFieldBuffers( + new ArrowFieldNode(child.getValueCount(), child.getNullCount()), + child.getFieldBuffers()); + } + List buffers = new ArrayList<>(source.getFieldBuffers()); + buffers.set(0, validity); + buffers.set(1, source.getOffsetBuffer().slice(width, (long) (count + 1) * width)); + delta.loadFieldBuffers(new ArrowFieldNode(count, 1), buffers); + try (FieldVector empty = source.getField().createVector(allocator)) { + empty.allocateNew(); + incoming.dictionary(empty, false); + } + incoming.dictionary(delta, true); + incoming.dictionary(delta, true); + incoming.batch(count * 2 - 1, 0, count + 1, null); + assertTrue(incoming.stream.next()); + assertEquals( + list + ? Arrays.asList("[12,13]", "[10,11]", null, null) + : Arrays.asList("gamma", "beta", null, null), + incoming.decode()); + assertEquals( + list + ? Arrays.asList( + "[10,11]", null, "[]", "[12,13]", "[10,11]", null, "[]", "[12,13]") + : Arrays.asList("beta", null, "gamma", "beta", null, "gamma"), + values(incoming.stream.getDictionaryProvider().lookup(7).getVector())); + incoming.observer.onCompleted(); + assertFalse(incoming.stream.next()); + } + } + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @Test + public void undefinedDictionaryFailsAndClosesMessages() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, false, "alpha", "beta"); + incoming.batch(1, 0); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("beta", "alpha"), incoming.decode()); + incoming.dictionary(8, true, "gamma"); + incoming.batch(0); + assertEquals( + "Dictionary not defined in schema: ID 8", + assertThrows(IllegalArgumentException.class, incoming.stream::next).getMessage()); + } + } + + @Test + public void invalidDictionaryIndexFailsAfterDelta() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, false, "alpha", "beta"); + incoming.dictionary(7, true, "gamma"); + incoming.batch(2, 0, null); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("gamma", "alpha", null), incoming.decode()); + incoming.batch(3); + assertTrue(incoming.stream.next()); + assertThrows(IllegalArgumentException.class, incoming::decode); + } + } + + @Test + public void cancellationDrainsDeltaAndReleasesBuffers() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, true)) { + incoming.dictionary(7, false, "alpha", "beta"); + incoming.batch(0, 1); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("alpha", "beta"), incoming.decode()); + incoming.dictionary(7, true, "gamma"); + incoming.batch(2, null); + incoming.stream.cancel("cancel dictionary stream", null); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("gamma", null), incoming.decode()); + final FlightRuntimeException failure = + assertThrows(FlightRuntimeException.class, incoming.stream::next); + assertEquals(FlightStatusCode.CANCELLED, failure.status().code()); + assertEquals("cancel dictionary stream", failure.status().description()); + assertEquals(1, incoming.cancellations); + } + } + + @Test + public void lateErrorPreservesCapturedValuesAndReleasesBuffers() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, false, "alpha", "beta"); + incoming.batch(0, 1); + assertTrue(incoming.stream.next()); + final List captured = incoming.decode(); + assertEquals(Arrays.asList("alpha", "beta"), captured); + incoming.dictionary(7, true, "gamma"); + incoming.batch(2); + incoming.observer.onError( + CallStatus.INTERNAL.withDescription("late failure").toRuntimeException()); + assertTrue(incoming.stream.next()); + assertEquals(Collections.singletonList("gamma"), incoming.decode()); + final FlightRuntimeException failure = + assertThrows(FlightRuntimeException.class, incoming.stream::next); + assertEquals(FlightStatusCode.INTERNAL, failure.status().code()); + assertEquals("late failure", failure.status().description()); + assertEquals(Arrays.asList("alpha", "beta"), captured); + } + } + + @Test + public void closeReleasesQueuedDictionaryAndRecordMessages() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, false, "alpha", "beta"); + incoming.dictionary(7, true, "gamma"); + incoming.batch(2, 0, null); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("gamma", "alpha", null), incoming.decode()); + incoming.dictionary(7, true, "delta"); + incoming.batch(3); + } + } + + @Test + public void dictionaryOwnershipSurvivesStreamClose() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, false, "alpha", "beta"); + incoming.dictionary(7, true, "gamma"); + incoming.batch(2, 0); + incoming.observer.onCompleted(); + assertTrue(incoming.stream.next()); + assertEquals(Arrays.asList("gamma", "alpha"), incoming.decode()); + assertFalse(incoming.stream.next()); + final DictionaryProvider borrowed = incoming.stream.getDictionaryProvider(); + try (DictionaryProvider.MapDictionaryProvider owned = + (DictionaryProvider.MapDictionaryProvider) incoming.stream.takeDictionaryOwnership()) { + assertSame(borrowed, owned); + assertThrows(IllegalStateException.class, incoming.stream::takeDictionaryOwnership); + assertThrows(IllegalStateException.class, incoming.stream::getDictionaryProvider); + incoming.stream.close(); + assertEquals(Arrays.asList("alpha", "beta", "gamma"), values(owned.lookup(7).getVector())); + } + assertEquals(0, incoming.allocator.getAllocatedMemory()); + } + } + + @Test + public void zeroRowBatchAndSchemaOnlyStream() throws Exception { + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.dictionary(7, false, "alpha", "beta"); + incoming.dictionary(7, true); + incoming.batch(); + incoming.observer.onCompleted(); + assertTrue(incoming.stream.next()); + assertEquals(MEMORY_SCHEMA, incoming.stream.getSchema()); + assertEquals(Collections.emptyList(), incoming.decode()); + assertEquals( + Arrays.asList("alpha", "beta"), + values(incoming.stream.getDictionaryProvider().lookup(7).getVector())); + assertFalse(incoming.stream.next()); + } + try (TestStream incoming = new TestStream(WIRE_SCHEMA, false)) { + incoming.observer.onCompleted(); + assertEquals(MEMORY_SCHEMA, incoming.stream.getSchema()); + assertFalse(incoming.stream.next()); + assertEquals( + Collections.emptyList(), + values(incoming.stream.getDictionaryProvider().lookup(7).getVector())); + } + } + + @Test + public void nonDictionaryStreamKeepsSchemaAndValues() throws Exception { + final Schema plainSchema = schema(new ArrowType.Int(32, true), null); + try (TestStream incoming = new TestStream(plainSchema, false)) { + incoming.batch(2, 0, null); + incoming.observer.onCompleted(); + assertTrue(incoming.stream.next()); + assertEquals(plainSchema, incoming.stream.getSchema()); + assertEquals(Arrays.asList("2", "0", null), values(incoming.stream.getRoot().getVector(0))); + assertEquals( + Collections.emptySet(), incoming.stream.getDictionaryProvider().getDictionaryIds()); + assertFalse(incoming.stream.next()); + } + } + + private static Schema schema(ArrowType type, DictionaryEncoding encoding) { + return new Schema( + Collections.singletonList( + new Field( + "word", + new FieldType(true, type, encoding, Map.of("field", "words")), + Collections.emptyList())), + Map.of("stream", "dictionary-delta")); + } + + private static List values(ValueVector vector) { + final List values = new ArrayList<>(); + for (int i = 0; i < vector.getValueCount(); i++) { + final Object value = vector.getObject(i); + values.add(value == null ? null : value.toString()); + } + return values; + } + + private static final class TestStream implements AutoCloseable { + private final RootAllocator allocator = new RootAllocator(1 << 20); + private final FlightStream stream; + private final StreamObserver observer; + private final List receivedBuffers = new ArrayList<>(); + private int cancellations; + + private TestStream(Schema schema, boolean cancellable) throws Exception { + stream = new FlightStream(allocator, 2, cancellable ? this::cancel : null, count -> {}); + observer = stream.asObserver(); + push(new ArrowMessage(null, schema, IpcOption.DEFAULT)); + } + + private void cancel(String message, Throwable exception) { + cancellations++; + observer.onError( + CallStatus.CANCELLED.withDescription(message).withCause(exception).toRuntimeException()); + } + + private void dictionary(long id, boolean delta, String... words) throws Exception { + try (VarCharVector vector = new VarCharVector("dictionary", allocator); + VectorSchemaRoot root = VectorSchemaRoot.of(vector)) { + vector.allocateNew(); + for (int i = 0; i < words.length; i++) { + if (words[i] == null) { + vector.setNull(i); + } else { + vector.setSafe(i, words[i].getBytes(StandardCharsets.UTF_8)); + } + } + root.setRowCount(words.length); + try (ArrowDictionaryBatch batch = + new ArrowDictionaryBatch(id, new VectorUnloader(root).getRecordBatch(), delta)) { + push(new ArrowMessage(batch, IpcOption.DEFAULT)); + } + } + } + + private void batch(Integer... indices) throws Exception { + try (IntVector vector = new IntVector("word", allocator); + VectorSchemaRoot root = VectorSchemaRoot.of(vector)) { + vector.allocateNew(indices.length); + for (int i = 0; i < indices.length; i++) { + if (indices[i] == null) { + vector.setNull(i); + } else { + vector.setSafe(i, indices[i]); + } + } + root.setRowCount(indices.length); + push( + new ArrowMessage( + new VectorUnloader(root).getRecordBatch(), null, false, IpcOption.DEFAULT)); + } + } + + private void nullDictionary(boolean delta, int count) throws Exception { + try (VectorSchemaRoot root = VectorSchemaRoot.of(new NullVector("dictionary"))) { + root.setRowCount(count); + try (ArrowDictionaryBatch batch = + new ArrowDictionaryBatch(7, new VectorUnloader(root).getRecordBatch(), delta)) { + push(new ArrowMessage(batch, IpcOption.DEFAULT)); + } + } + } + + private void longBatch(long... indices) throws Exception { + try (BigIntVector vector = new BigIntVector("word", allocator); + VectorSchemaRoot root = VectorSchemaRoot.of(vector)) { + vector.allocateNew(indices.length); + for (int i = 0; i < indices.length; i++) { + vector.setSafe(i, indices[i]); + } + root.setRowCount(indices.length); + push( + new ArrowMessage( + new VectorUnloader(root).getRecordBatch(), null, false, IpcOption.DEFAULT)); + } + } + + private void dictionary(FieldVector vector, boolean delta) throws Exception { + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + root.setRowCount(vector.getValueCount()); + try (ArrowDictionaryBatch batch = + new ArrowDictionaryBatch(7, new VectorUnloader(root).getRecordBatch(), delta)) { + push(new ArrowMessage(batch, IpcOption.DEFAULT)); + } + } + + private void push(ArrowMessage message) throws Exception { + final MethodDescriptor.Marshaller marshaller = + ArrowMessage.createMarshaller(allocator); + try (message; + InputStream wire = marshaller.stream(message)) { + final ArrowMessage received = marshaller.parse(wire); + received.getBufs().forEach(receivedBuffers::add); + observer.onNext(received); + } + } + + private List decode() { + try (ValueVector decoded = + DictionaryEncoder.decode( + stream.getRoot().getVector(0), stream.getDictionaryProvider().lookup(7))) { + return values(decoded); + } + } + + @Override + public void close() throws Exception { + try (allocator) { + stream.close(); + assertEquals( + receivedBuffers.stream() + .map(buffer -> buffer.capacity() == 0 ? 1 : 0) + .collect(Collectors.toList()), + receivedBuffers.stream() + .map(buffer -> buffer.getReferenceManager().getRefCount()) + .collect(Collectors.toList())); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/dictionary/DictionaryEncoder.java b/vector/src/main/java/org/apache/arrow/vector/dictionary/DictionaryEncoder.java index 4af1a8693f..a6d490c805 100644 --- a/vector/src/main/java/org/apache/arrow/vector/dictionary/DictionaryEncoder.java +++ b/vector/src/main/java/org/apache/arrow/vector/dictionary/DictionaryEncoder.java @@ -164,12 +164,12 @@ static void retrieveIndexVector( BaseIntVector indices, TransferPair transfer, int dictionaryCount, int start, int end) { for (int i = start; i < end; i++) { if (!indices.isNull(i)) { - int indexAsInt = (int) indices.getValueAsLong(i); - if (indexAsInt > dictionaryCount) { + long index = indices.getValueAsLong(i); + if (index < 0 || index >= dictionaryCount) { throw new IllegalArgumentException( - "Provided dictionary does not contain value for index " + indexAsInt); + "Provided dictionary does not contain value for index " + index); } - transfer.copyValueSafe(indexAsInt, i); + transfer.copyValueSafe((int) index, i); } } } diff --git a/vector/src/main/java/org/apache/arrow/vector/dictionary/StructSubfieldEncoder.java b/vector/src/main/java/org/apache/arrow/vector/dictionary/StructSubfieldEncoder.java index 8ff152fb1c..b183bc84c4 100644 --- a/vector/src/main/java/org/apache/arrow/vector/dictionary/StructSubfieldEncoder.java +++ b/vector/src/main/java/org/apache/arrow/vector/dictionary/StructSubfieldEncoder.java @@ -207,7 +207,8 @@ public static StructVector decode( TransferPair transfer = dictionary.getVector().makeTransferPair(decodedChildVector); BaseIntVector indices = (BaseIntVector) childVector; - DictionaryEncoder.retrieveIndexVector(indices, transfer, valueCount, 0, valueCount); + DictionaryEncoder.retrieveIndexVector( + indices, transfer, dictionary.getVector().getValueCount(), 0, valueCount); } } diff --git a/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java b/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java index e7c0d11cb9..d234b60d3e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java +++ b/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java @@ -122,13 +122,17 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) { int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); int targetDataSize = - targetVector - .getOffsetBuffer() - .getInt((long) targetVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); + targetVector.getOffsetBuffer().capacity() == 0 + ? 0 + : targetVector + .getOffsetBuffer() + .getInt((long) targetVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); + int deltaDataStart = deltaVector.getOffsetBuffer().getInt(0); int deltaDataSize = deltaVector - .getOffsetBuffer() - .getInt((long) deltaVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); + .getOffsetBuffer() + .getInt((long) deltaVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH) + - deltaDataStart; int newValueCapacity = targetDataSize + deltaDataSize; // make sure there is enough capacity @@ -149,7 +153,7 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) { // append data buffer MemoryUtil.copyMemory( - deltaVector.getDataBuffer().memoryAddress(), + deltaVector.getDataBuffer().memoryAddress() + deltaDataStart, targetVector.getDataBuffer().memoryAddress() + targetDataSize, deltaDataSize); @@ -172,7 +176,7 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) { .getOffsetBuffer() .setInt( (long) (targetVector.getValueCount() + 1 + i) * BaseVariableWidthVector.OFFSET_WIDTH, - oldOffset + targetDataSize); + oldOffset - deltaDataStart + targetDataSize); } ((BaseVariableWidthVector) targetVector).setLastSet(newValueCount - 1); targetVector.setValueCount(newValueCount); @@ -192,15 +196,20 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); long targetDataSize = - targetVector - .getOffsetBuffer() - .getLong( - (long) targetVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH); + targetVector.getOffsetBuffer().capacity() == 0 + ? 0 + : targetVector + .getOffsetBuffer() + .getLong( + (long) targetVector.getValueCount() + * BaseLargeVariableWidthVector.OFFSET_WIDTH); + long deltaDataStart = deltaVector.getOffsetBuffer().getLong(0); long deltaDataSize = deltaVector - .getOffsetBuffer() - .getLong( - (long) deltaVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH); + .getOffsetBuffer() + .getLong( + (long) deltaVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH) + - deltaDataStart; long newValueCapacity = targetDataSize + deltaDataSize; // make sure there is enough capacity @@ -221,7 +230,7 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { // append data buffer MemoryUtil.copyMemory( - deltaVector.getDataBuffer().memoryAddress(), + deltaVector.getDataBuffer().memoryAddress() + deltaDataStart, targetVector.getDataBuffer().memoryAddress() + targetDataSize, deltaDataSize); @@ -245,7 +254,7 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { .setLong( (long) (targetVector.getValueCount() + 1 + i) * BaseLargeVariableWidthVector.OFFSET_WIDTH, - oldOffset + targetDataSize); + oldOffset - deltaDataStart + targetDataSize); } ((BaseLargeVariableWidthVector) targetVector).setLastSet(newValueCount - 1); targetVector.setValueCount(newValueCount); @@ -328,19 +337,23 @@ public ValueVector visit(ListVector deltaVector, Void value) { int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); int targetListSize = - targetVector - .getOffsetBuffer() - .getInt((long) targetVector.getValueCount() * ListVector.OFFSET_WIDTH); - int deltaListSize = + targetVector.getOffsetBuffer().capacity() == 0 + ? 0 + : targetVector + .getOffsetBuffer() + .getInt((long) targetVector.getValueCount() * ListVector.OFFSET_WIDTH); + int deltaListStart = deltaVector.getOffsetBuffer().getInt(0); + int deltaListEnd = deltaVector .getOffsetBuffer() .getInt((long) deltaVector.getValueCount() * ListVector.OFFSET_WIDTH); + int deltaListSize = deltaListEnd - deltaListStart; + ListVector targetListVector = (ListVector) targetVector; // make sure the underlying vector has value count set targetListVector.getDataVector().setValueCount(targetListSize); - deltaVector.getDataVector().setValueCount(deltaListSize); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { @@ -372,13 +385,16 @@ public ValueVector visit(ListVector deltaVector, Void value) { .getOffsetBuffer() .setInt( (long) (targetVector.getValueCount() + 1 + i) * ListVector.OFFSET_WIDTH, - oldOffset + targetListSize); + oldOffset - deltaListStart + targetListSize); } targetListVector.setLastSet(newValueCount - 1); // append underlying vectors - VectorAppender innerAppender = new VectorAppender(targetListVector.getDataVector()); - deltaVector.getDataVector().accept(innerAppender, null); + appendDataVector( + targetListVector.getDataVector(), + deltaVector.getDataVector(), + deltaListStart, + deltaListSize); targetVector.setValueCount(newValueCount); return targetVector; @@ -397,20 +413,24 @@ public ValueVector visit(LargeListVector deltaVector, Void value) { int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); long targetListSize = - targetVector - .getOffsetBuffer() - .getLong((long) targetVector.getValueCount() * LargeListVector.OFFSET_WIDTH); - long deltaListSize = + targetVector.getOffsetBuffer().capacity() == 0 + ? 0 + : targetVector + .getOffsetBuffer() + .getLong((long) targetVector.getValueCount() * LargeListVector.OFFSET_WIDTH); + long deltaListStart = deltaVector.getOffsetBuffer().getLong(0); + long deltaListEnd = deltaVector .getOffsetBuffer() .getLong((long) deltaVector.getValueCount() * LargeListVector.OFFSET_WIDTH); - ListVector targetListVector = (ListVector) targetVector; + long deltaListSize = deltaListEnd - deltaListStart; + + LargeListVector targetListVector = (LargeListVector) targetVector; // make sure the underlying vector has value count set // todo recheck these casts when int64 vectors are supported targetListVector.getDataVector().setValueCount(checkedCastToInt(targetListSize)); - deltaVector.getDataVector().setValueCount(checkedCastToInt(deltaListSize)); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { @@ -427,10 +447,10 @@ public ValueVector visit(LargeListVector deltaVector, Void value) { // append offset buffer MemoryUtil.copyMemory( - deltaVector.getOffsetBuffer().memoryAddress() + ListVector.OFFSET_WIDTH, + deltaVector.getOffsetBuffer().memoryAddress() + LargeListVector.OFFSET_WIDTH, targetVector.getOffsetBuffer().memoryAddress() + (targetVector.getValueCount() + 1) * LargeListVector.OFFSET_WIDTH, - (long) deltaVector.getValueCount() * ListVector.OFFSET_WIDTH); + (long) deltaVector.getValueCount() * LargeListVector.OFFSET_WIDTH); // increase each offset from the second buffer for (int i = 0; i < deltaVector.getValueCount(); i++) { @@ -443,13 +463,16 @@ public ValueVector visit(LargeListVector deltaVector, Void value) { .getOffsetBuffer() .setLong( (long) (targetVector.getValueCount() + 1 + i) * LargeListVector.OFFSET_WIDTH, - oldOffset + targetListSize); + oldOffset - deltaListStart + targetListSize); } targetListVector.setLastSet(newValueCount - 1); // append underlying vectors - VectorAppender innerAppender = new VectorAppender(targetListVector.getDataVector()); - deltaVector.getDataVector().accept(innerAppender, null); + appendDataVector( + targetListVector.getDataVector(), + deltaVector.getDataVector(), + checkedCastToInt(deltaListStart), + checkedCastToInt(deltaListSize)); targetVector.setValueCount(newValueCount); return targetVector; @@ -693,6 +716,8 @@ public ValueVector visit(NullVector deltaVector, Void value) { Preconditions.checkArgument( targetVector.getField().getType().equals(deltaVector.getField().getType()), "The targetVector to append must have the same type as the targetVector being appended"); + targetVector.setValueCount( + Math.addExact(targetVector.getValueCount(), deltaVector.getValueCount())); return targetVector; } @@ -797,4 +822,19 @@ private void shiftAndAppendRunEndsDataBuffer( "Run-end vector and must be of type int with size 16, 32, or 64 bits."); } } + + private static void appendDataVector( + ValueVector targetDataVector, ValueVector deltaDataVector, int start, int length) { + VectorAppender innerAppender = new VectorAppender(targetDataVector); + if (start == 0 && length == deltaDataVector.getValueCount()) { + deltaDataVector.accept(innerAppender, null); + return; + } + TransferPair transferPair = + deltaDataVector.getTransferPair(deltaDataVector.getField(), deltaDataVector.getAllocator()); + try (ValueVector slicedDelta = transferPair.getTo()) { + transferPair.splitAndTransfer(start, length); + slicedDelta.accept(innerAppender, null); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java b/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java index 0945919b91..0c1ca13337 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java @@ -21,6 +21,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -31,6 +32,8 @@ import java.util.List; import java.util.Map; import java.util.function.ToIntBiFunction; +import java.util.stream.Collectors; +import java.util.stream.IntStream; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.FixedSizeListVector; @@ -54,6 +57,8 @@ 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.ValueSource; public class TestDictionaryVector { @@ -75,6 +80,78 @@ public void terminate() throws Exception { allocator.close(); } + @ParameterizedTest + @ValueSource(longs = {-1, 2, 4294967296L, Long.MAX_VALUE, Long.MIN_VALUE}) + public void testDecodeRejectsInvalidLongIndices(long index) { + try (VarCharVector values = newVarCharVector("dictionary", allocator); + BigIntVector indices = new BigIntVector("indices", allocator)) { + setVector(values, "foo", "bar"); + setVector(indices, 0L, 1L); + Dictionary dictionary = + new Dictionary(values, new DictionaryEncoding(7, false, new ArrowType.Int(64, true))); + try (ValueVector decoded = DictionaryEncoder.decode(indices, dictionary)) { + assertEquals( + Arrays.asList("foo", "bar"), + IntStream.range(0, decoded.getValueCount()) + .mapToObj(i -> decoded.getObject(i).toString()) + .collect(Collectors.toList())); + } + indices.setSafe(1, index); + long allocatedBefore = allocator.getAllocatedMemory(); + assertEquals( + "Provided dictionary does not contain value for index " + index, + assertThrows( + IllegalArgumentException.class, + () -> DictionaryEncoder.decode(indices, dictionary).close()) + .getMessage()); + assertEquals(allocatedBefore, allocator.getAllocatedMemory()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + @Test + public void testStructDecodeUsesDictionaryValueCount() { + try (StructVector validIndices = StructVector.empty("valid", allocator); + StructVector invalidIndices = StructVector.empty("invalid", allocator); + VarCharVector values = new VarCharVector("f0", allocator)) { + setVector(values, "aa", "bb"); + Dictionary dictionary = new Dictionary(values, new DictionaryEncoding(7, false, null)); + DictionaryProvider.MapDictionaryProvider provider = + new DictionaryProvider.MapDictionaryProvider(dictionary); + FieldType type = new FieldType(true, new ArrowType.Int(32, true), dictionary.getEncoding()); + validIndices.addOrGet("f0", type, IntVector.class); + invalidIndices.addOrGet("f0", type, IntVector.class); + NullableStructWriter writer = validIndices.getWriter(); + writer.allocate(); + writer.start(); + writer.integer("f0").writeInt(1); + writer.end(); + validIndices.setValueCount(1); + try (StructVector decoded = StructSubfieldEncoder.decode(validIndices, provider, allocator)) { + assertEquals( + List.of(Map.of("f0", new Text("bb"))), + IntStream.range(0, decoded.getValueCount()) + .mapToObj(decoded::getObject) + .collect(Collectors.toList())); + } + NullableStructWriter invalidWriter = invalidIndices.getWriter(); + invalidWriter.allocate(); + for (int i = 0; i < 5; i++) { + invalidWriter.start(); + invalidWriter.integer("f0").writeInt(i == 0 ? 2 : 0); + invalidWriter.end(); + } + invalidIndices.setValueCount(5); + assertEquals( + "Provided dictionary does not contain value for index 2", + assertThrows( + IllegalArgumentException.class, + () -> StructSubfieldEncoder.decode(invalidIndices, provider, allocator).close()) + .getMessage()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + @Test public void testEncodeStrings() { // Create a new value vector @@ -1053,7 +1130,7 @@ public void testStructNoMemoryLeak() { NullableStructWriter writer = indices.getWriter(); writer.allocate(); writer.start(); - writer.integer("f0").writeInt(1); + writer.integer("f0").writeInt(0); writer.integer("f1").writeInt(3); writer.end(); writer.setValueCount(1); diff --git a/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java index df5521a1ad..bdd7e71f1d 100644 --- a/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java +++ b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java @@ -17,6 +17,7 @@ package org.apache.arrow.vector.util; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -24,8 +25,10 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; +import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.CommonUtil; @@ -33,9 +36,11 @@ import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.Float4Vector; 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; @@ -53,6 +58,7 @@ import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableIntHolder; +import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -61,6 +67,8 @@ 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.EnumSource; /** Test cases for {@link VectorAppender}. */ public class TestVectorAppender { @@ -78,6 +86,252 @@ public void shutdown() { allocator.close(); } + @Test + public void testAppendNullVector() { + try (NullVector target = new NullVector("target"); + NullVector delta = new NullVector("delta", 2); + NullVector empty = new NullVector("empty")) { + VectorAppender appender = new VectorAppender(target); + assertSame(target, empty.accept(appender, null)); + assertEquals( + Collections.emptyList(), + IntStream.range(0, target.getValueCount()) + .mapToObj(target::getObject) + .collect(Collectors.toList())); + assertSame(target, delta.accept(appender, null)); + assertEquals( + Arrays.asList(null, null), + IntStream.range(0, target.getValueCount()) + .mapToObj(target::getObject) + .collect(Collectors.toList())); + VectorBatchAppender.batchAppend(target, empty, delta, delta); + assertEquals( + Arrays.asList(null, null, null, null, null, null), + IntStream.range(0, target.getValueCount()) + .mapToObj(target::getObject) + .collect(Collectors.toList())); + assertEquals( + Arrays.asList(null, null), + IntStream.range(0, delta.getValueCount()) + .mapToObj(delta::getObject) + .collect(Collectors.toList())); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @Test + public void testAppendNullVectorTypeMismatch() { + try (NullVector nulls = new NullVector("nulls", 1); + IntVector integers = new IntVector("integers", allocator)) { + ValueVectorDataPopulator.setVector(integers, 42); + assertThrows( + IllegalArgumentException.class, () -> nulls.accept(new VectorAppender(integers), null)); + assertThrows( + IllegalArgumentException.class, () -> integers.accept(new VectorAppender(nulls), null)); + assertEquals( + Collections.singletonList(42), + IntStream.range(0, integers.getValueCount()) + .mapToObj(integers::getObject) + .collect(Collectors.toList())); + assertEquals( + Collections.singletonList(null), + IntStream.range(0, nulls.getValueCount()) + .mapToObj(nulls::getObject) + .collect(Collectors.toList())); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + @Test + public void testAppendNullVectorCountBoundary() { + try (NullVector target = new NullVector("target", Integer.MAX_VALUE - 1); + NullVector delta = new NullVector("delta", 1); + NullVector empty = new NullVector("empty")) { + VectorBatchAppender.batchAppend(target, delta, empty); + assertEquals( + List.of(Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE), + List.of(target.getValueCount(), target.getNullCount(), target.getValueCapacity())); + assertThrows(ArithmeticException.class, () -> VectorBatchAppender.batchAppend(target, delta)); + assertEquals( + List.of(Integer.MAX_VALUE, 1, 0), + List.of(target.getValueCount(), delta.getValueCount(), empty.getValueCount())); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @Test + public void testAppendSlicedVariableWidthVector() { + try (VarCharVector target = new VarCharVector("target", allocator); + VarCharVector source = new VarCharVector("source", allocator); + VarCharVector delta = new VarCharVector("delta", allocator); + ArrowBuf validity = allocator.buffer(1)) { + ValueVectorDataPopulator.setVector(target, "alpha"); + ValueVectorDataPopulator.setVector(source, "skip", "beta", null, "gamma", "tail"); + validity.setByte(0, 5); + delta.loadFieldBuffers( + new ArrowFieldNode(3, 1), + Arrays.asList(validity, source.getOffsetBuffer().slice(4, 16), source.getDataBuffer())); + VectorBatchAppender.batchAppend(target, delta, delta); + assertEquals( + Arrays.asList("alpha", "beta", null, "gamma", "beta", null, "gamma"), + IntStream.range(0, target.getValueCount()) + .mapToObj(i -> target.isNull(i) ? null : target.getObject(i).toString()) + .collect(Collectors.toList())); + assertEquals( + Arrays.asList("skip", "beta", null, "gamma", "tail"), + IntStream.range(0, source.getValueCount()) + .mapToObj(i -> source.isNull(i) ? null : source.getObject(i).toString()) + .collect(Collectors.toList())); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + @Test + public void testAppendSlicedLargeVariableWidthVector() { + try (LargeVarCharVector target = new LargeVarCharVector("target", allocator); + LargeVarCharVector source = new LargeVarCharVector("source", allocator); + LargeVarCharVector delta = new LargeVarCharVector("delta", allocator); + ArrowBuf validity = allocator.buffer(1)) { + ValueVectorDataPopulator.setVector(target, "alpha"); + ValueVectorDataPopulator.setVector(source, "skip", "beta", null, "gamma", "tail"); + validity.setByte(0, 5); + delta.loadFieldBuffers( + new ArrowFieldNode(3, 1), + Arrays.asList(validity, source.getOffsetBuffer().slice(8, 32), source.getDataBuffer())); + VectorBatchAppender.batchAppend(target, delta, delta); + assertEquals( + Arrays.asList("alpha", "beta", null, "gamma", "beta", null, "gamma"), + IntStream.range(0, target.getValueCount()) + .mapToObj(i -> target.isNull(i) ? null : target.getObject(i).toString()) + .collect(Collectors.toList())); + assertEquals(5, source.getValueCount()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + @Test + public void testAppendSlicedListVector() { + try (ListVector target = ListVector.empty("target", allocator); + ListVector source = ListVector.empty("source", allocator); + ListVector delta = ListVector.empty("delta", allocator); + ArrowBuf validity = allocator.buffer(1)) { + ValueVectorDataPopulator.setVector(target, List.of(1)); + ValueVectorDataPopulator.setVector( + source, List.of(-1, -2), List.of(10, 11), null, List.of(), List.of(12, 13), List.of(99)); + delta.initializeChildrenFromFields(source.getField().getChildren()); + delta + .getDataVector() + .loadFieldBuffers( + new ArrowFieldNode(source.getDataVector().getValueCount(), 0), + source.getDataVector().getFieldBuffers()); + validity.setByte(0, 13); + delta.loadFieldBuffers( + new ArrowFieldNode(4, 1), Arrays.asList(validity, source.getOffsetBuffer().slice(4, 20))); + VectorBatchAppender.batchAppend(target, delta, delta); + assertEquals( + Arrays.asList( + List.of(1), + List.of(10, 11), + null, + List.of(), + List.of(12, 13), + List.of(10, 11), + null, + List.of(), + List.of(12, 13)), + IntStream.range(0, target.getValueCount()) + .mapToObj(target::getObject) + .collect(Collectors.toList())); + assertEquals( + Arrays.asList( + List.of(-1, -2), List.of(10, 11), null, List.of(), List.of(12, 13), List.of(99)), + IntStream.range(0, source.getValueCount()) + .mapToObj(source::getObject) + .collect(Collectors.toList())); + assertEquals(7, delta.getDataVector().getValueCount()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + @Test + public void testAppendSlicedLargeListVector() { + try (LargeListVector target = LargeListVector.empty("target", allocator); + LargeListVector source = LargeListVector.empty("source", allocator); + LargeListVector delta = LargeListVector.empty("delta", allocator); + ArrowBuf validity = allocator.buffer(1)) { + ValueVectorDataPopulator.setVector(target, List.of(1)); + ValueVectorDataPopulator.setVector( + source, List.of(-1, -2), List.of(10, 11), null, List.of(), List.of(12, 13), List.of(99)); + delta.initializeChildrenFromFields(source.getField().getChildren()); + delta + .getDataVector() + .loadFieldBuffers( + new ArrowFieldNode(source.getDataVector().getValueCount(), 0), + source.getDataVector().getFieldBuffers()); + validity.setByte(0, 13); + delta.loadFieldBuffers( + new ArrowFieldNode(4, 1), Arrays.asList(validity, source.getOffsetBuffer().slice(8, 40))); + VectorBatchAppender.batchAppend(target, delta, delta); + assertEquals( + Arrays.asList( + List.of(1), + List.of(10, 11), + null, + List.of(), + List.of(12, 13), + List.of(10, 11), + null, + List.of(), + List.of(12, 13)), + IntStream.range(0, target.getValueCount()) + .mapToObj(target::getObject) + .collect(Collectors.toList())); + assertEquals(7, delta.getDataVector().getValueCount()); + assertEquals(6, source.getValueCount()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + @ParameterizedTest + @EnumSource( + value = Types.MinorType.class, + names = {"VARCHAR", "LARGEVARCHAR", "LIST", "LARGELIST"}) + public void testAppendToUnallocatedVector(Types.MinorType type) { + try (FieldVector delta = + type.getNewVector( + new Field("delta", FieldType.nullable(type.getType()), Collections.emptyList()), + allocator, + null)) { + switch (type) { + case VARCHAR: + ValueVectorDataPopulator.setVector((VarCharVector) delta, "alpha", null, "beta"); + break; + case LARGEVARCHAR: + ValueVectorDataPopulator.setVector((LargeVarCharVector) delta, "alpha", null, "beta"); + break; + case LIST: + ValueVectorDataPopulator.setVector((ListVector) delta, List.of(42), null, List.of()); + break; + case LARGELIST: + ValueVectorDataPopulator.setVector((LargeListVector) delta, List.of(42), null, List.of()); + break; + default: + throw new AssertionError(type); + } + try (FieldVector target = delta.getField().createVector(allocator)) { + VectorBatchAppender.batchAppend(target, delta); + assertEquals( + type == Types.MinorType.LIST || type == Types.MinorType.LARGELIST + ? Arrays.asList("[42]", null, "[]") + : Arrays.asList("alpha", null, "beta"), + IntStream.range(0, target.getValueCount()) + .mapToObj(i -> target.isNull(i) ? null : target.getObject(i).toString()) + .collect(Collectors.toList())); + } + } + assertEquals(0, allocator.getAllocatedMemory()); + } + @Test public void testAppendFixedWidthVector() { final int length1 = 10;