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                    skipping_index: None,
1065                },
1066                JsonTypeHint {
1067                    path: vec!["commit".to_string(), "operation".to_string()],
1068                    data_type: ConcreteDataType::string_datatype(),
1069                    inverted_index: false,
1070                    skipping_index: None,
1071                },
1072                JsonTypeHint {
1073                    path: vec!["time_us".to_string()],
1074                    data_type: ConcreteDataType::int64_datatype(),
1075                    inverted_index: false,
1076                    skipping_index: None,
1077                },
1078            ],
1079            Some(0),
1080        )?;
1081        let mut builder = JsonVectorBuilder::with_settings(&settings, 2);
1082        assert!(matches!(
1083            &builder.state,
1084            JsonVectorBuilderState::ExplicitOnly { .. }
1085        ));
1086        let values = [
1087            json!({
1088                "kind": "record",
1089                "commit": {"operation": "create", "collection": "post"},
1090                "extra": 1,
1091                "time_us": 1
1092            }),
1093            json!({"kind": "other", "dynamic": true, "time_us": 2}),
1094        ];
1095        for value in values.clone() {
1096            let value = settings.encode(value)?;
1097            builder.try_push_value_ref(&value.as_value_ref())?;
1098        }
1099        let array = builder.to_vector().to_arrow_array();
1100        assert_eq!(&json2_physical_data_type(&settings), array.data_type());
1101        assert_eq!(0, builder.len());
1102        let structs = array.as_struct();
1103        assert_eq!(
1104            vec![JSON2_REMAINDER_FIELD_NAME, "commit", "kind", "time_us"],
1105            structs
1106                .fields()
1107                .iter()
1108                .map(|x| x.name().as_str())
1109                .collect::<Vec<_>>()
1110        );
1111        assert_eq!(
1112            vec![
1113                Some(json!({"commit": {"collection": "post"}, "extra": 1})),
1114                Some(json!({"commit": {"operation": null}, "dynamic": true})),
1115            ],
1116            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
1117        );
1118
1119        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
1120            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))),
1121        );
1122        let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?;
1123        let reconstructed = reconstructed.as_binary::<i32>();
1124        assert_eq!(
1125            values[0],
1126            decode_json_variant(reconstructed.value(0)).unwrap()
1127        );
1128        assert_eq!(
1129            json!({
1130                "kind": "other",
1131                "commit": {"operation": null},
1132                "dynamic": true,
1133                "time_us": 2
1134            }),
1135            decode_json_variant(reconstructed.value(1)).unwrap()
1136        );
1137
1138        let settings = JsonSettings::try_new(vec![], Some(0))?;
1139        let mut builder = JsonVectorBuilder::with_settings(&settings, 3);
1140        for value in [json!({}), json!({"x": 1})] {
1141            let value = settings.encode(value)?;
1142            builder.try_push_value_ref(&value.as_value_ref())?;
1143        }
1144        builder.push_null();
1145        let array = builder.to_vector().to_arrow_array();
1146        let structs = array.as_struct();
1147        assert_eq!(1, structs.num_columns());
1148        assert_eq!(
1149            vec![Some(json!({})), Some(json!({"x": 1})), None],
1150            variant_to_json_values(structs.column(0))?
1151        );
1152
1153        Ok(())
1154    }
1155
1156    #[test]
1157    fn test_finite_budget_selects_dynamic_paths() -> Result<()> {
1158        let builder = JsonVectorBuilder::with_settings(&JsonSettings::default(), 0);
1159        assert!(matches!(
1160            builder.state,
1161            JsonVectorBuilderState::AutoExpanding {
1162                max_auto_expanded_paths: JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS,
1163                ..
1164            }
1165        ));
1166
1167        let settings = JsonSettings::try_new(
1168            vec![JsonTypeHint {
1169                path: vec!["hint".to_string()],
1170                data_type: ConcreteDataType::string_datatype(),
1171                inverted_index: false,
1172                skipping_index: None,
1173            }],
1174            Some(2),
1175        )?;
1176        let values = [
1177            json!({
1178                "hint": "first",
1179                "conflict": 1,
1180                "popular": {"nested": 1},
1181                "tie_a": "a",
1182                "tie_b": true,
1183                "rare": 1
1184            }),
1185            json!({
1186                "hint": "second",
1187                "conflict": "string",
1188                "popular": {"nested": 2},
1189                "tie_a": "b",
1190                "tie_b": false
1191            }),
1192            json!({"hint": "third", "popular": "scalar"}),
1193        ];
1194        let mut builder = JsonVectorBuilder::with_settings(&settings, values.len());
1195        for value in values.clone() {
1196            let value = settings.encode(value)?;
1197            builder.try_push_value_ref(&value.as_value_ref())?;
1198        }
1199
1200        let array = builder.to_vector().to_arrow_array();
1201        let structs = array.as_struct();
1202        assert_eq!(
1203            vec![JSON2_REMAINDER_FIELD_NAME, "hint", "tie_a", "tie_b"],
1204            structs
1205                .fields()
1206                .iter()
1207                .map(|x| x.name().as_str())
1208                .collect::<Vec<_>>()
1209        );
1210        assert!(structs.column_by_name("popular").is_none());
1211        assert_eq!(
1212            vec![
1213                Some(json!({"conflict": 1, "popular": {"nested": 1}, "rare": 1})),
1214                Some(json!({"conflict": "string", "popular": {"nested": 2}})),
1215                Some(json!({"popular": "scalar"})),
1216            ],
1217            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
1218        );
1219
1220        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
1221            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))),
1222        );
1223        let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?;
1224        let reconstructed = reconstructed.as_binary::<i32>();
1225        assert_eq!(
1226            values[0],
1227            decode_json_variant(reconstructed.value(0)).unwrap()
1228        );
1229        assert_eq!(
1230            values[1],
1231            decode_json_variant(reconstructed.value(1)).unwrap()
1232        );
1233        assert_eq!(
1234            json!({
1235                "hint": "third",
1236                "popular": "scalar"
1237            }),
1238            decode_json_variant(reconstructed.value(2)).unwrap()
1239        );
1240
1241        Ok(())
1242    }
1243
1244    #[test]
1245    fn test_v2_builder_preserves_explicit_null_presence() -> Result<()> {
1246        let settings = JsonSettings::try_new(vec![], Some(1))?;
1247        let values = [json!({"value": 1}), json!({"value": null}), json!({})];
1248        let mut builder = JsonVectorBuilder::with_settings(&settings, values.len());
1249        for value in values.clone() {
1250            let value = settings.encode(value)?;
1251            builder.try_push_value_ref(&value.as_value_ref())?;
1252        }
1253        let array = builder.to_vector().to_arrow_array();
1254        let structs = array.as_struct();
1255        assert_eq!(
1256            vec![
1257                Some(json!({})),
1258                Some(json!({"value": null})),
1259                Some(json!({}))
1260            ],
1261            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
1262        );
1263        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
1264            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))),
1265        );
1266        let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?;
1267        let reconstructed = reconstructed.as_binary::<i32>();
1268
1269        assert_eq!(
1270            values[0],
1271            decode_json_variant(reconstructed.value(0)).unwrap()
1272        );
1273        assert_eq!(
1274            values[1],
1275            decode_json_variant(reconstructed.value(1)).unwrap()
1276        );
1277        assert_eq!(
1278            values[2],
1279            decode_json_variant(reconstructed.value(2)).unwrap()
1280        );
1281        Ok(())
1282    }
1283
1284    #[test]
1285    fn test_v2_builder_accepts_sql_null() -> Result<()> {
1286        let settings = JsonSettings::try_new(vec![], Some(1))?;
1287        let mut builder = JsonVectorBuilder::with_settings(&settings, 2);
1288        builder.try_push_value_ref(&ValueRef::Null)?;
1289        let value = settings.encode(json!({}))?;
1290        builder.try_push_value_ref(&value.as_value_ref())?;
1291
1292        let array = builder.to_vector().to_arrow_array();
1293        let structs = array.as_struct();
1294        assert_eq!(
1295            vec![None, Some(json!({}))],
1296            variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())?
1297        );
1298        Ok(())
1299    }
1300
1301    #[test]
1302    fn test_reconstruct_nested_remainder_only_value() -> Result<()> {
1303        let settings = JsonSettings::try_new(vec![], Some(1))?;
1304        let values = [
1305            json!({"a": {"hot": 1}}),
1306            json!({"a": {"hot": 2}}),
1307            json!({"a": {"cold": 3}}),
1308        ];
1309        let mut builder = JsonVectorBuilder::with_settings(&settings, values.len());
1310        for value in values.clone() {
1311            let value = settings.encode(value)?;
1312            builder.try_push_value_ref(&value.as_value_ref())?;
1313        }
1314
1315        let array = builder.to_vector().to_arrow_array();
1316        let field = Field::new("data", array.data_type().clone(), true).with_extension_type(
1317            Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))),
1318        );
1319        let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?;
1320        let reconstructed = reconstructed.as_binary::<i32>();
1321        assert_eq!(
1322            vec![values[0].clone(), values[1].clone(), values[2].clone()],
1323            (0..reconstructed.len())
1324                .map(|i| decode_json_variant(reconstructed.value(i)).unwrap())
1325                .collect::<Vec<_>>()
1326        );
1327
1328        Ok(())
1329    }
1330
1331    #[test]
1332    fn test_dynamic_paths_require_the_same_leaf_type() -> Result<()> {
1333        let settings = JsonSettings::try_new(
1334            vec![JsonTypeHint {
1335                path: vec!["nested".to_string(), "hinted".to_string()],
1336                data_type: ConcreteDataType::string_datatype(),
1337                inverted_index: false,
1338                skipping_index: None,
1339            }],
1340            Some(8),
1341        )?;
1342        let values = [
1343            json!({
1344                "branch": {},
1345                "different": [1],
1346                "empty": {},
1347                "nested": {},
1348                "reverse": "scalar",
1349                "same": [1]
1350            }),
1351            json!({
1352                "branch": {"leaf": 1},
1353                "different": ["x"],
1354                "empty": {},
1355                "nested": {"hinted": "x", "leaf": 1},
1356                "same": [2]
1357            }),
1358            json!({
1359                "branch": {},
1360                "different": [2],
1361                "empty": {},
1362                "nested": {},
1363                "reverse": {"leaf": 1},
1364                "same": [3]
1365            }),
1366        ];
1367        let mut builder = JsonVectorBuilder::with_settings(&settings, values.len());
1368        for value in values {
1369            let value = settings.encode(value)?;
1370            builder.try_push_value_ref(&value.as_value_ref())?;
1371        }
1372
1373        let JsonNativeType::Object(fields) = builder.state.native_type() else {
1374            unreachable!();
1375        };
1376        assert!(!fields.contains_key("different"));
1377        assert!(!fields.contains_key("empty"));
1378        assert!(!fields.contains_key("reverse"));
1379        assert!(matches!(fields.get("same"), Some(JsonNativeType::Array(_))));
1380        assert!(matches!(
1381            fields.get("branch"),
1382            Some(JsonNativeType::Object(fields)) if fields.contains_key("leaf")
1383        ));
1384        assert!(matches!(
1385            fields.get("nested"),
1386            Some(JsonNativeType::Object(fields))
1387                if fields.contains_key("hinted") && fields.contains_key("leaf")
1388        ));
1389
1390        Ok(())
1391    }
1392
1393    #[test]
1394    fn test_json_variant_into_struct_value() -> Result<()> {
1395        let struct_type = StructType::new(Arc::new(vec![StructField::new(
1396            "value".to_string(),
1397            ConcreteDataType::string_datatype(),
1398            true,
1399        )]));
1400        assert_eq!(
1401            json_variant_into_value(
1402                JsonVariant::Object(Default::default()),
1403                &ConcreteDataType::struct_datatype(struct_type.clone()),
1404                false,
1405            )?,
1406            Value::Null
1407        );
1408        assert_eq!(
1409            json_variant_into_value(
1410                JsonVariant::Object(Default::default()),
1411                &ConcreteDataType::struct_datatype(struct_type.clone()),
1412                true,
1413            )?,
1414            Value::Struct(StructValue::new(vec![Value::Null], struct_type))
1415        );
1416
1417        let item_type =
1418            ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![StructField::new(
1419                "id".to_string(),
1420                ConcreteDataType::int64_datatype(),
1421                true,
1422            )])));
1423        let struct_type = StructType::new(Arc::new(vec![
1424            StructField::new(
1425                "items".to_string(),
1426                ConcreteDataType::list_datatype(Arc::new(item_type.clone())),
1427                true,
1428            ),
1429            StructField::new(
1430                "meta".to_string(),
1431                ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![
1432                    StructField::new(
1433                        "name".to_string(),
1434                        ConcreteDataType::string_datatype(),
1435                        true,
1436                    ),
1437                ]))),
1438                true,
1439            ),
1440        ]));
1441        let variant = JsonObjectValue::from([
1442            (
1443                "items".to_string(),
1444                JsonVariant::Array(vec![
1445                    JsonVariant::from([("id", JsonVariant::from(1i64))]),
1446                    JsonVariant::from([("id", JsonVariant::from(2i64))]),
1447                ]),
1448            ),
1449            (
1450                "meta".to_string(),
1451                JsonVariant::from([("name", JsonVariant::from("foo"))]),
1452            ),
1453        ]);
1454        let value = Value::Struct(json_variant_into_struct_value(
1455            variant,
1456            struct_type.clone(),
1457            true,
1458        )?);
1459
1460        assert_eq!(
1461            value,
1462            Value::Struct(StructValue::new(
1463                vec![
1464                    Value::List(ListValue::new(
1465                        vec![
1466                            Value::Struct(StructValue::new(
1467                                vec![Value::Int64(1)],
1468                                StructType::new(Arc::new(vec![StructField::new(
1469                                    "id".to_string(),
1470                                    ConcreteDataType::int64_datatype(),
1471                                    true,
1472                                )]))
1473                            )),
1474                            Value::Struct(StructValue::new(
1475                                vec![Value::Int64(2)],
1476                                StructType::new(Arc::new(vec![StructField::new(
1477                                    "id".to_string(),
1478                                    ConcreteDataType::int64_datatype(),
1479                                    true,
1480                                )]))
1481                            )),
1482                        ],
1483                        Arc::new(item_type),
1484                    )),
1485                    Value::Struct(StructValue::new(
1486                        vec![Value::String("foo".into())],
1487                        StructType::new(Arc::new(vec![StructField::new(
1488                            "name".to_string(),
1489                            ConcreteDataType::string_datatype(),
1490                            true,
1491                        )])),
1492                    )),
1493                ],
1494                struct_type,
1495            ))
1496        );
1497        Ok(())
1498    }
1499}
datatypes/vectors/json/builder.rs

datatypes/vectors/json/
builder.rs
