Skip to content

Commit e2882db

Browse files
Add vector type support for Direct Access API inserts (#3733)
Enables inserting vector data types through the relational Direct Access API with full validation: - Type checking ensures vectors match column precision (FLOAT/16, HALF/32, DOUBLE/64) - Dimension validation verifies vector size matches schema definition - Null/empty vector handling for optional vector columns - Comprehensive test coverage via InsertVectorTest for all insert scenarios Also refactors vector parsing logic into VectorUtils.parseVector() to eliminate code duplication across MessageTuple and FRL. --------- Co-authored-by: Youssef Hatem <y_hatem@apple.com>
1 parent b62d08d commit e2882db

File tree

7 files changed

+295
-61
lines changed

7 files changed

+295
-61
lines changed

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/util/VectorUtils.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@ public static int getVectorPrecision(@Nonnull final RealVector vector) {
4747
@Nonnull
4848
public static RealVector parseVector(@Nonnull final ByteString byteString, @Nonnull final Type.Vector vectorType) {
4949
final var precision = vectorType.getPrecision();
50+
return parseVector(byteString, precision);
51+
}
52+
53+
@Nonnull
54+
public static RealVector parseVector(@Nonnull final ByteString byteString, int precision) {
5055
if (precision == 16) {
5156
return HalfRealVector.fromBytes(byteString.toByteArray());
5257
}
@@ -56,7 +61,7 @@ public static RealVector parseVector(@Nonnull final ByteString byteString, @Nonn
5661
if (precision == 64) {
5762
return DoubleRealVector.fromBytes(byteString.toByteArray());
5863
}
59-
throw new RecordCoreException("unexpected vector type " + vectorType);
64+
throw new RecordCoreException("unexpected vector precision " + precision);
6065
}
6166

6267

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/MessageTuple.java

Lines changed: 13 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,10 @@
2121
package com.apple.foundationdb.relational.recordlayer;
2222

2323
import com.apple.foundationdb.annotation.API;
24-
import com.apple.foundationdb.linear.DoubleRealVector;
25-
import com.apple.foundationdb.linear.FloatRealVector;
26-
import com.apple.foundationdb.linear.HalfRealVector;
27-
import com.apple.foundationdb.linear.RealVector;
28-
import com.apple.foundationdb.record.RecordCoreException;
2924
import com.apple.foundationdb.record.RecordMetaDataOptionsProto;
3025
import com.apple.foundationdb.record.TupleFieldsProto;
3126
import com.apple.foundationdb.record.metadata.expressions.TupleFieldsHelper;
27+
import com.apple.foundationdb.record.util.VectorUtils;
3228
import com.apple.foundationdb.relational.api.exceptions.InvalidColumnReferenceException;
3329
import com.google.protobuf.ByteString;
3430
import com.google.protobuf.DescriptorProtos;
@@ -61,9 +57,13 @@ public Object getObject(int position) throws InvalidColumnReferenceException {
6157
final var fieldOptions = fieldDescriptor.getOptions().getExtension(RecordMetaDataOptionsProto.field);
6258
final var fieldValue = message.getField(message.getDescriptorForType().getFields().get(position));
6359
if (fieldOptions.hasVectorOptions()) {
64-
final var precision = fieldOptions.getVectorOptions().getPrecision();
6560
final var byteStringFieldValue = (ByteString)fieldValue;
66-
return getVectorFromBytes(byteStringFieldValue, precision);
61+
if (byteStringFieldValue.isEmpty()) {
62+
return null;
63+
} else {
64+
final var precision = fieldOptions.getVectorOptions().getPrecision();
65+
return VectorUtils.parseVector(byteStringFieldValue, precision);
66+
}
6767
}
6868
if (fieldDescriptor.isRepeated()) {
6969
final var list = (List<?>) fieldValue;
@@ -76,21 +76,6 @@ public Object getObject(int position) throws InvalidColumnReferenceException {
7676
}
7777
}
7878

79-
@Nonnull
80-
private static RealVector getVectorFromBytes(@Nonnull final ByteString byteString, int precision) {
81-
final var bytes = byteString.toByteArray();
82-
if (precision == 64) {
83-
return DoubleRealVector.fromBytes(bytes);
84-
}
85-
if (precision == 32) {
86-
return FloatRealVector.fromBytes(bytes);
87-
}
88-
if (precision == 16) {
89-
return HalfRealVector.fromBytes(bytes);
90-
}
91-
throw new RecordCoreException("unexpected vector precision " + precision);
92-
}
93-
9479
public static Object sanitizeField(@Nonnull final Object field, @Nonnull final DescriptorProtos.FieldOptions fieldOptions) {
9580
if (field instanceof Message && ((Message) field).getDescriptorForType().equals(TupleFieldsProto.UUID.getDescriptor())) {
9681
return TupleFieldsHelper.fromProto((Message) field, TupleFieldsProto.UUID.getDescriptor());
@@ -102,8 +87,12 @@ public static Object sanitizeField(@Nonnull final Object field, @Nonnull final D
10287
final var byteString = (ByteString) field;
10388
final var fieldVectorOptionsMaybe = fieldOptions.getExtension(RecordMetaDataOptionsProto.field);
10489
if (fieldVectorOptionsMaybe.hasVectorOptions()) {
105-
final var precision = fieldVectorOptionsMaybe.getVectorOptions().getPrecision();
106-
return getVectorFromBytes(byteString, precision);
90+
if (byteString.isEmpty()) {
91+
return null;
92+
} else {
93+
final var precision = fieldVectorOptionsMaybe.getVectorOptions().getPrecision();
94+
return VectorUtils.parseVector(byteString, precision);
95+
}
10796
}
10897
return byteString.toByteArray();
10998
}

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

Lines changed: 49 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,14 @@
2121
package com.apple.foundationdb.relational.recordlayer;
2222

2323
import com.apple.foundationdb.annotation.API;
24+
import com.apple.foundationdb.linear.AbstractRealVector;
25+
import com.apple.foundationdb.linear.DoubleRealVector;
26+
import com.apple.foundationdb.linear.FloatRealVector;
27+
import com.apple.foundationdb.linear.HalfRealVector;
2428
import com.apple.foundationdb.record.RecordCoreException;
2529
import com.apple.foundationdb.record.RecordCursor;
2630
import com.apple.foundationdb.record.RecordMetaData;
31+
import com.apple.foundationdb.record.RecordMetaDataOptionsProto;
2732
import com.apple.foundationdb.record.TupleRange;
2833
import com.apple.foundationdb.record.metadata.MetaDataException;
2934
import com.apple.foundationdb.record.metadata.RecordType;
@@ -57,6 +62,7 @@
5762
import java.util.List;
5863
import java.util.Locale;
5964
import java.util.Map;
65+
import java.util.Optional;
6066
import java.util.Set;
6167
import java.util.TreeMap;
6268
import java.util.function.Function;
@@ -195,7 +201,7 @@ public static Message toDynamicMessage(RelationalStruct struct, Descriptors.Desc
195201
final var maybeEnumValue = struct.getString(i + 1);
196202
if (maybeEnumValue != null) {
197203
final var valueDescriptor = fd.getEnumType().findValueByName(maybeEnumValue);
198-
Assert.thatUnchecked(valueDescriptor != null, ErrorCode.CANNOT_CONVERT_TYPE, "Invalid enum value: %s", maybeEnumValue);
204+
Assert.that(valueDescriptor != null, ErrorCode.CANNOT_CONVERT_TYPE, "Invalid enum value: %s", maybeEnumValue);
199205
builder.setField(fd, valueDescriptor);
200206
}
201207
continue;
@@ -209,7 +215,11 @@ public static Message toDynamicMessage(RelationalStruct struct, Descriptors.Desc
209215
case STRING:
210216
final var obj = struct.getObject(i + 1);
211217
if (obj != null) {
212-
builder.setField(fd, obj);
218+
try {
219+
builder.setField(fd, obj);
220+
} catch (IllegalArgumentException ex) {
221+
throw new RelationalException("Unexpected Column type " + struct.getMetaData().getColumnTypeName(i) + " for column " + columnName, ErrorCode.CANNOT_CONVERT_TYPE, ex);
222+
}
213223
}
214224
break;
215225
case BYTES:
@@ -237,24 +247,49 @@ public static Message toDynamicMessage(RelationalStruct struct, Descriptors.Desc
237247
}
238248
}
239249
} else {
240-
if (fd.getType() == Descriptors.FieldDescriptor.Type.MESSAGE) {
241-
Assert.thatUnchecked(NullableArrayUtils.isWrappedArrayDescriptor(fd.getMessageType()));
242-
// wrap array in a struct and call toDynamicMessage again
243-
final var wrapper = new ImmutableRowStruct(new ArrayRow(array), RelationalStructMetaData.of(
244-
DataType.StructType.from("STRUCT", List.of(
245-
DataType.StructType.Field.from(NullableArrayUtils.REPEATED_FIELD_NAME, array.getMetaData().asRelationalType(), 0)
246-
), true)));
247-
builder.setField(fd, toDynamicMessage(wrapper, fd.getMessageType()));
248-
} else {
249-
Assert.failUnchecked("Field Type expected to be of Type ARRAY but is actually " + fd.getType());
250+
Assert.that(fd.getType() == Descriptors.FieldDescriptor.Type.MESSAGE, ErrorCode.CANNOT_CONVERT_TYPE,
251+
"Field Type expected to be of Type ARRAY but is actually " + fd.getType());
252+
Assert.that(NullableArrayUtils.isWrappedArrayDescriptor(fd.getMessageType()));
253+
// wrap array in a struct and call toDynamicMessage again
254+
final var wrapper = new ImmutableRowStruct(new ArrayRow(array), RelationalStructMetaData.of(
255+
DataType.StructType.from("STRUCT", List.of(
256+
DataType.StructType.Field.from(NullableArrayUtils.REPEATED_FIELD_NAME, array.getMetaData().asRelationalType(), 0)
257+
), true)));
258+
builder.setField(fd, toDynamicMessage(wrapper, fd.getMessageType()));
259+
}
260+
break;
261+
case VECTOR:
262+
final var vector = struct.getObject(i + 1);
263+
if (vector != null) {
264+
Assert.that(vector instanceof AbstractRealVector, ErrorCode.CANNOT_CONVERT_TYPE,
265+
"Field Type expected to be of Type VECTOR but is actually " + fd.getType());
266+
final var fieldOptionMaybe = Optional.ofNullable(fd.getOptions()).map(f -> f.getExtension(RecordMetaDataOptionsProto.field));
267+
Assert.that(fieldOptionMaybe.isPresent() && fieldOptionMaybe.get().hasVectorOptions(), ErrorCode.CANNOT_CONVERT_TYPE, "Cannot insert non vector type into vector column");
268+
final var vectorOptions = fieldOptionMaybe.get().getVectorOptions();
269+
Assert.that(vectorOptions.getDimensions() == ((AbstractRealVector)vector).getNumDimensions(), ErrorCode.CANNOT_CONVERT_TYPE, "Wrong number of dimension for vector");
270+
final int precision = vectorOptions.getPrecision();
271+
switch (precision) {
272+
case 16:
273+
Assert.that(vector instanceof HalfRealVector, ErrorCode.CANNOT_CONVERT_TYPE, "Wrong precision for vector");
274+
break;
275+
case 32:
276+
Assert.that(vector instanceof FloatRealVector, ErrorCode.CANNOT_CONVERT_TYPE, "Wrong precision for vector");
277+
break;
278+
case 64:
279+
Assert.that(vector instanceof DoubleRealVector, ErrorCode.CANNOT_CONVERT_TYPE, "Wrong precision for vector");
280+
break;
281+
default:
282+
Assert.fail(ErrorCode.INTERNAL_ERROR, "Unknown precision for vector");
283+
break;
250284
}
285+
builder.setField(fd, ((AbstractRealVector) vector).getRawData());
251286
}
252287
break;
253288
case NULL:
254289
break;
255290
default:
256-
Assert.failUnchecked(ErrorCode.INTERNAL_ERROR, (String.format(Locale.ROOT, "Unexpected Column type <%s> for column <%s>",
257-
struct.getMetaData().getColumnType(i), columnName)));
291+
Assert.fail(ErrorCode.INTERNAL_ERROR, (String.format(Locale.ROOT, "Unexpected Column type <%s> for column <%s>",
292+
struct.getMetaData().getColumnTypeName(i), columnName)));
258293
break;
259294
}
260295
}
Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
/*
2+
* InsertTest.java
3+
*
4+
* This source file is part of the FoundationDB open source project
5+
*
6+
* Copyright 2021-2025 Apple Inc. and the FoundationDB project authors
7+
*
8+
* Licensed under the Apache License, Version 2.0 (the "License");
9+
* you may not use this file except in compliance with the License.
10+
* You may obtain a copy of the License at
11+
*
12+
* http://www.apache.org/licenses/LICENSE-2.0
13+
*
14+
* Unless required by applicable law or agreed to in writing, software
15+
* distributed under the License is distributed on an "AS IS" BASIS,
16+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
* See the License for the specific language governing permissions and
18+
* limitations under the License.
19+
*/
20+
21+
package com.apple.foundationdb.relational.recordlayer;
22+
23+
import com.apple.foundationdb.linear.DoubleRealVector;
24+
import com.apple.foundationdb.linear.FloatRealVector;
25+
import com.apple.foundationdb.linear.HalfRealVector;
26+
import com.apple.foundationdb.relational.api.EmbeddedRelationalStruct;
27+
import com.apple.foundationdb.relational.api.KeySet;
28+
import com.apple.foundationdb.relational.api.Options;
29+
import com.apple.foundationdb.relational.api.RelationalConnection;
30+
import com.apple.foundationdb.relational.api.RelationalResultSet;
31+
import com.apple.foundationdb.relational.api.RelationalStatement;
32+
import com.apple.foundationdb.relational.api.RelationalStruct;
33+
import com.apple.foundationdb.relational.api.exceptions.ErrorCode;
34+
import com.apple.foundationdb.relational.utils.RelationalAssertions;
35+
import com.apple.foundationdb.relational.utils.ResultSetAssert;
36+
import com.apple.foundationdb.relational.utils.SimpleDatabaseRule;
37+
import org.junit.jupiter.api.Order;
38+
import org.junit.jupiter.api.Test;
39+
import org.junit.jupiter.api.extension.RegisterExtension;
40+
41+
import java.sql.DriverManager;
42+
import java.sql.SQLException;
43+
44+
public class InsertVectorTest {
45+
@RegisterExtension
46+
@Order(0)
47+
public final EmbeddedRelationalExtension relationalExtension = new EmbeddedRelationalExtension();
48+
49+
@RegisterExtension
50+
@Order(1)
51+
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))");
52+
53+
@Test
54+
void insertNulls() throws SQLException {
55+
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
56+
conn.setSchema("TEST_SCHEMA");
57+
try (RelationalStatement s = conn.createStatement()) {
58+
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).build();
59+
s.executeInsert("V", rec);
60+
try (RelationalResultSet rs = s.executeGet("V", new KeySet().setKeyColumn("PK", 0), Options.NONE)) {
61+
ResultSetAssert.assertThat(rs)
62+
.hasNextRow()
63+
.hasColumn("PK", 0)
64+
.hasColumn("V1", null)
65+
.hasColumn("V2", null)
66+
.hasColumn("V3", null)
67+
.hasNoNextRow();
68+
}
69+
}
70+
}
71+
}
72+
73+
@Test
74+
void partialInsert() throws SQLException {
75+
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
76+
conn.setSchema("TEST_SCHEMA");
77+
try (RelationalStatement s = conn.createStatement()) {
78+
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).addObject("V1", new FloatRealVector(new float[]{1f, 2f, 3f, 4f})).build();
79+
s.executeInsert("V", rec);
80+
try (RelationalResultSet rs = s.executeGet("V", new KeySet().setKeyColumn("PK", 0), Options.NONE)) {
81+
ResultSetAssert.assertThat(rs)
82+
.hasNextRow()
83+
.hasColumn("PK", 0)
84+
.hasColumn("V1", new FloatRealVector(new float[]{1f, 2f, 3f, 4f}))
85+
.hasColumn("V2", null)
86+
.hasColumn("V3", null)
87+
.hasNoNextRow();
88+
}
89+
}
90+
}
91+
}
92+
93+
@Test
94+
void fullInsert() throws SQLException {
95+
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
96+
conn.setSchema("TEST_SCHEMA");
97+
try (RelationalStatement s = conn.createStatement()) {
98+
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0)
99+
.addObject("V1", new FloatRealVector(new float[]{1f, 2f, 3f, 4f}))
100+
.addObject("V2", new HalfRealVector(new int[]{1, 2, 3}))
101+
.addObject("V3", new DoubleRealVector(new double[]{1d, 2d}))
102+
.build();
103+
s.executeInsert("V", rec);
104+
try (RelationalResultSet rs = s.executeGet("V", new KeySet().setKeyColumn("PK", 0), Options.NONE)) {
105+
ResultSetAssert.assertThat(rs)
106+
.hasNextRow()
107+
.hasColumn("PK", 0)
108+
.hasColumn("V1", new FloatRealVector(new float[]{1f, 2f, 3f, 4f}))
109+
.hasColumn("V2", new HalfRealVector(new int[]{1, 2, 3}))
110+
.hasColumn("V3", new DoubleRealVector(new double[]{1d, 2d}))
111+
.hasNoNextRow();
112+
}
113+
}
114+
}
115+
}
116+
117+
@Test
118+
void insertWrongDimensionFails() throws SQLException {
119+
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
120+
conn.setSchema("TEST_SCHEMA");
121+
try (RelationalStatement s = conn.createStatement()) {
122+
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).addObject("V1", new FloatRealVector(new float[] {1f, 2f, 3f, 4f, 5f})).build();
123+
RelationalAssertions.assertThrowsSqlException(
124+
() -> s.executeInsert("V", rec))
125+
.hasErrorCode(ErrorCode.CANNOT_CONVERT_TYPE);
126+
}
127+
}
128+
}
129+
130+
@Test
131+
void insertWrongPrecisionFails() throws SQLException {
132+
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
133+
conn.setSchema("TEST_SCHEMA");
134+
try (RelationalStatement s = conn.createStatement()) {
135+
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).addObject("V1", new DoubleRealVector(new double[] {1d, 2d, 3d, 4d})).build();
136+
RelationalAssertions.assertThrowsSqlException(
137+
() -> s.executeInsert("V", rec))
138+
.hasErrorCode(ErrorCode.CANNOT_CONVERT_TYPE);
139+
}
140+
}
141+
}
142+
143+
@Test
144+
void insertWrongTypeFails() throws SQLException {
145+
try (RelationalConnection conn = DriverManager.getConnection(database.getConnectionUri().toString()).unwrap(RelationalConnection.class)) {
146+
conn.setSchema("TEST_SCHEMA");
147+
try (RelationalStatement s = conn.createStatement()) {
148+
RelationalStruct rec = EmbeddedRelationalStruct.newBuilder().addInt("PK", 0).addInt("V1", 42).build();
149+
RelationalAssertions.assertThrowsSqlException(
150+
() -> s.executeInsert("V", rec))
151+
.hasErrorCode(ErrorCode.CANNOT_CONVERT_TYPE);
152+
}
153+
}
154+
}
155+
}

0 commit comments

Comments
 (0)