diff --git a/lance-spark-3.4_2.12/src/test/java/org/lance/spark/JsonColumnTest.java b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/JsonColumnTest.java new file mode 100644 index 000000000..de57a1472 --- /dev/null +++ b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/JsonColumnTest.java @@ -0,0 +1,16 @@ +/* + * 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 org.lance.spark; + +public class JsonColumnTest extends BaseJsonColumnTest {} diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/JsonColumnTest.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/JsonColumnTest.java new file mode 100644 index 000000000..de57a1472 --- /dev/null +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/JsonColumnTest.java @@ -0,0 +1,16 @@ +/* + * 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 org.lance.spark; + +public class JsonColumnTest extends BaseJsonColumnTest {} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/JsonUtils.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/JsonUtils.java new file mode 100644 index 000000000..29d07d5c4 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/JsonUtils.java @@ -0,0 +1,135 @@ +/* + * 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 org.lance.spark.utils; + +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StringType; +import org.apache.spark.sql.types.StructField; + +import java.util.Map; + +/** + * Helpers for Lance JSON columns, which Lance models as an Arrow extension type over UTF-8 storage. + * + *

Spark has no JSON type, so a JSON column surfaces as {@link StringType} carrying the extension + * name in the field metadata — the same approach {@link LargeVarCharUtils} uses for Arrow + * LargeUtf8. The JSON text itself is what Spark reads and writes; Lance encodes it to its internal + * JSONB form on write and decodes on read, so the connector never handles JSONB bytes. + * + *

Two spellings of the extension name are in play, and they are not interchangeable: + * + *

+ * + *

Both are recognized at the Arrow boundary. The connector normalizes either spelling to + * {@link #ARROW_JSON_EXTENSION_NAME} when constructing Spark metadata, so Spark schemas and all + * connector writes use the canonical Arrow name. + */ +public class JsonUtils { + + /** The canonical Arrow extension name, and the only one safe to write. */ + public static final String ARROW_JSON_EXTENSION_NAME = "arrow.json"; + + /** + * The internal spelling lance-core reports when a dataset is opened. Recognized, never written. + */ + public static final String LANCE_JSON_EXTENSION_NAME = "lance.json"; + + private JsonUtils() {} + + /** + * Checks whether an extension name denotes a JSON column, in either spelling. + * + * @param extensionName the value of the {@code ARROW:extension:name} key, may be null + * @return true if the name denotes a JSON column + */ + public static boolean isJsonExtensionName(String extensionName) { + return ARROW_JSON_EXTENSION_NAME.equals(extensionName) + || LANCE_JSON_EXTENSION_NAME.equals(extensionName); + } + + /** + * Checks whether an Arrow field is a JSON column. + * + *

The storage type is deliberately not checked. Lance reports JSON columns as LargeBinary + * because Arrow Java does not register the extension type, while other producers may present Utf8 + * or LargeUtf8; the extension name is the reliable signal in every case. + * + * @param field the Arrow field to check + * @return true if the field is a JSON column + */ + public static boolean hasJsonArrowExtension(Field field) { + if (field == null) { + return false; + } + + Map metadata = field.getMetadata(); + if (metadata == null) { + return false; + } + + return isJsonExtensionName(metadata.get(BlobUtils.ARROW_EXTENSION_NAME_KEY)); + } + + /** + * Checks whether an Arrow field uses Lance's physical JSON representation. + * + * @param field the Arrow field to check + * @return true if the field has the {@code lance.json} extension name + */ + public static boolean isLanceJsonField(Field field) { + if (field == null) { + return false; + } + + Map metadata = field.getMetadata(); + return metadata != null + && LANCE_JSON_EXTENSION_NAME.equals(metadata.get(BlobUtils.ARROW_EXTENSION_NAME_KEY)); + } + + /** + * Checks whether Spark metadata carries a JSON extension marker. + * + * @param metadata the Spark field metadata, may be null + * @return true if the metadata marks a JSON column + */ + public static boolean hasJsonMetadata(Metadata metadata) { + if (metadata == null || !metadata.contains(BlobUtils.ARROW_EXTENSION_NAME_KEY)) { + return false; + } + + return isJsonExtensionName(metadata.getString(BlobUtils.ARROW_EXTENSION_NAME_KEY)); + } + + /** + * Checks whether a Spark field is a JSON column. + * + * @param field the Spark struct field to check + * @return true if the field is a StringType column marked as JSON + */ + public static boolean isJsonSparkField(StructField field) { + if (field == null || !(field.dataType() instanceof StringType)) { + return false; + } + + return hasJsonMetadata(field.metadata()); + } +} diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/util/LanceArrowUtils.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/util/LanceArrowUtils.scala index 76cba616c..67ec52227 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/util/LanceArrowUtils.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/util/LanceArrowUtils.scala @@ -30,7 +30,7 @@ import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} import org.apache.spark.{SparkException, SparkUnsupportedOperationException} import org.apache.spark.sql.types._ import org.lance.spark.LanceConstant -import org.lance.spark.utils.{BlobUtils, DateMilliUtils, FixedSizeBinaryUtils, Float16Utils, LargeVarBinaryUtils, LargeVarCharUtils, ListChildUtils, VectorUtils} +import org.lance.spark.utils.{BlobUtils, DateMilliUtils, FixedSizeBinaryUtils, Float16Utils, JsonUtils, LargeVarBinaryUtils, LargeVarCharUtils, ListChildUtils, VectorUtils} import java.util.Locale import java.util.concurrent.atomic.AtomicInteger @@ -47,6 +47,7 @@ object LanceArrowUtils { val ARROW_EXT_NAME_KEY = BlobUtils.ARROW_EXTENSION_NAME_KEY val BLOB_V2_EXT_NAME = BlobUtils.ARROW_EXTENSION_BLOB_V2 val ARROW_LARGE_VAR_CHAR_KEY = LargeVarCharUtils.ARROW_LARGE_VAR_CHAR_KEY + val JSON_EXT_NAME = JsonUtils.ARROW_JSON_EXTENSION_NAME val ARROW_LARGE_VAR_BINARY_KEY = LargeVarBinaryUtils.ARROW_LARGE_VAR_BINARY_KEY val ARROW_DATE_MILLISECOND_KEY = DateMilliUtils.ARROW_DATE_MILLISECOND_KEY val ARROW_FIXED_SIZE_BINARY_BYTE_WIDTH_KEY = @@ -137,6 +138,9 @@ object LanceArrowUtils { // Lance returns LargeBinary in schema but Struct in data for blob columns // We need to handle this as binary to match the schema BinaryType + case _: ArrowType.LargeBinary if JsonUtils.hasJsonArrowExtension(field) => + // Lance stores JSON as LargeBinary JSONB, but Spark reads the decoded values as strings. + StringType case _: ArrowType.LargeUtf8 => // LargeUtf8 maps back to StringType in Spark StringType @@ -243,6 +247,10 @@ object LanceArrowUtils { if (Float16Utils.isFloat16ArrowField(field)) { builder.putString(ARROW_FLOAT16_KEY, Float16Utils.ARROW_FLOAT16_VALUE) } + case _: ArrowType.LargeBinary if JsonUtils.isLanceJsonField(field) => + // Dataset.getSchema exposes Lance's physical JSONB representation. Spark uses the + // logical Arrow extension name for the decoded StringType it presents to callers. + builder.putString(ARROW_EXT_NAME_KEY, JsonUtils.ARROW_JSON_EXTENSION_NAME) case _: ArrowType.LargeUtf8 => builder.putString(ARROW_LARGE_VAR_CHAR_KEY, LargeVarCharUtils.ARROW_LARGE_VAR_CHAR_VALUE) // Spark has a single BinaryType covering both Arrow Binary (32-bit offsets) and LargeBinary @@ -251,8 +259,10 @@ object LanceArrowUtils { // the schema (UPDATE, ADD COLUMNS FROM, or simply read -> transform -> write), and the // resulting write fails type validation against the existing Lance schema. // Blob columns are excluded: they are already LargeBinary-backed via the blob marker, which - // toArrowField honors on its own. - case _: ArrowType.LargeBinary if !isBlobField(field) => + // toArrowField honors on its own. JSON columns are excluded too: they surface as StringType, + // so a binary marker would contradict the Spark type and steer writeback to LargeBinary. + case _: ArrowType.LargeBinary + if !isBlobField(field) && !JsonUtils.hasJsonArrowExtension(field) => builder.putString( ARROW_LARGE_VAR_BINARY_KEY, LargeVarBinaryUtils.ARROW_LARGE_VAR_BINARY_VALUE) @@ -515,6 +525,15 @@ object LanceArrowUtils { toArrowField("uri", StringType, nullable = true, timeZoneId), arrowUInt64Field("position"), arrowUInt64Field("size")).asJava) + case _: StringType if JsonUtils.hasJsonMetadata(metadata) => + // Lance JSON column. Two things matter here. First, the storage must be UTF-8: Lance + // encodes the JSON text to its internal JSONB form itself, and rejects a LargeBinary + // array of pre-encoded bytes. The extension name must be the canonical `arrow.json`. + // Read metadata is normalized to that name, and forcing it here also supports callers + // that provide the physical `lance.json` spelling directly. + val jsonMeta = (meta + (ARROW_EXT_NAME_KEY -> JSON_EXT_NAME)).asJava + val jsonType = if (large) ArrowType.LargeUtf8.INSTANCE else ArrowType.Utf8.INSTANCE + new Field(name, new FieldType(nullable, jsonType, null, jsonMeta), Seq.empty[Field].asJava) case dataType => val fieldType = new FieldType(nullable, toArrowType(dataType, timeZoneId, large, name), null, meta.asJava) diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/BaseJsonColumnTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/BaseJsonColumnTest.java new file mode 100644 index 000000000..6673a712c --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/BaseJsonColumnTest.java @@ -0,0 +1,339 @@ +/* + * 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 org.lance.spark; + +import org.lance.spark.utils.JsonUtils; + +import org.apache.arrow.c.ArrowArrayStream; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +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.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SaveMode; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.MetadataBuilder; +import org.apache.spark.sql.types.StringType; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +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.assertTrue; + +/** + * Tests reading and writing Lance JSON columns through Spark. + * + *

Lance physically persists JSON as JSONB in a LargeBinary column carrying the {@code + * lance.json} extension name. Its Arrow read/write interface uses UTF-8 JSON text with the {@code + * arrow.json} extension name: Lance encodes the text to JSONB on write and decodes it on read. + * Spark has no JSON type, so the connector surfaces this logical representation as StringType + * carrying {@code arrow.json} in field metadata. + * + *

The datasets are created through the lance Java API rather than Spark DDL, matching what + * {@code BaseFixedSizeBinaryReadTest} does for a type Spark cannot express either. + * + *

Current lance-core exposes the physical {@code lance.json} schema from {@code + * Dataset.getSchema()}, so the connector translates that field to its logical UTF-8/{@code + * arrow.json} representation at the Spark boundary. Both extension names are recognized there to + * remain compatible if lance-core begins returning the logical Arrow schema. + */ +public abstract class BaseJsonColumnTest { + + private static SparkSession spark; + + @TempDir static Path tempDir; + + private static final String JSON_1 = "{\"ingested_at\":\"2026-09-04T17:00:00Z\",\"version\":7}"; + private static final String JSON_2 = "{\"ingested_at\":\"2026-09-05T09:30:00Z\",\"version\":8}"; + + @BeforeAll + static void setup() { + spark = SparkSession.builder().appName("json-column-test").master("local[*]").getOrCreate(); + } + + @AfterAll + static void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + /** + * Builds the Arrow field a producer would declare for a JSON column: Utf8 + the extension name. + */ + private static Field jsonField(String name) { + Map metadata = new HashMap<>(); + metadata.put("ARROW:extension:name", JsonUtils.ARROW_JSON_EXTENSION_NAME); + return new Field(name, new FieldType(true, ArrowType.Utf8.INSTANCE, null, metadata), null); + } + + /** + * Creates a lance dataset with an id column and a JSON column, optionally populated. + * + *

Goes through the Arrow C Data interface for the same reason {@code + * BaseFixedSizeBinaryReadTest} does: it keeps lance's JNI bridge and our buffers on the same + * allocator root. + */ + private static void createJsonDataset(String datasetUri, String... jsonValues) throws Exception { + Field idField = new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null); + Schema arrowSchema = new Schema(Arrays.asList(idField, jsonField("payload"))); + + BufferAllocator allocator = LanceRuntime.allocator(); + try (VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchema, allocator)) { + root.allocateNew(); + IntVector idVec = (IntVector) root.getVector("id"); + VarCharVector jsonVec = (VarCharVector) root.getVector("payload"); + + for (int i = 0; i < jsonValues.length; i++) { + idVec.setSafe(i, i + 1); + jsonVec.setSafe(i, jsonValues[i].getBytes(StandardCharsets.UTF_8)); + } + root.setRowCount(jsonValues.length); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, baos)) { + writer.start(); + writer.writeBatch(); + writer.end(); + } + + try (ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(baos.toByteArray()), allocator); + ArrowArrayStream arrowStream = ArrowArrayStream.allocateNew(allocator)) { + Data.exportArrayStream(allocator, reader, arrowStream); + org.lance.Dataset.write().stream(arrowStream).uri(datasetUri).execute().close(); + } + } + } + + @Test + public void testJsonColumnSurfacesAsStringType() throws Exception { + String datasetUri = tempDir.resolve("json_schema_test.lance").toString(); + createJsonDataset(datasetUri, JSON_1); + + Dataset df = spark.read().format(LanceDataSource.name).load(datasetUri); + StructField field = df.schema().apply("payload"); + + assertInstanceOf( + StringType.class, + field.dataType(), + "A JSON column should surface as StringType, not binary"); + assertEquals( + JsonUtils.ARROW_JSON_EXTENSION_NAME, + field.metadata().getString("ARROW:extension:name"), + "The Spark field must carry the canonical Arrow JSON extension name"); + } + + @Test + public void testJsonValuesAreReadAsText() throws Exception { + String datasetUri = tempDir.resolve("json_values_test.lance").toString(); + createJsonDataset(datasetUri, JSON_1, JSON_2); + + spark.read().format(LanceDataSource.name).load(datasetUri).createOrReplaceTempView("json_read"); + List rows = spark.sql("SELECT id, payload FROM json_read ORDER BY id").collectAsList(); + + assertEquals(2, rows.size()); + assertEquals(JSON_1, rows.get(0).getString(1)); + assertEquals(JSON_2, rows.get(1).getString(1)); + } + + /** The JSON text is an ordinary string to Spark, so Spark's own JSON functions apply to it. */ + @Test + public void testJsonColumnIsQueryableWithSparkJsonFunctions() throws Exception { + String datasetUri = tempDir.resolve("json_query_test.lance").toString(); + createJsonDataset(datasetUri, JSON_1, JSON_2); + + spark + .read() + .format(LanceDataSource.name) + .load(datasetUri) + .createOrReplaceTempView("json_query"); + List rows = + spark + .sql("SELECT id FROM json_query WHERE get_json_object(payload, '$.version') = '8'") + .collectAsList(); + + assertEquals(1, rows.size()); + assertEquals(2, rows.get(0).getInt(0)); + } + + /** + * Appends to a JSON column created by another producer. + * + *

The DataFrame is built from the logical schema Spark read back: StringType with the + * canonical {@code arrow.json} extension name. This verifies that the normalized schema can be + * written back to a table whose physical column is JSONB/LargeBinary with {@code lance.json}. + */ + @Test + public void testWriteToExistingJsonColumn() throws Exception { + String datasetUri = tempDir.resolve("json_append_test.lance").toString(); + createJsonDataset(datasetUri, JSON_1); + + StructType readSchema = spark.read().format(LanceDataSource.name).load(datasetUri).schema(); + Dataset toAppend = + spark.createDataFrame(Collections.singletonList(RowFactory.create(2, JSON_2)), readSchema); + + toAppend.write().format(LanceDataSource.name).mode(SaveMode.Append).save(datasetUri); + + spark + .read() + .format(LanceDataSource.name) + .load(datasetUri) + .createOrReplaceTempView("json_appended"); + List rows = spark.sql("SELECT id, payload FROM json_appended ORDER BY id").collectAsList(); + + assertEquals(2, rows.size(), "The appended row should be present"); + assertEquals(JSON_1, rows.get(0).getString(1)); + assertEquals(JSON_2, rows.get(1).getString(1), "The appended JSON should round-trip intact"); + } + + /** A JSON column must survive a full read-transform-write cycle without losing its type. */ + @Test + public void testJsonColumnRoundTripsThroughSparkWrite() throws Exception { + String sourceUri = tempDir.resolve("json_roundtrip_source.lance").toString(); + String targetUri = tempDir.resolve("json_roundtrip_target.lance").toString(); + createJsonDataset(sourceUri, JSON_1, JSON_2); + + spark + .read() + .format(LanceDataSource.name) + .load(sourceUri) + .write() + .format(LanceDataSource.name) + .mode(SaveMode.ErrorIfExists) + .save(targetUri); + + StructField field = + spark.read().format(LanceDataSource.name).load(targetUri).schema().apply("payload"); + assertInstanceOf( + StringType.class, field.dataType(), "The copied column should still be a string"); + assertTrue( + JsonUtils.isJsonSparkField(field), + "The copied column should still be a JSON column, not a plain string"); + } + + /** + * A plain string column must not become a JSON column, and a JSON marker must not be invented + * where the producer declared none. + */ + @Test + public void testPlainStringColumnIsUnaffected() throws Exception { + String datasetUri = tempDir.resolve("plain_string_test.lance").toString(); + + Field idField = new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field textField = new Field("payload", FieldType.nullable(ArrowType.Utf8.INSTANCE), null); + Schema arrowSchema = new Schema(Arrays.asList(idField, textField)); + + BufferAllocator allocator = LanceRuntime.allocator(); + try (VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchema, allocator)) { + root.allocateNew(); + ((IntVector) root.getVector("id")).setSafe(0, 1); + ((VarCharVector) root.getVector("payload")) + .setSafe(0, JSON_1.getBytes(StandardCharsets.UTF_8)); + root.setRowCount(1); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, baos)) { + writer.start(); + writer.writeBatch(); + writer.end(); + } + try (ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(baos.toByteArray()), allocator); + ArrowArrayStream arrowStream = ArrowArrayStream.allocateNew(allocator)) { + Data.exportArrayStream(allocator, reader, arrowStream); + org.lance.Dataset.write().stream(arrowStream).uri(datasetUri).execute().close(); + } + } + + StructField field = + spark.read().format(LanceDataSource.name).load(datasetUri).schema().apply("payload"); + assertInstanceOf(StringType.class, field.dataType()); + assertFalse( + JsonUtils.isJsonSparkField(field), + "A column declared without the extension name must not be treated as JSON"); + } + + /** Spark writes logical Arrow JSON fields as Lance's physical JSONB representation. */ + @Test + public void testSparkCanCreateJsonColumnViaMetadata() { + String datasetUri = tempDir.resolve("json_create_test.lance").toString(); + + StructType schema = + new StructType( + new StructField[] { + new StructField("id", DataTypes.IntegerType, true, Metadata.empty()), + new StructField( + "payload", + DataTypes.StringType, + true, + new MetadataBuilder() + .putString("ARROW:extension:name", JsonUtils.ARROW_JSON_EXTENSION_NAME) + .build()) + }); + + spark + .createDataFrame(Collections.singletonList(RowFactory.create(1, JSON_1)), schema) + .write() + .format(LanceDataSource.name) + .mode(SaveMode.ErrorIfExists) + .save(datasetUri); + + StructField field = + spark.read().format(LanceDataSource.name).load(datasetUri).schema().apply("payload"); + assertTrue( + JsonUtils.isJsonSparkField(field), + "A string column marked with the JSON extension name should be created as a JSON column"); + + try (org.lance.Dataset lanceDataset = + org.lance.Dataset.open().allocator(LanceRuntime.allocator()).uri(datasetUri).build()) { + Field physicalField = lanceDataset.getSchema().findField("payload"); + assertEquals( + ArrowType.LargeBinary.INSTANCE, + physicalField.getType(), + "Lance must persist JSON as LargeBinary JSONB"); + assertEquals( + JsonUtils.LANCE_JSON_EXTENSION_NAME, + physicalField.getMetadata().get("ARROW:extension:name"), + "Lance must persist the internal JSON extension name"); + } + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/utils/JsonUtilsTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/utils/JsonUtilsTest.java new file mode 100644 index 000000000..7da679401 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/utils/JsonUtilsTest.java @@ -0,0 +1,176 @@ +/* + * 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 org.lance.spark.utils; + +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.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.MetadataBuilder; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.LanceArrowUtils; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Unit tests for JSON column detection and the Arrow/Spark type mapping it drives. */ +public class JsonUtilsTest { + + private static final String EXT_KEY = "ARROW:extension:name"; + + private static Field arrowField(ArrowType type, String extensionName) { + Map metadata = new HashMap<>(); + if (extensionName != null) { + metadata.put(EXT_KEY, extensionName); + } + return new Field("payload", new FieldType(true, type, null, metadata), null); + } + + private static StructField sparkField(String extensionName) { + Metadata metadata = + extensionName == null + ? Metadata.empty() + : new MetadataBuilder().putString(EXT_KEY, extensionName).build(); + return new StructField("payload", DataTypes.StringType, true, metadata); + } + + @Test + public void testRecognizesBothExtensionNames() { + assertTrue(JsonUtils.isJsonExtensionName("arrow.json"), "the canonical declared name"); + assertTrue( + JsonUtils.isJsonExtensionName("lance.json"), + "the internal name lance-core reports when a dataset is reopened"); + assertFalse(JsonUtils.isJsonExtensionName("lance.blob.v2")); + assertFalse(JsonUtils.isJsonExtensionName(null)); + } + + /** + * Lance reports JSON columns as LargeBinary because Arrow Java does not register the extension, + * so detection must not depend on the storage type. + */ + @Test + public void testDetectsJsonRegardlessOfStorageType() { + assertTrue(JsonUtils.hasJsonArrowExtension(arrowField(ArrowType.Utf8.INSTANCE, "arrow.json"))); + assertTrue(JsonUtils.hasJsonArrowExtension(arrowField(ArrowType.LargeUtf8.INSTANCE, "arrow.json"))); + assertTrue( + JsonUtils.hasJsonArrowExtension(arrowField(ArrowType.LargeBinary.INSTANCE, "lance.json"))); + assertFalse(JsonUtils.hasJsonArrowExtension(arrowField(ArrowType.Utf8.INSTANCE, null))); + assertFalse(JsonUtils.hasJsonArrowExtension(null)); + } + + @Test + public void testDetectsPhysicalLanceJsonFields() { + assertTrue( + JsonUtils.isLanceJsonField(arrowField(ArrowType.LargeBinary.INSTANCE, "lance.json"))); + assertFalse(JsonUtils.isLanceJsonField(arrowField(ArrowType.Utf8.INSTANCE, "arrow.json"))); + assertFalse(JsonUtils.isLanceJsonField(null)); + } + + @Test + public void testSparkFieldDetectionRequiresStringType() { + assertTrue(JsonUtils.isJsonSparkField(sparkField("arrow.json"))); + assertFalse(JsonUtils.isJsonSparkField(sparkField(null))); + assertFalse( + JsonUtils.isJsonSparkField( + new StructField( + "payload", + DataTypes.BinaryType, + true, + new MetadataBuilder().putString(EXT_KEY, "arrow.json").build())), + "the marker alone must not make a binary column JSON"); + } + + @Test + public void testJsonArrowFieldMapsToStringType() { + assertEquals( + DataTypes.StringType, + LanceArrowUtils.fromArrowField(arrowField(ArrowType.LargeBinary.INSTANCE, "lance.json")), + "a JSON column reported as LargeBinary must still surface as a string"); + } + + /** A LargeBinary column with no JSON marker keeps its existing mapping. */ + @Test + public void testPlainLargeBinaryStillMapsToBinaryType() { + assertEquals( + DataTypes.BinaryType, + LanceArrowUtils.fromArrowField(arrowField(ArrowType.LargeBinary.INSTANCE, null))); + } + + /** + * The write path must emit the canonical name over UTF-8 storage. Echoing back the internal + * {@code lance.json} spelling is silently ignored by lance-core on a UTF-8 field, producing a + * plain string column and an append rejected as {@code should have type json}. + */ + @Test + public void testWriteTranslatesInternalNameToCanonicalName() { + Metadata readMetadata = new MetadataBuilder().putString(EXT_KEY, "lance.json").build(); + + Field written = + LanceArrowUtils.toArrowField( + "payload", DataTypes.StringType, true, null, readMetadata, false); + + assertEquals(ArrowType.Utf8.INSTANCE, written.getType(), "JSON storage must be UTF-8"); + assertEquals( + JsonUtils.ARROW_JSON_EXTENSION_NAME, + written.getMetadata().get(EXT_KEY), + "the internal name must be translated to the canonical one on write"); + } + + @Test + public void testWritePreservesCanonicalName() { + Metadata metadata = new MetadataBuilder().putString(EXT_KEY, "arrow.json").build(); + + Field written = + LanceArrowUtils.toArrowField("payload", DataTypes.StringType, true, null, metadata, false); + + assertEquals(ArrowType.Utf8.INSTANCE, written.getType()); + assertEquals(JsonUtils.ARROW_JSON_EXTENSION_NAME, written.getMetadata().get(EXT_KEY)); + } + + /** A JSON column read back and written again must not pick up a large-binary marker. */ + @Test + public void testJsonFieldDoesNotGetLargeBinaryMarker() { + StructType sparkSchema = + LanceArrowUtils.fromArrowSchema( + new Schema( + Collections.singletonList( + arrowField(ArrowType.LargeBinary.INSTANCE, "lance.json")))); + StructField sparkField = sparkSchema.apply("payload"); + + assertEquals(DataTypes.StringType, sparkField.dataType()); + assertEquals( + JsonUtils.ARROW_JSON_EXTENSION_NAME, + sparkField.metadata().getString(EXT_KEY), + "Spark metadata must use the canonical Arrow extension name"); + assertFalse( + sparkField.metadata().contains(LargeVarBinaryUtils.ARROW_LARGE_VAR_BINARY_KEY), + "a large-binary marker would contradict the StringType mapping and steer writeback " + + "back to LargeBinary"); + + Field written = + LanceArrowUtils.toArrowField( + "payload", DataTypes.StringType, true, null, sparkField.metadata(), false); + assertEquals(ArrowType.Utf8.INSTANCE, written.getType()); + assertEquals(JsonUtils.ARROW_JSON_EXTENSION_NAME, written.getMetadata().get(EXT_KEY)); + } +}