fix: support reading and writing JSON columns

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 {}

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 {}

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.
+ *
+ * <p>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.
+ *
+ * <p>Two spellings of the extension name are in play, and they are not interchangeable:
+ *
+ * <ul>
+ *   <li>{@code arrow.json} — the canonical Arrow extension name. This is what a producer declares,
+ *       and the only spelling lance-core recognizes when validating a write against an existing
+ *       schema.
+ *   <li>{@code lance.json} — the form lance-core reports back when a dataset is opened. It is an
+ *       internal label; declaring it on a UTF-8 field is silently ignored, leaving an ordinary
+ *       string column that merely looks like JSON.
+ * </ul>
+ *
+ * <p>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.
+   *
+   * <p>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<String, String> 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<String, String> 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());
+  }
+}

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)