mito2/sst/parquet/
json_align.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
15//! Shared JSON2 conversion for Parquet streams and in-memory query batches.
16//! Query types must agree before merge/dedup; compaction Rewrite retains physical layouts.
17
18use std::collections::HashMap;
19use std::pin::Pin;
20use std::sync::Arc;
21use std::task::{Context, Poll};
22
23use datafusion_common::cast_column;
24use datafusion_common::format::DEFAULT_CAST_OPTIONS;
25use datatypes::arrow::array::{ArrayRef, new_null_array};
26use datatypes::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef};
27use datatypes::arrow::record_batch::RecordBatch;
28use datatypes::extension::json::{JsonMetadata, is_json2_extension_type};
29use datatypes::json::{JsonSettings, TypeHintMismatchPolicy};
30use datatypes::vectors::json::array::JsonArray;
31use datatypes::vectors::json::json2_physical_data_type;
32use futures::Stream;
33use futures::stream::BoxStream;
34use serde_json::from_str;
35use snafu::{ResultExt, ensure};
36
37use crate::error::{
38    CastColumnSnafu, DataTypeMismatchSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu,
39};
40use crate::sst::parquet::Json2TargetLayout;
41
42pub(crate) type ProjectedRecordBatchStream = BoxStream<'static, Result<RecordBatch>>;
43
44/// Specifies how JSON columns in a record batch are aligned.
45#[derive(Debug)]
46pub(crate) enum AlignMode {
47    /// Aligns JSON columns to the logical fields in the output schema.
48    AlignToSchema,
49    /// Rewrites JSON columns to physical layouts, typically for compaction.
50    Rewrite {
51        /// Target layouts keyed by root column name, not nested field path.
52        ///
53        /// Only listed columns are rewritten. Other existing arrays are reused
54        /// unchanged and must already match their output field types.
55        /// An empty map therefore only fills missing roots.
56        columns: HashMap<String, Json2TargetLayout>,
57    },
58}
59
60/// Alignment mode with parsed rewrite metadata and validated target layouts.
61#[derive(Debug)]
62enum ResolvedAlignMode {
63    AlignToSchema,
64    Rewrite {
65        columns: HashMap<String, RewriteSettings>,
66    },
67}
68
69/// Adapts Parquet record batches to the output schema expected by the reader.
70///
71/// Nested projection can return only part of a JSON2 column, or omit its root
72/// entirely when no requested leaves are read. This stream restores missing
73/// roots with null arrays of the expected types.
74///
75/// Existing JSON2 columns are aligned to the logical schema inferred from
76/// type hints or rewritten to the specified JSON2 physical layout, as selected by
77/// [`AlignMode`].
78#[derive(derive_more::Debug)]
79pub struct JsonSchemaAligner<S> {
80    #[debug(skip)]
81    inner: S,
82    /// Output schema expected by the upper reader.
83    output_schema: SchemaRef,
84    /// Whether each projected root exists in the physical batch returned by
85    /// parquet.
86    projected_root_presence: Vec<bool>,
87    /// Number of columns expected from the physical batch returned by parquet.
88    expected_input_col_num: usize,
89    /// Whether all projected roots are present and the stream can pass batches
90    /// through.
91    all_roots_present: bool,
92    /// Alignment mode with parsed and validated rewrite settings.
93    mode: ResolvedAlignMode,
94    /// The cache for whether incoming batches already match output schema.
95    is_schema_matched: Option<bool>,
96}
97
98impl<S> JsonSchemaAligner<S>
99where
100    S: Stream<Item = Result<RecordBatch>>,
101{
102    /// Creates an aligner with a shared output schema and an explicit operation.
103    /// Parses rewrite metadata once and validates layouts against output field types.
104    pub(crate) fn new(
105        inner: S,
106        projected_root_presence: Vec<bool>,
107        output_schema: SchemaRef,
108        mode: AlignMode,
109    ) -> Result<JsonSchemaAligner<S>> {
110        ensure!(
111            projected_root_presence.len() == output_schema.fields().len(),
112            UnexpectedSnafu {
113                reason: format!(
114                    "JsonSchemaAligner projected root presence len {} does not match output schema columns {}",
115                    projected_root_presence.len(),
116                    output_schema.fields().len()
117                ),
118            }
119        );
120
121        let mode = resolve_align_mode(mode, output_schema.as_ref())?;
122
123        let expected_input_col_num = projected_root_presence
124            .iter()
125            .filter(|matched| **matched)
126            .count();
127        let all_roots_present = projected_root_presence.iter().all(|&m| m);
128        Ok(JsonSchemaAligner {
129            inner,
130            output_schema,
131            projected_root_presence,
132            expected_input_col_num,
133            all_roots_present,
134            mode,
135            is_schema_matched: None,
136        })
137    }
138}
139
140impl<S> Stream for JsonSchemaAligner<S>
141where
142    S: Stream<Item = Result<RecordBatch>> + Unpin,
143{
144    type Item = Result<RecordBatch>;
145
146    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
147        let this = self.get_mut();
148
149        match Pin::new(&mut this.inner).poll_next(cx) {
150            Poll::Ready(Some(Ok(rb))) => {
151                let is_schema_matched = matches!(this.mode, ResolvedAlignMode::AlignToSchema)
152                    && this.all_roots_present
153                    && *this
154                        .is_schema_matched
155                        .get_or_insert_with(|| rb.schema() == this.output_schema);
156
157                if is_schema_matched {
158                    Poll::Ready(Some(Ok(rb)))
159                } else {
160                    Poll::Ready(Some(align_projected_batch(
161                        rb,
162                        &this.output_schema,
163                        &this.projected_root_presence,
164                        this.expected_input_col_num,
165                        &this.mode,
166                    )))
167                }
168            }
169            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
170            Poll::Ready(None) => Poll::Ready(None),
171            Poll::Pending => Poll::Pending,
172        }
173    }
174}
175
176/// Aligns existing JSON2 columns to the Arrow types of matching fields in
177/// `query_schema`, matching columns by name.
178/// Preserves column order and skips columns absent from the batch.
179pub(crate) fn align_query_batch(
180    batch: RecordBatch,
181    query_schema: &SchemaRef,
182) -> Result<RecordBatch> {
183    let mut aligned = None;
184    for target in query_schema
185        .fields()
186        .iter()
187        .filter(|field| is_json2_extension_type(field))
188    {
189        let Ok(index) = batch.schema_ref().index_of(target.name()) else {
190            continue;
191        };
192        let source = batch.schema_ref().field(index);
193        if source.data_type() != target.data_type() {
194            let array = align_array(batch.column(index), source, target)?;
195            let (fields, columns) = aligned.get_or_insert_with(|| {
196                (
197                    batch.schema_ref().fields().to_vec(),
198                    batch.columns().to_vec(),
199                )
200            });
201            fields[index] = Arc::new(source.clone().with_data_type(target.data_type().clone()));
202            columns[index] = array;
203        }
204    }
205    let Some((fields, columns)) = aligned else {
206        return Ok(batch);
207    };
208    let schema = Arc::new(Schema::new_with_metadata(
209        fields,
210        batch.schema().metadata().clone(),
211    ));
212    RecordBatch::try_new(schema, columns).context(NewRecordBatchSnafu)
213}
214
215fn align_projected_batch(
216    rb: RecordBatch,
217    output_schema: &SchemaRef,
218    projected_root_presence: &[bool],
219    expected_input_col_num: usize,
220    mode: &ResolvedAlignMode,
221) -> Result<RecordBatch> {
222    ensure!(
223        rb.columns().len() == expected_input_col_num,
224        UnexpectedSnafu {
225            reason: format!(
226                "JsonSchemaAligner expected {} input columns but got {}",
227                expected_input_col_num,
228                rb.columns().len()
229            ),
230        }
231    );
232
233    let mut cols = Vec::with_capacity(projected_root_presence.len());
234    let mut idx = 0;
235    let input_schema = rb.schema_ref();
236
237    for (field, present) in output_schema.fields().iter().zip(projected_root_presence) {
238        if !present {
239            cols.push(new_null_array(field.data_type(), rb.num_rows()));
240            continue;
241        }
242
243        let array = match mode {
244            ResolvedAlignMode::AlignToSchema => {
245                align_array(rb.column(idx), input_schema.field(idx), field)?
246            }
247            ResolvedAlignMode::Rewrite { columns } => match columns.get(field.name()) {
248                Some(settings) => rewrite_array(rb.column(idx), input_schema.field(idx), settings)?,
249                None => rb.column(idx).clone(),
250            },
251        };
252        cols.push(array);
253        idx += 1;
254    }
255
256    RecordBatch::try_new(output_schema.clone(), cols).context(NewRecordBatchSnafu)
257}
258
259fn align_array(
260    source_array: &ArrayRef,
261    source_field: &Field,
262    target_field: &FieldRef,
263) -> Result<ArrayRef> {
264    if source_array.data_type() == target_field.data_type() {
265        return Ok(source_array.clone());
266    }
267
268    if is_json2_extension_type(target_field) {
269        if is_json2_extension_type(source_field) {
270            return JsonArray::from(source_array)
271                .project_to_v2(source_field, target_field.data_type())
272                .context(DataTypeMismatchSnafu);
273        }
274        return JsonArray::from(source_array)
275            .project_to(target_field.data_type())
276            .context(DataTypeMismatchSnafu);
277    }
278
279    if !matches!(target_field.data_type(), DataType::Struct(_)) {
280        return Ok(source_array.clone());
281    }
282
283    cast_column(
284        source_array,
285        target_field.data_type(),
286        &DEFAULT_CAST_OPTIONS,
287    )
288    .context(CastColumnSnafu)
289}
290
291fn rewrite_array(
292    source_array: &ArrayRef,
293    source_field: &Field,
294    settings: &RewriteSettings,
295) -> Result<ArrayRef> {
296    JsonArray::from(source_array)
297        .rewrite_to_v2_with_type_hint_mismatch_policy(
298            source_field,
299            &settings.logical_settings,
300            &settings.target_layout,
301            TypeHintMismatchPolicy::CoerceOrNull,
302        )
303        .context(DataTypeMismatchSnafu)
304}
305
306/// Resolved settings for rewriting one JSON2 column to a target physical layout.
307///
308/// Created from [`Json2TargetLayout`] when resolving the alignment mode, so
309/// extension metadata is parsed once and reused across batches.
310#[derive(Debug)]
311struct RewriteSettings {
312    /// Logical settings parsed from extension metadata and applied to JSON values
313    /// before encoding them into the target layout.
314    logical_settings: JsonSettings,
315    /// Settings defining the physical Arrow layout of the rewritten column.
316    target_layout: JsonSettings,
317}
318
319impl TryFrom<&Json2TargetLayout> for RewriteSettings {
320    type Error = crate::error::Error;
321
322    fn try_from(layout: &Json2TargetLayout) -> Result<Self> {
323        let metadata = from_str::<JsonMetadata>(&layout.extension_metadata).map_err(|e| {
324            UnexpectedSnafu {
325                reason: format!("invalid JSON2 extension metadata: {e}"),
326            }
327            .build()
328        })?;
329        Ok(Self {
330            logical_settings: metadata.into_json_settings(),
331            target_layout: layout.target_layout.clone(),
332        })
333    }
334}
335
336/// Parses rewrite metadata and validates target layouts against the output schema.
337fn resolve_align_mode(mode: AlignMode, output_schema: &Schema) -> Result<ResolvedAlignMode> {
338    let AlignMode::Rewrite { columns } = mode else {
339        return Ok(ResolvedAlignMode::AlignToSchema);
340    };
341
342    let mut rewrite_columns = HashMap::with_capacity(columns.len());
343    for (name, layout) in columns {
344        let settings = RewriteSettings::try_from(&layout)?;
345        let field = output_schema.field_with_name(&name).map_err(|_| {
346            UnexpectedSnafu {
347                reason: format!("JSON2 rewrite column '{name}' is missing from output schema"),
348            }
349            .build()
350        })?;
351        ensure!(
352            is_json2_extension_type(field)
353                && field.data_type() == &json2_physical_data_type(&settings.target_layout),
354            UnexpectedSnafu {
355                reason: format!(
356                    "JSON2 rewrite layout for column '{name}' does not match output field"
357                ),
358            }
359        );
360        rewrite_columns.insert(name, settings);
361    }
362
363    Ok(ResolvedAlignMode::Rewrite {
364        columns: rewrite_columns,
365    })
366}
367
368#[cfg(test)]
369mod tests {
370    use std::collections::HashMap;
371    use std::sync::Arc;
372
373    use super::*;
374    use crate::test_util;
375    use datatypes::arrow::array::{
376        Array, ArrayRef, BinaryArray, Int64Array, StringArray, StringViewArray, StructArray,
377    };
378    use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema};
379    use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
380    use datatypes::json::JsonTypeHint;
381    use datatypes::prelude::ConcreteDataType;
382    use datatypes::types::parse_string_to_jsonb;
383    use futures::{StreamExt, stream};
384
385    #[test]
386    fn test_aligner_resolves_json2_rewrite_settings()
387    -> std::result::Result<(), Box<dyn std::error::Error>> {
388        let logical_settings = JsonSettings::default();
389        let target_layout = JsonSettings::try_new(vec![], Some(0))?;
390        let rewrite_targets = HashMap::from([(
391            "j".to_string(),
392            Json2TargetLayout {
393                extension_metadata: serde_json::to_string(&JsonMetadata::new(
394                    logical_settings.clone(),
395                ))?,
396                target_layout: target_layout.clone(),
397            },
398        )]);
399        let aligner = JsonSchemaAligner::new(
400            stream::empty::<Result<RecordBatch>>(),
401            vec![false],
402            schema([
403                Field::new("j", json2_physical_data_type(&target_layout), true)
404                    .with_extension_type(Json2ExtensionType::default()),
405            ]),
406            AlignMode::Rewrite {
407                columns: rewrite_targets,
408            },
409        )?;
410        let ResolvedAlignMode::Rewrite { columns } = &aligner.mode else {
411            panic!("expected rewrite mode");
412        };
413        let settings = &columns["j"];
414        assert_eq!(logical_settings, settings.logical_settings);
415        assert_eq!(target_layout, settings.target_layout);
416        Ok(())
417    }
418
419    #[tokio::test]
420    async fn test_aligner_with_all_projected_roots_match() {
421        let output_schema = schema([
422            Field::new("a", DataType::Int64, true),
423            Field::new("b", DataType::Utf8, true),
424        ]);
425        let input = RecordBatch::try_new(
426            output_schema.clone(),
427            vec![int_array([1, 2, 3]), string_array(["x", "y", "z"])],
428        )
429        .unwrap();
430        let stream = stream::iter([Ok(input.clone())]);
431
432        let mut aligner = JsonSchemaAligner::new(
433            stream,
434            vec![true, true],
435            output_schema.clone(),
436            AlignMode::AlignToSchema,
437        )
438        .unwrap();
439        let output = aligner.next().await.unwrap().unwrap();
440
441        assert_eq!(input, output);
442        assert!(aligner.next().await.is_none());
443    }
444
445    #[tokio::test]
446    async fn test_aligner_with_fills_null_root_columns() {
447        let input_schema = schema([Field::new("a", DataType::Int64, true)]);
448        let output_schema = schema([
449            Field::new("a", DataType::Int64, true),
450            Field::new("missing", DataType::Utf8, true),
451            Field::new("c", DataType::Int64, true),
452        ]);
453        let input = RecordBatch::try_new(input_schema, vec![int_array([10, 20])]).unwrap();
454        let stream = stream::iter([Ok(input)]);
455
456        let mut aligner = JsonSchemaAligner::new(
457            stream,
458            vec![true, false, false],
459            output_schema.clone(),
460            AlignMode::AlignToSchema,
461        )
462        .unwrap();
463        let output = aligner.next().await.unwrap().unwrap();
464
465        assert_eq!(output_schema, output.schema());
466        assert_eq!(3, output.num_columns());
467        assert_eq!(
468            &[Some(10), Some(20)],
469            output
470                .column(0)
471                .as_any()
472                .downcast_ref::<Int64Array>()
473                .unwrap()
474                .iter()
475                .collect::<Vec<_>>()
476                .as_slice()
477        );
478        assert_eq!(DataType::Utf8, *output.column(1).data_type());
479        assert_eq!(output.num_rows(), output.column(1).null_count());
480        assert_eq!(DataType::Int64, *output.column(2).data_type());
481        assert_eq!(output.num_rows(), output.column(2).null_count());
482    }
483
484    #[tokio::test]
485    async fn test_aligner_with_fills_missing_struct_root_column() {
486        let input_schema = schema([Field::new("a", DataType::Int64, true)]);
487        let struct_type = DataType::Struct(Fields::from(vec![
488            Field::new("x", DataType::Int64, true),
489            Field::new("y", DataType::Utf8, true),
490        ]));
491        let output_schema = schema([
492            Field::new("a", DataType::Int64, true),
493            Field::new("missing_struct", struct_type.clone(), true),
494        ]);
495        let input = RecordBatch::try_new(input_schema, vec![int_array([10, 20])]).unwrap();
496        let stream = stream::iter([Ok(input)]);
497
498        let mut aligner = JsonSchemaAligner::new(
499            stream,
500            vec![true, false],
501            output_schema.clone(),
502            AlignMode::AlignToSchema,
503        )
504        .unwrap();
505        let output = aligner.next().await.unwrap().unwrap();
506
507        assert_eq!(output_schema, output.schema());
508        assert_eq!(2, output.num_columns());
509        assert_eq!(struct_type, output.column(1).data_type().clone());
510        assert_eq!(output.num_rows(), output.column(1).null_count());
511    }
512
513    #[tokio::test]
514    async fn test_aligner_reject_projection_len_mismatch() {
515        let output_schema = schema([Field::new("a", DataType::Int64, true)]);
516        let stream = stream::iter([]);
517
518        let err = match JsonSchemaAligner::new(
519            stream,
520            vec![true, false],
521            output_schema,
522            AlignMode::AlignToSchema,
523        ) {
524            Ok(_) => panic!("JsonSchemaAligner should reject projection length mismatch"),
525            Err(err) => err,
526        };
527
528        assert!(
529            err.to_string()
530                .contains("projected root presence len 2 does not match output schema columns 1")
531        );
532    }
533
534    #[tokio::test]
535    async fn test_aligner_reject_with_input_column_mismatch() {
536        let input_schema = schema([Field::new("a", DataType::Int64, true)]);
537        let output_schema = schema([
538            Field::new("a", DataType::Int64, true),
539            Field::new("b", DataType::Int64, true),
540            Field::new("missing", DataType::Int64, true),
541        ]);
542        let input = RecordBatch::try_new(input_schema, vec![int_array([1, 2])]).unwrap();
543        let stream = stream::iter([Ok(input)]);
544
545        let mut aligner = JsonSchemaAligner::new(
546            stream,
547            vec![true, true, false],
548            output_schema,
549            AlignMode::AlignToSchema,
550        )
551        .unwrap();
552        let err = aligner.next().await.unwrap().unwrap_err();
553
554        assert!(
555            err.to_string()
556                .contains("expected 2 input columns but got 1")
557        );
558    }
559
560    #[tokio::test]
561    async fn test_json_schema_aligner_aligns_struct_field() {
562        let output_schema = schema([Field::new(
563            "nested",
564            DataType::Struct(Fields::from(vec![
565                Field::new("x", DataType::Int64, true),
566                Field::new("y", DataType::Utf8, true),
567            ])),
568            true,
569        )]);
570        let input = RecordBatch::try_new(
571            schema([Field::new(
572                "nested",
573                DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int64, true)])),
574                true,
575            )]),
576            vec![Arc::new(StructArray::from(vec![(
577                Arc::new(Field::new("x", DataType::Int64, true)),
578                int_array([1, 2]),
579            )]))],
580        )
581        .unwrap();
582
583        let mut aligner = JsonSchemaAligner::new(
584            stream::iter([Ok(input)]),
585            vec![true],
586            output_schema.clone(),
587            AlignMode::AlignToSchema,
588        )
589        .unwrap();
590        let output = aligner.next().await.unwrap().unwrap();
591
592        assert_eq!(output_schema, output.schema());
593        let nested = output
594            .column(0)
595            .as_any()
596            .downcast_ref::<StructArray>()
597            .unwrap();
598        assert_eq!(2, nested.columns().len());
599        assert_eq!(2, nested.column(1).null_count());
600    }
601
602    #[tokio::test]
603    async fn test_json_schema_aligner_decodes_variant_to_struct() {
604        let source_values = [
605            Some(parse_string_to_jsonb("1").unwrap()),
606            Some(parse_string_to_jsonb(r#"{"b":2}"#).unwrap()),
607            None,
608        ];
609        let source = Arc::new(BinaryArray::from_iter(
610            source_values.iter().map(|value| value.as_deref()),
611        )) as ArrayRef;
612        let input_fields = Fields::from(vec![Arc::new(Field::new("a", DataType::Binary, true))]);
613        let input = RecordBatch::try_new(
614            schema([Field::new(
615                "j",
616                DataType::Struct(input_fields.clone()),
617                true,
618            )]),
619            vec![Arc::new(StructArray::new(input_fields, vec![source], None))],
620        )
621        .unwrap();
622
623        let output_schema = schema([Field::new(
624            "j",
625            DataType::Struct(Fields::from(vec![Arc::new(Field::new(
626                "a",
627                DataType::Struct(Fields::from(vec![
628                    Arc::new(Field::new("b", DataType::UInt64, true)),
629                    Arc::new(Field::new("c", DataType::Utf8View, true)),
630                ])),
631                true,
632            ))])),
633            true,
634        )
635        .with_extension_type(Json2ExtensionType::default())]);
636        let mut aligner = JsonSchemaAligner::new(
637            stream::iter([Ok(input)]),
638            vec![true],
639            output_schema.clone(),
640            AlignMode::AlignToSchema,
641        )
642        .unwrap();
643        let output = aligner.next().await.unwrap().unwrap();
644
645        assert_eq!(output_schema, output.schema());
646        let j = output
647            .column(0)
648            .as_any()
649            .downcast_ref::<StructArray>()
650            .unwrap();
651        let a = j.column(0).as_any().downcast_ref::<StructArray>().unwrap();
652        assert_eq!(
653            &[None, Some(2), None],
654            a.column(0)
655                .as_any()
656                .downcast_ref::<datatypes::arrow::array::UInt64Array>()
657                .unwrap()
658                .iter()
659                .collect::<Vec<_>>()
660                .as_slice()
661        );
662        assert_eq!(
663            &[None, None, None],
664            a.column(1)
665                .as_any()
666                .downcast_ref::<StringViewArray>()
667                .unwrap()
668                .iter()
669                .collect::<Vec<_>>()
670                .as_slice()
671        );
672    }
673
674    #[tokio::test]
675    async fn test_json_schema_aligner_preserves_struct_siblings() {
676        let source_values = [
677            Some(parse_string_to_jsonb(r#"{"x":1}"#).unwrap()),
678            Some(parse_string_to_jsonb(r#"{"x":2}"#).unwrap()),
679        ];
680        let source = Arc::new(BinaryArray::from_iter(
681            source_values.iter().map(|value| value.as_deref()),
682        )) as ArrayRef;
683        let c = Arc::new(Int64Array::from_iter_values([10, 20])) as ArrayRef;
684
685        let a_fields = Fields::from(vec![
686            Arc::new(Field::new("b", DataType::Binary, true)),
687            Arc::new(Field::new("c", DataType::Int64, true)),
688        ]);
689        let input_fields = Fields::from(vec![Arc::new(Field::new(
690            "a",
691            DataType::Struct(a_fields.clone()),
692            true,
693        ))]);
694        let input = RecordBatch::try_new(
695            schema([Field::new(
696                "j",
697                DataType::Struct(input_fields.clone()),
698                true,
699            )]),
700            vec![Arc::new(StructArray::new(
701                input_fields,
702                vec![Arc::new(StructArray::new(a_fields, vec![source, c], None))],
703                None,
704            ))],
705        )
706        .unwrap();
707
708        let output_schema = schema([Field::new(
709            "j",
710            DataType::Struct(Fields::from(vec![Arc::new(Field::new(
711                "a",
712                DataType::Struct(Fields::from(vec![
713                    Arc::new(Field::new(
714                        "b",
715                        DataType::Struct(Fields::from(vec![Arc::new(Field::new(
716                            "x",
717                            DataType::Int64,
718                            true,
719                        ))])),
720                        true,
721                    )),
722                    Arc::new(Field::new("c", DataType::Int64, true)),
723                ])),
724                true,
725            ))])),
726            true,
727        )
728        .with_extension_type(Json2ExtensionType::default())]);
729        let mut aligner = JsonSchemaAligner::new(
730            stream::iter([Ok(input)]),
731            vec![true],
732            output_schema.clone(),
733            AlignMode::AlignToSchema,
734        )
735        .unwrap();
736        let output = aligner.next().await.unwrap().unwrap();
737
738        assert_eq!(output_schema, output.schema());
739        let j = output
740            .column(0)
741            .as_any()
742            .downcast_ref::<StructArray>()
743            .unwrap();
744        let a = j.column(0).as_any().downcast_ref::<StructArray>().unwrap();
745        let b = a.column(0).as_any().downcast_ref::<StructArray>().unwrap();
746        assert_eq!(
747            &[Some(1), Some(2)],
748            b.column(0)
749                .as_any()
750                .downcast_ref::<Int64Array>()
751                .unwrap()
752                .iter()
753                .collect::<Vec<_>>()
754                .as_slice()
755        );
756        assert_eq!(
757            &[Some(10), Some(20)],
758            a.column(1)
759                .as_any()
760                .downcast_ref::<Int64Array>()
761                .unwrap()
762                .iter()
763                .collect::<Vec<_>>()
764                .as_slice()
765        );
766    }
767
768    #[tokio::test]
769    async fn test_rewrite_multiple_columns_and_fill_missing_roots() {
770        let logical_settings = JsonSettings::default();
771        let target_layout = JsonSettings::try_new(vec![], Some(0)).unwrap();
772        let target_type = json2_physical_data_type(&target_layout);
773        let output_schema = schema([
774            Field::new("j", target_type.clone(), true)
775                .with_extension_type(Json2ExtensionType::default()),
776            Field::new("missing", target_type.clone(), true)
777                .with_extension_type(Json2ExtensionType::default()),
778            Field::new("k", target_type.clone(), true)
779                .with_extension_type(Json2ExtensionType::default()),
780            Field::new("a", DataType::Int64, true),
781        ]);
782        let values = [Some(parse_string_to_jsonb(r#"{"x":1}"#).unwrap()), None];
783        let source = Arc::new(BinaryArray::from_iter(
784            values.iter().map(|value| value.as_deref()),
785        )) as ArrayRef;
786        let source_field = Field::new("j", DataType::Binary, true)
787            .with_extension_type(Json2ExtensionType::default());
788        let expected = JsonArray::from(&source)
789            .rewrite_to_v2(&source_field, &logical_settings, &target_layout)
790            .unwrap();
791        let input = RecordBatch::try_new(
792            schema([
793                source_field,
794                Field::new("k", DataType::Binary, true)
795                    .with_extension_type(Json2ExtensionType::default()),
796                Field::new("a", DataType::Int64, true),
797            ]),
798            vec![source.clone(), source, int_array([10, 20])],
799        )
800        .unwrap();
801        let columns = ["j", "missing", "k"]
802            .into_iter()
803            .map(|name| {
804                (
805                    name.to_string(),
806                    Json2TargetLayout {
807                        extension_metadata: serde_json::to_string(&JsonMetadata::new(
808                            logical_settings.clone(),
809                        ))
810                        .unwrap(),
811                        target_layout: target_layout.clone(),
812                    },
813                )
814            })
815            .collect();
816        let mut aligner = JsonSchemaAligner::new(
817            stream::iter([Ok(input)]),
818            vec![true, false, true, true],
819            output_schema.clone(),
820            AlignMode::Rewrite { columns },
821        )
822        .unwrap();
823        let output = aligner.next().await.unwrap().unwrap();
824        assert_eq!(output_schema, output.schema());
825        assert_eq!(expected.as_ref(), output.column(0).as_ref());
826        assert_eq!(expected.as_ref(), output.column(2).as_ref());
827        assert_eq!(&target_type, output.column(1).data_type());
828        assert_eq!(2, output.column(1).null_count());
829        assert_eq!(int_array([10, 20]).as_ref(), output.column(3).as_ref());
830    }
831
832    #[tokio::test]
833    async fn test_rewrite_keeps_rows_with_invalid_json2_settings() {
834        let settings = JsonSettings::try_new(
835            vec![JsonTypeHint {
836                path: vec!["kind".to_string()],
837                data_type: ConcreteDataType::string_datatype(),
838                inverted_index: false,
839                skipping_index: None,
840                fulltext_index: None,
841            }],
842            Some(0),
843        )
844        .unwrap();
845        let target_type = json2_physical_data_type(&settings);
846        let output_schema = schema([
847            Field::new("j", target_type.clone(), true).with_extension_type(
848                Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings.clone()))),
849            ),
850            Field::new("value", DataType::Int64, true),
851        ]);
852        let source = Arc::new(BinaryArray::from_iter([
853            Some(parse_string_to_jsonb(r#"{"kind":"valid"}"#).unwrap()),
854            Some(parse_string_to_jsonb(r#"{"kind":1}"#).unwrap()),
855        ])) as ArrayRef;
856        let input = RecordBatch::try_new(
857            schema([
858                Field::new("j", DataType::Binary, true)
859                    .with_extension_type(Json2ExtensionType::default()),
860                Field::new("value", DataType::Int64, true),
861            ]),
862            vec![source, int_array([10, 20])],
863        )
864        .unwrap();
865        let columns = HashMap::from([(
866            "j".to_string(),
867            Json2TargetLayout {
868                extension_metadata: serde_json::to_string(&JsonMetadata::new(settings.clone()))
869                    .unwrap(),
870                target_layout: settings,
871            },
872        )]);
873        let mut aligner = JsonSchemaAligner::new(
874            stream::iter([Ok(input)]),
875            vec![true, true],
876            output_schema,
877            AlignMode::Rewrite { columns },
878        )
879        .unwrap();
880
881        let output = aligner.next().await.unwrap().unwrap();
882        assert_eq!(2, output.num_rows());
883        assert_eq!(
884            10,
885            output
886                .column(1)
887                .as_any()
888                .downcast_ref::<Int64Array>()
889                .unwrap()
890                .value(0)
891        );
892        assert_eq!(
893            20,
894            output
895                .column(1)
896                .as_any()
897                .downcast_ref::<Int64Array>()
898                .unwrap()
899                .value(1)
900        );
901    }
902
903    #[test]
904    fn test_rewrite_rejects_mismatched_output_layout() {
905        let columns = HashMap::from([(
906            "j".to_string(),
907            Json2TargetLayout {
908                extension_metadata: serde_json::to_string(&JsonMetadata::new(
909                    JsonSettings::default(),
910                ))
911                .unwrap(),
912                target_layout: JsonSettings::try_new(vec![], Some(0)).unwrap(),
913            },
914        )]);
915        let result = JsonSchemaAligner::new(
916            stream::empty::<Result<RecordBatch>>(),
917            vec![false],
918            schema([Field::new("j", DataType::Binary, true)
919                .with_extension_type(Json2ExtensionType::default())]),
920            AlignMode::Rewrite { columns },
921        );
922        assert!(
923            result
924                .unwrap_err()
925                .to_string()
926                .contains("does not match output field")
927        );
928    }
929
930    #[tokio::test]
931    async fn test_empty_rewrite_only_fills_missing_roots() {
932        let source = int_array([10, 20]);
933        let input = RecordBatch::try_new(
934            schema([Field::new("a", DataType::Int64, true)]),
935            vec![source.clone()],
936        )
937        .unwrap();
938        let output_schema = schema([
939            Field::new("missing", DataType::Utf8, true),
940            Field::new("a", DataType::Int64, true),
941        ]);
942        let mut aligner = JsonSchemaAligner::new(
943            stream::iter([Ok(input)]),
944            vec![false, true],
945            output_schema.clone(),
946            AlignMode::Rewrite {
947                columns: HashMap::new(),
948            },
949        )
950        .unwrap();
951        let output = aligner.next().await.unwrap().unwrap();
952        assert_eq!(output_schema, output.schema());
953        assert_eq!(2, output.num_rows());
954        assert_eq!(2, output.column(0).null_count());
955        assert!(Arc::ptr_eq(&source, output.column(1)));
956    }
957
958    #[tokio::test]
959    async fn test_empty_rewrite_does_not_align_existing_struct() {
960        let source = Arc::new(StructArray::from(vec![(
961            Arc::new(Field::new("x", DataType::Int64, true)),
962            int_array([1, 2]),
963        )])) as ArrayRef;
964        let input = RecordBatch::try_new(
965            schema([Field::new("j", source.data_type().clone(), true)]),
966            vec![source],
967        )
968        .unwrap();
969        let output_schema = schema([Field::new(
970            "j",
971            DataType::Struct(Fields::from(vec![
972                Field::new("x", DataType::Int64, true),
973                Field::new("y", DataType::Utf8, true),
974            ])),
975            true,
976        )]);
977        let mut aligner = JsonSchemaAligner::new(
978            stream::iter([Ok(input)]),
979            vec![true],
980            output_schema,
981            AlignMode::Rewrite {
982                columns: HashMap::new(),
983            },
984        )
985        .unwrap();
986        assert!(matches!(
987            aligner.next().await.unwrap(),
988            Err(crate::error::Error::NewRecordBatch { .. })
989        ));
990    }
991
992    fn schema(fields: impl IntoIterator<Item = Field>) -> SchemaRef {
993        Arc::new(Schema::new(fields.into_iter().collect::<Vec<_>>()))
994    }
995
996    fn int_array(values: impl IntoIterator<Item = i64>) -> ArrayRef {
997        Arc::new(Int64Array::from_iter_values(values))
998    }
999
1000    fn string_array(values: impl IntoIterator<Item = &'static str>) -> ArrayRef {
1001        Arc::new(StringArray::from_iter_values(values))
1002    }
1003    #[test]
1004    fn test_query_alignment_preserves_source_columns_until_compatibility()
1005    -> std::result::Result<(), Box<dyn std::error::Error>> {
1006        let (_, batch) = test_util::sst_util::new_json2_record_batch()?;
1007        let values = int_array([10, 20]);
1008        let mut fields = vec![Arc::new(Field::new("value", DataType::Int64, true))];
1009        fields.extend(batch.schema_ref().fields().iter().cloned());
1010        let mut columns = vec![values.clone()];
1011        columns.extend_from_slice(batch.columns());
1012        let input = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?;
1013        let target = schema([
1014            Field::new("new_column", DataType::Utf8, true),
1015            Field::new("new_json", DataType::Binary, true)
1016                .with_extension_type(Json2ExtensionType::default()),
1017            Field::new("j", DataType::Binary, true)
1018                .with_extension_type(Json2ExtensionType::default()),
1019            Field::new("value", DataType::Int64, true),
1020        ]);
1021        let aligned = align_query_batch(input, &target)?;
1022        assert_eq!(batch.num_columns() + 1, aligned.num_columns());
1023        assert_eq!("value", aligned.schema_ref().field(0).name());
1024        assert!(Arc::ptr_eq(&values, aligned.column(0)));
1025        assert_eq!("j", aligned.schema_ref().field(1).name());
1026        assert_eq!(&DataType::Binary, aligned.column(1).data_type());
1027        let repeated = align_query_batch(aligned.clone(), &target)?;
1028        assert!(Arc::ptr_eq(aligned.column(1), repeated.column(1)));
1029        Ok(())
1030    }
1031}
mito2/sst/parquet/json_align.rs

mito2/sst/parquet/
json_align.rs
