Skip to content

Commit 6bff7c1

Browse files
LuQQiuLuQQiu
committed
refactor: Reorganize FixedSizeList tests into SQL and DataFrame tests
- Split FixedSizeListVectorTest into two focused test classes: - FixedSizeListSQLTest: Tests SQL CREATE TABLE and INSERT operations - FixedSizeListDataFrameTest: Tests DataFrame API write/read operations - Each test now validates both write and read paths - Added verification to VectorWriteTest to ensure data can be read back - Removed duplication and improved test organization
1 parent 9c254d5 commit 6bff7c1

3 files changed

Lines changed: 416 additions & 90 deletions

File tree

‎lance-spark-base_2.12/src/test/java/com/lancedb/lance/spark/FixedSizeListDataFrameTest.java‎

Lines changed: 312 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,312 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.lancedb.lance.spark;
15+
16+
import org.apache.spark.sql.Dataset;
17+
import org.apache.spark.sql.Row;
18+
import org.apache.spark.sql.RowFactory;
19+
import org.apache.spark.sql.SparkSession;
20+
import org.apache.spark.sql.types.*;
21+
import org.junit.jupiter.api.Test;
22+
import org.junit.jupiter.api.io.TempDir;
23+
24+
import java.nio.file.Path;
25+
import java.util.ArrayList;
26+
import java.util.List;
27+
28+
import static org.junit.jupiter.api.Assertions.*;
29+
30+
/**
31+
* Test for FixedSizeList support using DataFrame API.
32+
* Tests creating Lance tables with vector columns via DataFrame write operations
33+
* and validates both write and read paths.
34+
*/
35+
public class FixedSizeListDataFrameTest {
36+
37+
@TempDir Path tempDir;
38+
39+
@Test
40+
public void testDataFrameWriteAndReadWithFixedSizeList() {
41+
String catalogName = "lance_test";
42+
43+
SparkSession spark =
44+
SparkSession.builder()
45+
.appName("dataframe-fixedsizelist-test")
46+
.master("local[*]")
47+
.config(
48+
"spark.sql.catalog." + catalogName,
49+
"com.lancedb.lance.spark.LanceNamespaceSparkCatalog")
50+
.config("spark.sql.catalog." + catalogName + ".impl", "dir")
51+
.config("spark.sql.catalog." + catalogName + ".root", tempDir.toString())
52+
.getOrCreate();
53+
54+
try {
55+
String tableName = "df_vector_table";
56+
57+
// Create metadata for vector column - use Long value, not String
58+
Metadata vectorMetadata = Metadata.fromJson(
59+
"{\"arrow.fixed-size-list.size\":128}"
60+
);
61+
62+
// Create schema with vector column using DataFrame API
63+
StructType schema = new StructType(
64+
new StructField[] {
65+
DataTypes.createStructField("id", DataTypes.IntegerType, false),
66+
DataTypes.createStructField("text", DataTypes.StringType, true),
67+
new StructField(
68+
"embeddings",
69+
DataTypes.createArrayType(DataTypes.FloatType, false),
70+
false,
71+
vectorMetadata
72+
)
73+
}
74+
);
75+
76+
// Create test data
77+
List<Row> rows = new ArrayList<>();
78+
for (int i = 0; i < 10; i++) {
79+
float[] vector = new float[128];
80+
for (int j = 0; j < 128; j++) {
81+
vector[j] = i * 0.01f + j * 0.001f;
82+
}
83+
rows.add(RowFactory.create(i, "text_" + i, vector));
84+
}
85+
86+
Dataset<Row> df = spark.createDataFrame(rows, schema);
87+
88+
// Write to Lance table using DataFrame API
89+
df.writeTo(catalogName + ".default." + tableName)
90+
.using("lance")
91+
.createOrReplace();
92+
93+
// Read back and verify
94+
Dataset<Row> result = spark.table(catalogName + ".default." + tableName);
95+
assertEquals(10, result.count(), "Should have 10 rows");
96+
97+
// Verify the data was read correctly
98+
Row firstRow = result.first();
99+
assertEquals(0, firstRow.getInt(0));
100+
assertEquals("text_0", firstRow.getString(1));
101+
scala.collection.mutable.WrappedArray<Float> embeddings =
102+
(scala.collection.mutable.WrappedArray<Float>) firstRow.get(2);
103+
assertEquals(128, embeddings.size(), "Embeddings should have 128 elements");
104+
105+
// Verify values
106+
for (int i = 0; i < 10; i++) {
107+
float expected = i * 0.001f;
108+
assertEquals(expected, embeddings.apply(i), 0.0001f, "Value mismatch at index " + i);
109+
}
110+
111+
// Clean up
112+
spark.sql("DROP TABLE IF EXISTS " + catalogName + ".default." + tableName);
113+
114+
} finally {
115+
spark.stop();
116+
}
117+
}
118+
119+
@Test
120+
public void testDataFrameMultipleVectorColumns() {
121+
String catalogName = "lance_test";
122+
123+
SparkSession spark =
124+
SparkSession.builder()
125+
.appName("dataframe-multi-vector-test")
126+
.master("local[*]")
127+
.config(
128+
"spark.sql.catalog." + catalogName,
129+
"com.lancedb.lance.spark.LanceNamespaceSparkCatalog")
130+
.config("spark.sql.catalog." + catalogName + ".impl", "dir")
131+
.config("spark.sql.catalog." + catalogName + ".root", tempDir.toString())
132+
.getOrCreate();
133+
134+
try {
135+
String tableName = "df_multi_vector";
136+
137+
// Create metadata for different vector dimensions
138+
Metadata vec32Metadata = Metadata.fromJson("{\"arrow.fixed-size-list.size\":32}");
139+
Metadata vec128Metadata = Metadata.fromJson("{\"arrow.fixed-size-list.size\":128}");
140+
Metadata vec256Metadata = Metadata.fromJson("{\"arrow.fixed-size-list.size\":256}");
141+
142+
// Create schema with multiple vector columns
143+
StructType schema = new StructType(
144+
new StructField[] {
145+
DataTypes.createStructField("id", DataTypes.IntegerType, false),
146+
DataTypes.createStructField("name", DataTypes.StringType, true),
147+
new StructField(
148+
"small_embedding",
149+
DataTypes.createArrayType(DataTypes.FloatType, false),
150+
false,
151+
vec32Metadata
152+
),
153+
new StructField(
154+
"medium_embedding",
155+
DataTypes.createArrayType(DataTypes.FloatType, false),
156+
false,
157+
vec128Metadata
158+
),
159+
new StructField(
160+
"large_embedding",
161+
DataTypes.createArrayType(DataTypes.FloatType, false),
162+
false,
163+
vec256Metadata
164+
)
165+
}
166+
);
167+
168+
// Create test data
169+
List<Row> rows = new ArrayList<>();
170+
for (int i = 0; i < 5; i++) {
171+
float[] smallVec = new float[32];
172+
float[] mediumVec = new float[128];
173+
float[] largeVec = new float[256];
174+
175+
for (int j = 0; j < 32; j++) {
176+
smallVec[j] = i * 0.01f + j * 0.001f;
177+
}
178+
for (int j = 0; j < 128; j++) {
179+
mediumVec[j] = i * 0.005f + j * 0.0005f;
180+
}
181+
for (int j = 0; j < 256; j++) {
182+
largeVec[j] = i * 0.002f + j * 0.0002f;
183+
}
184+
185+
rows.add(RowFactory.create(i, "entity_" + i, smallVec, mediumVec, largeVec));
186+
}
187+
188+
Dataset<Row> df = spark.createDataFrame(rows, schema);
189+
190+
// Write to Lance table
191+
df.writeTo(catalogName + ".default." + tableName)
192+
.using("lance")
193+
.createOrReplace();
194+
195+
// Read back and verify
196+
Dataset<Row> result = spark.table(catalogName + ".default." + tableName);
197+
assertEquals(5, result.count(), "Should have 5 rows");
198+
199+
// Verify dimensions
200+
Row firstRow = result.first();
201+
scala.collection.mutable.WrappedArray<Float> smallEmb =
202+
(scala.collection.mutable.WrappedArray<Float>) firstRow.get(2);
203+
scala.collection.mutable.WrappedArray<Float> mediumEmb =
204+
(scala.collection.mutable.WrappedArray<Float>) firstRow.get(3);
205+
scala.collection.mutable.WrappedArray<Float> largeEmb =
206+
(scala.collection.mutable.WrappedArray<Float>) firstRow.get(4);
207+
208+
assertEquals(32, smallEmb.size(), "Small embedding should have 32 elements");
209+
assertEquals(128, mediumEmb.size(), "Medium embedding should have 128 elements");
210+
assertEquals(256, largeEmb.size(), "Large embedding should have 256 elements");
211+
212+
// Clean up
213+
spark.sql("DROP TABLE IF EXISTS " + catalogName + ".default." + tableName);
214+
215+
} finally {
216+
spark.stop();
217+
}
218+
}
219+
220+
@Test
221+
public void testDataFrameMixedPrecisionVectors() {
222+
String catalogName = "lance_test";
223+
224+
SparkSession spark =
225+
SparkSession.builder()
226+
.appName("dataframe-mixed-precision-test")
227+
.master("local[*]")
228+
.config(
229+
"spark.sql.catalog." + catalogName,
230+
"com.lancedb.lance.spark.LanceNamespaceSparkCatalog")
231+
.config("spark.sql.catalog." + catalogName + ".impl", "dir")
232+
.config("spark.sql.catalog." + catalogName + ".root", tempDir.toString())
233+
.getOrCreate();
234+
235+
try {
236+
String tableName = "df_mixed_precision";
237+
238+
// Create metadata
239+
Metadata floatVecMetadata = Metadata.fromJson("{\"arrow.fixed-size-list.size\":64}");
240+
Metadata doubleVecMetadata = Metadata.fromJson("{\"arrow.fixed-size-list.size\":64}");
241+
242+
// Create schema with float and double vectors
243+
StructType schema = new StructType(
244+
new StructField[] {
245+
DataTypes.createStructField("id", DataTypes.IntegerType, false),
246+
DataTypes.createStructField("label", DataTypes.StringType, true),
247+
new StructField(
248+
"float_embedding",
249+
DataTypes.createArrayType(DataTypes.FloatType, false),
250+
false,
251+
floatVecMetadata
252+
),
253+
new StructField(
254+
"double_embedding",
255+
DataTypes.createArrayType(DataTypes.DoubleType, false),
256+
false,
257+
doubleVecMetadata
258+
)
259+
}
260+
);
261+
262+
// Create test data
263+
List<Row> rows = new ArrayList<>();
264+
for (int i = 0; i < 5; i++) {
265+
float[] floatVec = new float[64];
266+
double[] doubleVec = new double[64];
267+
268+
for (int j = 0; j < 64; j++) {
269+
floatVec[j] = i * 0.1f + j * 0.01f;
270+
doubleVec[j] = i * 0.1 + j * 0.01;
271+
}
272+
273+
rows.add(RowFactory.create(i, "label_" + i, floatVec, doubleVec));
274+
}
275+
276+
Dataset<Row> df = spark.createDataFrame(rows, schema);
277+
278+
// Write to Lance table
279+
df.writeTo(catalogName + ".default." + tableName)
280+
.using("lance")
281+
.createOrReplace();
282+
283+
// Read back and verify
284+
Dataset<Row> result = spark.table(catalogName + ".default." + tableName);
285+
assertEquals(5, result.count(), "Should have 5 rows");
286+
287+
// Verify precision is maintained
288+
Row firstRow = result.first();
289+
scala.collection.mutable.WrappedArray<Float> floatEmb =
290+
(scala.collection.mutable.WrappedArray<Float>) firstRow.get(2);
291+
scala.collection.mutable.WrappedArray<Double> doubleEmb =
292+
(scala.collection.mutable.WrappedArray<Double>) firstRow.get(3);
293+
294+
assertEquals(64, floatEmb.size());
295+
assertEquals(64, doubleEmb.size());
296+
297+
// Check precision difference
298+
for (int i = 0; i < 10; i++) {
299+
float fVal = floatEmb.apply(i);
300+
double dVal = doubleEmb.apply(i);
301+
assertEquals(i * 0.01f, fVal, 0.0001f);
302+
assertEquals(i * 0.01, dVal, 0.0000001);
303+
}
304+
305+
// Clean up
306+
spark.sql("DROP TABLE IF EXISTS " + catalogName + ".default." + tableName);
307+
308+
} finally {
309+
spark.stop();
310+
}
311+
}
312+
}

0 commit comments

Comments
 (0)