From c7e75960188b74e6c4b150016acf80eba597d43e Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 21 Aug 2026 11:07:18 +0800 Subject: [PATCH 1/2] fix: read ARRAY through Arrow's Float8Vector readArrayData cast the child vector to the private inner Double8Vector wrapper instead of org.apache.arrow.vector.Float8Vector. The wrapper is never instantiated anywhere, so every read of an array-of-double column (List or FixedSizeList, both map to ARRAY) threw ClassCastException. The write path and the Float read branch already used the Arrow vector directly; this was a leftover from an unfinished refactor. RowDataConverterTest covers both list representations for double and float reads, plus empty/null arrays and a write batch that forces the child vector to reallocate (the production sink path); each guard was verified to redden under a targeted regression mutation. The tests allocate off-heap memory, so the root pom also sets the --add-opens=java.base/java.nio argLine that JDK 17 CI legs require for Arrow's MemoryUtil to initialize. --- pom.xml | 9 + .../lance/converter/RowDataConverter.java | 25 +- .../connector/lance/RowDataConverterTest.java | 236 ++++++++++++++++++ 3 files changed, 248 insertions(+), 22 deletions(-) create mode 100644 src/test/java/org/apache/flink/connector/lance/RowDataConverterTest.java diff --git a/pom.xml b/pom.xml index c2baa6c..1233fa6 100644 --- a/pom.xml +++ b/pom.xml @@ -18,6 +18,15 @@ 11 11 + + --add-opens=java.base/java.nio=ALL-UNNAMED + 0.1.0 diff --git a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java index 2c727ea..705bcee 100644 --- a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java +++ b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java @@ -277,13 +277,13 @@ private ArrayData readArrayData(FieldVector dataVector, int startIndex, int size } return new GenericArrayData(values); } else if (elementType instanceof DoubleType) { - Double8Vector double8Vector = (Double8Vector) dataVector; + Float8Vector float8Vector = (Float8Vector) dataVector; Double[] values = new Double[size]; for (int i = 0; i < size; i++) { - if (double8Vector.isNull(startIndex + i)) { + if (float8Vector.isNull(startIndex + i)) { values[i] = null; } else { - values[i] = double8Vector.get(startIndex + i); + values[i] = float8Vector.get(startIndex + i); } } return new GenericArrayData(values); @@ -326,25 +326,6 @@ private ArrayData readArrayData(FieldVector dataVector, int startIndex, int size "Unsupported array element type: " + elementType.getClass().getSimpleName()); } - /** - * Internal class for handling Double type Vector (alias for Float8Vector) - */ - private static class Double8Vector { - private final Float8Vector vector; - - Double8Vector(FieldVector vector) { - this.vector = (Float8Vector) vector; - } - - boolean isNull(int index) { - return vector.isNull(index); - } - - double get(int index) { - return vector.get(index); - } - } - /** * Read struct value */ diff --git a/src/test/java/org/apache/flink/connector/lance/RowDataConverterTest.java b/src/test/java/org/apache/flink/connector/lance/RowDataConverterTest.java new file mode 100644 index 0000000..f5b97a1 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/RowDataConverterTest.java @@ -0,0 +1,236 @@ +/* + * 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.flink.connector.lance; + +import org.apache.flink.connector.lance.converter.LanceTypeConverter; +import org.apache.flink.connector.lance.converter.RowDataConverter; +import org.apache.flink.table.data.ArrayData; +import org.apache.flink.table.data.GenericArrayData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.DoubleType; +import org.apache.flink.table.types.logical.FloatType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.RowType; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.FixedSizeListVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link RowDataConverter} array element handling, covering both + * Arrow representations the connector maps to Flink arrays: variable-size List + * (what the converter itself writes) and FixedSizeList (Lance vector columns). + * Reads are pinned for double and float elements; writes (including + * realloc-forcing batches) for double. + */ +class RowDataConverterTest { + + private BufferAllocator allocator; + + @BeforeEach + void setUp() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + void tearDown() { + allocator.close(); + } + + @Test + @DisplayName("Test ARRAY write/read round-trip via List vector") + void testWriteThenReadArrayOfDoubleRoundTrip() { + RowType rowType = RowType.of(new IntType(), new ArrayType(new DoubleType())); + RowDataConverter converter = new RowDataConverter(rowType); + + GenericRowData nullArrayRow = new GenericRowData(2); + nullArrayRow.setField(0, 4); + nullArrayRow.setField(1, null); + + List rows = + Arrays.asList( + row(1, new Double[] {1.5, 2.5, 3.5}), + row(2, new Double[] {4.5, null, 6.5}), + row(3, new Double[0]), + nullArrayRow); + + try (VectorSchemaRoot root = converter.createVectorSchemaRoot(allocator)) { + converter.toVectorSchemaRoot(rows, root); + + List readBack = converter.toRowDataList(root); + + assertThat(readBack).hasSize(4); + + assertThat(readBack.get(0).getInt(0)).isEqualTo(1); + ArrayData first = readBack.get(0).getArray(1); + assertThat(first.size()).isEqualTo(3); + assertThat(first.getDouble(0)).isEqualTo(1.5); + assertThat(first.getDouble(1)).isEqualTo(2.5); + assertThat(first.getDouble(2)).isEqualTo(3.5); + + assertThat(readBack.get(1).getInt(0)).isEqualTo(2); + ArrayData second = readBack.get(1).getArray(1); + assertThat(second.size()).isEqualTo(3); + assertThat(second.getDouble(0)).isEqualTo(4.5); + assertThat(second.isNullAt(1)).isTrue(); + assertThat(second.getDouble(2)).isEqualTo(6.5); + + assertThat(readBack.get(2).getInt(0)).isEqualTo(3); + assertThat(readBack.get(2).getArray(1).size()).isZero(); + + assertThat(readBack.get(3).getInt(0)).isEqualTo(4); + assertThat(readBack.get(3).isNullAt(1)).isTrue(); + } + } + + @Test + @DisplayName("Test FixedSizeList of double read (Lance float64 vector column)") + void testReadFixedSizeListOfDouble() { + Field embeddingField = + LanceTypeConverter.createFloat64VectorField("embedding", 2, true); + Schema schema = new Schema(Collections.singletonList(embeddingField)); + + ArrayType embeddingType = new ArrayType(new DoubleType()); + RowType rowType = + new RowType( + Collections.singletonList( + new RowType.RowField("embedding", embeddingType))); + RowDataConverter converter = new RowDataConverter(rowType); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + FixedSizeListVector listVector = (FixedSizeListVector) root.getVector("embedding"); + Float8Vector dataVector = (Float8Vector) listVector.getDataVector(); + dataVector.setSafe(0, 0.5); + dataVector.setSafe(1, 1.5); + dataVector.setSafe(2, 2.5); + dataVector.setNull(3); + listVector.setNotNull(0); + listVector.setNotNull(1); + listVector.setNull(2); + root.setRowCount(3); + + List readBack = converter.toRowDataList(root); + + assertThat(readBack).hasSize(3); + ArrayData first = readBack.get(0).getArray(0); + assertThat(first.size()).isEqualTo(2); + assertThat(first.getDouble(0)).isEqualTo(0.5); + assertThat(first.getDouble(1)).isEqualTo(1.5); + ArrayData second = readBack.get(1).getArray(0); + assertThat(second.getDouble(0)).isEqualTo(2.5); + assertThat(second.isNullAt(1)).isTrue(); + + assertThat(readBack.get(2).isNullAt(0)).isTrue(); + } + } + + @Test + @DisplayName("Test ARRAY write beyond the ListVector child's initial capacity") + void testWriteBeyondInitialListCapacity() { + RowType rowType = RowType.of(new IntType(), new ArrayType(new DoubleType())); + RowDataConverter converter = new RowDataConverter(rowType); + + // 200 rows x 3 elements = 600 elements against an initial child capacity of 4: + // every row past the first writes beyond that capacity, and the child doubles + // 4 -> 8 -> ... -> 1024 across the batch, so correctness depends on setSafe's + // reallocation copying prior data intact. + List rows = new ArrayList<>(200); + for (int i = 0; i < 200; i++) { + rows.add(row(i, new Double[] {i * 3.0, i * 3.0 + 1.0, i * 3.0 + 2.0})); + } + + try (VectorSchemaRoot root = converter.createVectorSchemaRoot(allocator)) { + ListVector listVector = (ListVector) root.getVector("f1"); + listVector.getDataVector().setInitialCapacity(4); + + converter.toVectorSchemaRoot(rows, root); + + List readBack = converter.toRowDataList(root); + assertThat(readBack).hasSize(200); + for (int i : new int[] {0, 1, 99, 100, 198, 199}) { + ArrayData array = readBack.get(i).getArray(1); + assertThat(array.size()).isEqualTo(3); + assertThat(array.getDouble(0)).isEqualTo(i * 3.0); + assertThat(array.getDouble(2)).isEqualTo(i * 3.0 + 2.0); + } + } + } + + @Test + @DisplayName("Test FixedSizeList of float read (Lance f32 vector column)") + void testReadFixedSizeListOfFloat() { + Field embeddingField = LanceTypeConverter.createVectorField("embedding", 2, true); + Schema schema = new Schema(Collections.singletonList(embeddingField)); + + ArrayType embeddingType = new ArrayType(new FloatType()); + RowType rowType = + new RowType( + Collections.singletonList( + new RowType.RowField("embedding", embeddingType))); + RowDataConverter converter = new RowDataConverter(rowType); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + FixedSizeListVector listVector = (FixedSizeListVector) root.getVector("embedding"); + Float4Vector dataVector = (Float4Vector) listVector.getDataVector(); + dataVector.setSafe(0, 0.5f); + dataVector.setSafe(1, 1.5f); + dataVector.setSafe(2, 2.5f); + dataVector.setNull(3); + listVector.setNotNull(0); + listVector.setNotNull(1); + root.setRowCount(2); + + List readBack = converter.toRowDataList(root); + + assertThat(readBack).hasSize(2); + ArrayData first = readBack.get(0).getArray(0); + assertThat(first.getFloat(0)).isEqualTo(0.5f); + assertThat(first.getFloat(1)).isEqualTo(1.5f); + ArrayData second = readBack.get(1).getArray(0); + assertThat(second.getFloat(0)).isEqualTo(2.5f); + assertThat(second.isNullAt(1)).isTrue(); + } + } + + private RowData row(int id, Double[] embedding) { + GenericRowData rowData = new GenericRowData(2); + rowData.setField(0, id); + rowData.setField(1, new GenericArrayData(embedding)); + return rowData; + } +} From 35bb30560fed28c13b965eb6aa7300582e5cfc43 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 22 Aug 2026 10:17:10 +0800 Subject: [PATCH 2/2] test: make the array guards pin what their comments claim Three things the guards did not actually establish. The realloc test's comment had every number wrong. setInitialCapacity(4) does not leave the child at 4: allocateNew rounds the 40-byte request to 64 and re-spreads it, landing on 7, and the child then grows 7 -> 15 -> 31 -> 63 -> 126 -> 252 -> 504 -> 1008, with the first realloc on row 2 rather than row 1. The comment now says that, and the test asserts the starting capacity is below 600 instead of taking it on faith. Without setInitialCapacity the default is 4032, which swallows all 600 with no realloc at all, so the test was one deleted line away from silently proving nothing; that line now reddens the assertion. setNull on a freshly allocated child is a no-op, since the validity bits start at zero. Both FixedSizeList tests wrote a value first, so the null slot sits over live data the way a real Arrow batch would, and row 2's own slots carry values so a drifted read fails on 4.5 instead of finding a conveniently null unwritten slot. The float test was missing two assertions the double test had: the list size, and a null parent row. It had no way to notice a broken parent-null short-circuit; removing that short-circuit now reddens it. Also asserts the middle element of each array in the realloc loop, which previously checked only the first and last. --- .../connector/lance/RowDataConverterTest.java | 43 ++++++++++++++++--- 1 file changed, 37 insertions(+), 6 deletions(-) diff --git a/src/test/java/org/apache/flink/connector/lance/RowDataConverterTest.java b/src/test/java/org/apache/flink/connector/lance/RowDataConverterTest.java index f5b97a1..076bee5 100644 --- a/src/test/java/org/apache/flink/connector/lance/RowDataConverterTest.java +++ b/src/test/java/org/apache/flink/connector/lance/RowDataConverterTest.java @@ -138,7 +138,18 @@ void testReadFixedSizeListOfDouble() { dataVector.setSafe(0, 0.5); dataVector.setSafe(1, 1.5); dataVector.setSafe(2, 2.5); + // Write index 3 before nulling it. On a freshly allocated child the validity + // bits are already zero, so setNull on its own would leave the slot merely + // unwritten rather than explicitly nulled. + dataVector.setSafe(3, 3.5); dataVector.setNull(3); + // Row 2 is nulled at the parent level, but its own two slots carry live + // values. That is not what pins the parent bit (an all-null array is still a + // non-null array, so the assertion below reddens either way); it makes the + // neighbouring assertions stricter, because a read that drifted into slots 4-5 + // now sees 4.5/5.5 instead of a conveniently null unwritten slot. + dataVector.setSafe(4, 4.5); + dataVector.setSafe(5, 5.5); listVector.setNotNull(0); listVector.setNotNull(1); listVector.setNull(2); @@ -165,10 +176,14 @@ void testWriteBeyondInitialListCapacity() { RowType rowType = RowType.of(new IntType(), new ArrayType(new DoubleType())); RowDataConverter converter = new RowDataConverter(rowType); - // 200 rows x 3 elements = 600 elements against an initial child capacity of 4: - // every row past the first writes beyond that capacity, and the child doubles - // 4 -> 8 -> ... -> 1024 across the batch, so correctness depends on setSafe's - // reallocation copying prior data intact. + // 200 rows x 3 elements = 600 elements. setInitialCapacity(4) does not literally + // leave the child at 4: allocateNew rounds the 40-byte request up to 64 and then + // re-spreads it over data plus validity, landing on 7. Rows 0 and 1 fit in that + // (indices 0-5), so the first reallocation happens on row 2, and the child then + // grows 7 -> 15 -> 31 -> 63 -> 126 -> 252 -> 504 -> 1008, seven times across the + // batch. Correctness therefore depends on setSafe's reallocation copying prior + // data intact. Without the setInitialCapacity call the default capacity is 4032, + // which swallows all 600 without a single realloc and would make this test vacuous. List rows = new ArrayList<>(200); for (int i = 0; i < 200; i++) { rows.add(row(i, new Double[] {i * 3.0, i * 3.0 + 1.0, i * 3.0 + 2.0})); @@ -178,6 +193,13 @@ void testWriteBeyondInitialListCapacity() { ListVector listVector = (ListVector) root.getVector("f1"); listVector.getDataVector().setInitialCapacity(4); + // Pin the premise rather than trusting the arithmetic above: allocate once and + // assert the child really does start below 600. If a future Arrow version + // rounds differently and the whole batch fits, this reddens instead of quietly + // turning the test into a no-op. + root.allocateNew(); + assertThat(listVector.getDataVector().getValueCapacity()).isLessThan(600); + converter.toVectorSchemaRoot(rows, root); List readBack = converter.toRowDataList(root); @@ -186,6 +208,7 @@ void testWriteBeyondInitialListCapacity() { ArrayData array = readBack.get(i).getArray(1); assertThat(array.size()).isEqualTo(3); assertThat(array.getDouble(0)).isEqualTo(i * 3.0); + assertThat(array.getDouble(1)).isEqualTo(i * 3.0 + 1.0); assertThat(array.getDouble(2)).isEqualTo(i * 3.0 + 2.0); } } @@ -210,20 +233,28 @@ void testReadFixedSizeListOfFloat() { dataVector.setSafe(0, 0.5f); dataVector.setSafe(1, 1.5f); dataVector.setSafe(2, 2.5f); + // See the double case: write before nulling, or the slot is only unwritten. + dataVector.setSafe(3, 3.5f); dataVector.setNull(3); + dataVector.setSafe(4, 4.5f); + dataVector.setSafe(5, 5.5f); listVector.setNotNull(0); listVector.setNotNull(1); - root.setRowCount(2); + listVector.setNull(2); + root.setRowCount(3); List readBack = converter.toRowDataList(root); - assertThat(readBack).hasSize(2); + assertThat(readBack).hasSize(3); ArrayData first = readBack.get(0).getArray(0); + assertThat(first.size()).isEqualTo(2); assertThat(first.getFloat(0)).isEqualTo(0.5f); assertThat(first.getFloat(1)).isEqualTo(1.5f); ArrayData second = readBack.get(1).getArray(0); assertThat(second.getFloat(0)).isEqualTo(2.5f); assertThat(second.isNullAt(1)).isTrue(); + + assertThat(readBack.get(2).isNullAt(0)).isTrue(); } }