diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/AggFunctionType.java b/fluss-common/src/main/java/org/apache/fluss/metadata/AggFunctionType.java index f0026fbba79..579c6952012 100644 --- a/fluss-common/src/main/java/org/apache/fluss/metadata/AggFunctionType.java +++ b/fluss-common/src/main/java/org/apache/fluss/metadata/AggFunctionType.java @@ -136,8 +136,20 @@ public DataTypeRoot[] getSupportedDataTypeRoots() { case LAST_VALUE_IGNORE_NULLS: case FIRST_VALUE: case FIRST_VALUE_IGNORE_NULLS: - // all data types are supported - return DataTypeRoot.values(); + // All data types are supported except VECTOR. VECTOR has no well-defined + // value-selection semantics in an aggregation context (dense vectors are + // typically compared by distance, not equality), so it is excluded here to + // prevent accidental misuse. + DataTypeRoot[] allRoots = DataTypeRoot.values(); + int vectorOrdinal = DataTypeRoot.VECTOR.ordinal(); + DataTypeRoot[] nonVectorRoots = new DataTypeRoot[allRoots.length - 1]; + int idx = 0; + for (int i = 0; i < allRoots.length; i++) { + if (i != vectorOrdinal) { + nonVectorRoots[idx++] = allRoots[i]; + } + } + return nonVectorRoots; default: throw new IllegalStateException("Unsupported aggregation function type: " + this); } diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/Schema.java b/fluss-common/src/main/java/org/apache/fluss/metadata/Schema.java index 92b675c428c..0e885ac5e92 100644 --- a/fluss-common/src/main/java/org/apache/fluss/metadata/Schema.java +++ b/fluss-common/src/main/java/org/apache/fluss/metadata/Schema.java @@ -812,6 +812,17 @@ private static List normalizeColumns( "The data type of auto increment column must be INT or BIGINT."); } + // VECTOR columns cannot be primary keys: equality comparisons on dense vectors are + // not supported (floating-point precision issues; semantically users want similarity + // distance, not bit-exact equality). See VECTOR design decision in the type system. + if (pkSet.contains(column.getName()) && column.getDataType().is(DataTypeRoot.VECTOR)) { + throw new IllegalArgumentException( + String.format( + "Column '%s' of type VECTOR cannot be used as a primary key. " + + "VECTOR columns do not support equality comparisons.", + column.getName())); + } + // primary key and auto increment column should not nullable if (pkSet.contains(column.getName()) && column.getDataType().isNullable()) { newColumns.add( diff --git a/fluss-common/src/main/java/org/apache/fluss/row/BinaryArray.java b/fluss-common/src/main/java/org/apache/fluss/row/BinaryArray.java index 35029bccaf2..8cbc76fa0a8 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/BinaryArray.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/BinaryArray.java @@ -90,6 +90,7 @@ public static int calculateFixLengthPartSize(DataType type) { case ARRAY: case MAP: case ROW: + case VECTOR: // long and double are 8 bytes; // otherwise it stores the length and offset of the variable-length part for types // such as is string, map, etc. diff --git a/fluss-common/src/main/java/org/apache/fluss/row/BinaryArrayWriter.java b/fluss-common/src/main/java/org/apache/fluss/row/BinaryArrayWriter.java index cf19c7ae360..de187870dc1 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/BinaryArrayWriter.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/BinaryArrayWriter.java @@ -220,6 +220,7 @@ public static NullSetter createNullSetter(DataType elementType) { case ARRAY: case MAP: case ROW: + case VECTOR: return BinaryArrayWriter::setNullLong; case BOOLEAN: return BinaryArrayWriter::setNullBoolean; diff --git a/fluss-common/src/main/java/org/apache/fluss/row/BinaryWriter.java b/fluss-common/src/main/java/org/apache/fluss/row/BinaryWriter.java index 1064750c1b1..41bb9c5179a 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/BinaryWriter.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/BinaryWriter.java @@ -24,6 +24,7 @@ import org.apache.fluss.row.serializer.RowSerializer; import org.apache.fluss.types.ArrayType; import org.apache.fluss.types.DataType; +import org.apache.fluss.types.FloatType; import org.apache.fluss.types.MapType; import org.apache.fluss.types.RowType; @@ -184,6 +185,12 @@ static BinaryWriter.ValueWriter createNotNullValueWriter( rowType.getFieldTypes().toArray(new DataType[0]), rowFormat); return (writer, pos, value) -> writer.writeRow(pos, (InternalRow) value, rowSerializer); + case VECTOR: + // VECTOR is serialized as an array of non-nullable FLOAT32 elements. + final ArraySerializer vectorSerializer = + new ArraySerializer(new FloatType(false), rowFormat); + return (writer, pos, value) -> + writer.writeArray(pos, (InternalArray) value, vectorSerializer); default: String msg = String.format( diff --git a/fluss-common/src/main/java/org/apache/fluss/row/InternalArray.java b/fluss-common/src/main/java/org/apache/fluss/row/InternalArray.java index 6575de1b398..2e9b97d612b 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/InternalArray.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/InternalArray.java @@ -143,6 +143,10 @@ static ElementGetter createElementGetter(DataType fieldType) { final int rowFieldCount = ((RowType) fieldType).getFieldCount(); elementGetter = (array, pos) -> array.getRow(pos, rowFieldCount); break; + case VECTOR: + // VECTOR values are represented as InternalArray of FLOAT32 elements. + elementGetter = InternalArray::getArray; + break; default: String msg = String.format( @@ -224,9 +228,17 @@ static ElementGetter createDeepElementGetter(DataType fieldType) { return genericRow; }; break; + case VECTOR: + elementGetter = + (array, pos) -> { + InternalArray inner = array.getArray(pos); + return new GenericArray(inner.toFloatArray()); + }; + break; default: // for primitive types, we can directly return the element getter elementGetter = createElementGetter(fieldType); + break; } if (!fieldType.isNullable()) { return elementGetter; diff --git a/fluss-common/src/main/java/org/apache/fluss/row/InternalRow.java b/fluss-common/src/main/java/org/apache/fluss/row/InternalRow.java index e19f42b2e0c..8a4484a6edd 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/InternalRow.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/InternalRow.java @@ -135,6 +135,9 @@ static Class getDataClass(DataType type) { return InternalMap.class; case ROW: return InternalRow.class; + case VECTOR: + // VECTOR values are stored as InternalArray of FLOAT32 elements. + return InternalArray.class; default: throw new IllegalArgumentException("Illegal type: " + type); } @@ -224,6 +227,10 @@ static FieldGetter createFieldGetter(DataType fieldType, int fieldPos) { final int numFields = ((RowType) fieldType).getFieldCount(); fieldGetter = row -> row.getRow(fieldPos, numFields); break; + case VECTOR: + // VECTOR values are InternalArray of FLOAT32 elements. + fieldGetter = row -> row.getArray(fieldPos); + break; default: throw new IllegalArgumentException("Illegal type: " + fieldType); } @@ -301,6 +308,13 @@ static FieldGetter createDeepFieldGetter(DataType fieldType, int fieldPos) { return genericRow; }; break; + case VECTOR: + fieldGetter = + row -> { + InternalArray array = row.getArray(fieldPos); + return new GenericArray(array.toFloatArray()); + }; + break; default: // for primitive types, use the normal field getter fieldGetter = createFieldGetter(fieldType, fieldPos); diff --git a/fluss-common/src/main/java/org/apache/fluss/row/arrow/ArrowWriter.java b/fluss-common/src/main/java/org/apache/fluss/row/arrow/ArrowWriter.java index a07329a0b12..c617c84ffc9 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/arrow/ArrowWriter.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/arrow/ArrowWriter.java @@ -30,6 +30,7 @@ import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.FieldVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorUnloader; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.ListVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.compression.CompressionCodec; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.compression.CompressionUtil; @@ -267,6 +268,9 @@ public int serializeToOutputView(AbstractPagedOutputView outputView) throws IOEx // update row count only when we try to write records to the output. root.setRowCount(recordsCount); + for (ArrowFieldWriter fieldWriter : fieldWriters) { + fieldWriter.finish(recordsCount); + } // update the uncompressed body size. int uncompressedBodySizeInBytes = getBodyLength(); @@ -328,6 +332,15 @@ private void initFieldVector(FieldVector fieldVector) { ((BaseFixedWidthVector) fieldVector).allocateNew(INITIAL_CAPACITY); } else if (fieldVector instanceof BaseVariableWidthVector) { ((BaseVariableWidthVector) fieldVector).allocateNew(INITIAL_CAPACITY); + } else if (fieldVector instanceof FixedSizeListVector) { + // FixedSizeListVector: allocate the top-level validity bitmap and then + // recursively initialize the child (Float32) data vector. + FixedSizeListVector fslv = (FixedSizeListVector) fieldVector; + fslv.allocateNew(); + FieldVector dataVector = fslv.getDataVector(); + if (dataVector != null) { + initFieldVector(dataVector); + } } else if (fieldVector instanceof ListVector) { ListVector listVector = (ListVector) fieldVector; listVector.allocateNew(); diff --git a/fluss-common/src/main/java/org/apache/fluss/row/arrow/vectors/ArrowVectorColumnVector.java b/fluss-common/src/main/java/org/apache/fluss/row/arrow/vectors/ArrowVectorColumnVector.java new file mode 100644 index 00000000000..494305910ce --- /dev/null +++ b/fluss-common/src/main/java/org/apache/fluss/row/arrow/vectors/ArrowVectorColumnVector.java @@ -0,0 +1,87 @@ +/* + * 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.fluss.row.arrow.vectors; + +import org.apache.fluss.annotation.Internal; +import org.apache.fluss.row.InternalArray; +import org.apache.fluss.row.columnar.ArrayColumnVector; +import org.apache.fluss.row.columnar.ColumnarArray; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.Float4Vector; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.FixedSizeListVector; + +import static org.apache.fluss.utils.Preconditions.checkNotNull; + +/** + * {@link org.apache.fluss.row.columnar.ColumnVector} backed by a shaded Arrow {@code + * FixedSizeListVector} for reading {@link org.apache.fluss.types.VectorType} columns. + * + *

Each row {@code i} maps to child Float32 elements at indices {@code [i*dimension, + * (i+1)*dimension)} in the child {@code Float4Vector}. The returned {@link InternalArray} is a + * {@link ColumnarArray} view over an {@link ArrowFloatColumnVector} wrapping the child vector. + */ +@Internal +public class ArrowVectorColumnVector implements ArrayColumnVector { + + /** The FixedSizeListVector holding per-row validity and the stride-based child data. */ + private final FixedSizeListVector vector; + + /** The fixed number of float elements per row (equals the declared VECTOR dimension). */ + private final int dimension; + + /** + * A ColumnVector view over the child Float4Vector, shared across all rows for zero-copy {@link + * ColumnarArray} slicing. + */ + private final ArrowFloatColumnVector elementVector; + + /** + * Creates a new {@link ArrowVectorColumnVector}. + * + * @param vector the {@code FixedSizeListVector} to read from + * @param dimension the declared VECTOR dimension (must equal {@code vector.getListSize()}) + */ + public ArrowVectorColumnVector(FixedSizeListVector vector, int dimension) { + this.vector = checkNotNull(vector); + this.dimension = dimension; + this.elementVector = + new ArrowFloatColumnVector((Float4Vector) checkNotNull(vector.getDataVector())); + } + + /** + * Returns the vector value at row {@code i} as an {@link InternalArray} of floats. + * + *

The returned array is a {@link ColumnarArray} window into the shared child vector, + * starting at element index {@code i * dimension} with length {@code dimension}. + * + * @param i row index (0-based) + * @return an {@link InternalArray} of {@code dimension} floats + */ + @Override + public InternalArray getArray(int i) { + if (vector.getDataVector().getValueCount() == 0 && vector.getValueCount() > 0) { + vector.getDataVector().setValueCount(vector.getValueCount() * dimension); + } + int start = i * dimension; + return new ColumnarArray(elementVector, start, dimension); + } + + @Override + public boolean isNullAt(int i) { + return vector.isNull(i); + } +} diff --git a/fluss-common/src/main/java/org/apache/fluss/row/arrow/writers/ArrowFieldWriter.java b/fluss-common/src/main/java/org/apache/fluss/row/arrow/writers/ArrowFieldWriter.java index 8dd16c64bbc..bc02d4499fb 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/arrow/writers/ArrowFieldWriter.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/arrow/writers/ArrowFieldWriter.java @@ -60,6 +60,14 @@ public void write(int rowIndex, DataGetters getters, int ordinal, boolean handle } } + /** + * Finishes writing the field vector for the given number of rows. Can be overridden by + * composite writers (e.g. {@link ArrowVectorWriter}) to finalize child vector state. + */ + public void finish(int recordsCount) { + // default no-op + } + /** Resets the state of the writer to write the next batch of fields. */ public void reset() { fieldVector.reset(); diff --git a/fluss-common/src/main/java/org/apache/fluss/row/arrow/writers/ArrowVectorWriter.java b/fluss-common/src/main/java/org/apache/fluss/row/arrow/writers/ArrowVectorWriter.java new file mode 100644 index 00000000000..1255e90373c --- /dev/null +++ b/fluss-common/src/main/java/org/apache/fluss/row/arrow/writers/ArrowVectorWriter.java @@ -0,0 +1,119 @@ +/* + * 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.fluss.row.arrow.writers; + +import org.apache.fluss.annotation.Internal; +import org.apache.fluss.row.DataGetters; +import org.apache.fluss.row.InternalArray; +import org.apache.fluss.row.arrow.ArrowWriter; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.FixedSizeListVector; + +/** + * {@link ArrowFieldWriter} for {@link org.apache.fluss.types.VectorType}, writing {@link + * InternalArray} of float values into a shaded Arrow {@code FixedSizeListVector}. + * + *

Unlike {@link ArrowArrayWriter} which uses {@code ListVector}'s {@code startNewValue/endValue} + * protocol, {@code FixedSizeListVector} uses stride-based child offsets: row {@code i} occupies + * child elements at indices {@code [i*listSize, (i+1)*listSize)}. Validity is set with {@code + * setNotNull(rowIndex)} and child elements are written via the delegated element writer. + */ +@Internal +public class ArrowVectorWriter extends ArrowFieldWriter { + + /** Writer for the Float32 child vector elements. */ + private final ArrowFieldWriter elementWriter; + + /** + * Running count of child element slots already written. Incremented by {@code listSize} for + * each row (including null rows, since FixedSizeListVector still allocates child slots for null + * rows). + */ + private int offset; + + /** + * Creates a new {@link ArrowVectorWriter}. + * + * @param vector the {@code FixedSizeListVector} to write into + * @param elementWriter writer for the Float32 child vector + */ + public ArrowVectorWriter(FixedSizeListVector vector, ArrowFieldWriter elementWriter) { + super(vector); + this.elementWriter = elementWriter; + this.offset = 0; + } + + @Override + public void doWrite(int rowIndex, DataGetters getters, int ordinal, boolean handleSafe) { + InternalArray array = getters.getArray(ordinal); + FixedSizeListVector listVector = (FixedSizeListVector) fieldVector; + int listSize = listVector.getListSize(); + if (array.size() != listSize) { + throw new IllegalArgumentException( + String.format( + "VECTOR dimension mismatch: expected %d elements but got %d.", + listSize, array.size())); + } + listVector.setNotNull(rowIndex); + for (int i = 0; i < listSize; i++) { + int elementIndex = offset + i; + // Use element-based index to determine handleSafe, not parent row count. + // When row count < INITIAL_CAPACITY but total elements > INITIAL_CAPACITY, + // we need safe mode for elements beyond the initial capacity. + boolean elementHandleSafe = elementIndex >= ArrowWriter.INITIAL_CAPACITY; + elementWriter.write(elementIndex, array, i, handleSafe || elementHandleSafe); + } + offset += listSize; + } + + /** + * Overrides the base {@link ArrowFieldWriter#write} to always advance the {@code offset} + * counter by {@code listSize}, even for null rows. + * + *

This is required because {@code FixedSizeListVector} uses stride-based child indexing: row + * {@code i}'s child elements always occupy positions {@code [i*listSize, (i+1)*listSize)}, + * regardless of whether the row is null. The base class short-circuits to {@code + * setNull(rowIndex)} without calling {@code doWrite}, so {@code offset} would never be + * incremented for null rows, causing subsequent non-null rows to write their child elements at + * the wrong positions. + */ + @Override + public void write(int rowIndex, DataGetters getters, int ordinal, boolean handleSafe) { + if (getters.isNullAt(ordinal)) { + fieldVector.setNull(rowIndex); + offset += ((FixedSizeListVector) fieldVector).getListSize(); + } else { + doWrite(rowIndex, getters, ordinal, handleSafe); + } + } + + @Override + public void finish(int recordsCount) { + ((FixedSizeListVector) fieldVector).getDataVector().setValueCount(offset); + } + + /** + * Resets the writer state for reuse (e.g. after batch serialization). The child element writer + * and offset counter are both reset to their initial state. + */ + @Override + public void reset() { + super.reset(); + elementWriter.reset(); + offset = 0; + } +} diff --git a/fluss-common/src/main/java/org/apache/fluss/row/compacted/CompactedRowReader.java b/fluss-common/src/main/java/org/apache/fluss/row/compacted/CompactedRowReader.java index e51f2ad6096..4bd24f06cb3 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/compacted/CompactedRowReader.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/compacted/CompactedRowReader.java @@ -30,6 +30,7 @@ import org.apache.fluss.row.map.CompactedMap; import org.apache.fluss.types.ArrayType; import org.apache.fluss.types.DataType; +import org.apache.fluss.types.FloatType; import org.apache.fluss.types.MapType; import org.apache.fluss.types.RowType; @@ -337,6 +338,10 @@ static FieldReader createFieldReader(DataType fieldType) { ((RowType) fieldType).getFieldTypes().toArray(new DataType[0]); fieldReader = (reader, pos) -> reader.readRow(nestedFieldTypes); break; + case VECTOR: + // VECTOR is stored as an array of non-nullable FLOAT32 elements. + fieldReader = (reader, pos) -> reader.readArray(new FloatType(false)); + break; default: throw new IllegalArgumentException( "Unsupported type for CompatedRow: " + fieldType); diff --git a/fluss-common/src/main/java/org/apache/fluss/row/indexed/IndexedRow.java b/fluss-common/src/main/java/org/apache/fluss/row/indexed/IndexedRow.java index 9a7b18f772f..dec78055641 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/indexed/IndexedRow.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/indexed/IndexedRow.java @@ -516,6 +516,7 @@ public static boolean isFixedLength(DataType dataType) { case ARRAY: case MAP: case ROW: + case VECTOR: return false; case DECIMAL: return Decimal.isCompact(((DecimalType) dataType).getPrecision()); diff --git a/fluss-common/src/main/java/org/apache/fluss/row/indexed/IndexedRowReader.java b/fluss-common/src/main/java/org/apache/fluss/row/indexed/IndexedRowReader.java index ff903cb7758..d0b5b3169f0 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/indexed/IndexedRowReader.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/indexed/IndexedRowReader.java @@ -31,6 +31,7 @@ import org.apache.fluss.row.map.IndexedMap; import org.apache.fluss.types.ArrayType; import org.apache.fluss.types.DataType; +import org.apache.fluss.types.FloatType; import org.apache.fluss.types.MapType; import org.apache.fluss.types.RowType; @@ -310,6 +311,10 @@ static FieldReader createFieldReader(DataType fieldType) { ((RowType) fieldType).getFieldTypes().toArray(new DataType[0]); fieldReader = (reader, pos) -> reader.readRow(nestedFieldTypes); break; + case VECTOR: + // VECTOR is stored as an array of non-nullable FLOAT32 elements. + fieldReader = (reader, pos) -> reader.readArray(new FloatType(false)); + break; default: throw new IllegalArgumentException("Unsupported type for IndexedRow: " + fieldType); } diff --git a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeChecks.java b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeChecks.java index ac598cee47f..34e1fb7a319 100644 --- a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeChecks.java +++ b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeChecks.java @@ -33,6 +33,8 @@ public final class DataTypeChecks { private static final FieldTypesExtractor FIELD_TYPES_EXTRACTOR = new FieldTypesExtractor(); + private static final DimensionExtractor DIMENSION_EXTRACTOR = new DimensionExtractor(); + public static int getLength(DataType dataType) { return dataType.accept(LENGTH_EXTRACTOR); } @@ -57,6 +59,14 @@ public static List getFieldTypes(DataType dataType) { return dataType.accept(FIELD_TYPES_EXTRACTOR); } + /** + * Returns the dimension (number of elements) of a {@link VectorType}. Throws {@link + * IllegalArgumentException} if called on any other type. + */ + public static int getDimension(DataType dataType) { + return dataType.accept(DIMENSION_EXTRACTOR); + } + /** Checks whether two data types are equal including field ids for row types. */ public static boolean equalsWithFieldId(DataType original, DataType that) { return that.accept(new DataTypeEqualsWithFieldId(original)); @@ -219,4 +229,11 @@ protected Boolean defaultMethod(DataType that) { return original.equals(that); } } + + private static class DimensionExtractor extends Extractor { + @Override + public Integer visit(VectorType vectorType) { + return vectorType.getDimension(); + } + } } diff --git a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeDefaultVisitor.java b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeDefaultVisitor.java index eb58c5a6c47..4c5ef493756 100644 --- a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeDefaultVisitor.java +++ b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeDefaultVisitor.java @@ -122,5 +122,10 @@ public R visit(RowType rowType) { return defaultMethod(rowType); } + @Override + public R visit(VectorType vectorType) { + return defaultMethod(vectorType); + } + protected abstract R defaultMethod(DataType dataType); } diff --git a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeFamily.java b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeFamily.java index 2c488bdd481..e858fa3cc63 100644 --- a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeFamily.java +++ b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeFamily.java @@ -54,5 +54,8 @@ public enum DataTypeFamily { COLLECTION, + /** Types that represent a fixed-size sequence of numeric values. */ + VECTOR, + EXTENSION } diff --git a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeParser.java b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeParser.java index 5068e02596a..ec2c4ea3f67 100644 --- a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeParser.java +++ b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeParser.java @@ -279,6 +279,7 @@ private enum Keyword { ARRAY, MAP, ROW, + VECTOR, NOT, NULL, } @@ -514,6 +515,8 @@ private DataType parseTypeByKeyword() { return parseMapType(); case ROW: return parseRowType(); + case VECTOR: + return parseVectorType(); default: throw parsingError("Unsupported type: " + token().value); } @@ -620,6 +623,14 @@ private int parseOptionalPrecision(int defaultPrecision) { return precision; } + private DataType parseVectorType() { + nextToken(TokenType.BEGIN_PARAMETER); + nextToken(TokenType.LITERAL_INT); + final int dimension = tokenAsInt(); + nextToken(TokenType.END_PARAMETER); + return new VectorType(dimension); + } + private DataType parseArrayType() { nextToken(TokenType.BEGIN_SUBTYPE); final DataType elementType = parseTypeWithNullability(); diff --git a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeRoot.java b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeRoot.java index f092e94661b..287509fc924 100644 --- a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeRoot.java +++ b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeRoot.java @@ -103,7 +103,9 @@ public enum DataTypeRoot { MAP(DataTypeFamily.CONSTRUCTED, DataTypeFamily.EXTENSION), - ROW(DataTypeFamily.CONSTRUCTED); + ROW(DataTypeFamily.CONSTRUCTED), + + VECTOR(DataTypeFamily.CONSTRUCTED, DataTypeFamily.VECTOR); private final Set families; diff --git a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeVisitor.java b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeVisitor.java index b3df5a12ba3..37eab873e8b 100644 --- a/fluss-common/src/main/java/org/apache/fluss/types/DataTypeVisitor.java +++ b/fluss-common/src/main/java/org/apache/fluss/types/DataTypeVisitor.java @@ -66,4 +66,6 @@ public interface DataTypeVisitor { R visit(MapType mapType); R visit(RowType rowType); + + R visit(VectorType vectorType); } diff --git a/fluss-common/src/main/java/org/apache/fluss/types/DataTypes.java b/fluss-common/src/main/java/org/apache/fluss/types/DataTypes.java index 758b717e19d..0e3235ec536 100644 --- a/fluss-common/src/main/java/org/apache/fluss/types/DataTypes.java +++ b/fluss-common/src/main/java/org/apache/fluss/types/DataTypes.java @@ -375,4 +375,16 @@ public static RowType ROW(DataField... fields) { public static RowType ROW(DataType... fieldTypes) { return RowType.builder().fields(fieldTypes).build(); } + + /** + * Data type of a fixed-dimension dense vector {@code VECTOR(n)} where {@code n} is the number + * of elements. {@code n} must be between 1 and {@link Integer#MAX_VALUE}. + * + *

Internally represented as an Arrow {@code FixedSizeList}. + * + * @see VectorType + */ + public static VectorType VECTOR(int dimension) { + return new VectorType(dimension); + } } diff --git a/fluss-common/src/main/java/org/apache/fluss/types/VectorElementType.java b/fluss-common/src/main/java/org/apache/fluss/types/VectorElementType.java new file mode 100644 index 00000000000..9634797296b --- /dev/null +++ b/fluss-common/src/main/java/org/apache/fluss/types/VectorElementType.java @@ -0,0 +1,50 @@ +/* + * 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.fluss.types; + +import org.apache.fluss.annotation.PublicEvolving; + +/** + * Enum representing the element precision of a {@link VectorType} column. + * + *

FLOAT32 is the baseline default. FLOAT16 and INT8 are reserved in the type descriptor for + * future scalar-quantization support and must not be used in writer/reader code until explicitly + * implemented. + * + * @since 0.7 + */ +@PublicEvolving +public enum VectorElementType { + + /** 32-bit IEEE 754 floating point. Fully supported. */ + FLOAT32, + + /** + * 16-bit half-precision floating point. + * + *

TODO: reserved for future scalar quantization — NOT YET IMPLEMENTED. + */ + FLOAT16, + + /** + * 8-bit signed integer (scalar quantized). + * + *

TODO: reserved for future scalar quantization — NOT YET IMPLEMENTED. + */ + INT8 +} diff --git a/fluss-common/src/main/java/org/apache/fluss/types/VectorType.java b/fluss-common/src/main/java/org/apache/fluss/types/VectorType.java new file mode 100644 index 00000000000..263fed58b26 --- /dev/null +++ b/fluss-common/src/main/java/org/apache/fluss/types/VectorType.java @@ -0,0 +1,132 @@ +/* + * 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.fluss.types; + +import org.apache.fluss.annotation.PublicEvolving; +import org.apache.fluss.row.InternalArray; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +import static org.apache.fluss.utils.Preconditions.checkArgument; + +/** + * Data type of a fixed-dimension dense vector {@code VECTOR(n)} where {@code n} is the number of + * elements. {@code n} must be between 1 and {@link Integer#MAX_VALUE}. + * + *

Internally represented as an Arrow {@code FixedSizeList} for zero-copy handoff to + * downstream lake storage (e.g. Lance). + * + *

Vector values are held in {@link InternalArray} with float elements at runtime. + * + *

Equality comparisons between VECTOR values (=, !=, IN, NOT IN) are not supported. + * + * @since 0.7 + */ +@PublicEvolving +public final class VectorType extends DataType { + + private static final long serialVersionUID = 1L; + + /** String format for serialization: {@code VECTOR(n)}. */ + public static final String FORMAT = "VECTOR(%d)"; + + /** Default element type when not specified. */ + public static final VectorElementType DEFAULT_ELEMENT_TYPE = VectorElementType.FLOAT32; + + private final int dimension; + private final VectorElementType elementType; + + /** + * Creates a {@link VectorType} with the given nullability, dimension, and element type. + * + * @param isNullable whether this type allows null values + * @param dimension the number of elements in the vector; must be positive + * @param elementType the precision of each element + */ + public VectorType(boolean isNullable, int dimension, VectorElementType elementType) { + super(isNullable, DataTypeRoot.VECTOR); + checkArgument(dimension > 0, "Dimension must be positive, got: %s", dimension); + if (elementType != VectorElementType.FLOAT32) { + throw new UnsupportedOperationException( + elementType + " is reserved for future use. Only FLOAT32 is supported."); + } + this.dimension = dimension; + this.elementType = elementType; + } + + /** + * Creates a nullable {@link VectorType} with {@link VectorElementType#FLOAT32} elements. + * + * @param dimension the number of elements in the vector; must be positive + */ + public VectorType(int dimension) { + this(true, dimension, DEFAULT_ELEMENT_TYPE); + } + + /** Returns the fixed number of elements in this vector type. */ + public int getDimension() { + return dimension; + } + + /** Returns the element precision of this vector type. */ + public VectorElementType getElementType() { + return elementType; + } + + @Override + public DataType copy(boolean isNullable) { + return new VectorType(isNullable, dimension, elementType); + } + + @Override + public String asSerializableString() { + return withNullability(FORMAT, dimension); + } + + @Override + public List getChildren() { + return Collections.emptyList(); + } + + @Override + public R accept(DataTypeVisitor visitor) { + return visitor.visit(this); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + if (!super.equals(o)) { + return false; + } + VectorType that = (VectorType) o; + return dimension == that.dimension && elementType == that.elementType; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), dimension, elementType); + } +} diff --git a/fluss-common/src/main/java/org/apache/fluss/utils/ArrowUtils.java b/fluss-common/src/main/java/org/apache/fluss/utils/ArrowUtils.java index 4219d45ffd8..1e1c1716aa3 100644 --- a/fluss-common/src/main/java/org/apache/fluss/utils/ArrowUtils.java +++ b/fluss-common/src/main/java/org/apache/fluss/utils/ArrowUtils.java @@ -41,6 +41,7 @@ import org.apache.fluss.row.arrow.vectors.ArrowTinyIntColumnVector; import org.apache.fluss.row.arrow.vectors.ArrowVarBinaryColumnVector; import org.apache.fluss.row.arrow.vectors.ArrowVarCharColumnVector; +import org.apache.fluss.row.arrow.vectors.ArrowVectorColumnVector; import org.apache.fluss.row.arrow.writers.ArrowArrayWriter; import org.apache.fluss.row.arrow.writers.ArrowBigIntWriter; import org.apache.fluss.row.arrow.writers.ArrowBinaryWriter; @@ -60,6 +61,7 @@ import org.apache.fluss.row.arrow.writers.ArrowTinyIntWriter; import org.apache.fluss.row.arrow.writers.ArrowVarBinaryWriter; import org.apache.fluss.row.arrow.writers.ArrowVarCharWriter; +import org.apache.fluss.row.arrow.writers.ArrowVectorWriter; import org.apache.fluss.row.columnar.ColumnVector; import org.apache.fluss.row.columnar.VectorizedColumnBatch; import org.apache.fluss.shaded.arrow.com.google.flatbuffers.FlatBufferBuilder; @@ -88,6 +90,7 @@ import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VarBinaryVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VarCharVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.ListVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.MapVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.StructVector; @@ -132,6 +135,7 @@ import org.apache.fluss.types.TimeType; import org.apache.fluss.types.TimestampType; import org.apache.fluss.types.TinyIntType; +import org.apache.fluss.types.VectorType; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -342,6 +346,11 @@ public static ArrowFieldWriter createArrowFieldWriter(FieldVector vector, DataTy FieldVector elementFieldVector = ((ListVector) vector).getDataVector(); return new ArrowArrayWriter( vector, ArrowUtils.createArrowFieldWriter(elementFieldVector, elementType)); + } else if (vector instanceof FixedSizeListVector && dataType instanceof VectorType) { + FieldVector childVector = ((FixedSizeListVector) vector).getDataVector(); + ArrowFieldWriter elementWriter = + ArrowUtils.createArrowFieldWriter(childVector, new FloatType(false)); + return new ArrowVectorWriter((FixedSizeListVector) vector, elementWriter); } else if (vector instanceof MapVector && dataType instanceof MapType) { MapType mapType = (MapType) dataType; MapVector mapVector = (MapVector) vector; @@ -412,6 +421,10 @@ public static ColumnVector createArrowColumnVector(ValueVector vector, DataType return new ArrowArrayColumnVector( listVector, ArrowUtils.createArrowColumnVector(listVector.getDataVector(), elementType)); + } else if (vector instanceof FixedSizeListVector && dataType instanceof VectorType) { + VectorType vectorType = (VectorType) dataType; + return new ArrowVectorColumnVector( + (FixedSizeListVector) vector, vectorType.getDimension()); } else if (vector instanceof MapVector && dataType instanceof MapType) { MapType mapType = (MapType) dataType; return new ArrowMapColumnVector( @@ -444,6 +457,9 @@ private static Field toArrowField(String fieldName, DataType logicalType) { children = Collections.singletonList( toArrowField("element", ((ArrayType) logicalType).getElementType())); + } else if (logicalType instanceof VectorType) { + // Child field: Float32 element, non-nullable (dense vector — no missing elements) + children = Collections.singletonList(toArrowField("element", new FloatType(false))); } else if (logicalType instanceof RowType) { RowType rowType = (RowType) logicalType; children = new ArrayList<>(rowType.getFieldCount()); @@ -582,6 +598,12 @@ public ArrowType visit(ArrayType arrayType) { return Types.MinorType.LIST.getType(); } + @Override + public ArrowType visit(VectorType vectorType) { + // VECTOR(n) maps to Arrow FixedSizeList with listSize = dimension + return new ArrowType.FixedSizeList(vectorType.getDimension()); + } + @Override public ArrowType visit(MapType mapType) { return new ArrowType.Map(false); diff --git a/fluss-common/src/main/java/org/apache/fluss/utils/InternalRowUtils.java b/fluss-common/src/main/java/org/apache/fluss/utils/InternalRowUtils.java index b2e4a2cbc25..90e3c5459f4 100644 --- a/fluss-common/src/main/java/org/apache/fluss/utils/InternalRowUtils.java +++ b/fluss-common/src/main/java/org/apache/fluss/utils/InternalRowUtils.java @@ -37,8 +37,10 @@ import org.apache.fluss.types.ArrayType; import org.apache.fluss.types.DataType; import org.apache.fluss.types.DataTypeRoot; +import org.apache.fluss.types.FloatType; import org.apache.fluss.types.MapType; import org.apache.fluss.types.RowType; +import org.apache.fluss.types.VectorType; import java.util.HashMap; import java.util.Map; @@ -81,6 +83,10 @@ public static InternalArray copyArray(InternalArray from, DataType eleType) { return new GenericArray(from.toDoubleArray()); } } + // VECTOR is always stored as float elements (FLOAT32, non-nullable elements) + if (eleType.getTypeRoot() == DataTypeRoot.VECTOR) { + return copyArray(from, new FloatType(false)); + } InternalArray.ElementGetter elementGetter = InternalArray.createElementGetter(eleType); Object[] newArray = new Object[from.size()]; @@ -123,6 +129,10 @@ private static Object copyValue(Object o, DataType type) { } else if (o instanceof InternalRow) { return copyRow((InternalRow) o, (RowType) type); } else if (o instanceof InternalArray) { + if (type instanceof VectorType) { + // VECTOR values are InternalArray of FLOAT32 elements; copy as float array + return copyArray((InternalArray) o, new FloatType(false)); + } return copyArray((InternalArray) o, ((ArrayType) type).getElementType()); } else if (o instanceof InternalMap) { return copyMap( diff --git a/fluss-common/src/main/java/org/apache/fluss/utils/json/DataTypeJsonSerde.java b/fluss-common/src/main/java/org/apache/fluss/utils/json/DataTypeJsonSerde.java index 12e6fb1c54b..95611f2683c 100644 --- a/fluss-common/src/main/java/org/apache/fluss/utils/json/DataTypeJsonSerde.java +++ b/fluss-common/src/main/java/org/apache/fluss/utils/json/DataTypeJsonSerde.java @@ -33,6 +33,8 @@ import org.apache.fluss.types.RowType; import org.apache.fluss.types.TimeType; import org.apache.fluss.types.TimestampType; +import org.apache.fluss.types.VectorElementType; +import org.apache.fluss.types.VectorType; import java.io.IOException; import java.util.ArrayList; @@ -73,6 +75,10 @@ public class DataTypeJsonSerde implements JsonSerializer, JsonDeserial static final String FIELD_NAME_FIELD_ID = "field_id"; static final String FIELD_NAME_FIELD_DESCRIPTION = "description"; + // VECTOR + static final String FIELD_NAME_DIMENSION = "dimension"; + static final String FIELD_NAME_VECTOR_ELEMENT_TYPE = "elementType"; + @Override public void serialize(DataType dataType, JsonGenerator generator) throws IOException { serializeTypeWithGenericSerialization(dataType, generator); @@ -139,6 +145,12 @@ private static void serializeTypeWithGenericSerialization( case ROW: serializeRow((RowType) dataType, jsonGenerator); break; + case VECTOR: + final VectorType vectorType = (VectorType) dataType; + jsonGenerator.writeNumberField(FIELD_NAME_DIMENSION, vectorType.getDimension()); + jsonGenerator.writeStringField( + FIELD_NAME_VECTOR_ELEMENT_TYPE, vectorType.getElementType().name()); + break; default: throw new UnsupportedOperationException( String.format( @@ -256,6 +268,18 @@ private static DataType deserializeFromRoot(JsonNode dataTypeNode) { return deserializeMap(dataTypeNode); case ROW: return deserializeRow(dataTypeNode); + case VECTOR: + final int vectorDimension = dataTypeNode.get(FIELD_NAME_DIMENSION).asInt(); + // elementType defaults to FLOAT32 for forward compatibility + final String elementTypeName = + dataTypeNode.has(FIELD_NAME_VECTOR_ELEMENT_TYPE) + ? dataTypeNode.get(FIELD_NAME_VECTOR_ELEMENT_TYPE).asText() + : VectorElementType.FLOAT32.name(); + if (!VectorElementType.FLOAT32.name().equals(elementTypeName)) { + throw new UnsupportedOperationException( + "Unsupported vector element type: " + elementTypeName); + } + return DataTypes.VECTOR(vectorDimension); default: throw new UnsupportedOperationException("Unsupported type root: " + typeRoot); } diff --git a/fluss-common/src/test/java/org/apache/fluss/row/arrow/ArrowReaderWriterTest.java b/fluss-common/src/test/java/org/apache/fluss/row/arrow/ArrowReaderWriterTest.java index 3190fa251bf..ea83ee8a24d 100644 --- a/fluss-common/src/test/java/org/apache/fluss/row/arrow/ArrowReaderWriterTest.java +++ b/fluss-common/src/test/java/org/apache/fluss/row/arrow/ArrowReaderWriterTest.java @@ -95,7 +95,8 @@ class ArrowReaderWriterTest { DataTypes.ROW( DataTypes.FIELD("i", DataTypes.INT()), DataTypes.FIELD("r", NESTED_DATA_TYPE), - DataTypes.FIELD("s", DataTypes.STRING()))); + DataTypes.FIELD("s", DataTypes.STRING())), + DataTypes.VECTOR(3)); private static final List TEST_DATA = Arrays.asList( @@ -137,7 +138,9 @@ class ArrowReaderWriterTest { GenericRow.of( 12, GenericRow.of(34, fromString("56"), 78L), - fromString("910"))), + fromString("910")), + // VECTOR(3) — row 0 + new GenericArray(new Float[] {0.1f, 0.2f, 0.3f})), GenericRow.of( false, (byte) 1, @@ -178,7 +181,9 @@ class ArrowReaderWriterTest { GenericRow.of( 12, GenericRow.of(34, fromString("56"), 78L), - fromString("910")))); + fromString("910")), + // VECTOR(3) — row 1 (null) + null)); @Test void testReaderWriter() throws IOException { @@ -407,4 +412,66 @@ void testMapWriterWithManyEntries() throws IOException { } } } + + /** + * Tests that VECTOR(3) columns write and read correctly via the full Arrow batch + * serialization/deserialization path. Verifies 100 rows where each row has a float vector with + * values derived from the row index. + */ + @Test + void testVectorReadWrite() throws IOException { + RowType rowType = + DataTypes.ROW( + DataTypes.FIELD("id", DataTypes.BIGINT()), + DataTypes.FIELD("embedding", DataTypes.VECTOR(3))); + + int numRows = 100; + InternalRow[] rows = new InternalRow[numRows]; + for (int i = 0; i < numRows; i++) { + rows[i] = + GenericRow.of( + (long) i, + new GenericArray( + new Float[] {(float) i, (float) i + 0.5f, (float) -i})); + } + + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot root = + VectorSchemaRoot.create(ArrowUtils.toArrowSchema(rowType), allocator); + ArrowWriterPool pool = new ArrowWriterPool(allocator); + ArrowWriter writer = + pool.getOrCreateWriter(1L, 1, Integer.MAX_VALUE, rowType, NO_COMPRESSION)) { + + for (InternalRow row : rows) { + writer.writeRow(row); + } + + AbstractPagedOutputView outputView = + new ManagedPagedOutputView(new TestingMemorySegmentPool(64 * 1024)); + int size = + writer.serializeToOutputView( + outputView, recordBatchHeaderSize(CURRENT_LOG_MAGIC_VALUE)); + assertThat(size).isGreaterThan(0); + + int heapSize = Math.max(size, writer.estimatedSizeInBytes()); + MemorySegment segment = MemorySegment.allocateHeapMemory(heapSize); + outputView + .getCurrentSegment() + .copyTo(recordBatchHeaderSize(CURRENT_LOG_MAGIC_VALUE), segment, 0, size); + + ArrowReader reader = + ArrowUtils.createArrowReader(segment, 0, size, root, allocator, rowType); + assertThat(reader.getRowCount()).isEqualTo(numRows); + + for (int i = 0; i < numRows; i++) { + ColumnarRow row = reader.read(i); + row.setRowId(i); + assertThat(row.getLong(0)).isEqualTo((long) i); + assertThat(row.getArray(1).size()).isEqualTo(3); + assertThat(row.getArray(1).getFloat(0)).isEqualTo((float) i); + assertThat(row.getArray(1).getFloat(1)).isEqualTo((float) i + 0.5f); + assertThat(row.getArray(1).getFloat(2)).isEqualTo((float) -i); + } + } + } } diff --git a/fluss-common/src/test/java/org/apache/fluss/row/arrow/writers/ArrowVectorWriterTest.java b/fluss-common/src/test/java/org/apache/fluss/row/arrow/writers/ArrowVectorWriterTest.java new file mode 100644 index 00000000000..03fa9ad24cc --- /dev/null +++ b/fluss-common/src/test/java/org/apache/fluss/row/arrow/writers/ArrowVectorWriterTest.java @@ -0,0 +1,298 @@ +/* + * 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.fluss.row.arrow.writers; + +import org.apache.fluss.memory.AbstractPagedOutputView; +import org.apache.fluss.memory.ManagedPagedOutputView; +import org.apache.fluss.memory.MemorySegment; +import org.apache.fluss.memory.TestingMemorySegmentPool; +import org.apache.fluss.row.GenericArray; +import org.apache.fluss.row.GenericRow; +import org.apache.fluss.row.InternalArray; +import org.apache.fluss.row.InternalRow; +import org.apache.fluss.row.arrow.ArrowReader; +import org.apache.fluss.row.arrow.ArrowWriter; +import org.apache.fluss.row.arrow.ArrowWriterPool; +import org.apache.fluss.row.columnar.ColumnarRow; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.RootAllocator; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.fluss.types.DataTypes; +import org.apache.fluss.types.RowType; +import org.apache.fluss.utils.ArrowUtils; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; + +import static org.apache.fluss.compression.ArrowCompressionInfo.NO_COMPRESSION; +import static org.apache.fluss.record.LogRecordBatch.CURRENT_LOG_MAGIC_VALUE; +import static org.apache.fluss.record.LogRecordBatchFormat.recordBatchHeaderSize; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Tests for {@link ArrowVectorWriter} and {@link + * org.apache.fluss.row.arrow.vectors.ArrowVectorColumnVector}. + */ +class ArrowVectorWriterTest { + + private static final int DIMENSION = 4; + + private static final RowType ROW_TYPE = + DataTypes.ROW( + DataTypes.FIELD("id", DataTypes.BIGINT()), + DataTypes.FIELD("embedding", DataTypes.VECTOR(DIMENSION))); + + // --------------------------------------------------------------------------- + // Helper: write rows and read them back, calling consumer inside resource scope + // --------------------------------------------------------------------------- + + @FunctionalInterface + private interface RowConsumer { + void accept(int rowIndex, ColumnarRow row) throws Exception; + } + + private void writeAndVerify(InternalRow[] rows, RowConsumer consumer) throws Exception { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot root = + VectorSchemaRoot.create(ArrowUtils.toArrowSchema(ROW_TYPE), allocator); + ArrowWriterPool pool = new ArrowWriterPool(allocator); + ArrowWriter writer = + pool.getOrCreateWriter( + 1L, 1, Integer.MAX_VALUE, ROW_TYPE, NO_COMPRESSION)) { + + for (InternalRow row : rows) { + writer.writeRow(row); + } + + AbstractPagedOutputView outputView = + new ManagedPagedOutputView(new TestingMemorySegmentPool(64 * 1024)); + int size = + writer.serializeToOutputView( + outputView, recordBatchHeaderSize(CURRENT_LOG_MAGIC_VALUE)); + int heapSize = Math.max(size, writer.estimatedSizeInBytes()); + MemorySegment segment = MemorySegment.allocateHeapMemory(heapSize); + outputView + .getCurrentSegment() + .copyTo(recordBatchHeaderSize(CURRENT_LOG_MAGIC_VALUE), segment, 0, size); + + ArrowReader reader = + ArrowUtils.createArrowReader(segment, 0, size, root, allocator, ROW_TYPE); + for (int i = 0; i < rows.length; i++) { + ColumnarRow row = reader.read(i); + row.setRowId(i); + try { + consumer.accept(i, row); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + } + } + + // --------------------------------------------------------------------------- + // Tests + // --------------------------------------------------------------------------- + + /** Write 3 non-null VECTOR(4) rows, read back and assert all float values. */ + @Test + void testWriteAndReadVectorRows() throws Exception { + float[] v0 = {1.0f, 2.0f, 3.0f, 4.0f}; + float[] v1 = {-1.0f, 0.5f, 100.0f, Float.MIN_VALUE}; + float[] v2 = {Float.MAX_VALUE, -Float.MAX_VALUE, 0.0f, -0.0f}; + + InternalRow[] rows = { + GenericRow.of(0L, new GenericArray(new Float[] {v0[0], v0[1], v0[2], v0[3]})), + GenericRow.of(1L, new GenericArray(new Float[] {v1[0], v1[1], v1[2], v1[3]})), + GenericRow.of(2L, new GenericArray(new Float[] {v2[0], v2[1], v2[2], v2[3]})) + }; + + float[][] expected = {v0, v1, v2}; + writeAndVerify( + rows, + (i, row) -> { + assertThat(row.isNullAt(1)).isFalse(); + InternalArray vec = row.getArray(1); + assertThat(vec.size()).isEqualTo(DIMENSION); + for (int j = 0; j < DIMENSION; j++) { + assertThat(vec.getFloat(j)).isEqualTo(expected[i][j]); + } + }); + } + + /** Write a null VECTOR row and verify isNullAt returns true for that row. */ + @Test + void testWriteNullVector() throws Exception { + InternalRow[] rows = { + GenericRow.of(0L, new GenericArray(new Float[] {1.0f, 2.0f, 3.0f, 4.0f})), + GenericRow.of(1L, (Object) null) // null embedding + }; + + writeAndVerify( + rows, + (i, row) -> { + if (i == 0) { + assertThat(row.isNullAt(1)).isFalse(); + InternalArray vec = row.getArray(1); + assertThat(vec.size()).isEqualTo(DIMENSION); + assertThat(vec.getFloat(0)).isEqualTo(1.0f); + assertThat(vec.getFloat(1)).isEqualTo(2.0f); + assertThat(vec.getFloat(2)).isEqualTo(3.0f); + assertThat(vec.getFloat(3)).isEqualTo(4.0f); + } else { + // Row 1 has null embedding + assertThat(row.isNullAt(1)).isTrue(); + } + }); + } + + /** + * Write more than INITIAL_CAPACITY (1024) rows with VECTOR(4). This exercises safe-mode element + * writes for child element indices beyond INITIAL_CAPACITY, analogous to the fix for + * ArrowArrayWriter. + */ + @Test + void testWriteBeyondInitialCapacity() throws Exception { + // 300 rows * 4 elements = 1200 child elements > INITIAL_CAPACITY (1024) + int numRows = 300; + InternalRow[] rows = new InternalRow[numRows]; + for (int i = 0; i < numRows; i++) { + float base = (float) i; + rows[i] = + GenericRow.of( + (long) i, + new GenericArray( + new Float[] {base, base + 0.1f, base + 0.2f, base + 0.3f})); + } + + writeAndVerify( + rows, + (i, row) -> { + assertThat(row.getLong(0)).isEqualTo((long) i); + assertThat(row.isNullAt(1)).isFalse(); + InternalArray vec = row.getArray(1); + assertThat(vec.size()).isEqualTo(DIMENSION); + float base = (float) i; + assertThat(vec.getFloat(0)).isEqualTo(base); + assertThat(vec.getFloat(1)).isEqualTo(base + 0.1f); + assertThat(vec.getFloat(2)).isEqualTo(base + 0.2f); + assertThat(vec.getFloat(3)).isEqualTo(base + 0.3f); + }); + } + + /** + * Write a batch, serialize it, then write a second batch and verify that the writer's internal + * offset counter is correctly reset between batches. + */ + @Test + void testResetAndRewrite() throws IOException { + float[] firstBatch = {9.0f, 8.0f, 7.0f, 6.0f}; + float[] secondBatch = {1.0f, 2.0f, 3.0f, 4.0f}; + + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot root = + VectorSchemaRoot.create(ArrowUtils.toArrowSchema(ROW_TYPE), allocator); + ArrowWriterPool pool = new ArrowWriterPool(allocator); + ArrowWriter writer = + pool.getOrCreateWriter( + 1L, 1, Integer.MAX_VALUE, ROW_TYPE, NO_COMPRESSION)) { + + // ---- First batch ---- + writer.writeRow( + GenericRow.of( + 0L, + new GenericArray( + new Float[] { + firstBatch[0], firstBatch[1], firstBatch[2], firstBatch[3] + }))); + + AbstractPagedOutputView outputView1 = + new ManagedPagedOutputView(new TestingMemorySegmentPool(64 * 1024)); + int size1 = + writer.serializeToOutputView( + outputView1, recordBatchHeaderSize(CURRENT_LOG_MAGIC_VALUE)); + int heapSize1 = Math.max(size1, writer.estimatedSizeInBytes()); + MemorySegment segment1 = MemorySegment.allocateHeapMemory(heapSize1); + outputView1 + .getCurrentSegment() + .copyTo(recordBatchHeaderSize(CURRENT_LOG_MAGIC_VALUE), segment1, 0, size1); + + ArrowReader reader1 = + ArrowUtils.createArrowReader(segment1, 0, size1, root, allocator, ROW_TYPE); + assertThat(reader1.getRowCount()).isEqualTo(1); + ColumnarRow row1 = reader1.read(0); + row1.setRowId(0); + InternalArray vec1 = row1.getArray(1); + assertThat(vec1.size()).isEqualTo(DIMENSION); + for (int i = 0; i < DIMENSION; i++) { + assertThat(vec1.getFloat(i)).isEqualTo(firstBatch[i]); + } + + // ---- Reset and second batch ---- + writer.reset(Integer.MAX_VALUE); + writer.writeRow( + GenericRow.of( + 1L, + new GenericArray( + new Float[] { + secondBatch[0], + secondBatch[1], + secondBatch[2], + secondBatch[3] + }))); + + AbstractPagedOutputView outputView2 = + new ManagedPagedOutputView(new TestingMemorySegmentPool(64 * 1024)); + int size2 = + writer.serializeToOutputView( + outputView2, recordBatchHeaderSize(CURRENT_LOG_MAGIC_VALUE)); + int heapSize2 = Math.max(size2, writer.estimatedSizeInBytes()); + MemorySegment segment2 = MemorySegment.allocateHeapMemory(heapSize2); + outputView2 + .getCurrentSegment() + .copyTo(recordBatchHeaderSize(CURRENT_LOG_MAGIC_VALUE), segment2, 0, size2); + + ArrowReader reader2 = + ArrowUtils.createArrowReader(segment2, 0, size2, root, allocator, ROW_TYPE); + assertThat(reader2.getRowCount()).isEqualTo(1); + ColumnarRow row2 = reader2.read(0); + row2.setRowId(0); + InternalArray vec2 = row2.getArray(1); + assertThat(vec2.size()).isEqualTo(DIMENSION); + for (int i = 0; i < DIMENSION; i++) { + assertThat(vec2.getFloat(i)).isEqualTo(secondBatch[i]); + } + } + } + + /** + * Writing an array with wrong number of elements into a VECTOR(4) writer must throw + * IllegalArgumentException immediately (before any child vector writes). + */ + @Test + void testMismatchedArraySize() { + // Write a 5-element array into a VECTOR(4) column — dimension mismatch + InternalRow[] rows = { + GenericRow.of(0L, new GenericArray(new Float[] {1.0f, 2.0f, 3.0f, 4.0f, 5.0f})) + }; + assertThatThrownBy(() -> writeAndVerify(rows, (i, row) -> {})) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("VECTOR dimension mismatch"); + } +} diff --git a/fluss-common/src/test/java/org/apache/fluss/testutils/InternalRowAssert.java b/fluss-common/src/test/java/org/apache/fluss/testutils/InternalRowAssert.java index 61903641482..269d92c2114 100644 --- a/fluss-common/src/test/java/org/apache/fluss/testutils/InternalRowAssert.java +++ b/fluss-common/src/test/java/org/apache/fluss/testutils/InternalRowAssert.java @@ -22,6 +22,7 @@ import org.apache.fluss.row.InternalRow; import org.apache.fluss.types.ArrayType; import org.apache.fluss.types.DataType; +import org.apache.fluss.types.FloatType; import org.apache.fluss.types.MapType; import org.apache.fluss.types.RowType; @@ -30,6 +31,7 @@ import static org.apache.fluss.types.DataTypeRoot.ARRAY; import static org.apache.fluss.types.DataTypeRoot.MAP; import static org.apache.fluss.types.DataTypeRoot.ROW; +import static org.apache.fluss.types.DataTypeRoot.VECTOR; import static org.assertj.core.api.Assertions.assertThat; /** Extend assertj assertions to easily assert {@link InternalRow}. */ @@ -78,6 +80,13 @@ public InternalRowAssert isEqualTo(InternalRow expected) { .withElementType(((ArrayType) fieldType).getElementType()) .as("InternalRow#get" + fieldType.getTypeRoot() + "(" + i + ")") .isEqualTo((InternalArray) expectedField); + } else if (fieldType.getTypeRoot() == VECTOR) { + // VECTOR is represented as InternalArray at runtime. + // Compare element-by-element as floats using FloatType(false) element type. + InternalArrayAssert.assertThatArray((InternalArray) actualField) + .withElementType(new FloatType(false)) + .as("InternalRow#get" + fieldType.getTypeRoot() + "(" + i + ")") + .isEqualTo((InternalArray) expectedField); } else if (fieldType.getTypeRoot() == MAP) { InternalMapAssert.assertThatMap((InternalMap) actualField) .withMapType((MapType) fieldType) diff --git a/fluss-common/src/test/java/org/apache/fluss/types/DataTypeParserVectorTest.java b/fluss-common/src/test/java/org/apache/fluss/types/DataTypeParserVectorTest.java new file mode 100644 index 00000000000..f9d9c8d5e56 --- /dev/null +++ b/fluss-common/src/test/java/org/apache/fluss/types/DataTypeParserVectorTest.java @@ -0,0 +1,117 @@ +/* + * 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.fluss.types; + +import org.apache.fluss.metadata.ValidationException; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for parsing {@link VectorType} via {@link DataTypeParser}. */ +public class DataTypeParserVectorTest { + + @Test + void testParseVectorBasic() { + DataType result = DataTypeParser.parse("VECTOR(1536)"); + assertThat(result).isInstanceOf(VectorType.class); + VectorType vectorType = (VectorType) result; + assertThat(vectorType.getDimension()).isEqualTo(1536); + assertThat(vectorType.getElementType()).isEqualTo(VectorElementType.FLOAT32); + assertThat(vectorType.isNullable()).isTrue(); + } + + @Test + void testParseVectorSmallDimension() { + DataType result = DataTypeParser.parse("VECTOR(1)"); + assertThat(result).isInstanceOf(VectorType.class); + assertThat(((VectorType) result).getDimension()).isEqualTo(1); + } + + @Test + void testParseVectorNotNull() { + DataType result = DataTypeParser.parse("VECTOR(1536) NOT NULL"); + assertThat(result).isInstanceOf(VectorType.class); + VectorType vectorType = (VectorType) result; + assertThat(vectorType.getDimension()).isEqualTo(1536); + assertThat(vectorType.isNullable()).isFalse(); + } + + @Test + void testParseVectorNull() { + DataType result = DataTypeParser.parse("VECTOR(768) NULL"); + assertThat(result).isInstanceOf(VectorType.class); + assertThat(result.isNullable()).isTrue(); + } + + @Test + void testParseVectorLowercaseKeyword() { + DataType result = DataTypeParser.parse("vector(4)"); + assertThat(result).isInstanceOf(VectorType.class); + assertThat(((VectorType) result).getDimension()).isEqualTo(4); + } + + @Test + void testParseVectorMixedCase() { + DataType result = DataTypeParser.parse("Vector(128)"); + assertThat(result).isInstanceOf(VectorType.class); + assertThat(((VectorType) result).getDimension()).isEqualTo(128); + } + + @Test + void testParseVectorMissingDimension() { + // VECTOR without parentheses should fail + assertThatThrownBy(() -> DataTypeParser.parse("VECTOR")) + .isInstanceOf(ValidationException.class); + } + + @Test + void testParseVectorZeroDimension() { + // VECTOR(0) should parse but fail at VectorType construction + assertThatThrownBy(() -> DataTypeParser.parse("VECTOR(0)")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Dimension must be positive"); + } + + @Test + void testRoundTrip() { + VectorType original = new VectorType(1536); + String serialized = original.asSerializableString(); + DataType parsed = DataTypeParser.parse(serialized); + assertThat(parsed).isEqualTo(original); + } + + @Test + void testRoundTripNotNull() { + VectorType original = new VectorType(false, 768, VectorElementType.FLOAT32); + String serialized = original.asSerializableString(); + DataType parsed = DataTypeParser.parse(serialized); + assertThat(parsed).isEqualTo(original); + } + + @Test + void testVectorInsideRowType() { + DataType result = DataTypeParser.parse("ROW(id BIGINT, embedding VECTOR(4))"); + assertThat(result).isInstanceOf(RowType.class); + RowType rowType = (RowType) result; + assertThat(rowType.getFieldCount()).isEqualTo(2); + assertThat(rowType.getTypeAt(1)).isInstanceOf(VectorType.class); + assertThat(((VectorType) rowType.getTypeAt(1)).getDimension()).isEqualTo(4); + } +} diff --git a/fluss-common/src/test/java/org/apache/fluss/types/VectorTypeTest.java b/fluss-common/src/test/java/org/apache/fluss/types/VectorTypeTest.java new file mode 100644 index 00000000000..e63f9e68c1f --- /dev/null +++ b/fluss-common/src/test/java/org/apache/fluss/types/VectorTypeTest.java @@ -0,0 +1,162 @@ +/* + * 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.fluss.types; + +import org.junit.jupiter.api.Test; + +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link VectorType}. */ +public class VectorTypeTest { + + @Test + void testConstructorValidation() { + // Valid construction + VectorType t = new VectorType(1536); + assertThat(t.getDimension()).isEqualTo(1536); + assertThat(t.getElementType()).isEqualTo(VectorElementType.FLOAT32); + assertThat(t.isNullable()).isTrue(); + + // Zero dimension should throw + assertThatThrownBy(() -> new VectorType(0)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Dimension must be positive"); + + // Negative dimension should throw + assertThatThrownBy(() -> new VectorType(-1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Dimension must be positive"); + } + + @Test + void testUnsupportedElementTypes() { + assertThatThrownBy(() -> new VectorType(true, 4, VectorElementType.FLOAT16)) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("reserved for future use"); + + assertThatThrownBy(() -> new VectorType(true, 4, VectorElementType.INT8)) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("reserved for future use"); + } + + @Test + void testNullability() { + VectorType nullable = new VectorType(4); + assertThat(nullable.isNullable()).isTrue(); + + VectorType notNull = (VectorType) nullable.copy(false); + assertThat(notNull.isNullable()).isFalse(); + assertThat(notNull.getDimension()).isEqualTo(4); + assertThat(notNull.getElementType()).isEqualTo(VectorElementType.FLOAT32); + } + + @Test + void testCopy() { + VectorType original = new VectorType(768); + VectorType copied = (VectorType) original.copy(); + assertThat(copied).isEqualTo(original); + assertThat(copied).isNotSameAs(original); + + VectorType notNullCopy = (VectorType) original.copy(false); + assertThat(notNullCopy.isNullable()).isFalse(); + assertThat(notNullCopy.getDimension()).isEqualTo(768); + } + + @Test + void testAsSerializableString() { + VectorType nullable = new VectorType(1536); + assertThat(nullable.asSerializableString()).isEqualTo("VECTOR(1536)"); + + VectorType notNull = new VectorType(false, 1536, VectorElementType.FLOAT32); + assertThat(notNull.asSerializableString()).isEqualTo("VECTOR(1536) NOT NULL"); + } + + @Test + void testAsSummaryString() { + VectorType t = new VectorType(256); + assertThat(t.asSummaryString()).isEqualTo("VECTOR(256)"); + } + + @Test + void testGetChildren() { + VectorType t = new VectorType(4); + assertThat(t.getChildren()).isEqualTo(Collections.emptyList()); + } + + @Test + void testAcceptVisitor() { + VectorType vectorType = new VectorType(4); + boolean[] visited = {false}; + DataTypeVisitor visitor = + new DataTypeDefaultVisitor() { + @Override + public Void visit(VectorType vt) { + visited[0] = true; + return null; + } + + @Override + protected Void defaultMethod(DataType dataType) { + return null; + } + }; + vectorType.accept(visitor); + assertThat(visited[0]).isTrue(); + } + + @Test + void testEqualsAndHashCode() { + VectorType a = new VectorType(4); + VectorType b = new VectorType(4); + VectorType c = new VectorType(8); + VectorType d = new VectorType(false, 4, VectorElementType.FLOAT32); + + assertThat(a).isEqualTo(b); + assertThat(a.hashCode()).isEqualTo(b.hashCode()); + + assertThat(a).isNotEqualTo(c); + assertThat(a).isNotEqualTo(d); // different nullability + assertThat(a).isNotEqualTo(new FloatType()); + } + + @Test + void testDataTypeRootFamily() { + VectorType t = new VectorType(4); + assertThat(t.getTypeRoot()).isEqualTo(DataTypeRoot.VECTOR); + assertThat(t.is(DataTypeRoot.VECTOR)).isTrue(); + assertThat(t.is(DataTypeFamily.VECTOR)).isTrue(); + assertThat(t.is(DataTypeFamily.CONSTRUCTED)).isTrue(); + assertThat(t.is(DataTypeFamily.COLLECTION)).isFalse(); + assertThat(t.is(DataTypeFamily.PREDEFINED)).isFalse(); + } + + @Test + void testGetDimensionViaDataTypeChecks() { + VectorType t = new VectorType(1024); + assertThat(DataTypeChecks.getDimension(t)).isEqualTo(1024); + } + + @Test + void testGetDimensionOnNonVectorTypeThrows() { + assertThatThrownBy(() -> DataTypeChecks.getDimension(new IntType())) + .isInstanceOf(IllegalArgumentException.class); + } +} diff --git a/fluss-common/src/test/java/org/apache/fluss/utils/json/DataTypeJsonSerdeVectorTest.java b/fluss-common/src/test/java/org/apache/fluss/utils/json/DataTypeJsonSerdeVectorTest.java new file mode 100644 index 00000000000..a80b3c15b1c --- /dev/null +++ b/fluss-common/src/test/java/org/apache/fluss/utils/json/DataTypeJsonSerdeVectorTest.java @@ -0,0 +1,122 @@ +/* + * 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.fluss.utils.json; + +import org.apache.fluss.types.DataType; +import org.apache.fluss.types.DataTypes; +import org.apache.fluss.types.VectorElementType; +import org.apache.fluss.types.VectorType; + +import org.junit.jupiter.api.Test; + +import java.nio.charset.StandardCharsets; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests for JSON serialization/deserialization of {@link VectorType} via {@link DataTypeJsonSerde}. + */ +public class DataTypeJsonSerdeVectorTest { + + @Test + void testSerializeVectorType() throws Exception { + VectorType vectorType = new VectorType(1536); + String json = + new String( + JsonSerdeUtils.writeValueAsBytes(vectorType, DataTypeJsonSerde.INSTANCE), + StandardCharsets.UTF_8); + assertThat(json).contains("\"type\":\"VECTOR\""); + assertThat(json).contains("\"dimension\":1536"); + assertThat(json).contains("\"elementType\":\"FLOAT32\""); + } + + @Test + void testSerializeVectorTypeNotNull() throws Exception { + VectorType vectorType = new VectorType(false, 768, VectorElementType.FLOAT32); + String json = + new String( + JsonSerdeUtils.writeValueAsBytes(vectorType, DataTypeJsonSerde.INSTANCE), + StandardCharsets.UTF_8); + assertThat(json).contains("\"type\":\"VECTOR\""); + assertThat(json).contains("\"nullable\":false"); + assertThat(json).contains("\"dimension\":768"); + } + + @Test + void testDeserializeVectorType() { + String json = "{\"type\":\"VECTOR\",\"dimension\":1536,\"elementType\":\"FLOAT32\"}"; + DataType result = + JsonSerdeUtils.readValue( + json.getBytes(StandardCharsets.UTF_8), DataTypeJsonSerde.INSTANCE); + assertThat(result).isInstanceOf(VectorType.class); + VectorType vectorType = (VectorType) result; + assertThat(vectorType.getDimension()).isEqualTo(1536); + assertThat(vectorType.getElementType()).isEqualTo(VectorElementType.FLOAT32); + assertThat(vectorType.isNullable()).isTrue(); + } + + @Test + void testDeserializeVectorTypeNotNull() { + String json = + "{\"type\":\"VECTOR\",\"nullable\":false,\"dimension\":768,\"elementType\":\"FLOAT32\"}"; + DataType result = + JsonSerdeUtils.readValue( + json.getBytes(StandardCharsets.UTF_8), DataTypeJsonSerde.INSTANCE); + assertThat(result).isInstanceOf(VectorType.class); + assertThat(result.isNullable()).isFalse(); + assertThat(((VectorType) result).getDimension()).isEqualTo(768); + } + + @Test + void testDeserializeMissingElementTypeDefaultsToFloat32() { + // elementType field absent — should default to FLOAT32 for backward compatibility + String json = "{\"type\":\"VECTOR\",\"dimension\":512}"; + DataType result = + JsonSerdeUtils.readValue( + json.getBytes(StandardCharsets.UTF_8), DataTypeJsonSerde.INSTANCE); + assertThat(result).isInstanceOf(VectorType.class); + assertThat(((VectorType) result).getElementType()).isEqualTo(VectorElementType.FLOAT32); + } + + @Test + void testRoundTrip() throws Exception { + VectorType original = new VectorType(1024); + byte[] serialized = JsonSerdeUtils.writeValueAsBytes(original, DataTypeJsonSerde.INSTANCE); + DataType deserialized = JsonSerdeUtils.readValue(serialized, DataTypeJsonSerde.INSTANCE); + assertThat(deserialized).isEqualTo(original); + } + + @Test + void testRoundTripNotNull() throws Exception { + VectorType original = new VectorType(false, 256, VectorElementType.FLOAT32); + byte[] serialized = JsonSerdeUtils.writeValueAsBytes(original, DataTypeJsonSerde.INSTANCE); + DataType deserialized = JsonSerdeUtils.readValue(serialized, DataTypeJsonSerde.INSTANCE); + assertThat(deserialized).isEqualTo(original); + } + + @Test + void testVectorEmbeddedInRowType() throws Exception { + DataType rowType = + DataTypes.ROW( + DataTypes.FIELD("id", DataTypes.BIGINT().copy(false)), + DataTypes.FIELD("embedding", DataTypes.VECTOR(4))); + byte[] serialized = JsonSerdeUtils.writeValueAsBytes(rowType, DataTypeJsonSerde.INSTANCE); + DataType deserialized = JsonSerdeUtils.readValue(serialized, DataTypeJsonSerde.INSTANCE); + assertThat(deserialized).isEqualTo(rowType); + } +} diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlussTypeToFlinkType.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlussTypeToFlinkType.java index a3b527966d8..2538bb74823 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlussTypeToFlinkType.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlussTypeToFlinkType.java @@ -38,6 +38,7 @@ import org.apache.fluss.types.TimeType; import org.apache.fluss.types.TimestampType; import org.apache.fluss.types.TinyIntType; +import org.apache.fluss.types.VectorType; import org.apache.flink.table.api.DataTypes; import org.apache.flink.table.types.DataType; @@ -168,6 +169,14 @@ public DataType visit(RowType rowType) { return withNullability(DataTypes.ROW(dataFields), rowType.isNullable()); } + @Override + public DataType visit(VectorType vectorType) { + // VECTOR is surfaced to Flink as ARRAY for query compatibility. + // The original VECTOR(n) descriptor is preserved in the Fluss schema store. + return withNullability( + DataTypes.ARRAY(DataTypes.FLOAT().notNull()), vectorType.isNullable()); + } + private DataType withNullability(DataType flinkType, boolean nullable) { if (flinkType.getLogicalType().isNullable() != nullable) { return nullable ? flinkType.nullable() : flinkType.notNull(); diff --git a/fluss-lake/fluss-lake-hudi/src/main/java/org/apache/fluss/lake/hudi/FlussDataTypeToHudiDataType.java b/fluss-lake/fluss-lake-hudi/src/main/java/org/apache/fluss/lake/hudi/FlussDataTypeToHudiDataType.java index 820e0ea2015..761575d0260 100644 --- a/fluss-lake/fluss-lake-hudi/src/main/java/org/apache/fluss/lake/hudi/FlussDataTypeToHudiDataType.java +++ b/fluss-lake/fluss-lake-hudi/src/main/java/org/apache/fluss/lake/hudi/FlussDataTypeToHudiDataType.java @@ -37,6 +37,7 @@ import org.apache.fluss.types.TimeType; import org.apache.fluss.types.TimestampType; import org.apache.fluss.types.TinyIntType; +import org.apache.fluss.types.VectorType; import org.apache.flink.table.api.DataTypes; import org.apache.flink.table.types.DataType; @@ -180,6 +181,12 @@ public DataType visit(RowType rowType) { return withNullability(DataTypes.ROW(fields), rowType.isNullable()); } + @Override + public DataType visit(VectorType vectorType) { + throw new UnsupportedOperationException( + "VECTOR type is not supported for Hudi lake format tiering."); + } + private DataType withNullability(DataType flinkDataType, boolean nullable) { if (flinkDataType.getLogicalType().isNullable() != nullable) { return nullable ? flinkDataType.nullable() : flinkDataType.notNull(); diff --git a/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/FlussDataTypeToIcebergDataType.java b/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/FlussDataTypeToIcebergDataType.java index fb8f8195748..3acb8170995 100644 --- a/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/FlussDataTypeToIcebergDataType.java +++ b/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/FlussDataTypeToIcebergDataType.java @@ -38,6 +38,7 @@ import org.apache.fluss.types.TimeType; import org.apache.fluss.types.TimestampType; import org.apache.fluss.types.TinyIntType; +import org.apache.fluss.types.VectorType; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; @@ -208,4 +209,10 @@ public Type visit(RowType rowType) { return Types.StructType.of(fields); } + + @Override + public Type visit(VectorType vectorType) { + throw new UnsupportedOperationException( + "VECTOR type is not supported for Iceberg lake format tiering."); + } } diff --git a/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/tiering/ShadedArrowBatchWriter.java b/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/tiering/ShadedArrowBatchWriter.java index 5ab4ca5470c..acce290d50d 100644 --- a/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/tiering/ShadedArrowBatchWriter.java +++ b/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/tiering/ShadedArrowBatchWriter.java @@ -24,6 +24,7 @@ import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.FieldVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.ListVector; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.StructVector; import org.apache.fluss.types.RowType; @@ -65,6 +66,9 @@ public void writeRow(InternalRow row) { public void finish() { shadedRoot.setRowCount(recordsCount); + for (ArrowFieldWriter fieldWriter : fieldWriters) { + fieldWriter.finish(recordsCount); + } } public void reset() { @@ -114,6 +118,13 @@ private void initFieldVector(FieldVector fieldVector) { for (FieldVector childVector : structVector.getChildrenFromFields()) { initFieldVector(childVector); } + } else if (fieldVector instanceof FixedSizeListVector) { + FixedSizeListVector fslv = (FixedSizeListVector) fieldVector; + fslv.allocateNew(); + FieldVector dataVector = fslv.getDataVector(); + if (dataVector != null) { + initFieldVector(dataVector); + } } else { fieldVector.allocateNew(); } diff --git a/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/utils/ArrowDataConverter.java b/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/utils/ArrowDataConverter.java index 2c1e7dcd659..8eceaf8025a 100644 --- a/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/utils/ArrowDataConverter.java +++ b/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/utils/ArrowDataConverter.java @@ -123,6 +123,24 @@ private static void copyVectorData( shadedVector.getClass().getSimpleName())); } + // Handle shaded FixedSizeListVector (for VECTOR type) — must come before the guard below + if (shadedVector + instanceof + org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.FixedSizeListVector) { + if (!(nonShadedVector instanceof FixedSizeListVector)) { + throw new IllegalArgumentException( + String.format( + "Expected non-shaded FixedSizeListVector for VECTOR type column, got: %s", + nonShadedVector.getClass().getSimpleName())); + } + copyFixedSizeListVectorData( + (org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex + .FixedSizeListVector) + shadedVector, + (FixedSizeListVector) nonShadedVector); + return; + } + if (nonShadedVector instanceof ListVector || nonShadedVector instanceof FixedSizeListVector) { throw new IllegalArgumentException( @@ -320,6 +338,51 @@ private static void copyChildDataWithOffsetRemapping( } } + /** + * Copies a shaded {@code FixedSizeListVector} (for the {@code VECTOR} type) to a non-shaded + * {@code FixedSizeListVector}. + * + *

Because both shaded and non-shaded Arrow use the same off-heap memory layout for + * FixedSizeList (validity buffer + stride-based child data), this is a direct bulk buffer copy + * — no offset remapping required (unlike the ListVector to FixedSizeListVector path). + */ + private static void copyFixedSizeListVectorData( + org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.FixedSizeListVector + shadedFSLV, + FixedSizeListVector nonShadedFSLV) { + + int valueCount = shadedFSLV.getValueCount(); + + // Recursively copy child (Float32) data vector + org.apache.fluss.shaded.arrow.org.apache.arrow.vector.FieldVector shadedDataVector = + shadedFSLV.getDataVector(); + FieldVector nonShadedDataVector = nonShadedFSLV.getDataVector(); + if (shadedDataVector != null && nonShadedDataVector != null) { + copyVectorData(shadedDataVector, nonShadedDataVector); + } + + // Copy the validity buffer (first field buffer) + List shadedBuffers = + shadedFSLV.getFieldBuffers(); + List nonShadedBuffers = nonShadedFSLV.getFieldBuffers(); + + if (!shadedBuffers.isEmpty() && !nonShadedBuffers.isEmpty()) { + org.apache.fluss.shaded.arrow.org.apache.arrow.memory.ArrowBuf shadedValidityBuf = + shadedBuffers.get(0); + ArrowBuf nonShadedValidityBuf = nonShadedBuffers.get(0); + + long size = Math.min(shadedValidityBuf.capacity(), nonShadedValidityBuf.capacity()); + if (size > 0) { + ByteBuffer srcBuffer = shadedValidityBuf.nioBuffer(0, (int) size); + srcBuffer.position(0); + srcBuffer.limit((int) Math.min(size, Integer.MAX_VALUE)); + nonShadedValidityBuf.setBytes(0, srcBuffer); + } + } + + nonShadedFSLV.setValueCount(valueCount); + } + private static void copyStructVectorData( org.apache.fluss.shaded.arrow.org.apache.arrow.vector.complex.StructVector shadedStructVector, diff --git a/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/utils/LanceArrowUtils.java b/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/utils/LanceArrowUtils.java index e0d8a0b3ab4..47d0d741f5e 100644 --- a/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/utils/LanceArrowUtils.java +++ b/fluss-lake/fluss-lake-lance/src/main/java/org/apache/fluss/lake/lance/utils/LanceArrowUtils.java @@ -37,6 +37,7 @@ import org.apache.fluss.types.TimeType; import org.apache.fluss.types.TimestampType; import org.apache.fluss.types.TinyIntType; +import org.apache.fluss.types.VectorType; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; @@ -118,6 +119,23 @@ private static Field toArrowField( } else { arrowType = toArrowType(logicalType); } + // VECTOR type: return early with explicit FixedSizeList(dimension) field and + // a non-nullable Float32 child field named "element". + if (logicalType instanceof VectorType) { + VectorType vectorType = (VectorType) logicalType; + ArrowType.FixedSizeList fslArrowType = + new ArrowType.FixedSizeList(vectorType.getDimension()); + FieldType fslFieldType = new FieldType(logicalType.isNullable(), fslArrowType, null); + Field childField = + new Field( + "element", + new FieldType( + false, + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), + null), + null); + return new Field(fieldName, fslFieldType, Collections.singletonList(childField)); + } FieldType fieldType = new FieldType(logicalType.isNullable(), arrowType, null); List children = null; if (logicalType instanceof ArrayType) { @@ -198,6 +216,9 @@ private static ArrowType toArrowType(DataType dataType) { } else { return new ArrowType.Timestamp(TimeUnit.NANOSECOND, null); } + } else if (dataType instanceof VectorType) { + VectorType vectorType = (VectorType) dataType; + return new ArrowType.FixedSizeList(vectorType.getDimension()); } else if (dataType instanceof ArrayType) { return ArrowType.List.INSTANCE; } else if (dataType instanceof RowType) { diff --git a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/testutils/FlinkLanceTieringTestBase.java b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/testutils/FlinkLanceTieringTestBase.java index 08c3ef22353..a0130c83d37 100644 --- a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/testutils/FlinkLanceTieringTestBase.java +++ b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/testutils/FlinkLanceTieringTestBase.java @@ -80,6 +80,7 @@ private static Configuration initConfig() { Configuration conf = new Configuration(); // not to clean snapshots for test purpose conf.set(ConfigOptions.KV_MAX_RETAINED_SNAPSHOTS, Integer.MAX_VALUE); + conf.setDouble("server.data-disk.write-limit-ratio", 1.0); conf.setString("datalake.format", "lance"); try { warehousePath = diff --git a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/tiering/LanceVectorTieringITCase.java b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/tiering/LanceVectorTieringITCase.java new file mode 100644 index 00000000000..9de1b51eb06 --- /dev/null +++ b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/tiering/LanceVectorTieringITCase.java @@ -0,0 +1,264 @@ +/* + * 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.fluss.lake.lance.tiering; + +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.lake.lance.LanceConfig; +import org.apache.fluss.lake.lance.testutils.FlinkLanceTieringTestBase; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.row.GenericArray; +import org.apache.fluss.row.GenericRow; +import org.apache.fluss.row.InternalRow; +import org.apache.fluss.types.DataTypes; + +import com.lancedb.lance.Dataset; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.FixedSizeListVector; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.flink.core.execution.JobClient; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * End-to-end integration test for tiering Fluss {@code VECTOR(n)} columns to Lance. + * + *

Verifies that VECTOR columns are written as Arrow {@code FixedSizeList(n)} in Lance, + * with correct float values and null handling. + */ +class LanceVectorTieringITCase extends FlinkLanceTieringTestBase { + + private static final int DIMENSION = 4; + private static final String DEFAULT_DB = "fluss"; + + private static Configuration lanceConf; + private static final RootAllocator allocator = new RootAllocator(); + + @BeforeAll + protected static void beforeAll() { + FlinkLanceTieringTestBase.beforeAll(); + lanceConf = Configuration.fromMap(getLanceCatalogConf()); + } + + /** + * Creates a log table with schema {@code (id BIGINT NOT NULL, embedding VECTOR(4))} backed by + * Lance, then writes 5 rows with known float values, triggers tiering, opens the resulting + * Lance dataset, and asserts: + * + *

+ */ + @Test + void testVectorTiering() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "vectorTable"); + + Schema schema = + Schema.newBuilder() + .column("id", DataTypes.BIGINT()) + .column("embedding", DataTypes.VECTOR(DIMENSION)) + .build(); + + TableDescriptor descriptor = + TableDescriptor.builder() + .schema(schema) + .distributedBy(1, "id") + .property(ConfigOptions.TABLE_DATALAKE_ENABLED.key(), "true") + .property(ConfigOptions.TABLE_DATALAKE_FRESHNESS, Duration.ofMillis(500)) + .build(); + + long tableId = createTable(tablePath, descriptor); + TableBucket tableBucket = new TableBucket(tableId, 0); + + // Write 5 rows: embedding = [i*1.0, i*1.1, i*1.2, i*1.3] + float[][] expectedData = { + {1.0f, 1.1f, 1.2f, 1.3f}, + {2.0f, 2.1f, 2.2f, 2.3f}, + {3.0f, 3.1f, 3.2f, 3.3f}, + {4.0f, 4.1f, 4.2f, 4.3f}, + {5.0f, 5.1f, 5.2f, 5.3f} + }; + + List rows = + Arrays.asList( + buildVectorRow(0L, expectedData[0]), + buildVectorRow(1L, expectedData[1]), + buildVectorRow(2L, expectedData[2]), + buildVectorRow(3L, expectedData[3]), + buildVectorRow(4L, expectedData[4])); + + writeRows(tablePath, rows, true); + + // Start tiering job and wait for replication + JobClient jobClient = buildTieringJob(execEnv); + assertReplicaStatus(tableBucket, 5); + + LanceConfig config = + LanceConfig.from( + lanceConf.toMap(), + Collections.emptyMap(), + tablePath.getDatabaseName(), + tablePath.getTableName()); + + try (Dataset dataset = + Dataset.open( + allocator, + config.getDatasetUri(), + LanceConfig.genReadOptionFromConfig(config))) { + + // Assert schema: embedding column must be FixedSizeList(DIMENSION) + org.apache.arrow.vector.types.pojo.Field embeddingField = + dataset.getSchema().findField("embedding"); + assertThat(embeddingField).isNotNull(); + assertThat(embeddingField.getType()).isInstanceOf(ArrowType.FixedSizeList.class); + assertThat(((ArrowType.FixedSizeList) embeddingField.getType()).getListSize()) + .isEqualTo(DIMENSION); + + // Assert child is Float32 + assertThat(embeddingField.getChildren()).hasSize(1); + assertThat(embeddingField.getChildren().get(0).getType()) + .isInstanceOf(ArrowType.FloatingPoint.class); + + // Read and assert data values + ArrowReader reader = dataset.newScan().scanBatches(); + VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot(); + assertThat(reader.loadNextBatch()).isTrue(); + + assertThat(readerRoot.getRowCount()).isEqualTo(5); + FixedSizeListVector embeddingVector = + (FixedSizeListVector) readerRoot.getVector("embedding"); + assertThat(embeddingVector.getListSize()).isEqualTo(DIMENSION); + + for (int i = 0; i < 5; i++) { + assertThat(embeddingVector.isNull(i)).isFalse(); + List values = embeddingVector.getObject(i); + assertThat(values).hasSize(DIMENSION); + for (int j = 0; j < DIMENSION; j++) { + assertThat((Float) values.get(j)) + .as("Row %d, element %d", i, j) + .isEqualTo(expectedData[i][j]); + } + } + } + + jobClient.cancel().get(); + } + + /** + * Creates a log table with a VECTOR(4) column, writes rows where some have null embeddings, + * triggers tiering, and asserts that the Lance dataset correctly represents null FixedSizeList + * rows (isNull returns true). + */ + @Test + void testNullVectorTiering() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "nullVectorTable"); + + Schema schema = + Schema.newBuilder() + .column("id", DataTypes.BIGINT()) + .column("embedding", DataTypes.VECTOR(DIMENSION)) + .build(); + + TableDescriptor descriptor = + TableDescriptor.builder() + .schema(schema) + .distributedBy(1, "id") + .property(ConfigOptions.TABLE_DATALAKE_ENABLED.key(), "true") + .property(ConfigOptions.TABLE_DATALAKE_FRESHNESS, Duration.ofMillis(500)) + .build(); + + long tableId = createTable(tablePath, descriptor); + TableBucket tableBucket = new TableBucket(tableId, 0); + + // Write 3 rows: non-null, null, non-null + GenericRow row0 = buildVectorRow(0L, new float[] {1.0f, 2.0f, 3.0f, 4.0f}); + GenericRow row1 = new GenericRow(2); // id=1, embedding=null + row1.setField(0, 1L); + row1.setField(1, null); + GenericRow row2 = buildVectorRow(2L, new float[] {5.0f, 6.0f, 7.0f, 8.0f}); + + writeRows(tablePath, Arrays.asList(row0, row1, row2), true); + + // Start tiering job and wait for replication + JobClient jobClient = buildTieringJob(execEnv); + assertReplicaStatus(tableBucket, 3); + + LanceConfig config = + LanceConfig.from( + lanceConf.toMap(), + Collections.emptyMap(), + tablePath.getDatabaseName(), + tablePath.getTableName()); + + try (Dataset dataset = + Dataset.open( + allocator, + config.getDatasetUri(), + LanceConfig.genReadOptionFromConfig(config))) { + + ArrowReader reader = dataset.newScan().scanBatches(); + VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot(); + assertThat(reader.loadNextBatch()).isTrue(); + + assertThat(readerRoot.getRowCount()).isEqualTo(3); + FixedSizeListVector embeddingVector = + (FixedSizeListVector) readerRoot.getVector("embedding"); + + // Row 0: non-null, [1.0, 2.0, 3.0, 4.0] + assertThat(embeddingVector.isNull(0)).isFalse(); + List v0 = embeddingVector.getObject(0); + assertThat((Float) v0.get(0)).isEqualTo(1.0f); + assertThat((Float) v0.get(1)).isEqualTo(2.0f); + assertThat((Float) v0.get(2)).isEqualTo(3.0f); + assertThat((Float) v0.get(3)).isEqualTo(4.0f); + + // Row 1: null embedding + assertThat(embeddingVector.isNull(1)).isTrue(); + + // Row 2: non-null, [5.0, 6.0, 7.0, 8.0] + assertThat(embeddingVector.isNull(2)).isFalse(); + List v2 = embeddingVector.getObject(2); + assertThat((Float) v2.get(0)).isEqualTo(5.0f); + assertThat((Float) v2.get(1)).isEqualTo(6.0f); + assertThat((Float) v2.get(2)).isEqualTo(7.0f); + assertThat((Float) v2.get(3)).isEqualTo(8.0f); + } + + jobClient.cancel().get(); + } + + private static GenericRow buildVectorRow(long id, float[] embedding) { + GenericRow row = new GenericRow(2); + row.setField(0, id); + row.setField(1, new GenericArray(embedding)); + return row; + } +} diff --git a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/utils/ArrowDataConverterTest.java b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/utils/ArrowDataConverterTest.java index a5face19aa9..99c2c390fae 100644 --- a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/utils/ArrowDataConverterTest.java +++ b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/utils/ArrowDataConverterTest.java @@ -194,4 +194,109 @@ void testConvertListToFixedSizeListWithNulls() { } } } + + /** + * Test converting a shaded FixedSizeListVector (VECTOR type, not the ARRAY+property legacy + * path) to a non-shaded FixedSizeListVector via {@link ArrowDataConverter#convertToNonShaded}. + * Verifies that 3 rows of VECTOR(4) are converted correctly with exact Float32 values. + */ + @Test + void testConvertFixedSizeListVectorData() { + int dimension = 4; + RowType rowType = DataTypes.ROW(DataTypes.FIELD("embedding", DataTypes.VECTOR(dimension))); + + Schema nonShadedSchema = LanceArrowUtils.toArrowSchema(rowType); + + float[][] data = { + {1.0f, 2.0f, 3.0f, 4.0f}, + {5.5f, 6.6f, 7.7f, 8.8f}, + {-1.0f, 0.0f, Float.MAX_VALUE, Float.MIN_VALUE} + }; + + try (ShadedArrowBatchWriter writer = new ShadedArrowBatchWriter(shadedAllocator, rowType)) { + for (float[] floats : data) { + GenericRow row = new GenericRow(1); + row.setField(0, new GenericArray(floats)); + writer.writeRow(row); + } + writer.finish(); + + try (VectorSchemaRoot nonShadedRoot = + ArrowDataConverter.convertToNonShaded( + writer.getShadedRoot(), nonShadedAllocator, nonShadedSchema)) { + assertThat(nonShadedRoot.getRowCount()).isEqualTo(3); + assertThat(nonShadedRoot.getVector("embedding")) + .isInstanceOf(FixedSizeListVector.class); + + FixedSizeListVector result = + (FixedSizeListVector) nonShadedRoot.getVector("embedding"); + assertThat(result.getListSize()).isEqualTo(dimension); + assertThat(result.getValueCount()).isEqualTo(3); + + // Verify exact Float32 values for each row + for (int i = 0; i < data.length; i++) { + assertThat(result.isNull(i)).isFalse(); + List values = result.getObject(i); + assertThat(values).hasSize(dimension); + for (int j = 0; j < dimension; j++) { + assertThat((Float) values.get(j)).isEqualTo(data[i][j]); + } + } + } + } + } + + /** Test that null VECTOR rows are correctly preserved in the converted non-shaded vector. */ + @Test + void testConvertFixedSizeListVectorDataWithNulls() { + int dimension = 3; + RowType rowType = DataTypes.ROW(DataTypes.FIELD("embedding", DataTypes.VECTOR(dimension))); + + Schema nonShadedSchema = LanceArrowUtils.toArrowSchema(rowType); + + try (ShadedArrowBatchWriter writer = new ShadedArrowBatchWriter(shadedAllocator, rowType)) { + // Row 0: non-null + GenericRow row0 = new GenericRow(1); + row0.setField(0, new GenericArray(new float[] {1.0f, 2.0f, 3.0f})); + writer.writeRow(row0); + + // Row 1: null embedding + GenericRow row1 = new GenericRow(1); + row1.setField(0, null); + writer.writeRow(row1); + + // Row 2: non-null + GenericRow row2 = new GenericRow(1); + row2.setField(0, new GenericArray(new float[] {4.0f, 5.0f, 6.0f})); + writer.writeRow(row2); + + writer.finish(); + + try (VectorSchemaRoot nonShadedRoot = + ArrowDataConverter.convertToNonShaded( + writer.getShadedRoot(), nonShadedAllocator, nonShadedSchema)) { + assertThat(nonShadedRoot.getRowCount()).isEqualTo(3); + + FixedSizeListVector result = + (FixedSizeListVector) nonShadedRoot.getVector("embedding"); + + // Row 0: non-null, values preserved exactly + assertThat(result.isNull(0)).isFalse(); + List v0 = result.getObject(0); + assertThat((Float) v0.get(0)).isEqualTo(1.0f); + assertThat((Float) v0.get(1)).isEqualTo(2.0f); + assertThat((Float) v0.get(2)).isEqualTo(3.0f); + + // Row 1: null + assertThat(result.isNull(1)).isTrue(); + + // Row 2: non-null, values preserved exactly + assertThat(result.isNull(2)).isFalse(); + List v2 = result.getObject(2); + assertThat((Float) v2.get(0)).isEqualTo(4.0f); + assertThat((Float) v2.get(1)).isEqualTo(5.0f); + assertThat((Float) v2.get(2)).isEqualTo(6.0f); + } + } + } } diff --git a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/utils/LanceArrowUtilsTest.java b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/utils/LanceArrowUtilsTest.java index ef224ca7ab8..9d615253c66 100644 --- a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/utils/LanceArrowUtilsTest.java +++ b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/utils/LanceArrowUtilsTest.java @@ -298,4 +298,60 @@ void testToArrowSchemaWithRowContainingArray() { assertThat(tagsChildren.get(0).getName()).isEqualTo("element"); assertThat(tagsChildren.get(0).getType()).isEqualTo(ArrowType.Utf8.INSTANCE); } + + @Test + void testVectorColumnToArrowSchema() { + // VECTOR(3) should produce a FixedSizeList(3) field with a non-nullable Float32 child + RowType rowType = DataTypes.ROW(DataTypes.FIELD("embedding", DataTypes.VECTOR(3))); + + Schema schema = LanceArrowUtils.toArrowSchema(rowType); + + assertThat(schema.getFields()).hasSize(1); + Field embeddingField = schema.findField("embedding"); + assertThat(embeddingField).isNotNull(); + assertThat(embeddingField.getType()).isInstanceOf(ArrowType.FixedSizeList.class); + assertThat(((ArrowType.FixedSizeList) embeddingField.getType()).getListSize()).isEqualTo(3); + + // Child field must be named "element" and be non-nullable Float32 + assertThat(embeddingField.getChildren()).hasSize(1); + Field childField = embeddingField.getChildren().get(0); + assertThat(childField.getName()).isEqualTo("element"); + assertThat(childField.getType()).isInstanceOf(ArrowType.FloatingPoint.class); + assertThat(childField.getFieldType().isNullable()).isFalse(); + } + + @Test + void testVectorNotNullColumn() { + // VECTOR(4) NOT NULL should produce a non-nullable top-level field + RowType rowType = + DataTypes.ROW(DataTypes.FIELD("embedding", DataTypes.VECTOR(4).copy(false))); + + Schema schema = LanceArrowUtils.toArrowSchema(rowType); + + Field embeddingField = schema.findField("embedding"); + assertThat(embeddingField).isNotNull(); + assertThat(embeddingField.getFieldType().isNullable()).isFalse(); + assertThat(embeddingField.getType()).isInstanceOf(ArrowType.FixedSizeList.class); + assertThat(((ArrowType.FixedSizeList) embeddingField.getType()).getListSize()).isEqualTo(4); + } + + @Test + void testVectorNoTablePropertyNeeded() { + // VECTOR type produces correct FixedSizeList schema without any table property, + // unlike the legacy ARRAY path which requires ".arrow.fixed-size-list.size" property. + RowType rowType = DataTypes.ROW(DataTypes.FIELD("embedding", DataTypes.VECTOR(8))); + + // No table properties passed + Schema schema = LanceArrowUtils.toArrowSchema(rowType, Collections.emptyMap()); + + Field embeddingField = schema.findField("embedding"); + assertThat(embeddingField).isNotNull(); + assertThat(embeddingField.getType()).isInstanceOf(ArrowType.FixedSizeList.class); + assertThat(((ArrowType.FixedSizeList) embeddingField.getType()).getListSize()).isEqualTo(8); + // Child must be Float32 + assertThat(embeddingField.getChildren()).hasSize(1); + assertThat(embeddingField.getChildren().get(0).getName()).isEqualTo("element"); + assertThat(embeddingField.getChildren().get(0).getType()) + .isInstanceOf(ArrowType.FloatingPoint.class); + } } diff --git a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/utils/FlussDataTypeToPaimonDataType.java b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/utils/FlussDataTypeToPaimonDataType.java index e97b4f67a69..b617dbc5a3e 100644 --- a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/utils/FlussDataTypeToPaimonDataType.java +++ b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/utils/FlussDataTypeToPaimonDataType.java @@ -38,6 +38,7 @@ import org.apache.fluss.types.TimeType; import org.apache.fluss.types.TimestampType; import org.apache.fluss.types.TinyIntType; +import org.apache.fluss.types.VectorType; import org.apache.paimon.types.DataType; import org.apache.paimon.types.DataTypes; @@ -160,6 +161,12 @@ public DataType visit(RowType rowType) { return withNullability(rowTypeBuilder.build(), rowType.isNullable()); } + @Override + public DataType visit(VectorType vectorType) { + throw new UnsupportedOperationException( + "VECTOR type is not supported for Paimon lake format tiering."); + } + private DataType withNullability(DataType paimon, boolean nullable) { if (paimon.isNullable() != nullable) { return nullable ? paimon.nullable() : paimon.notNull(); diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/util/PredicateMessageUtils.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/util/PredicateMessageUtils.java index f492e3957d6..737df62f309 100644 --- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/util/PredicateMessageUtils.java +++ b/fluss-rpc/src/main/java/org/apache/fluss/rpc/util/PredicateMessageUtils.java @@ -515,11 +515,12 @@ private enum DataTypeRootCode { TIMESTAMP_WITHOUT_TIME_ZONE(12, DataTypeRoot.TIMESTAMP_WITHOUT_TIME_ZONE), TIMESTAMP_WITH_LOCAL_TIME_ZONE(13, DataTypeRoot.TIMESTAMP_WITH_LOCAL_TIME_ZONE), BINARY(14, DataTypeRoot.BINARY), - BYTES(15, DataTypeRoot.BYTES); + BYTES(15, DataTypeRoot.BYTES), + VECTOR(16, DataTypeRoot.VECTOR); private final int value; private final DataTypeRoot dataTypeRoot; - private static final DataTypeRootCode[] VALUES = new DataTypeRootCode[16]; + private static final DataTypeRootCode[] VALUES = new DataTypeRootCode[17]; private static final Map ROOT_MAP = new HashMap<>(); static {