datatypes/vectors/json/
array.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use arrow::compute::{can_cast_types, cast};
18use arrow_array::cast::AsArray;
19use arrow_array::types::{
20    Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type,
21    UInt32Type, UInt64Type,
22};
23use arrow_array::{Array, ArrayRef, GenericListArray, StructArray, new_null_array};
24use arrow_schema::{DataType, Field};
25use serde_json::Value;
26use snafu::{OptionExt, ResultExt};
27
28use crate::arrow_array::{binary_array_value, string_array_value};
29use crate::data_type::{ConcreteDataType, DataType as _};
30use crate::error::{
31    AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result,
32};
33use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field};
34use crate::json::value::decode_json_variant;
35use crate::json::{JsonSettings, TypeHintMismatchPolicy, coerce_json_value_to_type};
36use crate::vectors::MutableVector;
37use crate::vectors::json::builder::{JsonVectorBuilder, json2_physical_data_type};
38use crate::vectors::json::variant::variant_to_json_values;
39
40pub struct JsonArray<'a> {
41    inner: &'a ArrayRef,
42}
43
44impl JsonArray<'_> {
45    /// Try to get the value (as a [Value]) at the index `i`.
46    pub fn try_get_value(&self, i: usize) -> Result<Value> {
47        let array = self.inner;
48        if array.is_null(i) {
49            return Ok(Value::Null);
50        }
51
52        let value = match array.data_type() {
53            DataType::Null => Value::Null,
54            DataType::Boolean => Value::Bool(array.as_boolean().value(i)),
55            DataType::Int8 => Value::from(array.as_primitive::<Int8Type>().value(i)),
56            DataType::Int16 => Value::from(array.as_primitive::<Int16Type>().value(i)),
57            DataType::Int32 => Value::from(array.as_primitive::<Int32Type>().value(i)),
58            DataType::Int64 => Value::from(array.as_primitive::<Int64Type>().value(i)),
59            DataType::UInt8 => Value::from(array.as_primitive::<UInt8Type>().value(i)),
60            DataType::UInt16 => Value::from(array.as_primitive::<UInt16Type>().value(i)),
61            DataType::UInt32 => Value::from(array.as_primitive::<UInt32Type>().value(i)),
62            DataType::UInt64 => Value::from(array.as_primitive::<UInt64Type>().value(i)),
63            DataType::Float32 => Value::from(array.as_primitive::<Float32Type>().value(i)),
64            DataType::Float64 => Value::from(array.as_primitive::<Float64Type>().value(i)),
65            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
66                Value::String(string_array_value(array, i).to_string())
67            }
68            DataType::Binary | DataType::LargeBinary | DataType::BinaryView => {
69                let bytes = binary_array_value(array, i);
70                decode_json_variant(bytes).map_err(|error| InvalidJsonbSnafu { error }.build())?
71            }
72            DataType::Struct(_) => {
73                let structs = array.as_struct();
74                let object = structs
75                    .fields()
76                    .iter()
77                    .zip(structs.columns())
78                    .map(|(field, column)| {
79                        JsonArray::from(column)
80                            .try_get_value(i)
81                            .map(|v| (field.name().clone(), v))
82                    })
83                    .collect::<Result<_>>()?;
84                Value::Object(object)
85            }
86            DataType::List(_) => {
87                let lists = array.as_list::<i32>();
88                let list = lists.value(i);
89                let list = JsonArray::from(&list);
90                let mut values = Vec::with_capacity(list.inner.len());
91                for i in 0..list.inner.len() {
92                    values.push(list.try_get_value(i)?);
93                }
94                Value::Array(values)
95            }
96            t => {
97                return InvalidJsonSnafu {
98                    value: format!("unknown JSON type {t}"),
99                }
100                .fail();
101            }
102        };
103        Ok(value)
104    }
105
106    /// Projects a physical JSON2 array to a logical query type.
107    ///
108    /// TODO(LFC) Supersede `project_to_v2` to `project_to`.
109    pub fn project_to_v2(&self, field: &Field, target: &DataType) -> Result<ArrayRef> {
110        if json2_remainder_field(field)?.is_some() {
111            project_json_values(self.json2_values()?, target)
112        } else {
113            self.project_to(target)
114        }
115    }
116
117    /// Rewrites a JSON2 array from the current physical layout into the specified
118    /// v2 physical layout.
119    pub fn rewrite_to_v2(
120        &self,
121        field: &Field,
122        logical_settings: &JsonSettings,
123        target_layout: &JsonSettings,
124    ) -> Result<ArrayRef> {
125        self.rewrite_to_v2_with_type_hint_mismatch_policy(
126            field,
127            logical_settings,
128            target_layout,
129            TypeHintMismatchPolicy::Reject,
130        )
131    }
132
133    /// Rewrites a JSON2 array to the specified v2 physical layout using the
134    /// given type hint mismatch policy.
135    pub fn rewrite_to_v2_with_type_hint_mismatch_policy(
136        &self,
137        field: &Field,
138        logical_settings: &JsonSettings,
139        target_layout: &JsonSettings,
140        policy: TypeHintMismatchPolicy,
141    ) -> Result<ArrayRef> {
142        let is_v2 = json2_remainder_field(field)?.is_some();
143        if is_v2 && self.inner.data_type() == &json2_physical_data_type(target_layout) {
144            return Ok(self.inner.clone());
145        }
146
147        let values = if is_v2 {
148            self.json2_values()?
149        } else {
150            (0..self.inner.len())
151                .map(|i| self.try_get_value(i))
152                .collect::<Result<Vec<_>>>()?
153        };
154        let mut builder = JsonVectorBuilder::with_settings(target_layout, values.len());
155        for value in values {
156            if value.is_null() {
157                builder.push_null();
158            } else {
159                let value =
160                    logical_settings.encode_with_type_hint_mismatch_policy(value, policy)?;
161                builder.try_push_value_ref(&value.as_value_ref())?;
162            }
163        }
164        Ok(builder.to_vector().to_arrow_array())
165    }
166
167    fn json2_values(&self) -> Result<Vec<Value>> {
168        let structs = self.inner.as_struct_opt().context(AlignJsonArraySnafu {
169            reason: "JSON2 layout v2 root array must be a struct",
170        })?;
171        let remainder = structs.column_by_name(JSON2_REMAINDER_FIELD_NAME);
172        let mut remainders = if let Some(remainder) = remainder {
173            variant_to_json_values(remainder)?
174        } else {
175            vec![None; structs.len()]
176        };
177        let mut values = Vec::with_capacity(structs.len());
178        let mut path = Vec::new();
179
180        for (i, remainder) in remainders.iter_mut().enumerate() {
181            if structs.is_null(i) {
182                values.push(Value::Null);
183                continue;
184            }
185
186            let mut object = match remainder.take() {
187                None => serde_json::Map::new(),
188                Some(Value::Object(object)) => object,
189                Some(value) => {
190                    return InvalidJsonSnafu {
191                        value: format!("JSON2 layout v2 remainder must be an object, got {value}"),
192                    }
193                    .fail();
194                }
195            };
196
197            for (child, column) in structs.fields().iter().zip(structs.columns()) {
198                if child.name() == JSON2_REMAINDER_FIELD_NAME {
199                    continue;
200                }
201                let mut value = JsonArray::from(column).try_get_value(i)?;
202                // Arrow child nulls cannot distinguish a missing path from an explicit JSON
203                // null. Builders preserve explicit null presence in the remainder, so nulls
204                // from the explicit branch must be discarded before merging both branches.
205                remove_null_object_fields(&mut value);
206                if value.is_null() {
207                    continue;
208                }
209                merge_explicit_value(&mut object, child.name().clone(), value, &mut path)?;
210            }
211            values.push(Value::Object(object));
212        }
213
214        Ok(values)
215    }
216
217    /// Projects this JSON array to `target` for query evaluation.
218    ///
219    /// Unlike [`Self::widen_to`], projection tolerates lossy conversions:
220    /// - source fields not present in `target` are discarded;
221    /// - fields missing from the source are filled with typed null arrays;
222    /// - values incompatible with the target type become NULL.
223    ///
224    /// Projection is applied recursively to structs and lists. Input nulls
225    /// remain NULL. Errors unrelated to type incompatibility, such as invalid
226    /// JSONB, are returned.
227    pub fn project_to(&self, target: &DataType) -> Result<ArrayRef> {
228        if self.inner.data_type() == target {
229            return Ok(self.inner.clone());
230        }
231
232        match (self.inner.data_type(), target) {
233            (DataType::Struct(_), DataType::Struct(target_fields)) => {
234                let struct_array = self.inner.as_struct();
235                let mut columns = Vec::with_capacity(target_fields.len());
236                for target_field in target_fields {
237                    let column = struct_array
238                        .column_by_name(target_field.name())
239                        .map(|column| JsonArray::from(column).project_to(target_field.data_type()))
240                        .transpose()?
241                        .unwrap_or_else(|| {
242                            new_null_array(target_field.data_type(), self.inner.len())
243                        });
244                    columns.push(column);
245                }
246                let projected = StructArray::try_new_with_length(
247                    target_fields.clone(),
248                    columns,
249                    struct_array.nulls().cloned(),
250                    struct_array.len(),
251                )
252                .context(ArrowComputeSnafu)?;
253                Ok(Arc::new(projected))
254            }
255            (DataType::List(_), DataType::List(target_item)) => {
256                let list_array = self.inner.as_list::<i32>();
257                let item_projected =
258                    JsonArray::from(list_array.values()).project_to(target_item.data_type())?;
259                Ok(Arc::new(
260                    GenericListArray::<i32>::try_new(
261                        target_item.clone(),
262                        list_array.offsets().clone(),
263                        item_projected,
264                        list_array.nulls().cloned(),
265                    )
266                    .context(ArrowComputeSnafu)?,
267                ))
268            }
269            _ => self.project_values_to(target),
270        }
271    }
272
273    fn project_values_to(&self, to_type: &DataType) -> Result<ArrayRef> {
274        let from_type = self.inner.data_type();
275        if can_fast_cast_types(from_type, to_type) {
276            return cast(self.inner.as_ref(), to_type).context(ArrowComputeSnafu);
277        }
278
279        let values = (0..self.inner.len())
280            .map(|i| self.try_get_value(i))
281            .collect::<Result<Vec<_>>>()?;
282        project_json_values(values, to_type)
283    }
284}
285
286fn merge_explicit_value(
287    remainder: &mut serde_json::Map<String, Value>,
288    key: String,
289    explicit: Value,
290    path: &mut Vec<String>,
291) -> Result<()> {
292    let Some(existing) = remainder.get_mut(&key) else {
293        remainder.insert(key, explicit);
294        return Ok(());
295    };
296    path.push(key);
297
298    let (Value::Object(remainder), Value::Object(explicit)) = (existing, explicit) else {
299        return InvalidJsonSnafu {
300            value: format!(
301                "cannot merge '{}' in explicit fields and remainder: not both objects",
302                path.join("."),
303            ),
304        }
305        .fail();
306    };
307    for (key, value) in explicit {
308        merge_explicit_value(remainder, key, value, path)?;
309    }
310    path.pop();
311    Ok(())
312}
313
314fn remove_null_object_fields(value: &mut Value) {
315    let Value::Object(object) = value else {
316        return;
317    };
318    object.retain(|_, value| {
319        remove_null_object_fields(value);
320        !value.is_null()
321    });
322}
323
324/// Returns whether Arrow can cast between the types without JSON-aware projection.
325/// Binary and nested types require JSONB decoding or recursive projection.
326fn can_fast_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
327    let is_scalar = |data_type: &DataType| {
328        data_type.is_numeric() || data_type.is_string() || data_type == &DataType::Boolean
329    };
330
331    is_scalar(from_type) && is_scalar(to_type) && can_cast_types(from_type, to_type)
332}
333
334fn project_json_values(values: Vec<Value>, to_type: &DataType) -> Result<ArrayRef> {
335    let concrete_type = ConcreteDataType::from_arrow_type(to_type);
336    let mut builder = concrete_type.create_mutable_vector(values.len());
337    for value in values {
338        let value = coerce_json_value_to_type(value, &concrete_type);
339        builder.try_push_value_ref(&value.as_value_ref())?;
340    }
341    Ok(builder.to_vector().to_arrow_array())
342}
343
344impl<'a> From<&'a ArrayRef> for JsonArray<'a> {
345    fn from(inner: &'a ArrayRef) -> Self {
346        Self { inner }
347    }
348}
349
350#[cfg(test)]
351mod test {
352    use std::sync::Arc;
353
354    use arrow_array::types::Int64Type;
355    use arrow_array::{
356        BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array,
357        Int64Array, ListArray, StringArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array,
358    };
359    use arrow_schema::{Field, Fields};
360    use serde_json::json;
361
362    use super::*;
363    use crate::extension::json::{Json2ExtensionType, JsonMetadata};
364    use crate::json::{JsonSettings, JsonTypeHint};
365    use crate::vectors::json::variant::{json_values_to_variant, variant_field};
366
367    #[test]
368    fn test_try_get_value() -> Result<()> {
369        let nulls = new_null_array(&DataType::Null, 2);
370        assert_eq!(JsonArray::from(&nulls).try_get_value(0)?, Value::Null);
371
372        let bools: ArrayRef = Arc::new(BooleanArray::from(vec![Some(true), None]));
373        assert_eq!(JsonArray::from(&bools).try_get_value(0)?, json!(true));
374        assert_eq!(JsonArray::from(&bools).try_get_value(1)?, Value::Null);
375
376        let ints: ArrayRef = Arc::new(Int64Array::from(vec![Some(-7), None]));
377        assert_eq!(JsonArray::from(&ints).try_get_value(0)?, json!(-7));
378        assert_eq!(JsonArray::from(&ints).try_get_value(1)?, Value::Null);
379
380        macro_rules! assert_number {
381            ($array:expr, $expected:expr) => {{
382                let array: ArrayRef = Arc::new($array);
383                assert_eq!(JsonArray::from(&array).try_get_value(0)?, json!($expected));
384            }};
385        }
386        assert_number!(Int8Array::from(vec![-8]), -8);
387        assert_number!(Int16Array::from(vec![-16]), -16);
388        assert_number!(Int32Array::from(vec![-32]), -32);
389        assert_number!(UInt8Array::from(vec![8]), 8);
390        assert_number!(UInt16Array::from(vec![16]), 16);
391        assert_number!(UInt32Array::from(vec![32]), 32);
392        assert_number!(Float32Array::from(vec![1.25]), 1.25);
393
394        let floats: ArrayRef = Arc::new(Float64Array::from(vec![Some(1.5)]));
395        assert_eq!(JsonArray::from(&floats).try_get_value(0)?, json!(1.5));
396
397        let strings: ArrayRef = Arc::new(StringArray::from(vec![Some("hello"), None]));
398        assert_eq!(JsonArray::from(&strings).try_get_value(0)?, json!("hello"));
399        assert_eq!(JsonArray::from(&strings).try_get_value(1)?, Value::Null);
400
401        let nested = jsonb::parse_value(br#"{"nested":[1,null,"x"]}"#)
402            .unwrap()
403            .to_vec();
404        let null = jsonb::parse_value(b"null").unwrap().to_vec();
405        let binaries: ArrayRef =
406            Arc::new(BinaryArray::from(vec![nested.as_slice(), null.as_slice()]));
407        assert_eq!(
408            JsonArray::from(&binaries).try_get_value(0)?,
409            json!({"nested": [1, null, "x"]})
410        );
411        assert_eq!(JsonArray::from(&binaries).try_get_value(1)?, Value::Null);
412
413        let lists: ArrayRef = Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
414            Some(vec![Some(1), None, Some(3)]),
415            None,
416        ]));
417        assert_eq!(
418            JsonArray::from(&lists).try_get_value(0)?,
419            json!([1, null, 3])
420        );
421        assert_eq!(JsonArray::from(&lists).try_get_value(1)?, Value::Null);
422
423        let structs: ArrayRef = Arc::new(StructArray::from(vec![
424            (
425                Arc::new(Field::new("flag", DataType::Boolean, true)),
426                Arc::new(BooleanArray::from(vec![Some(true), None])) as ArrayRef,
427            ),
428            (
429                Arc::new(Field::new_list(
430                    "items",
431                    Field::new_list_field(DataType::Int64, true),
432                    true,
433                )),
434                Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
435                    Some(vec![Some(1), None]),
436                    Some(vec![Some(2)]),
437                ])) as ArrayRef,
438            ),
439        ]));
440        assert_eq!(
441            JsonArray::from(&structs).try_get_value(0)?,
442            json!({"flag": true, "items": [1, null]})
443        );
444        assert_eq!(
445            JsonArray::from(&structs).try_get_value(1)?,
446            json!({"flag": null, "items": [2]})
447        );
448
449        Ok(())
450    }
451
452    #[test]
453    fn test_cast_variant_to_utf8_view_preserves_json_null() -> Result<()> {
454        let encode = |json: &[u8]| jsonb::parse_value(json).unwrap().to_vec();
455        let json_null = encode(b"null");
456        let object = encode(br#"{"value":1}"#);
457        let string = encode(br#""text""#);
458        let variants: ArrayRef = Arc::new(BinaryArray::from(vec![
459            Some(json_null.as_slice()),
460            Some(object.as_slice()),
461            Some(string.as_slice()),
462            None,
463        ]));
464
465        let casted = JsonArray::from(&variants).project_to(&DataType::Utf8View)?;
466        let casted = casted.as_string_view();
467        assert!(casted.is_null(0));
468        assert_eq!(casted.value(1), r#"{"value":1}"#);
469        assert_eq!(casted.value(2), "text");
470        assert!(casted.is_null(3));
471
472        Ok(())
473    }
474
475    #[test]
476    fn test_project_plain_scalars() -> Result<()> {
477        let integers: ArrayRef = Arc::new(Int64Array::from(vec![Some(42), Some(i64::MAX), None]));
478        let projected = JsonArray::from(&integers).project_to(&DataType::Int32)?;
479        let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(42), None, None]));
480        assert_eq!(&expected, &projected);
481
482        let booleans: ArrayRef = Arc::new(BooleanArray::from(vec![Some(true), Some(false), None]));
483        let projected = JsonArray::from(&booleans).project_to(&DataType::Float64)?;
484        let expected: ArrayRef = Arc::new(Float64Array::from(vec![Some(1.0), Some(0.0), None]));
485        assert_eq!(&expected, &projected);
486
487        let strings: ArrayRef = Arc::new(StringArray::from(vec![Some("42"), Some("bad"), None]));
488        let projected = JsonArray::from(&strings).project_to(&DataType::UInt64)?;
489        let expected: ArrayRef = Arc::new(UInt64Array::from(vec![Some(42), None, None]));
490        assert_eq!(&expected, &projected);
491
492        Ok(())
493    }
494
495    #[test]
496    fn test_align_variant_to_struct() -> Result<()> {
497        let encode = |json: &[u8]| jsonb::parse_value(json).unwrap().to_vec();
498        let object =
499            encode(br#"{"nested":{"flag":true,"items":[1,2],"raw":{"x":1},"text":42,"value":42}}"#);
500        let scalar = encode(b"1");
501        let variants: ArrayRef = Arc::new(BinaryArray::from(vec![
502            Some(object.as_slice()),
503            None,
504            Some(scalar.as_slice()),
505        ]));
506        let expected_type = DataType::Struct(Fields::from(vec![Field::new_struct(
507            "nested",
508            vec![
509                Field::new("flag", DataType::Boolean, true),
510                Field::new_list("items", Field::new_list_field(DataType::UInt64, true), true),
511                Field::new("raw", DataType::Binary, true),
512                Field::new("text", DataType::Utf8View, true),
513                Field::new("value", DataType::UInt64, true),
514            ],
515            true,
516        )]));
517
518        let aligned = JsonArray::from(&variants).project_to(&expected_type)?;
519        assert_eq!(&expected_type, aligned.data_type());
520        assert_eq!(
521            json!({
522                "nested": {
523                    "flag": true,
524                    "items": [1, 2],
525                    "raw": {"x": 1},
526                    "text": "42",
527                    "value": 42
528                }
529            }),
530            JsonArray::from(&aligned).try_get_value(0)?
531        );
532        assert!(aligned.is_null(1));
533        assert!(aligned.is_null(2));
534
535        Ok(())
536    }
537
538    #[test]
539    fn test_align_nested_variant_to_struct() -> Result<()> {
540        let object = jsonb::parse_value(br#"{"flag":true,"value":42}"#)
541            .unwrap()
542            .to_vec();
543        let variants: ArrayRef = Arc::new(BinaryArray::from(vec![Some(object.as_slice()), None]));
544        let input: ArrayRef = Arc::new(StructArray::from(vec![(
545            Arc::new(Field::new("nested", DataType::Binary, true)),
546            variants,
547        )]));
548        let expected_type = DataType::Struct(Fields::from(vec![Field::new_struct(
549            "nested",
550            vec![
551                Field::new("flag", DataType::Boolean, true),
552                Field::new("value", DataType::UInt64, true),
553            ],
554            true,
555        )]));
556
557        let aligned = JsonArray::from(&input).project_to(&expected_type)?;
558        assert_eq!(&expected_type, aligned.data_type());
559        assert_eq!(
560            json!({"nested": {"flag": true, "value": 42}}),
561            JsonArray::from(&aligned).try_get_value(0)?
562        );
563        assert_eq!(
564            json!({"nested": null}),
565            JsonArray::from(&aligned).try_get_value(1)?
566        );
567
568        Ok(())
569    }
570
571    #[test]
572    fn test_reconstruct_json2_v2_value() -> Result<()> {
573        let remainders = json_values_to_variant(&[
574            Some(json!({"cold": 1, "nested": {"right": true}})),
575            Some(json!({"!__remainder__!": "user value"})),
576        ])?;
577        let remainder = Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true));
578        let nested = Arc::new(Field::new_struct(
579            "nested",
580            [Arc::new(Field::new("left", DataType::Utf8, true))],
581            true,
582        ));
583        let nested_values: ArrayRef = Arc::new(StructArray::from(vec![(
584            Arc::new(Field::new("left", DataType::Utf8, true)),
585            Arc::new(StringArray::from(vec![Some("value"), None])) as ArrayRef,
586        )]));
587        let fields = Fields::from(vec![
588            remainder,
589            Arc::new(Field::new("count", DataType::Int64, true)),
590            nested,
591        ]);
592        let array: ArrayRef = Arc::new(StructArray::new(
593            fields.clone(),
594            vec![
595                remainders,
596                Arc::new(Int64Array::from(vec![Some(42), None])),
597                nested_values,
598            ],
599            None,
600        ));
601        let field = Field::new("data", DataType::Struct(fields), true).with_extension_type(
602            Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default()))),
603        );
604
605        assert_eq!(
606            json!({
607                "cold": 1,
608                "count": 42,
609                "nested": {"left": "value", "right": true}
610            }),
611            JsonArray::from(&array).json2_values()?[0]
612        );
613        assert_eq!(
614            json!({
615                "!__remainder__!": "user value",
616                "nested": {}
617            }),
618            JsonArray::from(&array).json2_values()?[1]
619        );
620        let target = DataType::Struct(
621            vec![
622                Arc::new(Field::new("cold", DataType::UInt64, true)),
623                Arc::new(Field::new("count", DataType::Int64, true)),
624            ]
625            .into(),
626        );
627        let projected = JsonArray::from(&array).project_to_v2(&field, &target)?;
628        assert_eq!(
629            json!({"cold": 1, "count": 42}),
630            JsonArray::from(&projected).try_get_value(0)?
631        );
632        assert_eq!(
633            json!({"cold": null, "count": null}),
634            JsonArray::from(&projected).try_get_value(1)?
635        );
636        Ok(())
637    }
638
639    #[test]
640    fn test_rewrite_to_v2_reuses_matching_layout() -> Result<()> {
641        let settings = JsonSettings::try_new(
642            vec![JsonTypeHint {
643                path: vec!["kind".to_string()],
644                data_type: ConcreteDataType::string_datatype(),
645                inverted_index: false,
646            }],
647            Some(0),
648        )?;
649        let value = settings.encode(json!({"kind": "access", "cold": 1}))?;
650        let mut builder = JsonVectorBuilder::with_settings(&settings, 1);
651        builder.try_push_value_ref(&value.as_value_ref())?;
652        let array = builder.to_vector().to_arrow_array();
653        let structs = array.as_struct();
654        assert!(structs.column_by_name("kind").is_some());
655        assert_eq!(
656            vec![Some(json!({"cold": 1}))],
657            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
658        );
659        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
660            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings.clone()))),
661        );
662
663        let rewritten = JsonArray::from(&array).rewrite_to_v2(&field, &settings, &settings)?;
664
665        assert!(Arc::ptr_eq(&array, &rewritten));
666        Ok(())
667    }
668
669    #[test]
670    fn test_project_partial_json2_v2_without_remainder() -> Result<()> {
671        let fields = Fields::from(vec![Arc::new(Field::new("hot", DataType::Int64, true))]);
672        let array: ArrayRef = Arc::new(StructArray::new(
673            fields.clone(),
674            vec![Arc::new(Int64Array::from(vec![1, 2]))],
675            None,
676        ));
677        let field = Field::new("data", DataType::Struct(fields), true).with_extension_type(
678            Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default()))),
679        );
680
681        let projected = JsonArray::from(&array).project_to_v2(&field, field.data_type())?;
682        assert!(Arc::ptr_eq(&array, &projected));
683        Ok(())
684    }
685
686    #[test]
687    fn test_reject_conflict_json2_v2_path() -> Result<()> {
688        let remainders = json_values_to_variant(&[Some(json!({"count": 1}))])?;
689        let fields = Fields::from(vec![
690            Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)),
691            Arc::new(Field::new("count", DataType::Int64, true)),
692        ]);
693        let array: ArrayRef = Arc::new(StructArray::new(
694            fields,
695            vec![remainders, Arc::new(Int64Array::from(vec![2]))],
696            None,
697        ));
698        let error = JsonArray::from(&array).json2_values().unwrap_err();
699        assert!(
700            error.to_string().contains(
701                "cannot merge 'count' in explicit fields and remainder: not both objects"
702            )
703        );
704
705        let Value::Object(mut remainder) = json!({"count": 1}) else {
706            unreachable!();
707        };
708        let error = merge_explicit_value(
709            &mut remainder,
710            "count".to_string(),
711            json!(1),
712            &mut Vec::new(),
713        )
714        .unwrap_err();
715        assert!(error.to_string().contains("cannot merge 'count'"));
716
717        let Value::Object(mut remainder) = json!({"nested": {"count": 1}}) else {
718            unreachable!();
719        };
720        let error = merge_explicit_value(
721            &mut remainder,
722            "nested".to_string(),
723            json!({"count": 2}),
724            &mut Vec::new(),
725        )
726        .unwrap_err();
727        assert!(error.to_string().contains("cannot merge 'nested.count'"));
728        Ok(())
729    }
730}
datatypes/vectors/json/array.rs

datatypes/vectors/json/
array.rs
