Skip to content

Commit d3716d1

Browse files
authored
[common] Support VECTOR elements in InternalArray accessors (#9773)
1 parent 2d0db9e commit d3716d1

4 files changed

Lines changed: 112 additions & 1 deletion

File tree

‎paimon-common/src/main/java/org/apache/paimon/data/BinaryArray.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ public static int calculateFixLengthPartSize(DataType type) {
8383
case TIMESTAMP_WITHOUT_TIME_ZONE:
8484
case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
8585
case ARRAY:
86+
case VECTOR:
8687
case MULTISET:
8788
case MAP:
8889
case ROW:
@@ -260,7 +261,8 @@ public InternalArray getArray(int pos) {
260261

261262
@Override
262263
public InternalVector getVector(int pos) {
263-
throw new IllegalArgumentException("Unsupported type: VectorType");
264+
assertIndexIsValid(pos);
265+
return MemorySegmentUtils.readVectorData(segments, offset, getLong(pos));
264266
}
265267

266268
@Override

‎paimon-common/src/main/java/org/apache/paimon/data/InternalArray.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,9 @@ static ElementGetter createElementGetter(DataType elementType) {
124124
case ARRAY:
125125
elementGetter = InternalArray::getArray;
126126
break;
127+
case VECTOR:
128+
elementGetter = InternalArray::getVector;
129+
break;
127130
case MULTISET:
128131
case MAP:
129132
elementGetter = InternalArray::getMap;

‎paimon-common/src/main/java/org/apache/paimon/data/InternalRow.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,8 @@ static Class<?> getDataClass(DataType type) {
150150
return Timestamp.class;
151151
case ARRAY:
152152
return InternalArray.class;
153+
case VECTOR:
154+
return InternalVector.class;
153155
case MULTISET:
154156
case MAP:
155157
return InternalMap.class;
Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.paimon.data;
20+
21+
import org.apache.paimon.data.serializer.InternalArraySerializer;
22+
import org.apache.paimon.io.DataInputViewStreamWrapper;
23+
import org.apache.paimon.io.DataOutputViewStreamWrapper;
24+
import org.apache.paimon.types.FloatType;
25+
import org.apache.paimon.types.VectorType;
26+
27+
import org.junit.jupiter.api.Test;
28+
29+
import java.io.ByteArrayInputStream;
30+
import java.io.ByteArrayOutputStream;
31+
32+
import static org.assertj.core.api.Assertions.assertThat;
33+
34+
/**
35+
* Tests that ARRAY&lt;VECTOR&gt; is supported by {@link InternalArray} accessors: the {@link
36+
* InternalArraySerializer} eagerly builds an element getter, so a missing VECTOR case fails
37+
* serializer construction for an accepted DDL type.
38+
*/
39+
class InternalArrayVectorGetterTest {
40+
41+
private static final VectorType VECTOR_TYPE = new VectorType(3, new FloatType());
42+
43+
@Test
44+
void elementGetterReadsVector() {
45+
BinaryVector vector = BinaryVector.fromPrimitiveArray(new float[] {1.0f, 2.0f, 3.0f});
46+
GenericArray array = new GenericArray(new Object[] {vector});
47+
48+
InternalArray.ElementGetter getter = InternalArray.createElementGetter(VECTOR_TYPE);
49+
50+
assertThat(getter.getElementOrNull(array, 0)).isEqualTo(vector);
51+
}
52+
53+
@Test
54+
void vectorArrayRoundTripsThroughSerializer() throws Exception {
55+
InternalArraySerializer serializer = new InternalArraySerializer(VECTOR_TYPE);
56+
// more than one element, with a null between them: only the binary form exercises the
57+
// 8-byte variable-length slot and the null bit, and the last vector then starts at an
58+
// offset the first one's payload decided
59+
GenericArray array =
60+
new GenericArray(
61+
new Object[] {
62+
BinaryVector.fromPrimitiveArray(new float[] {4.0f, 5.0f, 6.0f}),
63+
null,
64+
BinaryVector.fromPrimitiveArray(new float[] {7.0f, 8.0f, 9.0f})
65+
});
66+
67+
ByteArrayOutputStream out = new ByteArrayOutputStream();
68+
serializer.serialize(array, new DataOutputViewStreamWrapper(out));
69+
InternalArray readBack =
70+
serializer.deserialize(
71+
new DataInputViewStreamWrapper(
72+
new ByteArrayInputStream(out.toByteArray())));
73+
74+
assertThat(readBack.size()).isEqualTo(3);
75+
assertThat(readBack.isNullAt(1)).isTrue();
76+
assertThat(readBack.getVector(0).toFloatArray()).isEqualTo(new float[] {4.0f, 5.0f, 6.0f});
77+
assertThat(readBack.getVector(2).toFloatArray()).isEqualTo(new float[] {7.0f, 8.0f, 9.0f});
78+
79+
// toObjectArray and copy read the binary form through getDataClass and the element
80+
// getter; the projection and manifest paths use them, no test did
81+
InternalVector[] objects = ((BinaryArray) readBack).toObjectArray(VECTOR_TYPE);
82+
assertThat(objects).hasSize(3);
83+
assertThat(objects[1]).isNull();
84+
assertThat(objects[2].toFloatArray()).isEqualTo(new float[] {7.0f, 8.0f, 9.0f});
85+
86+
InternalArray copied = serializer.copy(readBack);
87+
assertThat(copied.isNullAt(1)).isTrue();
88+
assertThat(copied.getVector(2).toFloatArray()).isEqualTo(new float[] {7.0f, 8.0f, 9.0f});
89+
}
90+
91+
@Test
92+
void vectorArrayCopyResolvesDataClass() {
93+
// copy() of a GenericArray goes through InternalRow.getDataClass(VECTOR) to allocate the
94+
// element array; without the VECTOR case that throws "Illegal type", so this pins the
95+
// getDataClass fix (the round-trip test above copies raw bytes and never reaches it).
96+
InternalArraySerializer serializer = new InternalArraySerializer(VECTOR_TYPE);
97+
BinaryVector vector = BinaryVector.fromPrimitiveArray(new float[] {7.0f, 8.0f, 9.0f});
98+
GenericArray array = new GenericArray(new Object[] {vector});
99+
100+
InternalArray copied = serializer.copy(array);
101+
102+
assertThat(copied.getVector(0).toFloatArray()).isEqualTo(new float[] {7.0f, 8.0f, 9.0f});
103+
}
104+
}

0 commit comments

Comments
 (0)