datatypes/vectors/json/
builder.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::any::Any;
16use std::collections::{BTreeMap, HashMap};
17use std::sync::Arc;
18
19use arrow_array::cast::AsArray;
20use arrow_array::{Array, ArrayRef, StructArray};
21use arrow_schema::DataType;
22use parquet_variant_compute::VariantArrayBuilder;
23use snafu::{ResultExt, ensure};
24
25use crate::data_type::ConcreteDataType;
26use crate::error::{
27    ArrowComputeSnafu, Result, TryFromValueSnafu, UnexpectedSnafu, UnimplementedSnafu,
28    UnsupportedOperationSnafu,
29};
30use crate::extension::json::JSON2_REMAINDER_FIELD_NAME;
31use crate::json::value::{JsonNumber, JsonVariant, JsonVariantRef, encode_json_variant};
32use crate::json::{
33    JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JSON2_MAX_STRUCTURED_DEPTH, JsonSettings,
34};
35use crate::prelude::{ValueRef, Vector, VectorRef};
36use crate::types::StructType;
37use crate::types::json_type::{JsonNativeType, is_include};
38use crate::value::{ListValue, ListValueRef, StructValue, StructValueRef, Value};
39use crate::vectors::json::variant::{append_json_variant, append_json_variant_ref, variant_field};
40use crate::vectors::{Helper, MutableVector, NullVector, StructVectorBuilder};
41
42type JsonObjectValue = BTreeMap<String, JsonVariant>;
43
44/// Builds JSON2 vectors from object values.
45///
46/// Legacy mode merges all observed paths into the explicit Struct schema.
47/// Auto-expanding mode always materializes type-hinted paths, selects up to
48/// `max_auto_expanded_paths` compatible unhinted leaf paths by frequency, and stores
49/// conflicting or unselected paths in the Variant remainder field.
50pub(crate) struct JsonVectorBuilder {
51    state: JsonVectorBuilderState,
52}
53
54enum JsonVectorBuilderState {
55    Legacy {
56        merged_type: JsonNativeType,
57        values: Vec<JsonVariant>,
58    },
59    ExplicitOnly {
60        /// Paths declared by type hints and stored as dedicated Struct fields.
61        explicit_type: JsonNativeType,
62        /// Concrete Struct type used to append explicit values without buffering rows.
63        struct_type: StructType,
64        /// Builder for values selected by the explicit type hints.
65        explicit: StructVectorBuilder,
66        /// Builder for all values outside the explicit type hints.
67        remainder: VariantArrayBuilder,
68    },
69    AutoExpanding {
70        /// Paths declared by type hints and always stored as dedicated Struct fields.
71        explicit_type: JsonNativeType,
72        /// Maximum number of additional paths selected from buffered values.
73        max_auto_expanded_paths: u32,
74        /// Buffered values used to infer auto-expanded paths before building the vector.
75        values: Vec<JsonVariant>,
76    },
77}
78
79impl JsonVectorBuilderState {
80    fn native_type(&self) -> JsonNativeType {
81        match self {
82            Self::Legacy { merged_type, .. } => merged_type.clone(),
83            Self::ExplicitOnly { explicit_type, .. } => explicit_type.clone(),
84            Self::AutoExpanding {
85                explicit_type,
86                max_auto_expanded_paths,
87                values,
88            } => infer_expanded_type(explicit_type, *max_auto_expanded_paths, values),
89        }
90    }
91
92    fn len(&self) -> usize {
93        match self {
94            Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => values.len(),
95            Self::ExplicitOnly { explicit, .. } => explicit.len(),
96        }
97    }
98
99    fn try_build(&mut self) -> Result<VectorRef> {
100        match self {
101            Self::Legacy {
102                merged_type,
103                values,
104            } => build_legacy(values, merged_type),
105            Self::ExplicitOnly {
106                explicit,
107                remainder,
108                ..
109            } => {
110                let remainder = std::mem::replace(remainder, VariantArrayBuilder::new(0)).build();
111                finish_vector(explicit.to_vector(), ArrayRef::from(remainder))
112            }
113            Self::AutoExpanding {
114                explicit_type,
115                max_auto_expanded_paths,
116                values,
117            } => {
118                let expanded_type =
119                    infer_expanded_type(explicit_type, *max_auto_expanded_paths, values);
120                build_with_remainder(values, &expanded_type)
121            }
122        }
123    }
124
125    fn try_build_cloned(&self) -> Result<VectorRef> {
126        let mut state = match self {
127            Self::Legacy {
128                merged_type,
129                values,
130            } => Self::Legacy {
131                merged_type: merged_type.clone(),
132                values: values.clone(),
133            },
134            Self::AutoExpanding {
135                explicit_type,
136                max_auto_expanded_paths,
137                values,
138            } => Self::AutoExpanding {
139                explicit_type: explicit_type.clone(),
140                max_auto_expanded_paths: *max_auto_expanded_paths,
141                values: values.clone(),
142            },
143            // Only TimeSeriesMemtable requires a non-consuming snapshot, while JSON2 targets
144            // BulkMemtable. We've tried our best to support it above, but if this match arm does
145            // not, it's OK. The only reason it doesn't is because of `VariantArrayBuilder`. We'll
146            // track the upstream and see.
147            Self::ExplicitOnly { .. } => {
148                return UnimplementedSnafu {
149                    feat: "no auto expanded JSON2 array builder",
150                }
151                .fail();
152            }
153        };
154        state.try_build()
155    }
156
157    fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> {
158        if matches!(value, ValueRef::Null) {
159            self.push_null();
160            return Ok(());
161        }
162        let ValueRef::Json(value) = value else {
163            return TryFromValueSnafu {
164                reason: format!("expected JSON value, got {value:?}"),
165            }
166            .fail();
167        };
168        ensure!(
169            value.is_object() || value.is_null(),
170            TryFromValueSnafu {
171                reason: format!("expected JSON object value, got {value:?}"),
172            }
173        );
174        match self {
175            Self::Legacy {
176                merged_type,
177                values,
178            } => {
179                let json_type = value.json_type();
180                if !is_include(merged_type, json_type.as_ref()) {
181                    merged_type.merge(json_type.as_ref());
182                }
183                values.push(JsonVariant::from(value.variant()));
184            }
185            Self::ExplicitOnly {
186                explicit_type,
187                struct_type,
188                explicit,
189                remainder,
190            } => {
191                if value.is_null() {
192                    explicit.push_null();
193                    remainder.append_null();
194                } else {
195                    let (value, rest) =
196                        split_to_explicit_ref(value.variant(), explicit_type, struct_type)?;
197                    explicit.push_struct_value_ref(value)?;
198                    append_json_variant_ref(remainder, &rest).context(ArrowComputeSnafu)?;
199                }
200            }
201            Self::AutoExpanding { values, .. } => {
202                values.push(JsonVariant::from(value.variant()));
203            }
204        }
205        Ok(())
206    }
207
208    fn push_null(&mut self) {
209        match self {
210            Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => {
211                values.push(JsonVariant::Null)
212            }
213            Self::ExplicitOnly {
214                explicit,
215                remainder,
216                ..
217            } => {
218                explicit.push_null();
219                remainder.append_null();
220            }
221        }
222    }
223}
224
225/// Returns the fixed v2 Arrow physical type produced from `settings`.
226pub fn json2_physical_data_type(settings: &JsonSettings) -> DataType {
227    let DataType::Struct(fields) = explicit_type(settings).as_arrow_type() else {
228        unreachable!("JSON2 explicit type must map to Arrow Struct")
229    };
230    let mut fields = fields
231        .iter()
232        .cloned()
233        .chain(std::iter::once(Arc::new(variant_field(
234            JSON2_REMAINDER_FIELD_NAME,
235            true,
236        ))))
237        .collect::<Vec<_>>();
238    fields.sort_unstable_by(|x, y| x.name().cmp(y.name()));
239    DataType::Struct(fields.into())
240}
241
242fn explicit_type(settings: &JsonSettings) -> JsonNativeType {
243    let mut explicit_type = JsonNativeType::Object(Default::default());
244    for hint in settings.type_hints() {
245        insert_dynamic_type(&mut explicit_type, &hint.path, (&hint.data_type).into());
246    }
247    explicit_type
248}
249
250impl JsonVectorBuilder {
251    /// Creates a builder that merges all observed paths into the explicit schema.
252    pub(crate) fn new(initial_native_type: JsonNativeType, capacity: usize) -> Self {
253        debug_assert!(matches!(
254            initial_native_type,
255            JsonNativeType::Object(_) | JsonNativeType::Null
256        ));
257        Self {
258            state: JsonVectorBuilderState::Legacy {
259                merged_type: initial_native_type,
260                values: Vec::with_capacity(capacity),
261            },
262        }
263    }
264
265    /// Creates a builder bounded by the JSON settings and their type hints.
266    pub(crate) fn with_settings(settings: &JsonSettings, capacity: usize) -> Self {
267        let explicit_type = explicit_type(settings);
268        let state = if settings.max_auto_expanded_paths() == Some(0) {
269            let DataType::Struct(fields) = explicit_type.as_arrow_type() else {
270                unreachable!("JSON2 explicit type must map to Arrow Struct")
271            };
272            let struct_type = StructType::from(&fields);
273            JsonVectorBuilderState::ExplicitOnly {
274                explicit_type,
275                explicit: StructVectorBuilder::with_type_and_capacity(
276                    struct_type.clone(),
277                    capacity,
278                ),
279                struct_type,
280                remainder: VariantArrayBuilder::new(capacity),
281            }
282        } else {
283            JsonVectorBuilderState::AutoExpanding {
284                explicit_type,
285                max_auto_expanded_paths: settings
286                    .max_auto_expanded_paths()
287                    .unwrap_or(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS),
288                values: Vec::with_capacity(capacity),
289            }
290        };
291        Self { state }
292    }
293
294    fn try_build(&mut self) -> Result<VectorRef> {
295        self.state.try_build()
296    }
297}
298
299fn build_legacy(values: &mut Vec<JsonVariant>, merged_type: &JsonNativeType) -> Result<VectorRef> {
300    build_explicit(values, merged_type, false, |value| match value {
301        JsonVariant::Null => Ok(None),
302        JsonVariant::Object(value) => Ok(Some(value)),
303        _ => TryFromValueSnafu {
304            reason: "expected json object value".to_string(),
305        }
306        .fail(),
307    })
308}
309
310fn build_with_remainder(
311    values: &mut Vec<JsonVariant>,
312    expanded_type: &JsonNativeType,
313) -> Result<VectorRef> {
314    let mut remainder = VariantArrayBuilder::new(values.len());
315    let explicit = build_explicit(values, expanded_type, true, |value| {
316        if matches!(value, JsonVariant::Null) {
317            remainder.append_null();
318            return Ok(None);
319        }
320        let (value, rest) = split_to_explicit(value, expanded_type)?;
321        append_json_variant(&mut remainder, &JsonVariant::Object(rest))
322            .context(ArrowComputeSnafu)?;
323        Ok(Some(value))
324    })?;
325    finish_vector(explicit, ArrayRef::from(remainder.build()))
326}
327
328fn build_explicit(
329    values: &mut Vec<JsonVariant>,
330    explicit_type: &JsonNativeType,
331    // Temporary compatibility switch for the legacy storage layout. Once JSON2 fully switches to
332    // the v2 storage layout, empty objects should always be preserved instead of treated as null.
333    preserve_empty_structs: bool,
334    mut project: impl FnMut(JsonVariant) -> Result<Option<JsonObjectValue>>,
335) -> Result<VectorRef> {
336    let DataType::Struct(fields) = explicit_type.as_arrow_type() else {
337        return UnexpectedSnafu {
338            reason: "merged JSON2 type must map to Arrow Struct in JsonVectorBuilder",
339        }
340        .fail();
341    };
342    // TODO(LFC): Direct use Arrow's Struct datatype here.
343    let struct_type = StructType::from(&fields);
344
345    let mut builder =
346        StructVectorBuilder::with_type_and_capacity(struct_type.clone(), values.len());
347    for value in std::mem::take(values) {
348        let Some(value) = project(value)? else {
349            builder.push_null();
350            continue;
351        };
352        let value =
353            json_variant_into_struct_value(value, struct_type.clone(), preserve_empty_structs)?;
354        builder.push_struct_value_ref(StructValueRef::Ref(&value))?;
355    }
356    Ok(builder.to_vector())
357}
358
359fn finish_vector(explicit: VectorRef, remainder: ArrayRef) -> Result<VectorRef> {
360    let explicit = explicit.to_arrow_array();
361    let explicit = explicit.as_struct();
362    let mut children = explicit
363        .fields()
364        .iter()
365        .cloned()
366        .zip(explicit.columns().iter().cloned())
367        .chain(std::iter::once((
368            Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)),
369            remainder,
370        )))
371        .collect::<Vec<_>>();
372    children.sort_unstable_by(|(x, _), (y, _)| x.name().cmp(y.name()));
373    let (fields, columns): (Vec<_>, Vec<ArrayRef>) = children.into_iter().unzip();
374    let array: ArrayRef = Arc::new(StructArray::new(
375        fields.into(),
376        columns,
377        explicit.nulls().cloned(),
378    ));
379    Helper::try_into_vector(array)
380}
381
382fn infer_expanded_type(
383    explicit_type: &JsonNativeType,
384    max_auto_expanded_paths: u32,
385    values: &[JsonVariant],
386) -> JsonNativeType {
387    if max_auto_expanded_paths == 0 {
388        return explicit_type.clone();
389    }
390
391    let mut stats = HashMap::new();
392    let mut path = Vec::new();
393    init_explicit_path_stats(explicit_type, &mut path, &mut stats);
394    for value in values {
395        count_dynamic_paths(value, &mut path, &mut stats);
396    }
397    let mut candidates = stats
398        .iter()
399        // Explicit paths are already in the output schema and do not consume the dynamic
400        // expansion budget.
401        .filter(|(_, stat)| !stat.is_explicit && stat.is_leaf)
402        // Parquet cannot store empty structs, while widening an empty object to a non-empty struct
403        // loses its shape. Keep paths containing empty objects in the Variant remainder.
404        .filter(|(_, stat)| !stat.contains_empty_object)
405        // A leaf is eligible only when both itself and every object prefix have one stable
406        // role and type across all observed values.
407        .filter(|(path, _)| {
408            !(1..=path.len())
409                .any(|len| stats.get(&path[..len]).is_some_and(|stats| stats.conflicts))
410        })
411        .collect::<Vec<_>>();
412    candidates.sort_unstable_by(|(x_path, x), (y_path, y)| {
413        y.seen_count
414            .cmp(&x.seen_count)
415            .then_with(|| x_path.cmp(y_path))
416    });
417
418    let mut expanded_type = explicit_type.clone();
419    for (path, candidate) in candidates
420        .into_iter()
421        .take(max_auto_expanded_paths as usize)
422    {
423        insert_dynamic_type(
424            &mut expanded_type,
425            path,
426            candidate.expected_leaf_type.clone(),
427        );
428    }
429    expanded_type
430}
431
432/// Aggregated observations for one JSON path.
433///
434/// Schema inference first seeds the map with explicit paths, then walks all input values once.
435/// Objects, including empty objects, are non-leaf paths; every other non-null value is a leaf.
436/// The first dynamic observation fixes the path role and exact leaf type. A later role or type
437/// mismatch sets [`PathStats::conflicts`] permanently. Missing paths and null values do not affect
438/// the statistics.
439///
440/// After collection, dynamic leaves are ranked by [`PathStats::seen_count`]. A candidate is
441/// rejected when it or any parent path conflicts, so candidate selection never rescans the input
442/// values.
443struct PathStats {
444    /// Whether the path came from a type hint and is already part of the output schema.
445    is_explicit: bool,
446    /// Whether the path is a non-object value rather than an object prefix.
447    is_leaf: bool,
448    /// Exact type required for a leaf; unused non-leaf paths keep [`JsonNativeType::Null`].
449    expected_leaf_type: JsonNativeType,
450    /// Number of compatible observations used to rank dynamic leaves.
451    seen_count: usize,
452    /// Whether any observed value contains an empty object that requires lossless Variant storage.
453    contains_empty_object: bool,
454    /// Whether the path has ever had inconsistent roles or leaf types.
455    conflicts: bool,
456}
457
458/// Seeds path statistics from the configured explicit JSON shape.
459fn init_explicit_path_stats<'a>(
460    explicit_type: &'a JsonNativeType,
461    path: &mut Vec<&'a str>,
462    stats: &mut HashMap<Vec<&'a str>, PathStats>,
463) {
464    let JsonNativeType::Object(fields) = explicit_type else {
465        return;
466    };
467    for (name, data_type) in fields {
468        path.push(name);
469        let is_leaf = !matches!(data_type, JsonNativeType::Object(_));
470        let expected_leaf_type = if is_leaf {
471            data_type.clone()
472        } else {
473            JsonNativeType::default()
474        };
475        stats.insert(
476            path.clone(),
477            PathStats {
478                is_explicit: true,
479                is_leaf,
480                expected_leaf_type,
481                seen_count: 0,
482                contains_empty_object: false,
483                conflicts: false,
484            },
485        );
486        init_explicit_path_stats(data_type, path, stats);
487        path.pop();
488    }
489}
490
491/// Collects dynamic path statistics while traversing each input value once.
492fn count_dynamic_paths<'a>(
493    value: &'a JsonVariant,
494    path: &mut Vec<&'a str>,
495    stats: &mut HashMap<Vec<&'a str>, PathStats>,
496) {
497    if matches!(value, JsonVariant::Null) || path.len() > JSON2_MAX_STRUCTURED_DEPTH {
498        return;
499    }
500
501    if !path.is_empty() {
502        let is_leaf = !matches!(value, JsonVariant::Object(_));
503        let contains_empty_object = is_leaf && value.contains_empty_object();
504        let conflicts = if let Some(stats) = stats.get_mut(path.as_slice()) {
505            stats.contains_empty_object |= contains_empty_object;
506            if !stats.conflicts {
507                let role_conflict = stats.is_leaf != is_leaf;
508                let type_conflict = || match (&stats.expected_leaf_type, value) {
509                    // If both objects, they are compatible.
510                    (JsonNativeType::Null | JsonNativeType::Object(_), JsonVariant::Object(_)) => {
511                        false
512                    }
513                    _ => stats.expected_leaf_type != value.native_type(),
514                };
515                if role_conflict || type_conflict() {
516                    stats.conflicts = true;
517                } else {
518                    stats.seen_count += 1;
519                }
520            }
521            stats.conflicts
522        } else {
523            let expected_leaf_type = if is_leaf {
524                value.native_type()
525            } else {
526                JsonNativeType::default()
527            };
528            stats.insert(
529                path.clone(),
530                PathStats {
531                    is_explicit: false,
532                    is_leaf,
533                    expected_leaf_type,
534                    seen_count: 1,
535                    contains_empty_object,
536                    conflicts: false,
537                },
538            );
539            false
540        };
541        if conflicts {
542            return;
543        }
544    }
545
546    if let JsonVariant::Object(object) = value
547        && !object.is_empty()
548    {
549        for (name, value) in object {
550            path.push(name);
551            count_dynamic_paths(value, path, stats);
552            path.pop();
553        }
554    }
555}
556
557fn insert_dynamic_type<S: AsRef<str>>(
558    explicit_type: &mut JsonNativeType,
559    path: &[S],
560    data_type: JsonNativeType,
561) {
562    let JsonNativeType::Object(fields) = explicit_type else {
563        return;
564    };
565    let Some((name, path)) = path.split_first() else {
566        return;
567    };
568    let name = name.as_ref().to_string();
569    if path.is_empty() {
570        fields.insert(name, data_type);
571        return;
572    }
573    insert_dynamic_type(
574        fields
575            .entry(name)
576            .or_insert_with(|| JsonNativeType::Object(Default::default())),
577        path,
578        data_type,
579    )
580}
581
582fn split_to_explicit_ref<'a>(
583    value: &JsonVariantRef<'a>,
584    explicit_type: &JsonNativeType,
585    struct_type: &StructType,
586) -> Result<(StructValueRef<'a>, JsonVariantRef<'a>)> {
587    let JsonVariantRef::Object(object) = value else {
588        return TryFromValueSnafu {
589            reason: "expected json object value".to_string(),
590        }
591        .fail();
592    };
593    let explicit = json_object_ref_into_struct_value_ref(object, struct_type)?;
594    let remainder = remainder_ref(object, explicit_type)?;
595    Ok((explicit, JsonVariantRef::Object(remainder)))
596}
597
598fn json_object_ref_into_struct_value_ref<'a>(
599    object: &BTreeMap<&'a str, JsonVariantRef<'a>>,
600    struct_type: &StructType,
601) -> Result<StructValueRef<'a>> {
602    let mut values = Vec::with_capacity(struct_type.fields().len());
603    for field in struct_type.fields().iter() {
604        let value = match object.get(field.name()) {
605            Some(value) => json_variant_ref_into_value_ref(value, field.data_type())?,
606            None => ValueRef::Null,
607        };
608        values.push(value);
609    }
610    Ok(StructValueRef::RefList {
611        val: values,
612        fields: struct_type.clone(),
613    })
614}
615
616fn json_variant_ref_into_value_ref<'a>(
617    value: &JsonVariantRef<'a>,
618    expected_type: &ConcreteDataType,
619) -> Result<ValueRef<'a>> {
620    let value = match (value, expected_type) {
621        (JsonVariantRef::Null, _) | (_, ConcreteDataType::Null(_)) => ValueRef::Null,
622        (JsonVariantRef::Object(object), ConcreteDataType::Struct(struct_type)) => {
623            ValueRef::Struct(json_object_ref_into_struct_value_ref(object, struct_type)?)
624        }
625        (JsonVariantRef::Bool(x), ConcreteDataType::Boolean(_)) => ValueRef::Boolean(*x),
626        (JsonVariantRef::Number(x), ConcreteDataType::UInt64(_)) => {
627            let Some(x) = x.as_u64() else {
628                return TryFromValueSnafu {
629                    reason: format!("unable to convert {x:?} to UInt64"),
630                }
631                .fail();
632            };
633            ValueRef::UInt64(x)
634        }
635        (JsonVariantRef::Number(x), ConcreteDataType::Int64(_)) => {
636            let x = match x {
637                JsonNumber::PosInt(x) => i64::try_from(*x).ok(),
638                JsonNumber::NegInt(x) => Some(*x),
639                JsonNumber::Float(_) => None,
640            };
641            let Some(x) = x else {
642                return TryFromValueSnafu {
643                    reason: format!("unable to convert {x:?} to Int64"),
644                }
645                .fail();
646            };
647            ValueRef::Int64(x)
648        }
649        (JsonVariantRef::Number(JsonNumber::PosInt(x)), ConcreteDataType::Float64(_)) => {
650            ValueRef::Float64((*x as f64).into())
651        }
652        (JsonVariantRef::Number(JsonNumber::NegInt(x)), ConcreteDataType::Float64(_)) => {
653            ValueRef::Float64((*x as f64).into())
654        }
655        (JsonVariantRef::Number(JsonNumber::Float(x)), ConcreteDataType::Float64(_)) => {
656            ValueRef::Float64(*x)
657        }
658        (JsonVariantRef::String(x), ConcreteDataType::String(_)) => ValueRef::String(x),
659        (JsonVariantRef::Array(array), ConcreteDataType::List(list_type)) => {
660            let item_type = list_type.item_type().clone();
661            let values = array
662                .iter()
663                .map(|x| json_variant_ref_into_value_ref(x, &item_type))
664                .collect::<Result<Vec<_>>>()?;
665            ValueRef::List(ListValueRef::RefList {
666                val: values,
667                item_datatype: Arc::new(item_type),
668            })
669        }
670        (value, expected_type) => {
671            return TryFromValueSnafu {
672                reason: format!("unable to convert json value {value:?} to {expected_type}"),
673            }
674            .fail();
675        }
676    };
677    Ok(value)
678}
679
680fn remainder_ref<'a>(
681    object: &BTreeMap<&'a str, JsonVariantRef<'a>>,
682    explicit_type: &JsonNativeType,
683) -> Result<BTreeMap<&'a str, JsonVariantRef<'a>>> {
684    let JsonNativeType::Object(fields) = explicit_type else {
685        return UnexpectedSnafu {
686            reason: "JSON2 explicit type must be an object",
687        }
688        .fail();
689    };
690    let mut remainder = BTreeMap::new();
691    for (&name, value) in object {
692        // Preserve explicit JSON nulls in the remainder because Arrow child nulls cannot
693        // distinguish a present JSON null from a missing path.
694        if *value == JsonVariantRef::Null {
695            remainder.insert(name, JsonVariantRef::Null);
696            continue;
697        }
698
699        match fields.get(name) {
700            Some(data_type @ JsonNativeType::Object(_)) => match value {
701                JsonVariantRef::Object(object) => {
702                    let child = remainder_ref(object, data_type)?;
703                    if !child.is_empty() {
704                        remainder.insert(name, JsonVariantRef::Object(child));
705                    }
706                }
707                _ => {
708                    return TryFromValueSnafu {
709                        reason: "expected json object value".to_string(),
710                    }
711                    .fail();
712                }
713            },
714            // A non-object entry in the explicit type tree is an explicit leaf and is
715            // already written to the Struct builder. So here does nothing.
716            Some(_) => {}
717            None => {
718                remainder.insert(name, value.clone());
719            }
720        }
721    }
722    Ok(remainder)
723}
724
725fn split_to_explicit(
726    value: JsonVariant,
727    explicit_type: &JsonNativeType,
728) -> Result<(JsonObjectValue, JsonObjectValue)> {
729    let JsonVariant::Object(mut remainder) = value else {
730        return TryFromValueSnafu {
731            reason: "expected json object value".to_string(),
732        }
733        .fail();
734    };
735    let JsonNativeType::Object(fields) = explicit_type else {
736        return UnexpectedSnafu {
737            reason: "JSON2 explicit type must be an object",
738        }
739        .fail();
740    };
741    let mut explicit = JsonObjectValue::new();
742
743    for (name, data_type) in fields {
744        let Some(value) = remainder.remove(name) else {
745            continue;
746        };
747        if value == JsonVariant::Null {
748            explicit.insert(name.clone(), JsonVariant::Null);
749            // Preserve explicit JSON nulls in the remainder because Arrow child nulls cannot
750            // distinguish a present JSON null from a missing path.
751            remainder.insert(name.clone(), JsonVariant::Null);
752            continue;
753        }
754        if matches!(data_type, JsonNativeType::Object(_)) {
755            let (child_explicit, child_remainder) = split_to_explicit(value, data_type)?;
756            explicit.insert(name.clone(), JsonVariant::Object(child_explicit));
757            if !child_remainder.is_empty() {
758                remainder.insert(name.clone(), JsonVariant::Object(child_remainder));
759            }
760        } else {
761            explicit.insert(name.clone(), value);
762        }
763    }
764    Ok((explicit, remainder))
765}
766
767fn json_variant_into_struct_value(
768    object: JsonObjectValue,
769    struct_type: StructType,
770    preserve_empty_structs: bool,
771) -> Result<StructValue> {
772    let mut entries = object.into_iter();
773    let mut entry = entries.next();
774    let mut values = Vec::with_capacity(struct_type.fields().len());
775    for field in struct_type.fields().iter() {
776        let value = match entry.take() {
777            Some((name, value)) if name == field.name() => {
778                entry = entries.next();
779                json_variant_into_value(value, field.data_type(), preserve_empty_structs)?
780            }
781            Some((name, _)) if name.as_str() < field.name() => {
782                return TryFromValueSnafu {
783                    reason: format!("field {name} is missing from merged JSON type"),
784                }
785                .fail();
786            }
787            next => {
788                entry = next;
789                Value::Null
790            }
791        };
792        values.push(value);
793    }
794    if let Some((name, _)) = entry {
795        return TryFromValueSnafu {
796            reason: format!("field {name} is missing from merged JSON type"),
797        }
798        .fail();
799    }
800
801    Ok(StructValue::new(values, struct_type))
802}
803
804fn json_variant_into_value(
805    value: JsonVariant,
806    expected_type: &ConcreteDataType,
807    preserve_empty_structs: bool,
808) -> Result<Value> {
809    let value = match (value, expected_type) {
810        (JsonVariant::Null, _) | (_, ConcreteDataType::Null(_)) => Value::Null,
811        (JsonVariant::Object(object), _) if object.is_empty() && !preserve_empty_structs => {
812            Value::Null
813        }
814        (JsonVariant::Object(object), ConcreteDataType::Struct(struct_type)) => Value::Struct(
815            json_variant_into_struct_value(object, struct_type.clone(), preserve_empty_structs)?,
816        ),
817        (JsonVariant::Bool(x), ConcreteDataType::Boolean(_)) => Value::Boolean(x),
818        (JsonVariant::Number(x), ConcreteDataType::UInt64(_)) => {
819            let Some(x) = x.as_u64() else {
820                return TryFromValueSnafu {
821                    reason: format!("unable to convert {x:?} to UInt64"),
822                }
823                .fail();
824            };
825            Value::UInt64(x)
826        }
827        (JsonVariant::Number(x), ConcreteDataType::Int64(_)) => {
828            let x = match x {
829                JsonNumber::PosInt(x) => i64::try_from(x).ok(),
830                JsonNumber::NegInt(x) => Some(x),
831                JsonNumber::Float(_) => None,
832            };
833            let Some(x) = x else {
834                return TryFromValueSnafu {
835                    reason: format!("unable to convert {x:?} to Int64"),
836                }
837                .fail();
838            };
839            Value::Int64(x)
840        }
841        (JsonVariant::Number(JsonNumber::PosInt(x)), ConcreteDataType::Float64(_)) => {
842            Value::Float64((x as f64).into())
843        }
844        (JsonVariant::Number(JsonNumber::NegInt(x)), ConcreteDataType::Float64(_)) => {
845            Value::Float64((x as f64).into())
846        }
847        (JsonVariant::Number(JsonNumber::Float(x)), ConcreteDataType::Float64(_)) => {
848            Value::Float64(x)
849        }
850        (JsonVariant::String(x), ConcreteDataType::String(_)) => Value::String(x.into()),
851        (JsonVariant::Array(array), ConcreteDataType::List(list_type)) => {
852            let item_type = list_type.item_type().clone();
853            let values = array
854                .into_iter()
855                .map(|v| json_variant_into_value(v, &item_type, preserve_empty_structs))
856                .collect::<Result<Vec<_>>>()?;
857            Value::List(ListValue::new(values, Arc::new(item_type)))
858        }
859        (value, ConcreteDataType::Binary(_)) => Value::from(encode_json_variant(value)?),
860        (value, expected_type) => {
861            return TryFromValueSnafu {
862                reason: format!("unable to convert json value {value:?} to {expected_type}"),
863            }
864            .fail();
865        }
866    };
867    Ok(value)
868}
869
870impl MutableVector for JsonVectorBuilder {
871    fn data_type(&self) -> ConcreteDataType {
872        ConcreteDataType::json2(self.state.native_type())
873    }
874
875    fn len(&self) -> usize {
876        self.state.len()
877    }
878
879    fn as_any(&self) -> &dyn Any {
880        self
881    }
882
883    fn as_mut_any(&mut self) -> &mut dyn Any {
884        self
885    }
886
887    fn to_vector(&mut self) -> VectorRef {
888        self.try_build().unwrap_or_else(|e| {
889            // Just try to avoid panicking here.
890            common_telemetry::error!(e; "Unable to build JSON2 vector");
891            Arc::new(NullVector::new(self.len()))
892        })
893    }
894
895    fn to_vector_cloned(&self) -> VectorRef {
896        self.state.try_build_cloned().unwrap_or_else(|e| {
897            // Just try to avoid panicking here.
898            common_telemetry::error!(e; "Unable to build JSON2 vector");
899            Arc::new(NullVector::new(self.len()))
900        })
901    }
902
903    fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> {
904        self.state.try_push_value_ref(value)
905    }
906
907    fn push_null(&mut self) {
908        self.state.push_null()
909    }
910
911    fn extend_slice_of(&mut self, _: &dyn Vector, _: usize, _: usize) -> Result<()> {
912        UnsupportedOperationSnafu {
913            op: "extend_slice_of",
914            vector_type: "JsonVector",
915        }
916        .fail()
917    }
918}
919
920#[cfg(test)]
921mod tests {
922    use std::sync::Arc;
923
924    use arrow_array::cast::AsArray;
925    use arrow_schema::Field;
926    use common_base::bytes::Bytes;
927    use serde_json::json;
928
929    use super::*;
930    use crate::data_type::ConcreteDataType;
931    use crate::extension::json::{Json2ExtensionType, JsonMetadata};
932    use crate::json::JsonTypeHint;
933    use crate::json::value::decode_json_variant;
934    use crate::types::StructField;
935    use crate::types::json_type::JsonObjectType;
936    use crate::value::{ListValue, StructValue, Value, ValueRef};
937    use crate::vectors::json::array::JsonArray;
938    use crate::vectors::json::variant::variant_to_json_values;
939
940    #[test]
941    fn test_json_vector_builder() -> Result<()> {
942        fn parse_json_value(json: &str) -> Value {
943            let value: serde_json::Value = serde_json::from_str(json).unwrap();
944            Value::Json(Box::new(value.into()))
945        }
946
947        fn jsonb_bytes(json: &str) -> Bytes {
948            Bytes::from(jsonb::parse_value(json.as_bytes()).unwrap().to_vec())
949        }
950
951        // Object inputs should merge into a superset schema, preserve null rows,
952        // and project conflicting nested values into Variant payloads.
953        let mut builder = JsonVectorBuilder::new(JsonNativeType::Object(Default::default()), 3);
954        let first = parse_json_value(r#"{"id":1,"payload":{"name":"foo"}}"#);
955        let second = parse_json_value(r#"{"id":2,"extra":true,"payload":"raw"}"#);
956        builder.try_push_value_ref(&first.as_value_ref())?;
957        builder.push_null();
958        builder.try_push_value_ref(&second.as_value_ref())?;
959
960        let merged_type = JsonNativeType::Object(JsonObjectType::from([
961            ("extra".to_string(), JsonNativeType::Bool),
962            ("id".to_string(), JsonNativeType::i64()),
963            ("payload".to_string(), JsonNativeType::Variant),
964        ]));
965        assert_eq!(
966            builder.data_type(),
967            ConcreteDataType::json2(merged_type.clone())
968        );
969
970        let DataType::Struct(fields) = merged_type.as_arrow_type() else {
971            unreachable!()
972        };
973        let merged_struct_type = StructType::from(&fields);
974        let vector = builder.to_vector();
975        assert_eq!(vector.len(), 3);
976        assert_eq!(
977            vector.get(0),
978            Value::Struct(StructValue::new(
979                vec![
980                    Value::Null,
981                    Value::Int64(1),
982                    Value::Binary(jsonb_bytes(r#"{"name":"foo"}"#)),
983                ],
984                merged_struct_type.clone(),
985            ))
986        );
987        assert_eq!(vector.get(1), Value::Null);
988        assert_eq!(
989            vector.get(2),
990            Value::Struct(StructValue::new(
991                vec![
992                    Value::Boolean(true),
993                    Value::Int64(2),
994                    Value::Binary(jsonb_bytes(r#""raw""#)),
995                ],
996                merged_struct_type,
997            ))
998        );
999
1000        // A Null initial type represents an unknown JSON2 runtime type. The first
1001        // non-null value should set the concrete type instead of aligning all rows to Null.
1002        let mut inferred_builder = JsonVectorBuilder::new(JsonNativeType::Null, 2);
1003        let inferred_value = parse_json_value(r#"{"id":3}"#);
1004        inferred_builder.push_null();
1005        inferred_builder.try_push_value_ref(&inferred_value.as_value_ref())?;
1006
1007        let inferred_type = JsonNativeType::Object(JsonObjectType::from([(
1008            "id".to_string(),
1009            JsonNativeType::i64(),
1010        )]));
1011        assert_eq!(
1012            inferred_builder.data_type(),
1013            ConcreteDataType::json2(inferred_type.clone())
1014        );
1015
1016        let DataType::Struct(fields) = inferred_type.as_arrow_type() else {
1017            unreachable!()
1018        };
1019        let inferred_struct_type = StructType::from(&fields);
1020        let vector = inferred_builder.to_vector();
1021        assert_eq!(vector.get(0), Value::Null);
1022        assert_eq!(
1023            vector.get(1),
1024            Value::Struct(StructValue::new(
1025                vec![Value::Int64(3)],
1026                inferred_struct_type,
1027            ))
1028        );
1029
1030        // Non-object initial types are rejected by the builder invariant.
1031        let result = std::panic::catch_unwind(|| JsonVectorBuilder::new(JsonNativeType::Bool, 2));
1032        assert!(result.is_err());
1033
1034        // Non-object root values should be rejected at push time.
1035        let mut object_builder =
1036            JsonVectorBuilder::new(JsonNativeType::Object(Default::default()), 2);
1037        let object = parse_json_value(r#"{"k":1}"#);
1038        let boolean = parse_json_value("true");
1039        let err = object_builder
1040            .try_push_value_ref(&boolean.as_value_ref())
1041            .unwrap_err();
1042        assert!(err.to_string().contains("expected JSON object value"));
1043        object_builder.try_push_value_ref(&object.as_value_ref())?;
1044
1045        // Non-JSON values should be rejected at push time.
1046        let mut invalid_builder =
1047            JsonVectorBuilder::new(JsonNativeType::Object(Default::default()), 1);
1048        let err = invalid_builder
1049            .try_push_value_ref(&ValueRef::Boolean(true))
1050            .unwrap_err();
1051        assert!(err.to_string().contains("expected JSON value"));
1052
1053        Ok(())
1054    }
1055
1056    #[test]
1057    fn test_zero_budget_builder_uses_explicit_only_schema_and_remainder() -> Result<()> {
1058        let settings = JsonSettings::try_new(
1059            vec![
1060                JsonTypeHint {
1061                    path: vec!["kind".to_string()],
1062                    data_type: ConcreteDataType::string_datatype(),
1063                    inverted_index: false,
1064                },
1065                JsonTypeHint {
1066                    path: vec!["commit".to_string(), "operation".to_string()],
1067                    data_type: ConcreteDataType::string_datatype(),
1068                    inverted_index: false,
1069                },
1070                JsonTypeHint {
1071                    path: vec!["time_us".to_string()],
1072                    data_type: ConcreteDataType::int64_datatype(),
1073                    inverted_index: false,
1074                },
1075            ],
1076            Some(0),
1077        )?;
1078        let mut builder = JsonVectorBuilder::with_settings(&settings, 2);
1079        assert!(matches!(
1080            &builder.state,
1081            JsonVectorBuilderState::ExplicitOnly { .. }
1082        ));
1083        let values = [
1084            json!({
1085                "kind": "record",
1086                "commit": {"operation": "create", "collection": "post"},
1087                "extra": 1,
1088                "time_us": 1
1089            }),
1090            json!({"kind": "other", "dynamic": true, "time_us": 2}),
1091        ];
1092        for value in values.clone() {
1093            let value = settings.encode(value)?;
1094            builder.try_push_value_ref(&value.as_value_ref())?;
1095        }
1096        let array = builder.to_vector().to_arrow_array();
1097        assert_eq!(&json2_physical_data_type(&settings), array.data_type());
1098        assert_eq!(0, builder.len());
1099        let structs = array.as_struct();
1100        assert_eq!(
1101            vec![JSON2_REMAINDER_FIELD_NAME, "commit", "kind", "time_us"],
1102            structs
1103                .fields()
1104                .iter()
1105                .map(|x| x.name().as_str())
1106                .collect::<Vec<_>>()
1107        );
1108        assert_eq!(
1109            vec![
1110                Some(json!({"commit": {"collection": "post"}, "extra": 1})),
1111                Some(json!({"commit": {"operation": null}, "dynamic": true})),
1112            ],
1113            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
1114        );
1115
1116        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
1117            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))),
1118        );
1119        let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?;
1120        let reconstructed = reconstructed.as_binary::<i32>();
1121        assert_eq!(
1122            values[0],
1123            decode_json_variant(reconstructed.value(0)).unwrap()
1124        );
1125        assert_eq!(
1126            json!({
1127                "kind": "other",
1128                "commit": {"operation": null},
1129                "dynamic": true,
1130                "time_us": 2
1131            }),
1132            decode_json_variant(reconstructed.value(1)).unwrap()
1133        );
1134
1135        let settings = JsonSettings::try_new(vec![], Some(0))?;
1136        let mut builder = JsonVectorBuilder::with_settings(&settings, 3);
1137        for value in [json!({}), json!({"x": 1})] {
1138            let value = settings.encode(value)?;
1139            builder.try_push_value_ref(&value.as_value_ref())?;
1140        }
1141        builder.push_null();
1142        let array = builder.to_vector().to_arrow_array();
1143        let structs = array.as_struct();
1144        assert_eq!(1, structs.num_columns());
1145        assert_eq!(
1146            vec![Some(json!({})), Some(json!({"x": 1})), None],
1147            variant_to_json_values(structs.column(0))?
1148        );
1149
1150        Ok(())
1151    }
1152
1153    #[test]
1154    fn test_finite_budget_selects_dynamic_paths() -> Result<()> {
1155        let builder = JsonVectorBuilder::with_settings(&JsonSettings::default(), 0);
1156        assert!(matches!(
1157            builder.state,
1158            JsonVectorBuilderState::AutoExpanding {
1159                max_auto_expanded_paths: JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS,
1160                ..
1161            }
1162        ));
1163
1164        let settings = JsonSettings::try_new(
1165            vec![JsonTypeHint {
1166                path: vec!["hint".to_string()],
1167                data_type: ConcreteDataType::string_datatype(),
1168                inverted_index: false,
1169            }],
1170            Some(2),
1171        )?;
1172        let values = [
1173            json!({
1174                "hint": "first",
1175                "conflict": 1,
1176                "popular": {"nested": 1},
1177                "tie_a": "a",
1178                "tie_b": true,
1179                "rare": 1
1180            }),
1181            json!({
1182                "hint": "second",
1183                "conflict": "string",
1184                "popular": {"nested": 2},
1185                "tie_a": "b",
1186                "tie_b": false
1187            }),
1188            json!({"hint": "third", "popular": "scalar"}),
1189        ];
1190        let mut builder = JsonVectorBuilder::with_settings(&settings, values.len());
1191        for value in values.clone() {
1192            let value = settings.encode(value)?;
1193            builder.try_push_value_ref(&value.as_value_ref())?;
1194        }
1195
1196        let array = builder.to_vector().to_arrow_array();
1197        let structs = array.as_struct();
1198        assert_eq!(
1199            vec![JSON2_REMAINDER_FIELD_NAME, "hint", "tie_a", "tie_b"],
1200            structs
1201                .fields()
1202                .iter()
1203                .map(|x| x.name().as_str())
1204                .collect::<Vec<_>>()
1205        );
1206        assert!(structs.column_by_name("popular").is_none());
1207        assert_eq!(
1208            vec![
1209                Some(json!({"conflict": 1, "popular": {"nested": 1}, "rare": 1})),
1210                Some(json!({"conflict": "string", "popular": {"nested": 2}})),
1211                Some(json!({"popular": "scalar"})),
1212            ],
1213            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
1214        );
1215
1216        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
1217            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))),
1218        );
1219        let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?;
1220        let reconstructed = reconstructed.as_binary::<i32>();
1221        assert_eq!(
1222            values[0],
1223            decode_json_variant(reconstructed.value(0)).unwrap()
1224        );
1225        assert_eq!(
1226            values[1],
1227            decode_json_variant(reconstructed.value(1)).unwrap()
1228        );
1229        assert_eq!(
1230            json!({
1231                "hint": "third",
1232                "popular": "scalar"
1233            }),
1234            decode_json_variant(reconstructed.value(2)).unwrap()
1235        );
1236
1237        Ok(())
1238    }
1239
1240    #[test]
1241    fn test_v2_builder_preserves_explicit_null_presence() -> Result<()> {
1242        let settings = JsonSettings::try_new(vec![], Some(1))?;
1243        let values = [json!({"value": 1}), json!({"value": null}), json!({})];
1244        let mut builder = JsonVectorBuilder::with_settings(&settings, values.len());
1245        for value in values.clone() {
1246            let value = settings.encode(value)?;
1247            builder.try_push_value_ref(&value.as_value_ref())?;
1248        }
1249        let array = builder.to_vector().to_arrow_array();
1250        let structs = array.as_struct();
1251        assert_eq!(
1252            vec![
1253                Some(json!({})),
1254                Some(json!({"value": null})),
1255                Some(json!({}))
1256            ],
1257            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
1258        );
1259        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
1260            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))),
1261        );
1262        let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?;
1263        let reconstructed = reconstructed.as_binary::<i32>();
1264
1265        assert_eq!(
1266            values[0],
1267            decode_json_variant(reconstructed.value(0)).unwrap()
1268        );
1269        assert_eq!(
1270            values[1],
1271            decode_json_variant(reconstructed.value(1)).unwrap()
1272        );
1273        assert_eq!(
1274            values[2],
1275            decode_json_variant(reconstructed.value(2)).unwrap()
1276        );
1277        Ok(())
1278    }
1279
1280    #[test]
1281    fn test_v2_builder_accepts_sql_null() -> Result<()> {
1282        let settings = JsonSettings::try_new(vec![], Some(1))?;
1283        let mut builder = JsonVectorBuilder::with_settings(&settings, 2);
1284        builder.try_push_value_ref(&ValueRef::Null)?;
1285        let value = settings.encode(json!({}))?;
1286        builder.try_push_value_ref(&value.as_value_ref())?;
1287
1288        let array = builder.to_vector().to_arrow_array();
1289        let structs = array.as_struct();
1290        assert_eq!(
1291            vec![None, Some(json!({}))],
1292            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
1293        );
1294        Ok(())
1295    }
1296
1297    #[test]
1298    fn test_reconstruct_nested_remainder_only_value() -> Result<()> {
1299        let settings = JsonSettings::try_new(vec![], Some(1))?;
1300        let values = [
1301            json!({"a": {"hot": 1}}),
1302            json!({"a": {"hot": 2}}),
1303            json!({"a": {"cold": 3}}),
1304        ];
1305        let mut builder = JsonVectorBuilder::with_settings(&settings, values.len());
1306        for value in values.clone() {
1307            let value = settings.encode(value)?;
1308            builder.try_push_value_ref(&value.as_value_ref())?;
1309        }
1310
1311        let array = builder.to_vector().to_arrow_array();
1312        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
1313            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))),
1314        );
1315        let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?;
1316        let reconstructed = reconstructed.as_binary::<i32>();
1317        assert_eq!(
1318            vec![values[0].clone(), values[1].clone(), values[2].clone()],
1319            (0..reconstructed.len())
1320                .map(|i| decode_json_variant(reconstructed.value(i)).unwrap())
1321                .collect::<Vec<_>>()
1322        );
1323
1324        Ok(())
1325    }
1326
1327    #[test]
1328    fn test_dynamic_paths_require_the_same_leaf_type() -> Result<()> {
1329        let settings = JsonSettings::try_new(
1330            vec![JsonTypeHint {
1331                path: vec!["nested".to_string(), "hinted".to_string()],
1332                data_type: ConcreteDataType::string_datatype(),
1333                inverted_index: false,
1334            }],
1335            Some(8),
1336        )?;
1337        let values = [
1338            json!({
1339                "branch": {},
1340                "different": [1],
1341                "empty": {},
1342                "nested": {},
1343                "reverse": "scalar",
1344                "same": [1]
1345            }),
1346            json!({
1347                "branch": {"leaf": 1},
1348                "different": ["x"],
1349                "empty": {},
1350                "nested": {"hinted": "x", "leaf": 1},
1351                "same": [2]
1352            }),
1353            json!({
1354                "branch": {},
1355                "different": [2],
1356                "empty": {},
1357                "nested": {},
1358                "reverse": {"leaf": 1},
1359                "same": [3]
1360            }),
1361        ];
1362        let mut builder = JsonVectorBuilder::with_settings(&settings, values.len());
1363        for value in values {
1364            let value = settings.encode(value)?;
1365            builder.try_push_value_ref(&value.as_value_ref())?;
1366        }
1367
1368        let JsonNativeType::Object(fields) = builder.state.native_type() else {
1369            unreachable!();
1370        };
1371        assert!(!fields.contains_key("different"));
1372        assert!(!fields.contains_key("empty"));
1373        assert!(!fields.contains_key("reverse"));
1374        assert!(matches!(fields.get("same"), Some(JsonNativeType::Array(_))));
1375        assert!(matches!(
1376            fields.get("branch"),
1377            Some(JsonNativeType::Object(fields)) if fields.contains_key("leaf")
1378        ));
1379        assert!(matches!(
1380            fields.get("nested"),
1381            Some(JsonNativeType::Object(fields))
1382                if fields.contains_key("hinted") && fields.contains_key("leaf")
1383        ));
1384
1385        Ok(())
1386    }
1387
1388    #[test]
1389    fn test_json_variant_into_struct_value() -> Result<()> {
1390        let struct_type = StructType::new(Arc::new(vec![StructField::new(
1391            "value".to_string(),
1392            ConcreteDataType::string_datatype(),
1393            true,
1394        )]));
1395        assert_eq!(
1396            json_variant_into_value(
1397                JsonVariant::Object(Default::default()),
1398                &ConcreteDataType::struct_datatype(struct_type.clone()),
1399                false,
1400            )?,
1401            Value::Null
1402        );
1403        assert_eq!(
1404            json_variant_into_value(
1405                JsonVariant::Object(Default::default()),
1406                &ConcreteDataType::struct_datatype(struct_type.clone()),
1407                true,
1408            )?,
1409            Value::Struct(StructValue::new(vec![Value::Null], struct_type))
1410        );
1411
1412        let item_type =
1413            ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![StructField::new(
1414                "id".to_string(),
1415                ConcreteDataType::int64_datatype(),
1416                true,
1417            )])));
1418        let struct_type = StructType::new(Arc::new(vec![
1419            StructField::new(
1420                "items".to_string(),
1421                ConcreteDataType::list_datatype(Arc::new(item_type.clone())),
1422                true,
1423            ),
1424            StructField::new(
1425                "meta".to_string(),
1426                ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![
1427                    StructField::new(
1428                        "name".to_string(),
1429                        ConcreteDataType::string_datatype(),
1430                        true,
1431                    ),
1432                ]))),
1433                true,
1434            ),
1435        ]));
1436        let variant = JsonObjectValue::from([
1437            (
1438                "items".to_string(),
1439                JsonVariant::Array(vec![
1440                    JsonVariant::from([("id", JsonVariant::from(1i64))]),
1441                    JsonVariant::from([("id", JsonVariant::from(2i64))]),
1442                ]),
1443            ),
1444            (
1445                "meta".to_string(),
1446                JsonVariant::from([("name", JsonVariant::from("foo"))]),
1447            ),
1448        ]);
1449        let value = Value::Struct(json_variant_into_struct_value(
1450            variant,
1451            struct_type.clone(),
1452            true,
1453        )?);
1454
1455        assert_eq!(
1456            value,
1457            Value::Struct(StructValue::new(
1458                vec![
1459                    Value::List(ListValue::new(
1460                        vec![
1461                            Value::Struct(StructValue::new(
1462                                vec![Value::Int64(1)],
1463                                StructType::new(Arc::new(vec![StructField::new(
1464                                    "id".to_string(),
1465                                    ConcreteDataType::int64_datatype(),
1466                                    true,
1467                                )]))
1468                            )),
1469                            Value::Struct(StructValue::new(
1470                                vec![Value::Int64(2)],
1471                                StructType::new(Arc::new(vec![StructField::new(
1472                                    "id".to_string(),
1473                                    ConcreteDataType::int64_datatype(),
1474                                    true,
1475                                )]))
1476                            )),
1477                        ],
1478                        Arc::new(item_type),
1479                    )),
1480                    Value::Struct(StructValue::new(
1481                        vec![Value::String("foo".into())],
1482                        StructType::new(Arc::new(vec![StructField::new(
1483                            "name".to_string(),
1484                            ConcreteDataType::string_datatype(),
1485                            true,
1486                        )])),
1487                    )),
1488                ],
1489                struct_type,
1490            ))
1491        );
1492        Ok(())
1493    }
1494}
datatypes/vectors/json/builder.rs

datatypes/vectors/json/
builder.rs
