Skip to content

Add vector type support for Direct Access API inserts #3733

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Nov 7, 2025
Merged

Add vector type support for Direct Access API inserts #3733

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
9360f3a
Add vector type support for Direct Access API inserts
arnaud-lacurie Nov 7, 2025
b5eae70
Fix checks
arnaud-lacurie Nov 7, 2025
d0a43e8
Fix teamscale issues
arnaud-lacurie Nov 7, 2025
b1b0373
Add defensive assert
arnaud-lacurie Nov 7, 2025
a469ea2
Fix missing import
arnaud-lacurie Nov 7, 2025
73ab83a
Fix teamscale
arnaud-lacurie Nov 7, 2025
78e494e
Merge remote-tracking branch 'upstream/main' into vector
hatyo Nov 7, 2025
0c67274
Adapt test behavior.
hatyo Nov 7, 2025
b9b095b
Add tests for inserting `null` vector.
hatyo Nov 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/util/VectorUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,11 @@ public static int getVectorPrecision(@Nonnull final RealVector vector) {
@Nonnull
public static RealVector parseVector(@Nonnull final ByteString byteString, @Nonnull final Type.Vector vectorType) {
final var precision = vectorType.getPrecision();
return parseVector(byteString, precision);
}

@Nonnull
public static RealVector parseVector(@Nonnull final ByteString byteString, int precision) {
if (precision == 16) {
return HalfRealVector.fromBytes(byteString.toByteArray());
}
Expand All @@ -56,7 +61,7 @@ public static RealVector parseVector(@Nonnull final ByteString byteString, @Nonn
if (precision == 64) {
return DoubleRealVector.fromBytes(byteString.toByteArray());
}
throw new RecordCoreException("unexpected vector type " + vectorType);
throw new RecordCoreException("unexpected vector precision " + precision);
}


Expand Down
37 changes: 13 additions & 24 deletions ...tional-core/src/main/java/com/apple/foundationdb/relational/recordlayer/MessageTuple.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,10 @@
package com.apple.foundationdb.relational.recordlayer;

import com.apple.foundationdb.annotation.API;
import com.apple.foundationdb.linear.DoubleRealVector;
import com.apple.foundationdb.linear.FloatRealVector;
import com.apple.foundationdb.linear.HalfRealVector;
import com.apple.foundationdb.linear.RealVector;
import com.apple.foundationdb.record.RecordCoreException;
import com.apple.foundationdb.record.RecordMetaDataOptionsProto;
import com.apple.foundationdb.record.TupleFieldsProto;
import com.apple.foundationdb.record.metadata.expressions.TupleFieldsHelper;
import com.apple.foundationdb.record.util.VectorUtils;
import com.apple.foundationdb.relational.api.exceptions.InvalidColumnReferenceException;
import com.google.protobuf.ByteString;
import com.google.protobuf.DescriptorProtos;
Expand Down Expand Up @@ -61,9 +57,13 @@ public Object getObject(int position) throws InvalidColumnReferenceException {
final var fieldOptions = fieldDescriptor.getOptions().getExtension(RecordMetaDataOptionsProto.field);
final var fieldValue = message.getField(message.getDescriptorForType().getFields().get(position));
if (fieldOptions.hasVectorOptions()) {
final var precision = fieldOptions.getVectorOptions().getPrecision();
final var byteStringFieldValue = (ByteString)fieldValue;
return getVectorFromBytes(byteStringFieldValue, precision);
if (byteStringFieldValue.isEmpty()) {
return null;
} else {
final var precision = fieldOptions.getVectorOptions().getPrecision();
return VectorUtils.parseVector(byteStringFieldValue, precision);
}
}
if (fieldDescriptor.isRepeated()) {
final var list = (List<?>) fieldValue;
Expand All @@ -76,21 +76,6 @@ public Object getObject(int position) throws InvalidColumnReferenceException {
}
}

@Nonnull
private static RealVector getVectorFromBytes(@Nonnull final ByteString byteString, int precision) {
final var bytes = byteString.toByteArray();
if (precision == 64) {
return DoubleRealVector.fromBytes(bytes);
}
if (precision == 32) {
return FloatRealVector.fromBytes(bytes);
}
if (precision == 16) {
return HalfRealVector.fromBytes(bytes);
}
throw new RecordCoreException("unexpected vector precision " + precision);
}

public static Object sanitizeField(@Nonnull final Object field, @Nonnull final DescriptorProtos.FieldOptions fieldOptions) {
if (field instanceof Message && ((Message) field).getDescriptorForType().equals(TupleFieldsProto.UUID.getDescriptor())) {
return TupleFieldsHelper.fromProto((Message) field, TupleFieldsProto.UUID.getDescriptor());
Expand All @@ -102,8 +87,12 @@ public static Object sanitizeField(@Nonnull final Object field, @Nonnull final D
final var byteString = (ByteString) field;
final var fieldVectorOptionsMaybe = fieldOptions.getExtension(RecordMetaDataOptionsProto.field);
if (fieldVectorOptionsMaybe.hasVectorOptions()) {
final var precision = fieldVectorOptionsMaybe.getVectorOptions().getPrecision();
return getVectorFromBytes(byteString, precision);
if (byteString.isEmpty()) {
return null;
} else {
final var precision = fieldVectorOptionsMaybe.getVectorOptions().getPrecision();
return VectorUtils.parseVector(byteString, precision);
}
}
return byteString.toByteArray();
}
Expand Down
63 changes: 49 additions & 14 deletions ...nal-core/src/main/java/com/apple/foundationdb/relational/recordlayer/RecordTypeTable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,14 @@
package com.apple.foundationdb.relational.recordlayer;

import com.apple.foundationdb.annotation.API;
import com.apple.foundationdb.linear.AbstractRealVector;
import com.apple.foundationdb.linear.DoubleRealVector;
import com.apple.foundationdb.linear.FloatRealVector;
import com.apple.foundationdb.linear.HalfRealVector;
import com.apple.foundationdb.record.RecordCoreException;
import com.apple.foundationdb.record.RecordCursor;
import com.apple.foundationdb.record.RecordMetaData;
import com.apple.foundationdb.record.RecordMetaDataOptionsProto;
import com.apple.foundationdb.record.TupleRange;
import com.apple.foundationdb.record.metadata.MetaDataException;
import com.apple.foundationdb.record.metadata.RecordType;
Expand Down Expand Up @@ -57,6 +62,7 @@
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.TreeMap;
import java.util.function.Function;
Expand Down Expand Up @@ -195,7 +201,7 @@
final var maybeEnumValue = struct.getString(i + 1);
if (maybeEnumValue != null) {
final var valueDescriptor = fd.getEnumType().findValueByName(maybeEnumValue);
Assert.thatUnchecked(valueDescriptor != null, ErrorCode.CANNOT_CONVERT_TYPE, "Invalid enum value: %s", maybeEnumValue);
Assert.that(valueDescriptor != null, ErrorCode.CANNOT_CONVERT_TYPE, "Invalid enum value: %s", maybeEnumValue);
builder.setField(fd, valueDescriptor);
}
continue;
Expand All @@ -209,7 +215,11 @@
case STRING:
final var obj = struct.getObject(i + 1);
if (obj != null) {
builder.setField(fd, obj);
try {
builder.setField(fd, obj);

Check warning on line 219 in fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/RecordTypeTable.java

View check run for this annotation

fdb.teamscale.io / Teamscale | Findings

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/RecordTypeTable.java#L219 

This method is quite long [0] and deeply nested in multiple places [1, 2, 3] which can make it hard to understand and maintain. Consider extracting helper methods or reducing the nesting by using early breaks or returns.

[0] https://fdb.teamscale.io/findings/details/foundationdb-fdb-record-layer?t=FORK_MR%2F3733%2Farnaud-lacurie%2Fvector%3AHEAD&id=91DE432808F3C9A882BBCED7300CA61E
[1] https://fdb.teamscale.io/findings/details/foundationdb-fdb-record-layer?t=FORK_MR%2F3733%2Farnaud-lacurie%2Fvector%3AHEAD&id=A3C2C13102F3BF22EB2D7739FC7B3843
[2] https://fdb.teamscale.io/findings/details/foundationdb-fdb-record-layer?t=FORK_MR%2F3733%2Farnaud-lacurie%2Fvector%3AHEAD&id=7EAD693963C48407D4FE9B303B0F2CE4
[3] https://fdb.teamscale.io/findings/details/foundationdb-fdb-record-layer?t=FORK_MR%2F3733%2Farnaud-lacurie%2Fvector%3AHEAD&id=B99191BCC27D59E978995FD56E93EC95
} catch (IllegalArgumentException ex) {
throw new RelationalException("Unexpected Column type " + struct.getMetaData().getColumnTypeName(i) + " for column " + columnName, ErrorCode.CANNOT_CONVERT_TYPE, ex);
}
}
break;
case BYTES:
Expand Down Expand Up @@ -237,24 +247,49 @@
}
}
} else {
if (fd.getType() == Descriptors.FieldDescriptor.Type.MESSAGE) {
Assert.thatUnchecked(NullableArrayUtils.isWrappedArrayDescriptor(fd.getMessageType()));
// wrap array in a struct and call toDynamicMessage again
final var wrapper = new ImmutableRowStruct(new ArrayRow(array), RelationalStructMetaData.of(
DataType.StructType.from("STRUCT", List.of(
DataType.StructType.Field.from(NullableArrayUtils.REPEATED_FIELD_NAME, array.getMetaData().asRelationalType(), 0)
), true)));
builder.setField(fd, toDynamicMessage(wrapper, fd.getMessageType()));
} else {
Assert.failUnchecked("Field Type expected to be of Type ARRAY but is actually " + fd.getType());
Assert.that(fd.getType() == Descriptors.FieldDescriptor.Type.MESSAGE, ErrorCode.CANNOT_CONVERT_TYPE,
"Field Type expected to be of Type ARRAY but is actually " + fd.getType());
Assert.that(NullableArrayUtils.isWrappedArrayDescriptor(fd.getMessageType()));
// wrap array in a struct and call toDynamicMessage again
final var wrapper = new ImmutableRowStruct(new ArrayRow(array), RelationalStructMetaData.of(
DataType.StructType.from("STRUCT", List.of(
DataType.StructType.Field.from(NullableArrayUtils.REPEATED_FIELD_NAME, array.getMetaData().asRelationalType(), 0)
), true)));
builder.setField(fd, toDynamicMessage(wrapper, fd.getMessageType()));
}
break;
case VECTOR:
final var vector = struct.getObject(i + 1);
if (vector != null) {
Assert.that(vector instanceof AbstractRealVector, ErrorCode.CANNOT_CONVERT_TYPE,
"Field Type expected to be of Type VECTOR but is actually " + fd.getType());
final var fieldOptionMaybe = Optional.ofNullable(fd.getOptions()).map(f -> f.getExtension(RecordMetaDataOptionsProto.field));
Assert.that(fieldOptionMaybe.isPresent() && fieldOptionMaybe.get().hasVectorOptions(), ErrorCode.CANNOT_CONVERT_TYPE, "Cannot insert non vector type into vector column");
final var vectorOptions = fieldOptionMaybe.get().getVectorOptions();
Assert.that(vectorOptions.getDimensions() == ((AbstractRealVector)vector).getNumDimensions(), ErrorCode.CANNOT_CONVERT_TYPE, "Wrong number of dimension for vector");
final int precision = vectorOptions.getPrecision();
switch (precision) {
case 16:
Assert.that(vector instanceof HalfRealVector, ErrorCode.CANNOT_CONVERT_TYPE, "Wrong precision for vector");
break;
case 32:
Assert.that(vector instanceof FloatRealVector, ErrorCode.CANNOT_CONVERT_TYPE, "Wrong precision for vector");
break;
case 64:
Assert.that(vector instanceof DoubleRealVector, ErrorCode.CANNOT_CONVERT_TYPE, "Wrong precision for vector");
break;
default:
Assert.fail(ErrorCode.INTERNAL_ERROR, "Unknown precision for vector");
break;
}
builder.setField(fd, ((AbstractRealVector) vector).getRawData());
}
break;
case NULL:
break;
default:
Assert.failUnchecked(ErrorCode.INTERNAL_ERROR, (String.format(Locale.ROOT, "Unexpected Column type <%s> for column <%s>",
struct.getMetaData().getColumnType(i), columnName)));
Assert.fail(ErrorCode.INTERNAL_ERROR, (String.format(Locale.ROOT, "Unexpected Column type <%s> for column <%s>",
struct.getMetaData().getColumnTypeName(i), columnName)));
break;
}
}
Expand Down
155 changes: 155 additions & 0 deletions ...al-core/src/test/java/com/apple/foundationdb/relational/recordlayer/InsertVectorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,155 @@
/*
* InsertTest.java
*
* This source file is part of the FoundationDB open source project
*
* Copyright 2021-2025 Apple Inc. and the FoundationDB project authors
*
* Licensed 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 com.apple.foundationdb.relational.recordlayer;

import com.apple.foundationdb.linear.DoubleRealVector;
import com.apple.foundationdb.linear.FloatRealVector;
import com.apple.foundationdb.linear.HalfRealVector;
import com.apple.foundationdb.relational.api.EmbeddedRelationalStruct;
import com.apple.foundationdb.relational.api.KeySet;
import com.apple.foundationdb.relational.api.Options;
import com.apple.foundationdb.relational.api.RelationalConnection;
import com.apple.foundationdb.relational.api.RelationalResultSet;
import com.apple.foundationdb.relational.api.RelationalStatement;
import com.apple.foundationdb.relational.api.RelationalStruct;
import com.apple.foundationdb.relational.api.exceptions.ErrorCode;
import com.apple.foundationdb.relational.utils.RelationalAssertions;
import com.apple.foundationdb.relational.utils.ResultSetAssert;
import com.apple.foundationdb.relational.utils.SimpleDatabaseRule;
import org.junit.jupiter.api.Order;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;

import java.sql.DriverManager;
import java.sql.SQLException;

public class InsertVectorTest {
@RegisterExtension
@Order(0)
public final EmbeddedRelationalExtension relationalExtension = new EmbeddedRelationalExtension();

@RegisterExtension
@Order(1)
public final SimpleDatabaseRule database = new SimpleDatabaseRule(InsertVectorTest.class, "CREATE TABLE V(PK INTEGER, V1 VECTOR(4, FLOAT), V2 VECTOR(3, HALF), V3 VECTOR(2, DOUBLE), PRIMARY KEY(PK))");

@Test
void insertNulls() throws SQLException {
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
conn.setSchema("TEST_SCHEMA");
try (RelationalStatement s = conn.createStatement()) {
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).build();
s.executeInsert("V", rec);
try (RelationalResultSet rs = s.executeGet("V", new KeySet().setKeyColumn("PK", 0), Options.NONE)) {
ResultSetAssert.assertThat(rs)
.hasNextRow()
.hasColumn("PK", 0)
.hasColumn("V1", null)
.hasColumn("V2", null)
.hasColumn("V3", null)
.hasNoNextRow();
}
}
}
}

@Test
void partialInsert() throws SQLException {
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
conn.setSchema("TEST_SCHEMA");
try (RelationalStatement s = conn.createStatement()) {
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).addObject("V1", new FloatRealVector(new float[]{1f, 2f, 3f, 4f})).build();
s.executeInsert("V", rec);
try (RelationalResultSet rs = s.executeGet("V", new KeySet().setKeyColumn("PK", 0), Options.NONE)) {
ResultSetAssert.assertThat(rs)
.hasNextRow()
.hasColumn("PK", 0)
.hasColumn("V1", new FloatRealVector(new float[]{1f, 2f, 3f, 4f}))
.hasColumn("V2", null)
.hasColumn("V3", null)
.hasNoNextRow();
}
}
}
}

@Test
void fullInsert() throws SQLException {
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
conn.setSchema("TEST_SCHEMA");
try (RelationalStatement s = conn.createStatement()) {
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0)
.addObject("V1", new FloatRealVector(new float[]{1f, 2f, 3f, 4f}))
.addObject("V2", new HalfRealVector(new int[]{1, 2, 3}))
.addObject("V3", new DoubleRealVector(new double[]{1d, 2d}))
.build();
s.executeInsert("V", rec);
try (RelationalResultSet rs = s.executeGet("V", new KeySet().setKeyColumn("PK", 0), Options.NONE)) {
ResultSetAssert.assertThat(rs)
.hasNextRow()
.hasColumn("PK", 0)
.hasColumn("V1", new FloatRealVector(new float[]{1f, 2f, 3f, 4f}))
.hasColumn("V2", new HalfRealVector(new int[]{1, 2, 3}))
.hasColumn("V3", new DoubleRealVector(new double[]{1d, 2d}))
.hasNoNextRow();
}
}
}
}

@Test
void insertWrongDimensionFails() throws SQLException {
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
conn.setSchema("TEST_SCHEMA");
try (RelationalStatement s = conn.createStatement()) {
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).addObject("V1", new FloatRealVector(new float[] {1f, 2f, 3f, 4f, 5f})).build();
RelationalAssertions.assertThrowsSqlException(
() -> s.executeInsert("V", rec))
.hasErrorCode(ErrorCode.CANNOT_CONVERT_TYPE);
}
}
}

@Test
void insertWrongPrecisionFails() throws SQLException {
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
conn.setSchema("TEST_SCHEMA");
try (RelationalStatement s = conn.createStatement()) {
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).addObject("V1", new DoubleRealVector(new double[] {1d, 2d, 3d, 4d})).build();
RelationalAssertions.assertThrowsSqlException(
() -> s.executeInsert("V", rec))
.hasErrorCode(ErrorCode.CANNOT_CONVERT_TYPE);
}
}
}

@Test
void insertWrongTypeFails() throws SQLException {
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
conn.setSchema("TEST_SCHEMA");
try (RelationalStatement s = conn.createStatement()) {
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).addInt("V1", 42).build();
RelationalAssertions.assertThrowsSqlException(
() -> s.executeInsert("V", rec))
.hasErrorCode(ErrorCode.CANNOT_CONVERT_TYPE);
}
}
}
}
Loading
Loading