lance_arrow/
json.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! JSON support for Apache Arrow.
5
6use std::convert::TryFrom;
7use std::sync::Arc;
8
9use arrow_array::builder::{GenericStringBuilder, LargeBinaryBuilder};
10use arrow_array::cast::AsArray;
11use arrow_array::{
12    Array, ArrayRef, FixedSizeListArray, GenericStringArray, LargeBinaryArray, LargeListArray,
13    LargeStringArray, ListArray, MapArray, OffsetSizeTrait, RecordBatch, StringArray, StructArray,
14};
15use arrow_schema::{ArrowError, DataType, Field as ArrowField, Fields, Schema};
16use jsonb::OwnedJsonb;
17use jsonb::jsonpath::{JsonPath, Selector};
18
19use crate::ARROW_EXT_NAME_KEY;
20
21/// Arrow extension type name for JSON data (Lance internal)
22pub const JSON_EXT_NAME: &str = "lance.json";
23
24/// Arrow extension type name for JSON data (Arrow official)
25pub const ARROW_JSON_EXT_NAME: &str = "arrow.json";
26
27/// Check if a field is a JSON extension field (Lance internal JSONB storage)
28pub fn is_json_field(field: &ArrowField) -> bool {
29    field.data_type() == &DataType::LargeBinary
30        && field
31            .metadata()
32            .get(ARROW_EXT_NAME_KEY)
33            .map(|name| name == JSON_EXT_NAME)
34            .unwrap_or_default()
35}
36
37/// Check if a field is an Arrow JSON extension field (PyArrow pa.json() type)
38pub fn is_arrow_json_field(field: &ArrowField) -> bool {
39    // Arrow JSON extension type uses Utf8 or LargeUtf8 as storage type
40    (field.data_type() == &DataType::Utf8 || field.data_type() == &DataType::LargeUtf8)
41        && field
42            .metadata()
43            .get(ARROW_EXT_NAME_KEY)
44            .map(|name| name == ARROW_JSON_EXT_NAME)
45            .unwrap_or_default()
46}
47
48/// Check if a field or any of its descendants is a JSON field
49pub fn has_json_fields(field: &ArrowField) -> bool {
50    if is_json_field(field) {
51        return true;
52    }
53
54    match field.data_type() {
55        DataType::Struct(fields) => fields.iter().any(|f| has_json_fields(f)),
56        DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => {
57            has_json_fields(f)
58        }
59        DataType::Map(f, _) => has_json_fields(f),
60        _ => false,
61    }
62}
63
64/// Check if a field or any of its descendants is an Arrow JSON field
65pub fn has_arrow_json_fields(field: &ArrowField) -> bool {
66    if is_arrow_json_field(field) {
67        return true;
68    }
69
70    match field.data_type() {
71        DataType::Struct(fields) => fields.iter().any(|f| has_arrow_json_fields(f)),
72        DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => {
73            has_arrow_json_fields(f)
74        }
75        DataType::Map(f, _) => has_arrow_json_fields(f),
76        _ => false,
77    }
78}
79
80/// The Arrow representation holding a JSON field's values.
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub enum JsonEncoding {
83    /// JSONB bytes: [`JSON_EXT_NAME`] over `LargeBinary`. Lance stores JSON
84    /// this way.
85    Jsonb,
86    /// JSON text: [`ARROW_JSON_EXT_NAME`] over `Utf8` or `LargeUtf8`. Callers
87    /// write JSON this way, and reads return it this way by default.
88    Text,
89}
90
91impl JsonEncoding {
92    /// The encoding of `field`'s values, or `None` when `field` does not hold
93    /// JSON.
94    pub fn of_field(field: &ArrowField) -> Option<Self> {
95        if is_json_field(field) {
96            Some(Self::Jsonb)
97        } else if is_arrow_json_field(field) {
98            Some(Self::Text)
99        } else {
100            None
101        }
102    }
103}
104
105/// The values of a JSON column, readable as JSONB or as text whichever
106/// encoding they arrived in.
107///
108/// Code that computes on JSON, such as tokenizers and indices, reads values
109/// through this type, so its behavior does not depend on whether they come
110/// from storage or from a caller.
111#[derive(Debug, Clone, Copy)]
112pub struct JsonValues<'a> {
113    encoding: JsonEncoding,
114    array: &'a ArrayRef,
115}
116
117impl<'a> JsonValues<'a> {
118    /// Read `array` as the values of `field`, or `None` when `field` does not
119    /// hold JSON.
120    pub fn try_new(field: &ArrowField, array: &'a ArrayRef) -> Option<Self> {
121        JsonEncoding::of_field(field).map(|encoding| Self { encoding, array })
122    }
123
124    pub fn encoding(&self) -> JsonEncoding {
125        self.encoding
126    }
127
128    /// The values as JSONB. Text values are parsed, so invalid JSON fails.
129    pub fn to_jsonb(&self) -> Result<LargeBinaryArray, ArrowError> {
130        match self.encoding {
131            JsonEncoding::Jsonb => Ok(self.array.as_binary::<i64>().clone()),
132            JsonEncoding::Text => Ok(JsonArray::try_from(self.array.clone())?.into_inner()),
133        }
134    }
135
136    /// The values as JSON text: the input array when it already holds text,
137    /// otherwise `LargeUtf8` decoded from JSONB.
138    pub fn to_text(&self) -> ArrayRef {
139        match self.encoding {
140            JsonEncoding::Jsonb => Arc::new(decode_jsonb_array::<i64>(self.array.as_binary())),
141            JsonEncoding::Text => self.array.clone(),
142        }
143    }
144}
145
146/// Create a JSON field with the appropriate extension metadata
147pub fn json_field(name: &str, nullable: bool) -> ArrowField {
148    let mut field = ArrowField::new(name, DataType::LargeBinary, nullable);
149    let mut metadata = std::collections::HashMap::new();
150    metadata.insert(ARROW_EXT_NAME_KEY.to_string(), JSON_EXT_NAME.to_string());
151    field.set_metadata(metadata);
152    field
153}
154
155/// A specialized array for JSON data stored as JSONB binary format
156#[derive(Debug, Clone)]
157pub struct JsonArray {
158    inner: LargeBinaryArray,
159}
160
161impl JsonArray {
162    /// Create a new JsonArray from an iterator of JSON strings
163    pub fn try_from_iter<I, S>(iter: I) -> Result<Self, ArrowError>
164    where
165        I: IntoIterator<Item = Option<S>>,
166        S: AsRef<str>,
167    {
168        let mut builder = LargeBinaryBuilder::new();
169
170        for json_str in iter {
171            match json_str {
172                Some(s) => {
173                    let encoded = encode_json(s.as_ref()).map_err(|e| {
174                        ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e))
175                    })?;
176                    builder.append_value(&encoded);
177                }
178                None => builder.append_null(),
179            }
180        }
181
182        Ok(Self {
183            inner: builder.finish(),
184        })
185    }
186
187    /// Get the underlying LargeBinaryArray
188    pub fn into_inner(self) -> LargeBinaryArray {
189        self.inner
190    }
191
192    /// Get a reference to the underlying LargeBinaryArray
193    pub fn inner(&self) -> &LargeBinaryArray {
194        &self.inner
195    }
196
197    /// Get the value at index i as decoded JSON string
198    pub fn value(&self, i: usize) -> Result<String, ArrowError> {
199        if self.inner.is_null(i) {
200            return Err(ArrowError::InvalidArgumentError(
201                "Value is null".to_string(),
202            ));
203        }
204
205        let jsonb_bytes = self.inner.value(i);
206        Ok(decode_json(jsonb_bytes))
207    }
208
209    /// Get the value at index i as raw JSONB bytes
210    pub fn value_bytes(&self, i: usize) -> &[u8] {
211        self.inner.value(i)
212    }
213
214    /// Get JSONPath value from the JSON at index i
215    pub fn json_path(&self, i: usize, path: &str) -> Result<Option<String>, ArrowError> {
216        if self.inner.is_null(i) {
217            return Ok(None);
218        }
219
220        let path = parse_json_path(path)?;
221        Ok(select_json_path(self.inner.value(i), &path)?.map(|value| value.to_string()))
222    }
223
224    /// Convert to Arrow string array (JSON as UTF-8)
225    pub fn to_arrow_json(&self) -> ArrayRef {
226        Arc::new(decode_jsonb_array::<i32>(&self.inner))
227    }
228
229    pub fn len(&self) -> usize {
230        self.inner.len()
231    }
232
233    pub fn is_empty(&self) -> bool {
234        self.inner.is_empty()
235    }
236
237    pub fn is_null(&self, i: usize) -> bool {
238        self.inner.is_null(i)
239    }
240}
241
242// TryFrom implementations for string arrays
243impl TryFrom<StringArray> for JsonArray {
244    type Error = ArrowError;
245
246    fn try_from(array: StringArray) -> Result<Self, Self::Error> {
247        Self::try_from(&array)
248    }
249}
250
251impl TryFrom<&StringArray> for JsonArray {
252    type Error = ArrowError;
253
254    fn try_from(array: &StringArray) -> Result<Self, Self::Error> {
255        encode_json_text_array(array)
256    }
257}
258
259impl TryFrom<LargeStringArray> for JsonArray {
260    type Error = ArrowError;
261
262    fn try_from(array: LargeStringArray) -> Result<Self, Self::Error> {
263        Self::try_from(&array)
264    }
265}
266
267impl TryFrom<&LargeStringArray> for JsonArray {
268    type Error = ArrowError;
269
270    fn try_from(array: &LargeStringArray) -> Result<Self, Self::Error> {
271        encode_json_text_array(array)
272    }
273}
274
275fn encode_json_text_array<O: OffsetSizeTrait>(
276    array: &GenericStringArray<O>,
277) -> Result<JsonArray, ArrowError> {
278    let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len());
279    for value in array {
280        match value {
281            Some(json_str) => {
282                let encoded = encode_json(json_str).map_err(|e| {
283                    ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e))
284                })?;
285                builder.append_value(&encoded);
286            }
287            None => builder.append_null(),
288        }
289    }
290    Ok(JsonArray {
291        inner: builder.finish(),
292    })
293}
294
295fn decode_jsonb_array<O: OffsetSizeTrait>(array: &LargeBinaryArray) -> GenericStringArray<O> {
296    let mut builder =
297        GenericStringBuilder::<O>::with_capacity(array.len(), array.value_data().len());
298    for value in array {
299        match value {
300            Some(jsonb_bytes) => builder.append_value(decode_json(jsonb_bytes)),
301            None => builder.append_null(),
302        }
303    }
304    builder.finish()
305}
306
307impl TryFrom<ArrayRef> for JsonArray {
308    type Error = ArrowError;
309
310    fn try_from(array_ref: ArrayRef) -> Result<Self, Self::Error> {
311        match array_ref.data_type() {
312            DataType::Utf8 => {
313                // Downcast is guaranteed to succeed after matching on DataType::Utf8
314                let string_array = array_ref
315                    .as_any()
316                    .downcast_ref::<StringArray>()
317                    .expect("DataType::Utf8 array must be StringArray");
318                Self::try_from(string_array)
319            }
320            DataType::LargeUtf8 => {
321                // Downcast is guaranteed to succeed after matching on DataType::LargeUtf8
322                let large_string_array = array_ref
323                    .as_any()
324                    .downcast_ref::<LargeStringArray>()
325                    .expect("DataType::LargeUtf8 array must be LargeStringArray");
326                Self::try_from(large_string_array)
327            }
328            dt => Err(ArrowError::InvalidArgumentError(format!(
329                "Unsupported array type for JSON: {:?}. Expected Utf8 or LargeUtf8",
330                dt
331            ))),
332        }
333    }
334}
335
336/// Encode JSON string to JSONB format
337pub fn encode_json(json_str: &str) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
338    let value = jsonb::parse_value(json_str.as_bytes())?;
339    Ok(value.to_vec())
340}
341
342/// Decode JSONB bytes to JSON string
343pub fn decode_json(jsonb_bytes: &[u8]) -> String {
344    let raw_jsonb = jsonb::RawJsonb::new(jsonb_bytes);
345    raw_jsonb.to_string()
346}
347
348/// Parse a JSONPath expression.
349pub fn parse_json_path(path: &str) -> Result<JsonPath<'_>, ArrowError> {
350    jsonb::jsonpath::parse_json_path(path.as_bytes())
351        .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid JSONPath '{path}': {e}")))
352}
353
354/// Select the value at `path` in a JSONB value.
355///
356/// A path matching several values selects them together as one JSON array,
357/// so no match is silently dropped.
358pub fn select_json_path(
359    jsonb_bytes: &[u8],
360    path: &JsonPath<'_>,
361) -> Result<Option<OwnedJsonb>, ArrowError> {
362    Selector::new(jsonb::RawJsonb::new(jsonb_bytes))
363        .select_value(path)
364        .map_err(|e| {
365            ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}"))
366        })
367}
368
369/// Whether `path` matches any value in a JSONB value.
370pub fn json_path_exists(jsonb_bytes: &[u8], path: &JsonPath<'_>) -> Result<bool, ArrowError> {
371    Selector::new(jsonb::RawJsonb::new(jsonb_bytes))
372        .exists(path)
373        .map_err(|e| {
374            ArrowError::InvalidArgumentError(format!("Failed to match JSONPath {path}: {e}"))
375        })
376}
377
378/// Select every value matched by `path` in a JSONB value.
379pub fn select_json_path_values(
380    jsonb_bytes: &[u8],
381    path: &JsonPath<'_>,
382) -> Result<Vec<OwnedJsonb>, ArrowError> {
383    Selector::new(jsonb::RawJsonb::new(jsonb_bytes))
384        .select_values(path)
385        .map_err(|e| {
386            ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}"))
387        })
388}
389
390/// Convert an Arrow JSON field to Lance JSON field (with JSONB storage)
391pub fn arrow_json_to_lance_json(field: &ArrowField) -> ArrowField {
392    if is_arrow_json_field(field) {
393        return field_with_extension(field, DataType::LargeBinary, JSON_EXT_NAME);
394    }
395
396    let data_type = match field.data_type() {
397        DataType::Struct(fields) => {
398            let fields = fields
399                .iter()
400                .map(|field| Arc::new(arrow_json_to_lance_json(field)))
401                .collect::<Vec<_>>();
402            DataType::Struct(Fields::from(fields))
403        }
404        DataType::List(item) => DataType::List(Arc::new(arrow_json_to_lance_json(item))),
405        DataType::LargeList(item) => DataType::LargeList(Arc::new(arrow_json_to_lance_json(item))),
406        DataType::FixedSizeList(item, size) => {
407            DataType::FixedSizeList(Arc::new(arrow_json_to_lance_json(item)), *size)
408        }
409        DataType::Map(entries, keys_sorted) => {
410            DataType::Map(Arc::new(arrow_json_to_lance_json(entries)), *keys_sorted)
411        }
412        _ => return field.clone(),
413    };
414
415    field_with_data_type(field, data_type)
416}
417
418/// Convert a Lance JSON field to Arrow JSON field.
419pub fn lance_json_to_arrow_json(field: &ArrowField) -> ArrowField {
420    if is_json_field(field) {
421        return field_with_extension(field, DataType::Utf8, ARROW_JSON_EXT_NAME);
422    }
423
424    let data_type = match field.data_type() {
425        DataType::Struct(fields) => {
426            let fields = fields
427                .iter()
428                .map(|field| Arc::new(lance_json_to_arrow_json(field)))
429                .collect::<Vec<_>>();
430            DataType::Struct(Fields::from(fields))
431        }
432        DataType::List(item) => DataType::List(Arc::new(lance_json_to_arrow_json(item))),
433        DataType::LargeList(item) => DataType::LargeList(Arc::new(lance_json_to_arrow_json(item))),
434        DataType::FixedSizeList(item, size) => {
435            DataType::FixedSizeList(Arc::new(lance_json_to_arrow_json(item)), *size)
436        }
437        DataType::Map(entries, keys_sorted) => {
438            DataType::Map(Arc::new(lance_json_to_arrow_json(entries)), *keys_sorted)
439        }
440        _ => return field.clone(),
441    };
442
443    field_with_data_type(field, data_type)
444}
445
446fn field_with_data_type(field: &ArrowField, data_type: DataType) -> ArrowField {
447    ArrowField::new(field.name(), data_type, field.is_nullable())
448        .with_metadata(field.metadata().clone())
449}
450
451fn field_with_extension(
452    field: &ArrowField,
453    data_type: DataType,
454    extension_name: &str,
455) -> ArrowField {
456    let mut metadata = field.metadata().clone();
457    metadata.insert(ARROW_EXT_NAME_KEY.to_string(), extension_name.to_string());
458    ArrowField::new(field.name(), data_type, field.is_nullable()).with_metadata(metadata)
459}
460
461fn convert_json_array<F>(
462    field: &ArrowField,
463    array: &ArrayRef,
464    convert_leaf: &F,
465) -> Result<(ArrowField, ArrayRef, bool), ArrowError>
466where
467    F: Fn(&ArrowField, &ArrayRef) -> Result<Option<(ArrowField, ArrayRef)>, ArrowError>,
468{
469    if let Some((field, array)) = convert_leaf(field, array)? {
470        return Ok((field, array, true));
471    }
472
473    match field.data_type() {
474        DataType::Struct(fields) => {
475            let struct_array = array.as_struct();
476            let mut new_fields = Vec::with_capacity(fields.len());
477            let mut new_columns = Vec::with_capacity(fields.len());
478            let mut changed = false;
479
480            for (field, column) in fields.iter().zip(struct_array.columns()) {
481                let (new_field, new_column, field_changed) =
482                    convert_json_array(field, column, convert_leaf)?;
483                changed |= field_changed;
484                new_fields.push(Arc::new(new_field));
485                new_columns.push(new_column);
486            }
487
488            if changed {
489                let fields = Fields::from(new_fields);
490                let new_field = field_with_data_type(field, DataType::Struct(fields.clone()));
491                let new_array =
492                    StructArray::new(fields, new_columns, struct_array.nulls().cloned());
493                Ok((new_field, Arc::new(new_array) as ArrayRef, true))
494            } else {
495                Ok((field.clone(), array.clone(), false))
496            }
497        }
498        DataType::List(item) => {
499            let list_array: &ListArray = array.as_list();
500            let (new_item, new_values, changed) =
501                convert_json_array(item, list_array.values(), convert_leaf)?;
502            if changed {
503                let new_field =
504                    field_with_data_type(field, DataType::List(Arc::new(new_item.clone())));
505                let new_array = ListArray::new(
506                    Arc::new(new_item),
507                    list_array.offsets().clone(),
508                    new_values,
509                    list_array.nulls().cloned(),
510                );
511                Ok((new_field, Arc::new(new_array) as ArrayRef, true))
512            } else {
513                Ok((field.clone(), array.clone(), false))
514            }
515        }
516        DataType::LargeList(item) => {
517            let list_array: &LargeListArray = array.as_list();
518            let (new_item, new_values, changed) =
519                convert_json_array(item, list_array.values(), convert_leaf)?;
520            if changed {
521                let new_field =
522                    field_with_data_type(field, DataType::LargeList(Arc::new(new_item.clone())));
523                let new_array = LargeListArray::new(
524                    Arc::new(new_item),
525                    list_array.offsets().clone(),
526                    new_values,
527                    list_array.nulls().cloned(),
528                );
529                Ok((new_field, Arc::new(new_array) as ArrayRef, true))
530            } else {
531                Ok((field.clone(), array.clone(), false))
532            }
533        }
534        DataType::FixedSizeList(item, size) => {
535            let list_array: &FixedSizeListArray = array.as_fixed_size_list();
536            let (new_item, new_values, changed) =
537                convert_json_array(item, list_array.values(), convert_leaf)?;
538            if changed {
539                let new_field = field_with_data_type(
540                    field,
541                    DataType::FixedSizeList(Arc::new(new_item.clone()), *size),
542                );
543                let new_array = FixedSizeListArray::try_new_with_length(
544                    Arc::new(new_item),
545                    *size,
546                    new_values,
547                    list_array.nulls().cloned(),
548                    list_array.len(),
549                )?;
550                Ok((new_field, Arc::new(new_array) as ArrayRef, true))
551            } else {
552                Ok((field.clone(), array.clone(), false))
553            }
554        }
555        DataType::Map(entries, keys_sorted) => {
556            let map_array = array
557                .as_any()
558                .downcast_ref::<MapArray>()
559                .expect("DataType::Map array must be MapArray");
560            let entries_array = Arc::new(map_array.entries().clone()) as ArrayRef;
561            let (new_entries, new_entries_array, changed) =
562                convert_json_array(entries, &entries_array, convert_leaf)?;
563            if changed {
564                let entries_struct = new_entries_array
565                    .as_any()
566                    .downcast_ref::<StructArray>()
567                    .expect("Map entries must be StructArray")
568                    .clone();
569                let new_field = field_with_data_type(
570                    field,
571                    DataType::Map(Arc::new(new_entries.clone()), *keys_sorted),
572                );
573                let new_array = MapArray::new(
574                    Arc::new(new_entries),
575                    map_array.offsets().clone(),
576                    entries_struct,
577                    map_array.nulls().cloned(),
578                    *keys_sorted,
579                );
580                Ok((new_field, Arc::new(new_array) as ArrayRef, true))
581            } else {
582                Ok((field.clone(), array.clone(), false))
583            }
584        }
585        _ => Ok((field.clone(), array.clone(), false)),
586    }
587}
588
589fn convert_arrow_json_array(
590    field: &ArrowField,
591    array: &ArrayRef,
592) -> Result<(ArrowField, ArrayRef, bool), ArrowError> {
593    convert_json_array(field, array, &|field, array| {
594        if is_arrow_json_field(field) {
595            let json_array = JsonArray::try_from(array.clone())?;
596            Ok(Some((
597                arrow_json_to_lance_json(field),
598                Arc::new(json_array.into_inner()) as ArrayRef,
599            )))
600        } else {
601            Ok(None)
602        }
603    })
604}
605
606fn convert_lance_json_array(
607    field: &ArrowField,
608    array: &ArrayRef,
609) -> Result<(ArrowField, ArrayRef, bool), ArrowError> {
610    convert_json_array(field, array, &|field, array| {
611        if is_json_field(field) {
612            Ok(Some((
613                lance_json_to_arrow_json(field),
614                Arc::new(decode_jsonb_array::<i32>(array.as_binary())) as ArrayRef,
615            )))
616        } else {
617            Ok(None)
618        }
619    })
620}
621
622/// Convert a RecordBatch with Lance JSON columns (JSONB) back to Arrow JSON format (strings)
623pub fn convert_lance_json_to_arrow(
624    batch: &arrow_array::RecordBatch,
625) -> Result<arrow_array::RecordBatch, ArrowError> {
626    let schema = batch.schema();
627    let mut needs_conversion = false;
628    let mut new_fields = Vec::with_capacity(schema.fields().len());
629    let mut new_columns = Vec::with_capacity(batch.num_columns());
630
631    for (i, field) in schema.fields().iter().enumerate() {
632        let column = batch.column(i);
633        let (new_field, new_column, changed) = convert_lance_json_array(field, column)?;
634
635        needs_conversion |= changed;
636        new_fields.push(new_field);
637        new_columns.push(new_column);
638    }
639
640    if needs_conversion {
641        let new_schema = Arc::new(Schema::new_with_metadata(
642            new_fields,
643            schema.metadata().clone(),
644        ));
645        RecordBatch::try_new(new_schema, new_columns)
646    } else {
647        // No conversion needed, return original batch
648        Ok(batch.clone())
649    }
650}
651
652/// Convert a RecordBatch with Arrow JSON columns to Lance JSON format (JSONB)
653pub fn convert_json_columns(
654    batch: &arrow_array::RecordBatch,
655) -> Result<arrow_array::RecordBatch, ArrowError> {
656    let schema = batch.schema();
657    let mut needs_conversion = false;
658    let mut new_fields = Vec::with_capacity(schema.fields().len());
659    let mut new_columns = Vec::with_capacity(batch.num_columns());
660
661    for (i, field) in schema.fields().iter().enumerate() {
662        let column = batch.column(i);
663        let (new_field, new_column, changed) = convert_arrow_json_array(field, column)?;
664
665        needs_conversion |= changed;
666        new_fields.push(new_field);
667        new_columns.push(new_column);
668    }
669
670    if needs_conversion {
671        let new_schema = Arc::new(Schema::new_with_metadata(
672            new_fields,
673            schema.metadata().clone(),
674        ));
675        RecordBatch::try_new(new_schema, new_columns)
676    } else {
677        // No conversion needed, return original batch
678        Ok(batch.clone())
679    }
680}
681
682#[cfg(test)]
683mod tests {
684    use super::*;
685
686    #[test]
687    fn test_json_field_creation() {
688        let field = json_field("data", true);
689        assert_eq!(field.name(), "data");
690        assert_eq!(field.data_type(), &DataType::LargeBinary);
691        assert!(field.is_nullable());
692        assert!(is_json_field(&field));
693    }
694
695    #[test]
696    fn test_json_array_from_strings() {
697        let json_strings = vec![
698            Some(r#"{"name": "Alice", "age": 30}"#),
699            None,
700            Some(r#"{"name": "Bob", "age": 25}"#),
701        ];
702
703        let array = JsonArray::try_from_iter(json_strings).unwrap();
704        assert_eq!(array.len(), 3);
705        assert!(!array.is_null(0));
706        assert!(array.is_null(1));
707        assert!(!array.is_null(2));
708
709        let decoded = array.value(0).unwrap();
710        assert!(decoded.contains("Alice"));
711    }
712
713    #[test]
714    fn test_json_array_from_string_array() {
715        let string_array = StringArray::from(vec![
716            Some(r#"{"name": "Alice"}"#),
717            Some(r#"{"name": "Bob"}"#),
718            None,
719        ]);
720
721        let json_array = JsonArray::try_from(string_array).unwrap();
722        assert_eq!(json_array.len(), 3);
723        assert!(!json_array.is_null(0));
724        assert!(!json_array.is_null(1));
725        assert!(json_array.is_null(2));
726    }
727
728    #[test]
729    fn test_json_path_extraction() {
730        let json_array = JsonArray::try_from_iter(vec![
731            Some(r#"{"user": {"name": "Alice", "age": 30}}"#),
732            Some(r#"{"user": {"name": "Bob"}}"#),
733        ])
734        .unwrap();
735
736        let name = json_array.json_path(0, "$.user.name").unwrap();
737        assert_eq!(name, Some("\"Alice\"".to_string()));
738
739        let age = json_array.json_path(1, "$.user.age").unwrap();
740        assert_eq!(age, None);
741
742        // Several matches are selected together, as the SQL JSON functions do.
743        let json_array =
744            JsonArray::try_from_iter(vec![Some(r#"{"a": [{"b": 1}, {"b": 2}]}"#)]).unwrap();
745        let bs = json_array.json_path(0, "$.a[*].b").unwrap();
746        assert_eq!(bs, Some("[1,2]".to_string()));
747    }
748
749    #[test]
750    fn test_json_values_read_either_encoding() {
751        let docs = vec![Some(r#"{"a":1}"#), None];
752        let text: ArrayRef = Arc::new(StringArray::from(docs.clone()));
753        let jsonb: ArrayRef = Arc::new(JsonArray::try_from_iter(docs).unwrap().into_inner());
754        let mut text_field = ArrowField::new("j", DataType::Utf8, true);
755        text_field.set_metadata(std::collections::HashMap::from([(
756            ARROW_EXT_NAME_KEY.to_string(),
757            ARROW_JSON_EXT_NAME.to_string(),
758        )]));
759        let jsonb_field = json_field("j", true);
760
761        let from_text = JsonValues::try_new(&text_field, &text).unwrap();
762        let from_jsonb = JsonValues::try_new(&jsonb_field, &jsonb).unwrap();
763        assert_eq!(from_text.encoding(), JsonEncoding::Text);
764        assert_eq!(from_jsonb.encoding(), JsonEncoding::Jsonb);
765
766        assert_eq!(&from_text.to_jsonb().unwrap(), jsonb.as_binary::<i64>());
767        assert_eq!(&from_jsonb.to_jsonb().unwrap(), jsonb.as_binary::<i64>());
768        assert_eq!(&from_text.to_text(), &text);
769        let decoded = from_jsonb.to_text();
770        let decoded = decoded.as_string::<i64>();
771        assert_eq!(decoded.value(0), r#"{"a":1}"#);
772        assert!(decoded.is_null(1));
773
774        let plain = ArrowField::new("j", DataType::Utf8, true);
775        assert!(JsonValues::try_new(&plain, &text).is_none());
776    }
777
778    #[test]
779    fn test_convert_json_columns() {
780        // Create a batch with Arrow JSON column
781        let json_strings = vec![Some(r#"{"name": "Alice"}"#), Some(r#"{"name": "Bob"}"#)];
782        let json_arr = StringArray::from(json_strings);
783
784        // Create field with arrow.json extension
785        let mut field = ArrowField::new("data", DataType::Utf8, false);
786        let mut metadata = std::collections::HashMap::new();
787        metadata.insert(
788            ARROW_EXT_NAME_KEY.to_string(),
789            ARROW_JSON_EXT_NAME.to_string(),
790        );
791        field.set_metadata(metadata);
792
793        let schema = Arc::new(Schema::new(vec![field]));
794        let batch = RecordBatch::try_new(schema, vec![Arc::new(json_arr) as ArrayRef]).unwrap();
795
796        // Convert the batch
797        let converted = convert_json_columns(&batch).unwrap();
798
799        // Check the converted schema
800        assert_eq!(converted.num_columns(), 1);
801        let converted_schema = converted.schema();
802        let converted_field = converted_schema.field(0);
803        assert_eq!(converted_field.data_type(), &DataType::LargeBinary);
804        assert_eq!(
805            converted_field.metadata().get(ARROW_EXT_NAME_KEY),
806            Some(&JSON_EXT_NAME.to_string())
807        );
808
809        // Check the data was converted
810        let converted_column = converted.column(0);
811        assert_eq!(converted_column.data_type(), &DataType::LargeBinary);
812        assert_eq!(converted_column.len(), 2);
813
814        // Verify the data is valid JSONB
815        let binary_array = converted_column
816            .as_any()
817            .downcast_ref::<LargeBinaryArray>()
818            .unwrap();
819        for i in 0..binary_array.len() {
820            let jsonb_bytes = binary_array.value(i);
821            let decoded = decode_json(jsonb_bytes);
822            assert!(decoded.contains("name"));
823        }
824    }
825
826    #[test]
827    fn test_convert_nested_json_columns() {
828        use arrow_buffer::{OffsetBuffer, ScalarBuffer};
829
830        let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, false));
831        let mut metadata = std::collections::HashMap::new();
832        metadata.insert(
833            ARROW_EXT_NAME_KEY.to_string(),
834            ARROW_JSON_EXT_NAME.to_string(),
835        );
836        let extra_field =
837            Arc::new(ArrowField::new("extra", DataType::Utf8, true).with_metadata(metadata));
838        let item_fields = Fields::from(vec![uri_field, extra_field]);
839
840        let values = StructArray::new(
841            item_fields.clone(),
842            vec![
843                Arc::new(StringArray::from(vec![Some("a.jpg"), Some("b.jpg")])) as ArrayRef,
844                Arc::new(StringArray::from(vec![
845                    Some(r#"{"codec":"h264"}"#),
846                    None::<&str>,
847                ])) as ArrayRef,
848            ],
849            None,
850        );
851        let item = Arc::new(ArrowField::new("item", DataType::Struct(item_fields), true));
852        let media = ListArray::new(
853            item,
854            OffsetBuffer::new(ScalarBuffer::from(vec![0, 1, 2])),
855            Arc::new(values),
856            None,
857        );
858        let schema = Arc::new(Schema::new(vec![ArrowField::new(
859            "media",
860            media.data_type().clone(),
861            true,
862        )]));
863        let batch = RecordBatch::try_new(schema, vec![Arc::new(media) as ArrayRef]).unwrap();
864
865        assert!(has_arrow_json_fields(batch.schema().field(0)));
866
867        let converted = convert_json_columns(&batch).unwrap();
868        let converted_schema = converted.schema();
869        let DataType::List(item) = converted_schema.field(0).data_type() else {
870            panic!("expected list field");
871        };
872        let DataType::Struct(fields) = item.data_type() else {
873            panic!("expected struct item");
874        };
875        assert!(is_json_field(&fields[1]));
876
877        let list_array: &ListArray = converted.column(0).as_list();
878        let values = list_array.values().as_struct();
879        let extra = values
880            .column(1)
881            .as_any()
882            .downcast_ref::<LargeBinaryArray>()
883            .unwrap();
884        assert!(decode_json(extra.value(0)).contains("h264"));
885        assert!(extra.is_null(1));
886
887        let logical = convert_lance_json_to_arrow(&converted).unwrap();
888        let logical_schema = logical.schema();
889        let DataType::List(item) = logical_schema.field(0).data_type() else {
890            panic!("expected list field");
891        };
892        let DataType::Struct(fields) = item.data_type() else {
893            panic!("expected struct item");
894        };
895        assert!(is_arrow_json_field(&fields[1]));
896
897        let list_array: &ListArray = logical.column(0).as_list();
898        let values = list_array.values().as_struct();
899        let extra = values
900            .column(1)
901            .as_any()
902            .downcast_ref::<StringArray>()
903            .unwrap();
904        assert!(extra.value(0).contains("h264"));
905        assert!(extra.is_null(1));
906    }
907
908    #[test]
909    fn test_convert_fixed_size_list_zero_json_preserves_length() {
910        let mut metadata = std::collections::HashMap::new();
911        metadata.insert(
912            ARROW_EXT_NAME_KEY.to_string(),
913            ARROW_JSON_EXT_NAME.to_string(),
914        );
915        let item = Arc::new(ArrowField::new("item", DataType::Utf8, true).with_metadata(metadata));
916        let values = Arc::new(StringArray::from(Vec::<Option<&str>>::new())) as ArrayRef;
917        let lists = FixedSizeListArray::try_new_with_length(item, 0, values, None, 3).unwrap();
918        let schema = Arc::new(Schema::new(vec![ArrowField::new(
919            "lists",
920            lists.data_type().clone(),
921            true,
922        )]));
923        let batch = RecordBatch::try_new(schema, vec![Arc::new(lists) as ArrayRef]).unwrap();
924
925        let converted = convert_json_columns(&batch).unwrap();
926        assert_eq!(converted.num_rows(), 3);
927        assert_eq!(converted.column(0).len(), 3);
928
929        let converted_schema = converted.schema();
930        let DataType::FixedSizeList(item, size) = converted_schema.field(0).data_type() else {
931            panic!("expected fixed size list field");
932        };
933        assert_eq!(*size, 0);
934        assert!(is_json_field(item));
935
936        let logical = convert_lance_json_to_arrow(&converted).unwrap();
937        assert_eq!(logical.num_rows(), 3);
938        assert_eq!(logical.column(0).len(), 3);
939
940        let logical_schema = logical.schema();
941        let DataType::FixedSizeList(item, size) = logical_schema.field(0).data_type() else {
942            panic!("expected fixed size list field");
943        };
944        assert_eq!(*size, 0);
945        assert!(is_arrow_json_field(item));
946    }
947
948    #[test]
949    fn test_has_json_fields() {
950        // Test direct JSON field
951        let json_f = json_field("data", true);
952        assert!(has_json_fields(&json_f));
953
954        // Test non-JSON field
955        let non_json = ArrowField::new("data", DataType::Utf8, true);
956        assert!(!has_json_fields(&non_json));
957
958        // Test struct containing JSON field
959        let struct_field = ArrowField::new(
960            "struct",
961            DataType::Struct(vec![json_field("nested_json", true)].into()),
962            true,
963        );
964        assert!(has_json_fields(&struct_field));
965
966        // Test struct without JSON field
967        let struct_no_json = ArrowField::new(
968            "struct",
969            DataType::Struct(vec![ArrowField::new("text", DataType::Utf8, true)].into()),
970            true,
971        );
972        assert!(!has_json_fields(&struct_no_json));
973
974        // Test List containing JSON field
975        let list_field = ArrowField::new(
976            "list",
977            DataType::List(Arc::new(json_field("item", true))),
978            true,
979        );
980        assert!(has_json_fields(&list_field));
981
982        // Test LargeList containing JSON field
983        let large_list_field = ArrowField::new(
984            "large_list",
985            DataType::LargeList(Arc::new(json_field("item", true))),
986            true,
987        );
988        assert!(has_json_fields(&large_list_field));
989
990        // Test FixedSizeList containing JSON field
991        let fixed_list_field = ArrowField::new(
992            "fixed_list",
993            DataType::FixedSizeList(Arc::new(json_field("item", true)), 3),
994            true,
995        );
996        assert!(has_json_fields(&fixed_list_field));
997
998        // Test Map containing JSON field
999        let map_field = ArrowField::new(
1000            "map",
1001            DataType::Map(
1002                Arc::new(ArrowField::new(
1003                    "entries",
1004                    DataType::Struct(
1005                        vec![
1006                            ArrowField::new("key", DataType::Utf8, false),
1007                            json_field("value", true),
1008                        ]
1009                        .into(),
1010                    ),
1011                    false,
1012                )),
1013                false,
1014            ),
1015            true,
1016        );
1017        assert!(has_json_fields(&map_field));
1018    }
1019
1020    #[test]
1021    fn test_json_array_inner() {
1022        let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap();
1023        let inner = json_array.inner();
1024        assert_eq!(inner.len(), 1);
1025    }
1026
1027    #[test]
1028    fn test_json_array_value_null_error() {
1029        let json_array = JsonArray::try_from_iter(vec![None::<&str>]).unwrap();
1030        let result = json_array.value(0);
1031        assert!(result.is_err());
1032        assert!(result.unwrap_err().to_string().contains("null"));
1033    }
1034
1035    #[test]
1036    fn test_json_array_value_bytes() {
1037        let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap();
1038        let bytes = json_array.value_bytes(0);
1039        assert!(!bytes.is_empty());
1040    }
1041
1042    #[test]
1043    fn test_json_path_with_null() {
1044        let json_array =
1045            JsonArray::try_from_iter(vec![Some(r#"{"user": {"name": "Alice"}}"#), None::<&str>])
1046                .unwrap();
1047
1048        let result = json_array.json_path(1, "$.user.name").unwrap();
1049        assert_eq!(result, None);
1050    }
1051
1052    #[test]
1053    fn test_to_arrow_json() {
1054        let json_array = JsonArray::try_from_iter(vec![
1055            Some(r#"{"name": "Alice"}"#),
1056            None::<&str>,
1057            Some(r#"{"name": "Bob"}"#),
1058        ])
1059        .unwrap();
1060
1061        let arrow_json = json_array.to_arrow_json();
1062        assert_eq!(arrow_json.len(), 3);
1063        assert!(!arrow_json.is_null(0));
1064        assert!(arrow_json.is_null(1));
1065        assert!(!arrow_json.is_null(2));
1066
1067        let string_array = arrow_json.as_any().downcast_ref::<StringArray>().unwrap();
1068        assert!(string_array.value(0).contains("Alice"));
1069        assert!(string_array.value(2).contains("Bob"));
1070    }
1071
1072    #[test]
1073    fn test_json_array_trait_methods() {
1074        let json_array =
1075            JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#), Some(r#"{"b": 2}"#)]).unwrap();
1076
1077        // Wrapper methods
1078        assert_eq!(json_array.len(), 2);
1079        assert!(!json_array.is_empty());
1080        assert!(!json_array.is_null(0));
1081
1082        // Underlying Arrow array
1083        assert_eq!(json_array.inner().data_type(), &DataType::LargeBinary);
1084        assert_eq!(json_array.inner().len(), 2);
1085    }
1086
1087    #[test]
1088    fn test_json_array_empty() {
1089        let json_array = JsonArray::try_from_iter(Vec::<Option<&str>>::new()).unwrap();
1090        assert!(json_array.is_empty());
1091        assert_eq!(json_array.len(), 0);
1092    }
1093
1094    #[test]
1095    fn test_try_from_large_string_array() {
1096        let large_string_array = LargeStringArray::from(vec![
1097            Some(r#"{"name": "Alice"}"#),
1098            Some(r#"{"name": "Bob"}"#),
1099            None,
1100        ]);
1101
1102        // Test TryFrom<&LargeStringArray>
1103        let json_array = JsonArray::try_from(&large_string_array).unwrap();
1104        assert_eq!(json_array.len(), 3);
1105        assert!(!json_array.is_null(0));
1106        assert!(!json_array.is_null(1));
1107        assert!(json_array.is_null(2));
1108
1109        // Test TryFrom<LargeStringArray> (owned)
1110        let large_string_array2 = LargeStringArray::from(vec![Some(r#"{"x": 1}"#)]);
1111        let json_array2 = JsonArray::try_from(large_string_array2).unwrap();
1112        assert_eq!(json_array2.len(), 1);
1113    }
1114
1115    #[test]
1116    fn test_try_from_array_ref() {
1117        // Test with Utf8
1118        let string_array: ArrayRef = Arc::new(StringArray::from(vec![
1119            Some(r#"{"a": 1}"#),
1120            Some(r#"{"b": 2}"#),
1121        ]));
1122        let json_array = JsonArray::try_from(string_array).unwrap();
1123        assert_eq!(json_array.len(), 2);
1124
1125        // Test with LargeUtf8
1126        let large_string_array: ArrayRef = Arc::new(LargeStringArray::from(vec![
1127            Some(r#"{"c": 3}"#),
1128            Some(r#"{"d": 4}"#),
1129        ]));
1130        let json_array2 = JsonArray::try_from(large_string_array).unwrap();
1131        assert_eq!(json_array2.len(), 2);
1132
1133        // Test with unsupported type
1134        let int_array: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![1, 2, 3]));
1135        let result = JsonArray::try_from(int_array);
1136        assert!(result.is_err());
1137        assert!(result.unwrap_err().to_string().contains("Unsupported"));
1138    }
1139
1140    #[test]
1141    fn test_arrow_json_to_lance_json_non_json_field() {
1142        // Test that non-JSON fields are returned unchanged
1143        let field = ArrowField::new("text", DataType::Utf8, true);
1144        let converted = arrow_json_to_lance_json(&field);
1145        assert_eq!(converted.data_type(), &DataType::Utf8);
1146        assert_eq!(converted.name(), "text");
1147    }
1148
1149    #[test]
1150    fn test_convert_lance_json_to_arrow() {
1151        // Create a batch with Lance JSON column (JSONB)
1152        let json_array = JsonArray::try_from_iter(vec![
1153            Some(r#"{"name": "Alice"}"#),
1154            None::<&str>,
1155            Some(r#"{"name": "Bob"}"#),
1156        ])
1157        .unwrap();
1158
1159        let lance_json_field = json_field("data", true);
1160        let schema = Arc::new(Schema::new(vec![lance_json_field]));
1161        let batch =
1162            RecordBatch::try_new(schema, vec![Arc::new(json_array.into_inner()) as ArrayRef])
1163                .unwrap();
1164
1165        // Convert back to Arrow JSON
1166        let converted = convert_lance_json_to_arrow(&batch).unwrap();
1167
1168        // Check schema
1169        let converted_schema = converted.schema();
1170        let converted_field = converted_schema.field(0);
1171        assert_eq!(converted_field.data_type(), &DataType::Utf8);
1172        assert_eq!(
1173            converted_field.metadata().get(ARROW_EXT_NAME_KEY),
1174            Some(&ARROW_JSON_EXT_NAME.to_string())
1175        );
1176
1177        // Check data
1178        let string_array = converted
1179            .column(0)
1180            .as_any()
1181            .downcast_ref::<StringArray>()
1182            .unwrap();
1183        assert!(!string_array.is_null(0));
1184        assert!(string_array.is_null(1));
1185        assert!(!string_array.is_null(2));
1186        assert!(string_array.value(0).contains("Alice"));
1187        assert!(string_array.value(2).contains("Bob"));
1188    }
1189
1190    #[test]
1191    fn test_convert_lance_json_to_arrow_empty_batch() {
1192        // Create an empty batch with Lance JSON column
1193        let lance_json_field = json_field("data", true);
1194        let schema = Arc::new(Schema::new(vec![lance_json_field]));
1195        let empty_binary = LargeBinaryBuilder::new().finish();
1196        let batch = RecordBatch::try_new(schema, vec![Arc::new(empty_binary) as ArrayRef]).unwrap();
1197
1198        // Convert back to Arrow JSON
1199        let converted = convert_lance_json_to_arrow(&batch).unwrap();
1200        assert_eq!(converted.num_rows(), 0);
1201        assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8);
1202    }
1203
1204    #[test]
1205    fn test_convert_lance_json_to_arrow_no_json_columns() {
1206        // Create a batch without JSON columns
1207        let field = ArrowField::new("text", DataType::Utf8, true);
1208        let schema = Arc::new(Schema::new(vec![field]));
1209        let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
1210        let batch = RecordBatch::try_new(schema, vec![Arc::new(string_array) as ArrayRef]).unwrap();
1211
1212        // Convert - should return the same batch
1213        let converted = convert_lance_json_to_arrow(&batch).unwrap();
1214        assert_eq!(converted.num_columns(), 1);
1215        assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8);
1216    }
1217
1218    #[test]
1219    fn test_convert_json_columns_empty_batch() {
1220        // Create an empty batch with Arrow JSON column
1221        let mut field = ArrowField::new("data", DataType::Utf8, false);
1222        let mut metadata = std::collections::HashMap::new();
1223        metadata.insert(
1224            ARROW_EXT_NAME_KEY.to_string(),
1225            ARROW_JSON_EXT_NAME.to_string(),
1226        );
1227        field.set_metadata(metadata);
1228
1229        let schema = Arc::new(Schema::new(vec![field]));
1230        let empty_strings = arrow_array::builder::StringBuilder::new().finish();
1231        let batch =
1232            RecordBatch::try_new(schema, vec![Arc::new(empty_strings) as ArrayRef]).unwrap();
1233
1234        let converted = convert_json_columns(&batch).unwrap();
1235        assert_eq!(converted.num_rows(), 0);
1236        assert_eq!(
1237            converted.schema().field(0).data_type(),
1238            &DataType::LargeBinary
1239        );
1240    }
1241
1242    #[test]
1243    fn test_convert_json_columns_large_string() {
1244        // Create a batch with Arrow JSON column using LargeUtf8
1245        let json_strings = LargeStringArray::from(vec![
1246            Some(r#"{"name": "Alice"}"#),
1247            Some(r#"{"name": "Bob"}"#),
1248        ]);
1249
1250        let mut field = ArrowField::new("data", DataType::LargeUtf8, false);
1251        let mut metadata = std::collections::HashMap::new();
1252        metadata.insert(
1253            ARROW_EXT_NAME_KEY.to_string(),
1254            ARROW_JSON_EXT_NAME.to_string(),
1255        );
1256        field.set_metadata(metadata);
1257
1258        let schema = Arc::new(Schema::new(vec![field]));
1259        let batch = RecordBatch::try_new(schema, vec![Arc::new(json_strings) as ArrayRef]).unwrap();
1260
1261        let converted = convert_json_columns(&batch).unwrap();
1262        assert_eq!(converted.num_columns(), 1);
1263        assert_eq!(
1264            converted.schema().field(0).data_type(),
1265            &DataType::LargeBinary
1266        );
1267        assert_eq!(converted.num_rows(), 2);
1268    }
1269
1270    #[test]
1271    fn test_convert_json_columns_no_json_columns() {
1272        // Create a batch without JSON columns
1273        let field = ArrowField::new("text", DataType::Utf8, true);
1274        let schema = Arc::new(Schema::new(vec![field]));
1275        let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
1276        let batch = RecordBatch::try_new(schema, vec![Arc::new(string_array) as ArrayRef]).unwrap();
1277
1278        // Convert - should return the same batch
1279        let converted = convert_json_columns(&batch).unwrap();
1280        assert_eq!(converted.num_columns(), 1);
1281        assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8);
1282    }
1283
1284    #[test]
1285    fn test_convert_json_columns_mixed_columns() {
1286        // Create a batch with both JSON and non-JSON columns
1287        let json_strings = StringArray::from(vec![
1288            Some(r#"{"name": "Alice"}"#),
1289            Some(r#"{"name": "Bob"}"#),
1290        ]);
1291        let text_strings = StringArray::from(vec![Some("hello"), Some("world")]);
1292
1293        let mut json_field = ArrowField::new("json_data", DataType::Utf8, false);
1294        let mut metadata = std::collections::HashMap::new();
1295        metadata.insert(
1296            ARROW_EXT_NAME_KEY.to_string(),
1297            ARROW_JSON_EXT_NAME.to_string(),
1298        );
1299        json_field.set_metadata(metadata);
1300
1301        let text_field = ArrowField::new("text_data", DataType::Utf8, true);
1302
1303        let schema = Arc::new(Schema::new(vec![json_field, text_field]));
1304        let batch = RecordBatch::try_new(
1305            schema,
1306            vec![
1307                Arc::new(json_strings) as ArrayRef,
1308                Arc::new(text_strings) as ArrayRef,
1309            ],
1310        )
1311        .unwrap();
1312
1313        let converted = convert_json_columns(&batch).unwrap();
1314        assert_eq!(converted.num_columns(), 2);
1315        assert_eq!(
1316            converted.schema().field(0).data_type(),
1317            &DataType::LargeBinary
1318        );
1319        assert_eq!(converted.schema().field(1).data_type(), &DataType::Utf8);
1320    }
1321
1322    #[test]
1323    fn test_is_arrow_json_field_large_utf8() {
1324        // Test with LargeUtf8 storage type
1325        let mut field = ArrowField::new("data", DataType::LargeUtf8, true);
1326        let mut metadata = std::collections::HashMap::new();
1327        metadata.insert(
1328            ARROW_EXT_NAME_KEY.to_string(),
1329            ARROW_JSON_EXT_NAME.to_string(),
1330        );
1331        field.set_metadata(metadata);
1332
1333        assert!(is_arrow_json_field(&field));
1334    }
1335
1336    #[test]
1337    fn test_encode_json_invalid() {
1338        // Test encoding invalid JSON
1339        let result = encode_json("not valid json {");
1340        assert!(result.is_err());
1341    }
1342
1343    #[test]
1344    fn test_json_array_from_invalid_json() {
1345        // Test creating JsonArray from invalid JSON strings
1346        let result = JsonArray::try_from_iter(vec![Some("invalid json {")]);
1347        assert!(result.is_err());
1348        assert!(result.unwrap_err().to_string().contains("Failed to encode"));
1349    }
1350
1351    #[test]
1352    fn test_try_from_string_array_invalid_json() {
1353        let string_array = StringArray::from(vec![Some("invalid json {")]);
1354        let result = JsonArray::try_from(string_array);
1355        assert!(result.is_err());
1356    }
1357
1358    #[test]
1359    fn test_try_from_large_string_array_invalid_json() {
1360        let large_string_array = LargeStringArray::from(vec![Some("invalid json {")]);
1361        let result = JsonArray::try_from(large_string_array);
1362        assert!(result.is_err());
1363    }
1364
1365    #[test]
1366    fn test_convert_lance_json_to_arrow_mixed_columns() {
1367        // Create a batch with both JSON and non-JSON columns
1368        let json_array = JsonArray::try_from_iter(vec![
1369            Some(r#"{"name": "Alice"}"#),
1370            Some(r#"{"name": "Bob"}"#),
1371        ])
1372        .unwrap();
1373        let text_strings = StringArray::from(vec![Some("hello"), Some("world")]);
1374
1375        let json_f = json_field("json_data", true);
1376        let text_field = ArrowField::new("text_data", DataType::Utf8, true);
1377
1378        let schema = Arc::new(Schema::new(vec![json_f, text_field]));
1379        let batch = RecordBatch::try_new(
1380            schema,
1381            vec![
1382                Arc::new(json_array.into_inner()) as ArrayRef,
1383                Arc::new(text_strings) as ArrayRef,
1384            ],
1385        )
1386        .unwrap();
1387
1388        let converted = convert_lance_json_to_arrow(&batch).unwrap();
1389        assert_eq!(converted.num_columns(), 2);
1390        assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8);
1391        assert_eq!(converted.schema().field(1).data_type(), &DataType::Utf8);
1392    }
1393
1394    #[test]
1395    fn test_json_path_invalid_path() {
1396        let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap();
1397        // Invalid JSONPath syntax should return error
1398        let result = json_array.json_path(0, "invalid path without $");
1399        assert!(result.is_err());
1400        assert!(result.unwrap_err().to_string().contains("Invalid JSONPath"));
1401    }
1402
1403    #[test]
1404    fn test_convert_json_columns_invalid_storage_type() {
1405        // Create a batch with Arrow JSON field but wrong storage type (Int32 instead of Utf8)
1406        let int_array = arrow_array::Int32Array::from(vec![1, 2, 3]);
1407
1408        let mut field = ArrowField::new("data", DataType::Int32, false);
1409        let mut metadata = std::collections::HashMap::new();
1410        metadata.insert(
1411            ARROW_EXT_NAME_KEY.to_string(),
1412            ARROW_JSON_EXT_NAME.to_string(),
1413        );
1414        field.set_metadata(metadata);
1415
1416        let schema = Arc::new(Schema::new(vec![field]));
1417        let batch = RecordBatch::try_new(schema, vec![Arc::new(int_array) as ArrayRef]).unwrap();
1418
1419        // This should succeed since Int32 doesn't match is_arrow_json_field check
1420        // (is_arrow_json_field requires Utf8 or LargeUtf8)
1421        let result = convert_json_columns(&batch);
1422        assert!(result.is_ok());
1423    }
1424
1425    #[test]
1426    fn test_is_json_field_wrong_extension() {
1427        // LargeBinary field without the correct extension metadata
1428        let field = ArrowField::new("data", DataType::LargeBinary, true);
1429        assert!(!is_json_field(&field));
1430
1431        // LargeBinary field with wrong extension name
1432        let mut field2 = ArrowField::new("data", DataType::LargeBinary, true);
1433        let mut metadata = std::collections::HashMap::new();
1434        metadata.insert(
1435            ARROW_EXT_NAME_KEY.to_string(),
1436            "other.extension".to_string(),
1437        );
1438        field2.set_metadata(metadata);
1439        assert!(!is_json_field(&field2));
1440    }
1441
1442    #[test]
1443    fn test_is_arrow_json_field_wrong_extension() {
1444        // Utf8 field without extension metadata
1445        let field = ArrowField::new("data", DataType::Utf8, true);
1446        assert!(!is_arrow_json_field(&field));
1447
1448        // Utf8 field with wrong extension name
1449        let mut field2 = ArrowField::new("data", DataType::Utf8, true);
1450        let mut metadata = std::collections::HashMap::new();
1451        metadata.insert(
1452            ARROW_EXT_NAME_KEY.to_string(),
1453            "other.extension".to_string(),
1454        );
1455        field2.set_metadata(metadata);
1456        assert!(!is_arrow_json_field(&field2));
1457
1458        // Wrong type entirely
1459        let field3 = ArrowField::new("data", DataType::Int32, true);
1460        assert!(!is_arrow_json_field(&field3));
1461    }
1462
1463    #[test]
1464    fn test_convert_json_columns_invalid_json_utf8() {
1465        // Test error propagation when converting invalid JSON (Utf8)
1466        let invalid_json = StringArray::from(vec![Some("invalid json {")]);
1467
1468        let mut field = ArrowField::new("data", DataType::Utf8, false);
1469        let mut metadata = std::collections::HashMap::new();
1470        metadata.insert(
1471            ARROW_EXT_NAME_KEY.to_string(),
1472            ARROW_JSON_EXT_NAME.to_string(),
1473        );
1474        field.set_metadata(metadata);
1475
1476        let schema = Arc::new(Schema::new(vec![field]));
1477        let batch = RecordBatch::try_new(schema, vec![Arc::new(invalid_json) as ArrayRef]).unwrap();
1478
1479        let result = convert_json_columns(&batch);
1480        assert!(result.is_err());
1481    }
1482
1483    #[test]
1484    fn test_convert_json_columns_invalid_json_large_utf8() {
1485        // Test error propagation when converting invalid JSON (LargeUtf8)
1486        let invalid_json = LargeStringArray::from(vec![Some("invalid json {")]);
1487
1488        let mut field = ArrowField::new("data", DataType::LargeUtf8, false);
1489        let mut metadata = std::collections::HashMap::new();
1490        metadata.insert(
1491            ARROW_EXT_NAME_KEY.to_string(),
1492            ARROW_JSON_EXT_NAME.to_string(),
1493        );
1494        field.set_metadata(metadata);
1495
1496        let schema = Arc::new(Schema::new(vec![field]));
1497        let batch = RecordBatch::try_new(schema, vec![Arc::new(invalid_json) as ArrayRef]).unwrap();
1498
1499        let result = convert_json_columns(&batch);
1500        assert!(result.is_err());
1501    }
1502
1503    #[test]
1504    fn test_json_path_on_corrupted_jsonb() {
1505        // Create corrupted JSONB bytes directly
1506        let corrupted_bytes: &[u8] = &[0xFF, 0xFE, 0x00, 0x01, 0x02];
1507        let corrupted_binary = LargeBinaryArray::from(vec![Some(corrupted_bytes)]);
1508
1509        // Wrap in JsonArray
1510        let corrupted_json = JsonArray {
1511            inner: corrupted_binary,
1512        };
1513
1514        // Try to use json_path on corrupted data - the selector might fail or return unexpected results
1515        // This exercises the code path but may not produce an error depending on jsonb library behavior
1516        let _result = corrupted_json.json_path(0, "$.a");
1517        // We don't assert on the result as the behavior depends on the jsonb library
1518    }
1519
1520    #[test]
1521    fn test_decode_json_on_various_inputs() {
1522        // Test decode_json with various inputs
1523        let valid_jsonb = encode_json(r#"{"key": "value"}"#).unwrap();
1524        let decoded = decode_json(&valid_jsonb);
1525        assert!(decoded.contains("key"));
1526
1527        // Empty bytes - jsonb library handles this gracefully
1528        let decoded_empty = decode_json(&[]);
1529        // Just verify it doesn't panic
1530        let _ = decoded_empty;
1531
1532        // Random bytes - jsonb library handles this gracefully
1533        let decoded_random = decode_json(&[0xFF, 0xFE, 0x00]);
1534        // Just verify it doesn't panic
1535        let _ = decoded_random;
1536    }
1537}
lance_arrow/json.rs

lance_arrow/
json.rs
