This is an automated email from the ASF dual-hosted git repository.
danny0405 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/master by this push:
new 5e58dd363aed feat(flink): support vector columns in Lance writer
(#19831)
5e58dd363aed is described below
commit 5e58dd363aed21aa9c048b573023fff94bd798f6
Author: Shuo Cheng <[email protected]>
AuthorDate: Fri Sep 4 17:08:16 2026 +0800
feat(flink): support vector columns in Lance writer (#19831)
---
.../row/HoodieRowDataFileWriterFactory.java | 5 +-
.../io/storage/row/HoodieRowDataLanceWriter.java | 22 +++--
.../row/lance/HoodieFlinkLanceArrowUtils.java | 53 +++++++++-
.../io/storage/row/lance/LanceRowDataWriter.java | 86 +++++++++++++++-
.../row/TestHoodieFlinkLanceArrowUtils.java | 66 ++++++++++++-
.../storage/row/TestHoodieRowDataLanceWriter.java | 103 ++++++++++++++++++++
.../io/storage/row/TestLanceRowDataWriter.java | 108 ++++++++++++++++++---
7 files changed, 408 insertions(+), 35 deletions(-)
diff --git
a/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/HoodieRowDataFileWriterFactory.java
b/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/HoodieRowDataFileWriterFactory.java
index 77054f98bf0f..7b15cfaca90d 100644
---
a/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/HoodieRowDataFileWriterFactory.java
+++
b/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/HoodieRowDataFileWriterFactory.java
@@ -36,9 +36,7 @@ import org.apache.hudi.storage.HoodieStorage;
import org.apache.hudi.storage.StorageConfiguration;
import org.apache.hudi.storage.StoragePath;
import org.apache.hudi.storage.hadoop.HadoopStorageConfiguration;
-import org.apache.hudi.util.HoodieSchemaConverter;
-import org.apache.flink.table.types.logical.RowType;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
@@ -126,10 +124,9 @@ public class HoodieRowDataFileWriterFactory extends
HoodieFileWriterFactory {
boolean withOperation =
config.getBooleanOrDefault(HoodieWriteConfig.ALLOW_OPERATION_METADATA_FIELD);
Option<org.apache.hudi.common.bloom.BloomFilter> bloomFilter =
enableBloomFilter(metaFieldsMode, config)
? Option.of(createBloomFilter(config)) : Option.empty();
- RowType rowType = HoodieSchemaConverter.convertToRowType(schema);
return new HoodieRowDataLanceWriter(
path,
- rowType,
+ schema,
instantTime,
taskContextSupplier,
bloomFilter,
diff --git
a/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/HoodieRowDataLanceWriter.java
b/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/HoodieRowDataLanceWriter.java
index 0990e3d6527a..238974cf3221 100644
---
a/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/HoodieRowDataLanceWriter.java
+++
b/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/HoodieRowDataLanceWriter.java
@@ -23,6 +23,7 @@ import org.apache.hudi.common.bloom.BloomFilter;
import org.apache.hudi.common.engine.TaskContextSupplier;
import org.apache.hudi.common.model.HoodieKey;
import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.schema.HoodieSchema;
import org.apache.hudi.common.util.Option;
import org.apache.hudi.common.util.ValidationUtils;
import org.apache.hudi.io.lance.HoodieBaseLanceWriter;
@@ -33,9 +34,10 @@ import org.apache.hudi.storage.StoragePath;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.flink.table.data.RowData;
-import org.apache.flink.table.types.logical.RowType;
import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
import java.util.function.Function;
/**
@@ -47,7 +49,7 @@ public class HoodieRowDataLanceWriter extends
HoodieBaseLanceWriter<RowData, Str
private static final long MIN_RECORDS_FOR_SIZE_CHECK = 100L;
private static final long MAX_RECORDS_FOR_SIZE_CHECK = 10000L;
- private final RowType rowType;
+ private final HoodieSchema hoodieSchema;
private final Schema arrowSchema;
private final String fileName;
private final String instantTime;
@@ -60,7 +62,7 @@ public class HoodieRowDataLanceWriter extends
HoodieBaseLanceWriter<RowData, Str
public HoodieRowDataLanceWriter(
StoragePath file,
- RowType rowType,
+ HoodieSchema hoodieSchema,
String instantTime,
TaskContextSupplier taskContextSupplier,
Option<BloomFilter> bloomFilterOpt,
@@ -78,8 +80,8 @@ public class HoodieRowDataLanceWriter extends
HoodieBaseLanceWriter<RowData, Str
ValidationUtils.checkArgument(flushByteWatermark < allocatorSize,
"flushByteWatermark (" + flushByteWatermark + ") must be less than
allocatorSize ("
+ allocatorSize + ")");
- this.rowType = rowType;
- this.arrowSchema = HoodieFlinkLanceArrowUtils.toArrowSchema(rowType);
+ this.hoodieSchema = hoodieSchema.getNonNullType();
+ this.arrowSchema =
HoodieFlinkLanceArrowUtils.toArrowSchema(this.hoodieSchema);
this.fileName = file.getName();
this.instantTime = instantTime;
this.maxFileSize = maxFileSize;
@@ -134,6 +136,14 @@ public class HoodieRowDataLanceWriter extends
HoodieBaseLanceWriter<RowData, Str
return arrowSchema;
}
+ @Override
+ protected Map<String, String> additionalSchemaMetadata() {
+ String value = HoodieSchema.buildVectorColumnsMetadataValue(hoodieSchema);
+ return value.isEmpty()
+ ? Collections.emptyMap()
+ : Collections.singletonMap(HoodieSchema.VECTOR_COLUMNS_METADATA_KEY,
value);
+ }
+
private class RowDataArrowWriter implements ArrowWriter<RowData> {
private final VectorSchemaRoot root;
private final LanceRowDataWriter writer;
@@ -141,7 +151,7 @@ public class HoodieRowDataLanceWriter extends
HoodieBaseLanceWriter<RowData, Str
private RowDataArrowWriter(VectorSchemaRoot root) {
this.root = root;
- this.writer = new LanceRowDataWriter(rowType, root.getFieldVectors(),
utcTimestamp);
+ this.writer = new LanceRowDataWriter(hoodieSchema,
root.getFieldVectors(), utcTimestamp);
}
@Override
diff --git
a/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/lance/HoodieFlinkLanceArrowUtils.java
b/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/lance/HoodieFlinkLanceArrowUtils.java
index 6c4e4ac29cb0..9806f5e9ad1e 100644
---
a/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/lance/HoodieFlinkLanceArrowUtils.java
+++
b/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/lance/HoodieFlinkLanceArrowUtils.java
@@ -18,8 +18,12 @@
package org.apache.hudi.io.storage.row.lance;
+import org.apache.hudi.common.schema.HoodieSchema;
+import org.apache.hudi.common.schema.HoodieSchemaField;
+import org.apache.hudi.common.schema.HoodieSchemaType;
import org.apache.hudi.common.util.ValidationUtils;
import org.apache.hudi.exception.HoodieNotSupportedException;
+import org.apache.hudi.util.HoodieSchemaConverter;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
@@ -72,6 +76,7 @@ import org.apache.flink.table.types.logical.VarCharType;
import java.math.BigDecimal;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
/**
@@ -88,10 +93,22 @@ public final class HoodieFlinkLanceArrowUtils {
private HoodieFlinkLanceArrowUtils() {
}
- public static Schema toArrowSchema(RowType rowType) {
+ /**
+ * Converts a Hoodie write schema to the Arrow schema used by Lance.
+ *
+ * <p>This method distinguishes a Hoodie VECTOR from an ordinary Flink ARRAY
and encodes
+ * top-level FLOAT/DOUBLE vectors as Arrow FixedSizeList fields.
+ */
+ public static Schema toArrowSchema(HoodieSchema hoodieSchema) {
+ HoodieSchema recordSchema = hoodieSchema.getNonNullType();
+ RowType rowType = HoodieSchemaConverter.convertToRowType(recordSchema);
+ List<HoodieSchemaField> hoodieFields = recordSchema.getFields();
+
List<Field> fields = new ArrayList<>(rowType.getFieldCount());
- for (RowType.RowField field : rowType.getFields()) {
- fields.add(toArrowField(field.getName(), field.getType()));
+ for (int i = 0; i < rowType.getFieldCount(); i++) {
+ RowType.RowField rowField = rowType.getFields().get(i);
+ fields.add(toArrowField(
+ rowField.getName(), rowField.getType(),
hoodieFields.get(i).schema().getNonNullType()));
}
return new Schema(fields);
}
@@ -180,6 +197,36 @@ public final class HoodieFlinkLanceArrowUtils {
return new Field(name, new FieldType(type.isNullable(), toArrowType(type),
null), children);
}
+ private static Field toArrowField(String name, LogicalType type,
HoodieSchema hoodieSchema) {
+ if (hoodieSchema.getType() != HoodieSchemaType.VECTOR) {
+ return toArrowField(name, type);
+ }
+
+ HoodieSchema.Vector vectorSchema = (HoodieSchema.Vector) hoodieSchema;
+ validateLanceVector(name, vectorSchema);
+ ValidationUtils.checkArgument(type instanceof ArrayType,
+ "VECTOR column '" + name + "' must map to a Flink ARRAY type");
+ ArrayType arrayType = (ArrayType) type;
+ Field elementField = toArrowField("element", arrayType.getElementType());
+ return new Field(
+ name,
+ new FieldType(
+ type.isNullable(),
+ new ArrowType.FixedSizeList(vectorSchema.getDimension()),
+ null),
+ Collections.singletonList(elementField));
+ }
+
+ private static void validateLanceVector(String fieldName,
HoodieSchema.Vector vectorSchema) {
+ HoodieSchema.Vector.VectorElementType elementType =
vectorSchema.getVectorElementType();
+ if (elementType != HoodieSchema.Vector.VectorElementType.FLOAT
+ && elementType != HoodieSchema.Vector.VectorElementType.DOUBLE) {
+ throw new HoodieNotSupportedException(
+ "Lance base-file format currently supports FLOAT/DOUBLE VECTOR
columns only; "
+ + "got element type " + elementType + " for field '" + fieldName
+ "'");
+ }
+ }
+
private static ArrowType toArrowType(LogicalType type) {
switch (type.getTypeRoot()) {
case BOOLEAN:
diff --git
a/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/lance/LanceRowDataWriter.java
b/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/lance/LanceRowDataWriter.java
index 41c2481d5987..7d85b887fb05 100644
---
a/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/lance/LanceRowDataWriter.java
+++
b/hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/io/storage/row/lance/LanceRowDataWriter.java
@@ -18,7 +18,12 @@
package org.apache.hudi.io.storage.row.lance;
+import org.apache.hudi.common.schema.HoodieSchema;
+import org.apache.hudi.common.schema.HoodieSchemaField;
+import org.apache.hudi.common.schema.HoodieSchemaType;
+import org.apache.hudi.common.schema.HoodieSchemaUtils;
import org.apache.hudi.exception.HoodieNotSupportedException;
+import org.apache.hudi.util.HoodieSchemaConverter;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
@@ -34,6 +39,7 @@ import org.apache.arrow.vector.TimeStampVector;
import org.apache.arrow.vector.TinyIntVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.complex.FixedSizeListVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.flink.table.data.ArrayData;
@@ -64,10 +70,16 @@ public class LanceRowDataWriter {
private final FieldWriter[] fieldWriters;
- public LanceRowDataWriter(RowType rowType, List<FieldVector> vectors,
boolean utcTimestamp) {
+ public LanceRowDataWriter(
+ HoodieSchema hoodieSchema, List<FieldVector> vectors, boolean
utcTimestamp) {
+ HoodieSchema recordSchema = hoodieSchema.getNonNullType();
+ RowType rowType = HoodieSchemaConverter.convertToRowType(recordSchema);
+ List<HoodieSchemaField> hoodieFields = recordSchema.getFields();
this.fieldWriters = new FieldWriter[rowType.getFieldCount()];
for (int i = 0; i < fieldWriters.length; i++) {
- fieldWriters[i] = createWriter(rowType.getTypeAt(i), vectors.get(i),
utcTimestamp);
+ HoodieSchema fieldSchema = hoodieFields.get(i).schema().getNonNullType();
+ fieldWriters[i] = createWriter(
+ rowType.getTypeAt(i), fieldSchema, vectors.get(i), utcTimestamp,
hoodieFields.get(i).name());
}
}
@@ -77,7 +89,12 @@ public class LanceRowDataWriter {
}
}
- private static FieldWriter createWriter(LogicalType type, FieldVector
vector, boolean utcTimestamp) {
+ private static FieldWriter createWriter(
+ LogicalType type,
+ HoodieSchema hoodieSchema,
+ FieldVector vector,
+ boolean utcTimestamp,
+ String fieldPath) {
switch (type.getTypeRoot()) {
case BOOLEAN:
return new BooleanWriter((BitVector) vector);
@@ -113,15 +130,36 @@ public class LanceRowDataWriter {
StructVector structVector = (StructVector) vector;
FieldWriter[] fieldWriters = new FieldWriter[rowType.getFieldCount()];
for (int i = 0; i < fieldWriters.length; i++) {
+ String childName = rowType.getFieldNames().get(i);
fieldWriters[i] = createWriter(
- rowType.getTypeAt(i), (FieldVector)
structVector.getChildByOrdinal(i), utcTimestamp);
+ rowType.getTypeAt(i),
+ HoodieSchemaUtils.getFieldSchema(hoodieSchema,
childName).getNonNullType(),
+ (FieldVector) structVector.getChildByOrdinal(i),
+ utcTimestamp,
+ fieldPath + "." + childName);
}
return new RowWriter(structVector, fieldWriters);
case ARRAY:
ArrayType arrayType = (ArrayType) type;
+ if (hoodieSchema != null && hoodieSchema.getType() ==
HoodieSchemaType.VECTOR) {
+ HoodieSchema.Vector vectorSchema = (HoodieSchema.Vector)
hoodieSchema;
+ FixedSizeListVector fixedSizeListVector = (FixedSizeListVector)
vector;
+ FieldWriter vectorElementWriter = createWriter(
+ arrayType.getElementType(),
+ null,
+ fixedSizeListVector.getDataVector(),
+ utcTimestamp,
+ fieldPath + "[]");
+ return new VectorWriter(
+ fixedSizeListVector, vectorElementWriter,
vectorSchema.getDimension(), fieldPath);
+ }
ListVector listVector = (ListVector) vector;
FieldWriter elementWriter = createWriter(
- arrayType.getElementType(), listVector.getDataVector(),
utcTimestamp);
+ arrayType.getElementType(),
+ hoodieSchema.getElementType().getNonNullType(),
+ listVector.getDataVector(),
+ utcTimestamp,
+ fieldPath + "[]");
return new ArrayWriter(listVector, elementWriter);
default:
throw unsupported(type);
@@ -468,6 +506,44 @@ public class LanceRowDataWriter {
}
}
+ private static class VectorWriter extends FieldWriter {
+ private final FixedSizeListVector vector;
+ private final FieldWriter elementWriter;
+ private final int dimension;
+ private final String fieldPath;
+
+ private VectorWriter(
+ FixedSizeListVector vector, FieldWriter elementWriter, int dimension,
String fieldPath) {
+ super(vector);
+ this.vector = vector;
+ this.elementWriter = elementWriter;
+ this.dimension = dimension;
+ this.fieldPath = fieldPath;
+ }
+
+ @Override
+ void writeNonNull(RowData row, int ordinal, int rowId) {
+ writeVector(row.getArray(ordinal), rowId);
+ }
+
+ @Override
+ void writeNonNull(ArrayData array, int ordinal, int rowId) {
+ writeVector(array.getArray(ordinal), rowId);
+ }
+
+ private void writeVector(ArrayData array, int rowId) {
+ if (array.size() != dimension) {
+ throw new IllegalArgumentException(
+ "VECTOR column '" + fieldPath + "' requires dimension " + dimension
+ + " but row contains " + array.size() + " elements");
+ }
+ int startIndex = vector.startNewValue(rowId);
+ for (int i = 0; i < dimension; i++) {
+ elementWriter.write(array, i, startIndex + i);
+ }
+ }
+ }
+
private static long timestampToMicros(TimestampData timestampData, int
precision, boolean utcTimestamp) {
long millis = utcTimestamp ? timestampData.getMillisecond() :
timestampData.toTimestamp().getTime();
return precision > 3 && utcTimestamp
diff --git
a/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestHoodieFlinkLanceArrowUtils.java
b/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestHoodieFlinkLanceArrowUtils.java
index 007148a4df97..dc0b20839c58 100644
---
a/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestHoodieFlinkLanceArrowUtils.java
+++
b/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestHoodieFlinkLanceArrowUtils.java
@@ -18,30 +18,38 @@
package org.apache.hudi.io.storage.row;
+import org.apache.hudi.common.schema.HoodieSchema;
import org.apache.hudi.exception.HoodieNotSupportedException;
import org.apache.hudi.io.storage.row.lance.HoodieFlinkLanceArrowUtils;
+import org.apache.hudi.util.HoodieSchemaConverter;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.TimeStampMicroVector;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.data.TimestampData;
import org.apache.flink.table.types.logical.ArrayType;
+import org.apache.flink.table.types.logical.DoubleType;
+import org.apache.flink.table.types.logical.FloatType;
import org.apache.flink.table.types.logical.IntType;
import org.apache.flink.table.types.logical.LocalZonedTimestampType;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.MapType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.TimestampType;
+import org.apache.flink.table.types.logical.TinyIntType;
import org.apache.flink.table.types.logical.VarCharType;
import org.junit.jupiter.api.Test;
import java.util.Collections;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -56,9 +64,10 @@ public class TestHoodieFlinkLanceArrowUtils {
RowType rowType = RowType.of(
new LogicalType[] {new TimestampType(6), new
LocalZonedTimestampType(6)},
new String[] {"timestamp", "local_timestamp"});
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(rowType,
"record");
RowType roundTripped = HoodieFlinkLanceArrowUtils.toRowType(
- HoodieFlinkLanceArrowUtils.toArrowSchema(rowType));
+ HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema));
assertInstanceOf(TimestampType.class, roundTripped.getTypeAt(0));
assertInstanceOf(LocalZonedTimestampType.class, roundTripped.getTypeAt(1));
@@ -86,22 +95,69 @@ public class TestHoodieFlinkLanceArrowUtils {
@Test
public void testNestedSchemaRoundTrip() {
RowType rowType = nestedRowType();
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(rowType,
"record");
RowType roundTripped = HoodieFlinkLanceArrowUtils.toRowType(
- HoodieFlinkLanceArrowUtils.toArrowSchema(rowType));
+ HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema));
assertEquals(rowType, roundTripped);
}
@Test
public void testRejectsMapTypeWhenWritingSchema() {
- HoodieNotSupportedException exception =
assertThrows(HoodieNotSupportedException.class,
- () -> HoodieFlinkLanceArrowUtils.toArrowSchema(RowType.of(
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(
+ RowType.of(
new LogicalType[] {new MapType(new VarCharType(), new IntType())},
- new String[] {"attributes"})));
+ new String[] {"attributes"}),
+ "record");
+ HoodieNotSupportedException exception =
assertThrows(HoodieNotSupportedException.class,
+ () -> HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema));
assertTrue(exception.getMessage().contains("Flink Lance base-file support
currently supports primitive, ROW, and ARRAY columns;"));
}
+ @Test
+ public void testVectorColumnsUseFixedSizeListAndOrdinaryArraysRemainLists() {
+ RowType rowType = RowType.of(
+ new LogicalType[] {
+ new ArrayType(new FloatType()),
+ new ArrayType(new FloatType(false)),
+ new ArrayType(new DoubleType(false))},
+ new String[] {"values", "embedding", "scores"});
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(
+ rowType, "vector_record", "embedding:2,scores:3");
+
+ Schema arrowSchema =
HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema);
+
+ assertInstanceOf(ArrowType.List.class,
arrowSchema.findField("values").getType());
+ assertFixedSizeList(arrowSchema.findField("embedding"), 2);
+ assertFixedSizeList(arrowSchema.findField("scores"), 3);
+ }
+
+ @Test
+ public void testRejectsInt8VectorForLance() {
+ RowType rowType = RowType.of(
+ new LogicalType[] {new ArrayType(new TinyIntType(false))},
+ new String[] {"embedding"});
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(
+ rowType, "vector_record", "embedding:2");
+
+ HoodieNotSupportedException exception = assertThrows(
+ HoodieNotSupportedException.class,
+ () -> HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema));
+
+ assertTrue(exception.getMessage().contains("FLOAT/DOUBLE VECTOR columns
only"));
+ assertTrue(exception.getMessage().contains("embedding"));
+ }
+
+ private static void assertFixedSizeList(Field field, int dimension) {
+ ArrowType.FixedSizeList fixedSizeList =
+ assertInstanceOf(ArrowType.FixedSizeList.class, field.getType());
+ assertEquals(dimension, fixedSizeList.getListSize());
+ assertEquals("element", field.getChildren().get(0).getName());
+ assertFalse(field.getChildren().get(0).isNullable());
+ assertTrue(field.getMetadata().isEmpty());
+ }
+
private static RowType nestedRowType() {
RowType profileType = RowType.of(
new LogicalType[] {new VarCharType(), new IntType()},
diff --git
a/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestHoodieRowDataLanceWriter.java
b/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestHoodieRowDataLanceWriter.java
new file mode 100644
index 000000000000..b7303a61e93f
--- /dev/null
+++
b/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestHoodieRowDataLanceWriter.java
@@ -0,0 +1,103 @@
+/*
+ * 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.hudi.io.storage.row;
+
+import org.apache.hudi.common.engine.TaskContextSupplier;
+import org.apache.hudi.common.schema.HoodieSchema;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.storage.StoragePath;
+import org.apache.hudi.util.HoodieSchemaConverter;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.Float4Vector;
+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.table.data.GenericArrayData;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.types.logical.ArrayType;
+import org.apache.flink.table.types.logical.FloatType;
+import org.apache.flink.table.types.logical.IntType;
+import org.apache.flink.table.types.logical.LogicalType;
+import org.apache.flink.table.types.logical.RowType;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.lance.file.LanceFileReader;
+
+import java.nio.file.Path;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+
+/** Tests for {@link HoodieRowDataLanceWriter}. */
+public class TestHoodieRowDataLanceWriter {
+
+ @TempDir
+ Path tempDir;
+
+ @Test
+ public void testWritesVectorDataAndFooterMetadata() throws Exception {
+ RowType rowType = RowType.of(
+ new LogicalType[] {new IntType(false), new ArrayType(new
FloatType(false))},
+ new String[] {"id", "embedding"});
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(
+ rowType, "vector_record", "embedding:2");
+ StoragePath path = new
StoragePath(tempDir.resolve("vector.lance").toUri());
+
+ try (HoodieRowDataLanceWriter writer = new HoodieRowDataLanceWriter(
+ path,
+ hoodieSchema,
+ "001",
+ mock(TaskContextSupplier.class),
+ Option.empty(),
+ 128 * 1024 * 1024L,
+ 64 * 1024 * 1024L,
+ 16 * 1024 * 1024L,
+ true,
+ false,
+ false)) {
+ writer.writeRow("key1", GenericRowData.of(
+ 1, new GenericArrayData(new Object[] {1.25F, 2.5F})));
+ }
+
+ try (BufferAllocator allocator = new RootAllocator();
+ LanceFileReader reader = LanceFileReader.open(path.toString(),
allocator);
+ ArrowReader arrowReader = reader.readAll(null, null,
Integer.MAX_VALUE)) {
+ assertEquals(1, reader.numRows());
+ Map<String, String> metadata = reader.schema().getCustomMetadata();
+ assertEquals("embedding:VECTOR(2)",
+ metadata.get(HoodieSchema.VECTOR_COLUMNS_METADATA_KEY));
+ ArrowType.FixedSizeList fixedSizeList = assertInstanceOf(
+ ArrowType.FixedSizeList.class,
reader.schema().findField("embedding").getType());
+ assertEquals(2, fixedSizeList.getListSize());
+
+ assertTrue(arrowReader.loadNextBatch());
+ VectorSchemaRoot root = arrowReader.getVectorSchemaRoot();
+ FixedSizeListVector vector = (FixedSizeListVector)
root.getVector("embedding");
+ Float4Vector elements = (Float4Vector) vector.getDataVector();
+ assertEquals(1.25F, elements.get(0));
+ assertEquals(2.5F, elements.get(1));
+ }
+ }
+}
diff --git
a/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestLanceRowDataWriter.java
b/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestLanceRowDataWriter.java
index d43c2ad6d9b1..19e28f345224 100644
---
a/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestLanceRowDataWriter.java
+++
b/hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/io/storage/row/TestLanceRowDataWriter.java
@@ -18,15 +18,20 @@
package org.apache.hudi.io.storage.row;
+import org.apache.hudi.common.schema.HoodieSchema;
import org.apache.hudi.exception.HoodieNotSupportedException;
import org.apache.hudi.io.storage.row.lance.HoodieFlinkLanceArrowUtils;
import org.apache.hudi.io.storage.row.lance.LanceRowDataWriter;
+import org.apache.hudi.util.HoodieSchemaConverter;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.Float4Vector;
+import org.apache.arrow.vector.Float8Vector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.TimeStampMicroVector;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.complex.FixedSizeListVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.types.TimeUnit;
@@ -65,6 +70,8 @@ import java.util.Collections;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -75,14 +82,15 @@ public class TestLanceRowDataWriter {
@Test
public void testPrimitiveValuesAndNulls() {
- RowType rowType = primitiveRowType();
+ HoodieSchema hoodieSchema =
HoodieSchemaConverter.convertToSchema(primitiveRowType(), "record");
+ RowType rowType =
HoodieSchemaConverter.convertToRowType(hoodieSchema.getNonNullType());
TimestampData timestamp = TimestampData.fromEpochMillis(1234L, 567000);
TimestampData localTimestamp = TimestampData.fromEpochMillis(5678L,
123000);
DecimalData decimal = DecimalData.fromBigDecimal(new
BigDecimal("12345.67"), 10, 2);
GenericRowData values = GenericRowData.of(
true,
- (byte) 12,
- (short) 1234,
+ 12,
+ 1234,
123456,
20000,
12345678,
@@ -100,9 +108,9 @@ public class TestLanceRowDataWriter {
try (BufferAllocator allocator = new RootAllocator();
VectorSchemaRoot root = VectorSchemaRoot.create(
- HoodieFlinkLanceArrowUtils.toArrowSchema(rowType), allocator)) {
+ HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema),
allocator)) {
root.allocateNew();
- LanceRowDataWriter writer = new LanceRowDataWriter(rowType,
root.getFieldVectors(), true);
+ LanceRowDataWriter writer = new LanceRowDataWriter(hoodieSchema,
root.getFieldVectors(), true);
writer.write(values, 0);
writer.write(nulls, 1);
root.getFieldVectors().forEach(vector -> vector.setValueCount(2));
@@ -110,8 +118,8 @@ public class TestLanceRowDataWriter {
RowData actual = HoodieFlinkLanceArrowUtils.toRowData(rowType,
root.getFieldVectors(), 0);
assertEquals(true, actual.getBoolean(0));
- assertEquals((byte) 12, actual.getByte(1));
- assertEquals((short) 1234, actual.getShort(2));
+ assertEquals(12, actual.getInt(1));
+ assertEquals(1234, actual.getInt(2));
assertEquals(123456, actual.getInt(3));
assertEquals(20000, actual.getInt(4));
assertEquals(12345678, actual.getInt(5));
@@ -144,10 +152,11 @@ public class TestLanceRowDataWriter {
FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND,
null)),
allocator)) {
RowType rowType = RowType.of(new LogicalType[] {new TimestampType(6)},
new String[] {"ts"});
- new LanceRowDataWriter(rowType, Collections.singletonList(vector),
true).write(rowData, 0);
+ HoodieSchema hoodieSchema =
HoodieSchemaConverter.convertToSchema(rowType, "record");
+ new LanceRowDataWriter(hoodieSchema, Collections.singletonList(vector),
true).write(rowData, 0);
assertEquals(1234567L, vector.get(0));
- new LanceRowDataWriter(rowType, Collections.singletonList(vector),
false).write(rowData, 1);
+ new LanceRowDataWriter(hoodieSchema, Collections.singletonList(vector),
false).write(rowData, 1);
assertEquals(timestampData.toTimestamp().getTime() * 1000L,
vector.get(1));
}
}
@@ -155,6 +164,7 @@ public class TestLanceRowDataWriter {
@Test
public void testNestedValueRoundTripAndListValueCount() {
RowType rowType = nestedRowType();
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(rowType,
"record");
GenericRowData first = GenericRowData.of(
GenericRowData.of(
StringData.fromString("alice"),
@@ -180,9 +190,9 @@ public class TestLanceRowDataWriter {
try (BufferAllocator allocator = new RootAllocator();
VectorSchemaRoot root = VectorSchemaRoot.create(
- HoodieFlinkLanceArrowUtils.toArrowSchema(rowType), allocator)) {
+ HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema),
allocator)) {
root.allocateNew();
- LanceRowDataWriter writer = new LanceRowDataWriter(rowType,
root.getFieldVectors(), true);
+ LanceRowDataWriter writer = new LanceRowDataWriter(hoodieSchema,
root.getFieldVectors(), true);
writer.write(first, 0);
writer.write(second, 1);
writer.write(third, 2);
@@ -205,18 +215,92 @@ public class TestLanceRowDataWriter {
@Test
public void testRejectsMapType() {
MapType mapType = new MapType(new VarCharType(), new IntType());
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(
+ RowType.of(new LogicalType[] {mapType}, new String[] {"attributes"}),
"record");
try (BufferAllocator allocator = new RootAllocator();
IntVector vector = new IntVector("attributes", allocator)) {
HoodieNotSupportedException exception =
assertThrows(HoodieNotSupportedException.class,
() -> new LanceRowDataWriter(
- RowType.of(new LogicalType[] {mapType}, new String[]
{"attributes"}),
+ hoodieSchema,
Collections.singletonList(vector),
true));
assertTrue(exception.getMessage().contains(mapUnsupportedMessage()));
}
}
+ @Test
+ public void testWritesFloatAndDoubleVectors() {
+ RowType rowType = RowType.of(
+ new LogicalType[] {
+ new ArrayType(new FloatType(false)),
+ new ArrayType(new DoubleType(false)),
+ new ArrayType(new IntType())},
+ new String[] {"embedding", "scores", "values"});
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(
+ rowType, "vector_record", "embedding:2,scores:3");
+ GenericRowData first = GenericRowData.of(
+ new GenericArrayData(new Object[] {1.25F, 2.5F}),
+ new GenericArrayData(new Object[] {3.0D, 4.0D, 5.0D}),
+ new GenericArrayData(new Object[] {6, null}));
+ GenericRowData second = GenericRowData.of(
+ null,
+ new GenericArrayData(new Object[] {7.0D, 8.0D, 9.0D}),
+ new GenericArrayData(new Object[0]));
+
+ try (BufferAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root = VectorSchemaRoot.create(
+ HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema),
allocator)) {
+ root.allocateNew();
+ LanceRowDataWriter writer = new LanceRowDataWriter(
+ hoodieSchema, root.getFieldVectors(), true);
+ writer.write(first, 0);
+ writer.write(second, 1);
+ root.getFieldVectors().forEach(vector -> vector.setValueCount(2));
+ root.setRowCount(2);
+
+ FixedSizeListVector floatVector = (FixedSizeListVector)
root.getVector("embedding");
+ Float4Vector floatElements = (Float4Vector) floatVector.getDataVector();
+ assertFalse(floatVector.isNull(0));
+ assertTrue(floatVector.isNull(1));
+ assertEquals(1.25F, floatElements.get(0));
+ assertEquals(2.5F, floatElements.get(1));
+
+ FixedSizeListVector doubleVector = (FixedSizeListVector)
root.getVector("scores");
+ Float8Vector doubleElements = (Float8Vector)
doubleVector.getDataVector();
+ assertEquals(3, doubleVector.getListSize());
+ assertEquals(3.0D, doubleElements.get(0));
+ assertEquals(5.0D, doubleElements.get(2));
+ assertEquals(7.0D, doubleElements.get(3));
+ assertEquals(9.0D, doubleElements.get(5));
+
+ assertInstanceOf(ListVector.class, root.getVector("values"));
+ }
+ }
+
+ @Test
+ public void testRejectsInvalidVectorDimension() {
+ RowType rowType = RowType.of(
+ new LogicalType[] {new ArrayType(new FloatType(false))},
+ new String[] {"embedding"});
+ HoodieSchema hoodieSchema = HoodieSchemaConverter.convertToSchema(
+ rowType, "vector_record", "embedding:2");
+
+ try (BufferAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root = VectorSchemaRoot.create(
+ HoodieFlinkLanceArrowUtils.toArrowSchema(hoodieSchema),
allocator)) {
+ root.allocateNew();
+ LanceRowDataWriter writer = new LanceRowDataWriter(
+ hoodieSchema, root.getFieldVectors(), true);
+
+ IllegalArgumentException dimensionError = assertThrows(
+ IllegalArgumentException.class,
+ () -> writer.write(GenericRowData.of(
+ new GenericArrayData(new Object[] {1.0F})), 0));
+ assertTrue(dimensionError.getMessage().contains("requires dimension 2"));
+ }
+ }
+
private static String mapUnsupportedMessage() {
return "Flink Lance base-file support currently supports primitive, ROW,
and ARRAY columns;";
}
