arrow_integration_test/
lib.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Partial support for the [Apache Arrow JSON test data format](https://github.com/apache/arrow/blob/master/docs/source/format/Integration.rst#json-test-data-format)
19//!
20//! These utilities define structs that read the integration JSON format for integration testing purposes.
21//!
22//! This is not a canonical format, but provides a human-readable way of verifying language implementations
23//!
24//! <div class="warning">
25//!
26//! This crate is **only intended for integration testing the
27//! [Arrow project](https://github.com/apache/arrow-rs)**. It is not [intended for usage outside of
28//! this context](https://github.com/apache/arrow-rs/issues/8684#issuecomment-3433193158).
29//!
30//! </div>
31
32#![doc(
33    html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg",
34    html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg"
35)]
36#![cfg_attr(docsrs, feature(doc_cfg))]
37#![warn(missing_docs)]
38use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, ScalarBuffer};
39use hex::decode;
40use num_bigint::BigInt;
41use num_traits::Signed;
42use serde::{Deserialize, Serialize};
43use serde_json::{Map as SJMap, Value};
44use std::collections::HashMap;
45use std::sync::Arc;
46
47use arrow::array::*;
48use arrow::buffer::{Buffer, MutableBuffer};
49use arrow::datatypes::*;
50use arrow::error::{ArrowError, Result};
51use arrow::util::bit_util;
52
53mod datatype;
54mod field;
55mod schema;
56
57pub use datatype::*;
58pub use field::*;
59pub use schema::*;
60
61/// A struct that represents an Arrow file with a schema and record batches
62///
63/// See <https://github.com/apache/arrow/blob/master/docs/source/format/Integration.rst#json-test-data-format>
64#[derive(Deserialize, Serialize, Debug)]
65pub struct ArrowJson {
66    /// The Arrow schema for JSON file
67    pub schema: ArrowJsonSchema,
68    /// The `RecordBatch`es in the JSON file
69    pub batches: Vec<ArrowJsonBatch>,
70    /// The dictionaries in the JSON file
71    #[serde(skip_serializing_if = "Option::is_none")]
72    pub dictionaries: Option<Vec<ArrowJsonDictionaryBatch>>,
73}
74
75/// A struct that partially reads the Arrow JSON schema.
76///
77/// Fields are left as JSON `Value` as they vary by `DataType`
78#[derive(Deserialize, Serialize, Debug)]
79pub struct ArrowJsonSchema {
80    /// An array of JSON fields
81    pub fields: Vec<ArrowJsonField>,
82    /// An array of metadata key-value pairs
83    #[serde(skip_serializing_if = "Option::is_none")]
84    pub metadata: Option<Vec<HashMap<String, String>>>,
85}
86
87/// Fields are left as JSON `Value` as they vary by `DataType`
88#[derive(Deserialize, Serialize, Debug)]
89pub struct ArrowJsonField {
90    /// The name of the field
91    pub name: String,
92    /// The data type of the field,
93    /// can be any valid JSON value
94    #[serde(rename = "type")]
95    pub field_type: Value,
96    /// Whether the field is nullable
97    pub nullable: bool,
98    /// The children fields
99    pub children: Vec<ArrowJsonField>,
100    /// The dictionary for the field
101    #[serde(skip_serializing_if = "Option::is_none")]
102    pub dictionary: Option<ArrowJsonFieldDictionary>,
103    /// The metadata for the field, if any
104    #[serde(skip_serializing_if = "Option::is_none")]
105    pub metadata: Option<Value>,
106}
107
108impl From<&FieldRef> for ArrowJsonField {
109    fn from(value: &FieldRef) -> Self {
110        Self::from(value.as_ref())
111    }
112}
113
114impl From<&Field> for ArrowJsonField {
115    fn from(field: &Field) -> Self {
116        let metadata_value = match field.metadata().is_empty() {
117            false => {
118                let mut array = Vec::new();
119                for (k, v) in field.metadata() {
120                    let mut kv_map = SJMap::new();
121                    kv_map.insert(k.clone(), Value::String(v.clone()));
122                    array.push(Value::Object(kv_map));
123                }
124                if !array.is_empty() {
125                    Some(Value::Array(array))
126                } else {
127                    None
128                }
129            }
130            _ => None,
131        };
132
133        Self {
134            name: field.name().clone(),
135            field_type: data_type_to_json(field.data_type()),
136            nullable: field.is_nullable(),
137            children: vec![],
138            dictionary: None, // TODO: not enough info
139            metadata: metadata_value,
140        }
141    }
142}
143
144/// Represents a dictionary-encoded field in the Arrow JSON format
145#[derive(Deserialize, Serialize, Debug)]
146pub struct ArrowJsonFieldDictionary {
147    /// A unique identifier for the dictionary
148    pub id: i64,
149    /// The type of the dictionary index
150    #[serde(rename = "indexType")]
151    pub index_type: DictionaryIndexType,
152    /// Whether the dictionary is ordered
153    #[serde(rename = "isOrdered")]
154    pub is_ordered: bool,
155}
156
157/// Type of an index for a dictionary-encoded field in the Arrow JSON format
158#[derive(Deserialize, Serialize, Debug)]
159pub struct DictionaryIndexType {
160    /// The name of the dictionary index type
161    pub name: String,
162    /// Whether the dictionary index type is signed
163    #[serde(rename = "isSigned")]
164    pub is_signed: bool,
165    /// The bit width of the dictionary index type
166    #[serde(rename = "bitWidth")]
167    pub bit_width: i64,
168}
169
170/// A struct that partially reads the Arrow JSON record batch
171#[derive(Deserialize, Serialize, Debug, Clone)]
172pub struct ArrowJsonBatch {
173    count: usize,
174    /// The columns in the record batch
175    pub columns: Vec<ArrowJsonColumn>,
176}
177
178/// A struct that partially reads the Arrow JSON dictionary batch
179#[derive(Deserialize, Serialize, Debug, Clone)]
180pub struct ArrowJsonDictionaryBatch {
181    /// The unique identifier for the dictionary
182    pub id: i64,
183    /// The data for the dictionary
184    pub data: ArrowJsonBatch,
185}
186
187/// A struct that partially reads the Arrow JSON column/array
188#[derive(Deserialize, Serialize, Clone, Debug)]
189pub struct ArrowJsonColumn {
190    name: String,
191    /// The number of elements in the column
192    pub count: usize,
193    /// The validity bitmap to determine null values
194    #[serde(rename = "VALIDITY")]
195    pub validity: Option<Vec<u8>>,
196    /// The data values in the column
197    #[serde(rename = "DATA")]
198    pub data: Option<Vec<Value>>,
199    /// The offsets for variable-sized data types
200    #[serde(rename = "OFFSET")]
201    pub offset: Option<Vec<Value>>, // leaving as Value as 64-bit offsets are strings
202    /// The type id for union types
203    #[serde(rename = "TYPE_ID")]
204    pub type_id: Option<Vec<i8>>,
205    /// The sizes for ListView/LargeListView types
206    #[serde(rename = "SIZE")]
207    pub size: Option<Vec<Value>>,
208    /// The views for BinaryView/Utf8View types
209    #[serde(rename = "VIEWS")]
210    pub views: Option<Vec<Value>>,
211    /// The variadic data buffers for BinaryView/Utf8View types
212    #[serde(rename = "VARIADIC_DATA_BUFFERS")]
213    pub variadic_data_buffers: Option<Vec<String>>,
214    /// The children columns for nested types
215    pub children: Option<Vec<ArrowJsonColumn>>,
216}
217
218impl ArrowJson {
219    /// Compare the Arrow JSON with a record batch reader
220    pub fn equals_reader(&self, reader: &mut dyn RecordBatchReader) -> Result<bool> {
221        if !self.schema.equals_schema(&reader.schema()) {
222            return Ok(false);
223        }
224
225        for json_batch in self.get_record_batches()? {
226            let batch = reader.next();
227            match batch {
228                Some(Ok(batch)) => {
229                    if json_batch != batch {
230                        println!("json: {json_batch:?}");
231                        println!("batch: {batch:?}");
232                        return Ok(false);
233                    }
234                }
235                Some(Err(e)) => return Err(e),
236                None => return Ok(false),
237            }
238        }
239
240        Ok(true)
241    }
242
243    /// Convert the stored dictionaries to `Vec[RecordBatch]`
244    pub fn get_record_batches(&self) -> Result<Vec<RecordBatch>> {
245        let schema = self.schema.to_arrow_schema()?;
246
247        let mut dictionaries = HashMap::new();
248        self.dictionaries.iter().for_each(|dict_batches| {
249            dict_batches.iter().for_each(|d| {
250                dictionaries.insert(d.id, d.clone());
251            });
252        });
253
254        let batches: Result<Vec<_>> = self
255            .batches
256            .iter()
257            .map(|col| record_batch_from_json(&schema, col.clone(), Some(&dictionaries)))
258            .collect();
259
260        batches
261    }
262}
263
264impl ArrowJsonSchema {
265    /// Compare the Arrow JSON schema with the Arrow `Schema`
266    fn equals_schema(&self, schema: &Schema) -> bool {
267        let field_len = self.fields.len();
268        if field_len != schema.fields().len() {
269            return false;
270        }
271        for i in 0..field_len {
272            let json_field = &self.fields[i];
273            let field = schema.field(i);
274            if !json_field.equals_field(field) {
275                return false;
276            }
277        }
278        true
279    }
280
281    fn to_arrow_schema(&self) -> Result<Schema> {
282        let arrow_fields: Result<Vec<_>> = self
283            .fields
284            .iter()
285            .map(|field| field.to_arrow_field())
286            .collect();
287
288        if let Some(metadatas) = &self.metadata {
289            let mut metadata: HashMap<String, String> = HashMap::new();
290
291            metadatas.iter().for_each(|pair| {
292                let key = pair.get("key").unwrap();
293                let value = pair.get("value").unwrap();
294                metadata.insert(key.clone(), value.clone());
295            });
296
297            Ok(Schema::new_with_metadata(arrow_fields?, metadata))
298        } else {
299            Ok(Schema::new(arrow_fields?))
300        }
301    }
302}
303
304impl ArrowJsonField {
305    /// Compare the Arrow JSON field with the Arrow `Field`
306    fn equals_field(&self, field: &Field) -> bool {
307        // convert to a field
308        match self.to_arrow_field() {
309            Ok(self_field) => {
310                assert_eq!(&self_field, field, "Arrow fields not the same");
311                true
312            }
313            Err(e) => {
314                eprintln!("Encountered error while converting JSON field to Arrow field: {e:?}");
315                false
316            }
317        }
318    }
319
320    /// Convert to an Arrow Field
321    /// TODO: convert to use an Into
322    fn to_arrow_field(&self) -> Result<Field> {
323        // a bit regressive, but we have to convert the field to JSON in order to convert it
324        let field =
325            serde_json::to_value(self).map_err(|error| ArrowError::JsonError(error.to_string()))?;
326        field_from_json(&field)
327    }
328}
329
330/// Generates a [`RecordBatch`] from an Arrow JSON batch, given a schema
331pub fn record_batch_from_json(
332    schema: &Schema,
333    json_batch: ArrowJsonBatch,
334    json_dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
335) -> Result<RecordBatch> {
336    let mut columns = vec![];
337
338    for (field, json_col) in schema.fields().iter().zip(json_batch.columns) {
339        let col = array_from_json(field, json_col, json_dictionaries)?;
340        columns.push(col);
341    }
342
343    RecordBatch::try_new(Arc::new(schema.clone()), columns)
344}
345
346/// Construct an Arrow array from a partially typed JSON column
347pub fn array_from_json(
348    field: &Field,
349    json_col: ArrowJsonColumn,
350    dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
351) -> Result<ArrayRef> {
352    match field.data_type() {
353        DataType::Null => Ok(Arc::new(NullArray::new(json_col.count))),
354        DataType::Boolean => {
355            let mut b = BooleanBuilder::with_capacity(json_col.count);
356            for (is_valid, value) in json_col
357                .validity
358                .as_ref()
359                .unwrap()
360                .iter()
361                .zip(json_col.data.unwrap())
362            {
363                match is_valid {
364                    1 => b.append_value(value.as_bool().unwrap()),
365                    _ => b.append_null(),
366                }
367            }
368            Ok(Arc::new(b.finish()))
369        }
370        DataType::Int8 => {
371            let mut b = Int8Builder::with_capacity(json_col.count);
372            for (is_valid, value) in json_col
373                .validity
374                .as_ref()
375                .unwrap()
376                .iter()
377                .zip(json_col.data.unwrap())
378            {
379                match is_valid {
380                    1 => b.append_value(value.as_i64().ok_or_else(|| {
381                        ArrowError::JsonError(format!("Unable to get {value:?} as int64"))
382                    })? as i8),
383                    _ => b.append_null(),
384                }
385            }
386            Ok(Arc::new(b.finish()))
387        }
388        DataType::Int16 => {
389            let mut b = Int16Builder::with_capacity(json_col.count);
390            for (is_valid, value) in json_col
391                .validity
392                .as_ref()
393                .unwrap()
394                .iter()
395                .zip(json_col.data.unwrap())
396            {
397                match is_valid {
398                    1 => b.append_value(value.as_i64().unwrap() as i16),
399                    _ => b.append_null(),
400                }
401            }
402            Ok(Arc::new(b.finish()))
403        }
404        DataType::Int32 | DataType::Date32 | DataType::Time32(_) => {
405            let mut b = Int32Builder::with_capacity(json_col.count);
406            for (is_valid, value) in json_col
407                .validity
408                .as_ref()
409                .unwrap()
410                .iter()
411                .zip(json_col.data.unwrap())
412            {
413                match is_valid {
414                    1 => b.append_value(value.as_i64().unwrap() as i32),
415                    _ => b.append_null(),
416                }
417            }
418            let array = Arc::new(b.finish()) as ArrayRef;
419            arrow::compute::cast(&array, field.data_type())
420        }
421        DataType::Interval(IntervalUnit::YearMonth) => {
422            let mut b = IntervalYearMonthBuilder::with_capacity(json_col.count);
423            for (is_valid, value) in json_col
424                .validity
425                .as_ref()
426                .unwrap()
427                .iter()
428                .zip(json_col.data.unwrap())
429            {
430                match is_valid {
431                    1 => b.append_value(value.as_i64().unwrap() as i32),
432                    _ => b.append_null(),
433                }
434            }
435            Ok(Arc::new(b.finish()))
436        }
437        DataType::Int64
438        | DataType::Date64
439        | DataType::Time64(_)
440        | DataType::Timestamp(_, _)
441        | DataType::Duration(_) => {
442            let mut b = Int64Builder::with_capacity(json_col.count);
443            for (is_valid, value) in json_col
444                .validity
445                .as_ref()
446                .unwrap()
447                .iter()
448                .zip(json_col.data.unwrap())
449            {
450                match is_valid {
451                    1 => b.append_value(match value {
452                        Value::Number(n) => n.as_i64().unwrap(),
453                        Value::String(s) => s.parse().expect("Unable to parse string as i64"),
454                        _ => panic!("Unable to parse {value:?} as number"),
455                    }),
456                    _ => b.append_null(),
457                }
458            }
459            let array = Arc::new(b.finish()) as ArrayRef;
460            arrow::compute::cast(&array, field.data_type())
461        }
462        DataType::Interval(IntervalUnit::DayTime) => {
463            let mut b = IntervalDayTimeBuilder::with_capacity(json_col.count);
464            for (is_valid, value) in json_col
465                .validity
466                .as_ref()
467                .unwrap()
468                .iter()
469                .zip(json_col.data.unwrap())
470            {
471                match is_valid {
472                    1 => b.append_value(match value {
473                        Value::Object(ref map)
474                            if map.contains_key("days") && map.contains_key("milliseconds") =>
475                        {
476                            match field.data_type() {
477                                DataType::Interval(IntervalUnit::DayTime) => {
478                                    let days = map.get("days").unwrap();
479                                    let milliseconds = map.get("milliseconds").unwrap();
480
481                                    match (days, milliseconds) {
482                                        (Value::Number(d), Value::Number(m)) => {
483                                            let days = d.as_i64().unwrap() as _;
484                                            let millis = m.as_i64().unwrap() as _;
485                                            IntervalDayTime::new(days, millis)
486                                        }
487                                        _ => {
488                                            panic!("Unable to parse {value:?} as interval daytime")
489                                        }
490                                    }
491                                }
492                                _ => panic!("Unable to parse {value:?} as interval daytime"),
493                            }
494                        }
495                        _ => panic!("Unable to parse {value:?} as number"),
496                    }),
497                    _ => b.append_null(),
498                }
499            }
500            Ok(Arc::new(b.finish()))
501        }
502        DataType::UInt8 => {
503            let mut b = UInt8Builder::with_capacity(json_col.count);
504            for (is_valid, value) in json_col
505                .validity
506                .as_ref()
507                .unwrap()
508                .iter()
509                .zip(json_col.data.unwrap())
510            {
511                match is_valid {
512                    1 => b.append_value(value.as_u64().unwrap() as u8),
513                    _ => b.append_null(),
514                }
515            }
516            Ok(Arc::new(b.finish()))
517        }
518        DataType::UInt16 => {
519            let mut b = UInt16Builder::with_capacity(json_col.count);
520            for (is_valid, value) in json_col
521                .validity
522                .as_ref()
523                .unwrap()
524                .iter()
525                .zip(json_col.data.unwrap())
526            {
527                match is_valid {
528                    1 => b.append_value(value.as_u64().unwrap() as u16),
529                    _ => b.append_null(),
530                }
531            }
532            Ok(Arc::new(b.finish()))
533        }
534        DataType::UInt32 => {
535            let mut b = UInt32Builder::with_capacity(json_col.count);
536            for (is_valid, value) in json_col
537                .validity
538                .as_ref()
539                .unwrap()
540                .iter()
541                .zip(json_col.data.unwrap())
542            {
543                match is_valid {
544                    1 => b.append_value(value.as_u64().unwrap() as u32),
545                    _ => b.append_null(),
546                }
547            }
548            Ok(Arc::new(b.finish()))
549        }
550        DataType::UInt64 => {
551            let mut b = UInt64Builder::with_capacity(json_col.count);
552            for (is_valid, value) in json_col
553                .validity
554                .as_ref()
555                .unwrap()
556                .iter()
557                .zip(json_col.data.unwrap())
558            {
559                match is_valid {
560                    1 => {
561                        if value.is_string() {
562                            b.append_value(
563                                value
564                                    .as_str()
565                                    .unwrap()
566                                    .parse()
567                                    .expect("Unable to parse string as u64"),
568                            )
569                        } else if value.is_number() {
570                            b.append_value(value.as_u64().expect("Unable to read number as u64"))
571                        } else {
572                            panic!("Unable to parse value {value:?} as u64")
573                        }
574                    }
575                    _ => b.append_null(),
576                }
577            }
578            Ok(Arc::new(b.finish()))
579        }
580        DataType::Interval(IntervalUnit::MonthDayNano) => {
581            let mut b = IntervalMonthDayNanoBuilder::with_capacity(json_col.count);
582            for (is_valid, value) in json_col
583                .validity
584                .as_ref()
585                .unwrap()
586                .iter()
587                .zip(json_col.data.unwrap())
588            {
589                match is_valid {
590                    1 => b.append_value(match value {
591                        Value::Object(v) => {
592                            let months = v.get("months").unwrap();
593                            let days = v.get("days").unwrap();
594                            let nanoseconds = v.get("nanoseconds").unwrap();
595                            match (months, days, nanoseconds) {
596                                (
597                                    Value::Number(months),
598                                    Value::Number(days),
599                                    Value::Number(nanoseconds),
600                                ) => {
601                                    let months = months.as_i64().unwrap() as i32;
602                                    let days = days.as_i64().unwrap() as i32;
603                                    let nanoseconds = nanoseconds.as_i64().unwrap();
604                                    IntervalMonthDayNano::new(months, days, nanoseconds)
605                                }
606                                (_, _, _) => {
607                                    panic!("Unable to parse {v:?} as MonthDayNano")
608                                }
609                            }
610                        }
611                        _ => panic!("Unable to parse {value:?} as MonthDayNano"),
612                    }),
613                    _ => b.append_null(),
614                }
615            }
616            Ok(Arc::new(b.finish()))
617        }
618        DataType::Float32 => {
619            let mut b = Float32Builder::with_capacity(json_col.count);
620            for (is_valid, value) in json_col
621                .validity
622                .as_ref()
623                .unwrap()
624                .iter()
625                .zip(json_col.data.unwrap())
626            {
627                match is_valid {
628                    1 => b.append_value(value.as_f64().unwrap() as f32),
629                    _ => b.append_null(),
630                }
631            }
632            Ok(Arc::new(b.finish()))
633        }
634        DataType::Float64 => {
635            let mut b = Float64Builder::with_capacity(json_col.count);
636            for (is_valid, value) in json_col
637                .validity
638                .as_ref()
639                .unwrap()
640                .iter()
641                .zip(json_col.data.unwrap())
642            {
643                match is_valid {
644                    1 => b.append_value(value.as_f64().unwrap()),
645                    _ => b.append_null(),
646                }
647            }
648            Ok(Arc::new(b.finish()))
649        }
650        DataType::Binary => {
651            let mut b = BinaryBuilder::with_capacity(json_col.count, 1024);
652            for (is_valid, value) in json_col
653                .validity
654                .as_ref()
655                .unwrap()
656                .iter()
657                .zip(json_col.data.unwrap())
658            {
659                match is_valid {
660                    1 => {
661                        let v = decode(value.as_str().unwrap()).unwrap();
662                        b.append_value(&v)
663                    }
664                    _ => b.append_null(),
665                }
666            }
667            Ok(Arc::new(b.finish()))
668        }
669        DataType::LargeBinary => {
670            let mut b = LargeBinaryBuilder::with_capacity(json_col.count, 1024);
671            for (is_valid, value) in json_col
672                .validity
673                .as_ref()
674                .unwrap()
675                .iter()
676                .zip(json_col.data.unwrap())
677            {
678                match is_valid {
679                    1 => {
680                        let v = decode(value.as_str().unwrap()).unwrap();
681                        b.append_value(&v)
682                    }
683                    _ => b.append_null(),
684                }
685            }
686            Ok(Arc::new(b.finish()))
687        }
688        DataType::Utf8 => {
689            let mut b = StringBuilder::with_capacity(json_col.count, 1024);
690            for (is_valid, value) in json_col
691                .validity
692                .as_ref()
693                .unwrap()
694                .iter()
695                .zip(json_col.data.unwrap())
696            {
697                match is_valid {
698                    1 => b.append_value(value.as_str().unwrap()),
699                    _ => b.append_null(),
700                }
701            }
702            Ok(Arc::new(b.finish()))
703        }
704        DataType::LargeUtf8 => {
705            let mut b = LargeStringBuilder::with_capacity(json_col.count, 1024);
706            for (is_valid, value) in json_col
707                .validity
708                .as_ref()
709                .unwrap()
710                .iter()
711                .zip(json_col.data.unwrap())
712            {
713                match is_valid {
714                    1 => b.append_value(value.as_str().unwrap()),
715                    _ => b.append_null(),
716                }
717            }
718            Ok(Arc::new(b.finish()))
719        }
720        DataType::FixedSizeBinary(len) => {
721            let mut b = FixedSizeBinaryBuilder::with_capacity(json_col.count, *len);
722            for (is_valid, value) in json_col
723                .validity
724                .as_ref()
725                .unwrap()
726                .iter()
727                .zip(json_col.data.unwrap())
728            {
729                match is_valid {
730                    1 => {
731                        let v = hex::decode(value.as_str().unwrap()).unwrap();
732                        b.append_value(&v)?
733                    }
734                    _ => b.append_null(),
735                }
736            }
737            Ok(Arc::new(b.finish()))
738        }
739        DataType::List(child_field) => {
740            let null_buf = create_null_buf(&json_col);
741            let children = json_col.children.clone().unwrap();
742            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
743            let offsets: Vec<i32> = json_col
744                .offset
745                .unwrap()
746                .iter()
747                .map(|v| v.as_i64().unwrap() as i32)
748                .collect();
749            let list_data = ArrayData::builder(field.data_type().clone())
750                .len(json_col.count)
751                .offset(0)
752                .add_buffer(Buffer::from(offsets.to_byte_slice()))
753                .add_child_data(child_array.into_data())
754                .null_bit_buffer(Some(null_buf))
755                .build()
756                .unwrap();
757            Ok(Arc::new(ListArray::from(list_data)))
758        }
759        DataType::LargeList(child_field) => {
760            let null_buf = create_null_buf(&json_col);
761            let children = json_col.children.clone().unwrap();
762            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
763            let offsets: Vec<i64> = json_col
764                .offset
765                .unwrap()
766                .iter()
767                .map(|v| match v {
768                    Value::Number(n) => n.as_i64().unwrap(),
769                    Value::String(s) => s.parse::<i64>().unwrap(),
770                    _ => panic!("64-bit offset must be either string or number"),
771                })
772                .collect();
773            let list_data = ArrayData::builder(field.data_type().clone())
774                .len(json_col.count)
775                .offset(0)
776                .add_buffer(Buffer::from(offsets.to_byte_slice()))
777                .add_child_data(child_array.into_data())
778                .null_bit_buffer(Some(null_buf))
779                .build()
780                .unwrap();
781            Ok(Arc::new(LargeListArray::from(list_data)))
782        }
783        DataType::ListView(child_field) => {
784            let null_buf = create_null_buf(&json_col);
785            let children = json_col.children.clone().unwrap();
786            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
787            let offsets: Vec<i32> = json_col
788                .offset
789                .unwrap()
790                .iter()
791                .map(|v| v.as_i64().unwrap() as i32)
792                .collect();
793            let sizes: Vec<i32> = json_col
794                .size
795                .unwrap()
796                .iter()
797                .map(|v| v.as_i64().unwrap() as i32)
798                .collect();
799            let list_data = ArrayData::builder(field.data_type().clone())
800                .len(json_col.count)
801                .add_buffer(Buffer::from(offsets.to_byte_slice()))
802                .add_buffer(Buffer::from(sizes.to_byte_slice()))
803                .add_child_data(child_array.into_data())
804                .null_bit_buffer(Some(null_buf))
805                .build()
806                .unwrap();
807            Ok(Arc::new(ListViewArray::from(list_data)))
808        }
809        DataType::LargeListView(child_field) => {
810            let null_buf = create_null_buf(&json_col);
811            let children = json_col.children.clone().unwrap();
812            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
813            let offsets: Vec<i64> = json_col
814                .offset
815                .unwrap()
816                .iter()
817                .map(|v| match v {
818                    Value::Number(n) => n.as_i64().unwrap(),
819                    Value::String(s) => s.parse::<i64>().unwrap(),
820                    _ => panic!("64-bit offset must be either string or number"),
821                })
822                .collect();
823            let sizes: Vec<i64> = json_col
824                .size
825                .unwrap()
826                .iter()
827                .map(|v| match v {
828                    Value::Number(n) => n.as_i64().unwrap(),
829                    Value::String(s) => s.parse::<i64>().unwrap(),
830                    _ => panic!("64-bit size must be either string or number"),
831                })
832                .collect();
833            let list_data = ArrayData::builder(field.data_type().clone())
834                .len(json_col.count)
835                .add_buffer(Buffer::from(offsets.to_byte_slice()))
836                .add_buffer(Buffer::from(sizes.to_byte_slice()))
837                .add_child_data(child_array.into_data())
838                .null_bit_buffer(Some(null_buf))
839                .build()
840                .unwrap();
841            Ok(Arc::new(LargeListViewArray::from(list_data)))
842        }
843        DataType::FixedSizeList(child_field, _) => {
844            let children = json_col.children.clone().unwrap();
845            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
846            let null_buf = create_null_buf(&json_col);
847            let list_data = ArrayData::builder(field.data_type().clone())
848                .len(json_col.count)
849                .add_child_data(child_array.into_data())
850                .null_bit_buffer(Some(null_buf))
851                .build()
852                .unwrap();
853            Ok(Arc::new(FixedSizeListArray::from(list_data)))
854        }
855        DataType::Struct(fields) => {
856            // construct struct with null data
857            let null_buf = create_null_buf(&json_col);
858            let mut array_data = ArrayData::builder(field.data_type().clone())
859                .len(json_col.count)
860                .null_bit_buffer(Some(null_buf));
861
862            for (field, col) in fields.iter().zip(json_col.children.unwrap()) {
863                let array = array_from_json(field, col, dictionaries)?;
864                array_data = array_data.add_child_data(array.into_data());
865            }
866
867            let array = StructArray::from(array_data.build().unwrap());
868            Ok(Arc::new(array))
869        }
870        DataType::Dictionary(key_type, value_type) => {
871            #[expect(deprecated)]
872            let dict_id = field.dict_id().ok_or_else(|| {
873                ArrowError::JsonError(format!("Unable to find dict_id for field {field}"))
874            })?;
875            // find dictionary
876            let dictionary = dictionaries
877                .ok_or_else(|| {
878                    ArrowError::JsonError(format!(
879                        "Unable to find any dictionaries for field {field}"
880                    ))
881                })?
882                .get(&dict_id);
883            match dictionary {
884                Some(dictionary) => dictionary_array_from_json(
885                    field,
886                    json_col,
887                    key_type,
888                    value_type,
889                    dictionary,
890                    dictionaries,
891                ),
892                None => Err(ArrowError::JsonError(format!(
893                    "Unable to find dictionary for field {field}"
894                ))),
895            }
896        }
897        DataType::Decimal32(precision, scale) => {
898            let mut b = Decimal32Builder::with_capacity(json_col.count);
899            for (is_valid, value) in json_col
900                .validity
901                .as_ref()
902                .unwrap()
903                .iter()
904                .zip(json_col.data.unwrap())
905            {
906                match is_valid {
907                    1 => b.append_value(value.as_str().unwrap().parse::<i32>().unwrap()),
908                    _ => b.append_null(),
909                }
910            }
911            Ok(Arc::new(
912                b.finish().with_precision_and_scale(*precision, *scale)?,
913            ))
914        }
915        DataType::Decimal64(precision, scale) => {
916            let mut b = Decimal64Builder::with_capacity(json_col.count);
917            for (is_valid, value) in json_col
918                .validity
919                .as_ref()
920                .unwrap()
921                .iter()
922                .zip(json_col.data.unwrap())
923            {
924                match is_valid {
925                    1 => b.append_value(value.as_str().unwrap().parse::<i64>().unwrap()),
926                    _ => b.append_null(),
927                }
928            }
929            Ok(Arc::new(
930                b.finish().with_precision_and_scale(*precision, *scale)?,
931            ))
932        }
933        DataType::Decimal128(precision, scale) => {
934            let mut b = Decimal128Builder::with_capacity(json_col.count);
935            for (is_valid, value) in json_col
936                .validity
937                .as_ref()
938                .unwrap()
939                .iter()
940                .zip(json_col.data.unwrap())
941            {
942                match is_valid {
943                    1 => b.append_value(value.as_str().unwrap().parse::<i128>().unwrap()),
944                    _ => b.append_null(),
945                }
946            }
947            Ok(Arc::new(
948                b.finish().with_precision_and_scale(*precision, *scale)?,
949            ))
950        }
951        DataType::Decimal256(precision, scale) => {
952            let mut b = Decimal256Builder::with_capacity(json_col.count);
953            for (is_valid, value) in json_col
954                .validity
955                .as_ref()
956                .unwrap()
957                .iter()
958                .zip(json_col.data.unwrap())
959            {
960                match is_valid {
961                    1 => {
962                        let str = value.as_str().unwrap();
963                        let integer = BigInt::parse_bytes(str.as_bytes(), 10).unwrap();
964                        let integer_bytes = integer.to_signed_bytes_le();
965                        let mut bytes = if integer.is_positive() {
966                            [0_u8; 32]
967                        } else {
968                            [255_u8; 32]
969                        };
970                        bytes[0..integer_bytes.len()].copy_from_slice(integer_bytes.as_slice());
971                        b.append_value(i256::from_le_bytes(bytes));
972                    }
973                    _ => b.append_null(),
974                }
975            }
976            Ok(Arc::new(
977                b.finish().with_precision_and_scale(*precision, *scale)?,
978            ))
979        }
980        DataType::Map(child_field, _) => {
981            let null_buf = create_null_buf(&json_col);
982            let children = json_col.children.clone().unwrap();
983            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
984            let offsets: Vec<i32> = json_col
985                .offset
986                .unwrap()
987                .iter()
988                .map(|v| v.as_i64().unwrap() as i32)
989                .collect();
990            let array_data = ArrayData::builder(field.data_type().clone())
991                .len(json_col.count)
992                .add_buffer(Buffer::from(offsets.to_byte_slice()))
993                .add_child_data(child_array.into_data())
994                .null_bit_buffer(Some(null_buf))
995                .build()
996                .unwrap();
997
998            let array = MapArray::from(array_data);
999            Ok(Arc::new(array))
1000        }
1001        DataType::Union(fields, _) => {
1002            let Some(type_ids) = json_col.type_id else {
1003                return Err(ArrowError::JsonError(
1004                    "Cannot find expected type_id in json column".to_string(),
1005                ));
1006            };
1007
1008            let offset: Option<ScalarBuffer<i32>> = json_col
1009                .offset
1010                .map(|offsets| offsets.iter().map(|v| v.as_i64().unwrap() as i32).collect());
1011
1012            let mut children = Vec::with_capacity(fields.len());
1013            for ((_, field), col) in fields.iter().zip(json_col.children.unwrap()) {
1014                let array = array_from_json(field, col, dictionaries)?;
1015                children.push(array);
1016            }
1017
1018            let array =
1019                UnionArray::try_new(fields.clone(), type_ids.into(), offset, children).unwrap();
1020            Ok(Arc::new(array))
1021        }
1022        DataType::Utf8View => {
1023            let views = json_col.views.ok_or_else(|| {
1024                ArrowError::JsonError("Utf8View requires VIEWS field".to_string())
1025            })?;
1026            let variadic_buffers = json_col.variadic_data_buffers.unwrap_or_default();
1027            let validity = json_col.validity.as_ref();
1028
1029            let mut builder = StringViewBuilder::new();
1030            for (i, view) in views.iter().enumerate() {
1031                let is_valid = validity.map_or(1, |v| v[i]);
1032                if is_valid == 0 {
1033                    builder.append_null();
1034                } else {
1035                    let view_obj = view.as_object().unwrap();
1036                    let size = view_obj["SIZE"].as_u64().unwrap() as usize;
1037                    // Check for INLINED key presence - inlined if SIZE <= 12
1038                    if let Some(inlined) = view_obj.get("INLINED") {
1039                        builder.append_value(inlined.as_str().unwrap());
1040                    } else {
1041                        // Reference to variadic buffer
1042                        let buffer_index = view_obj["BUFFER_INDEX"].as_u64().unwrap() as usize;
1043                        let offset = view_obj["OFFSET"].as_u64().unwrap() as usize;
1044                        let buffer_data = hex::decode(&variadic_buffers[buffer_index]).unwrap();
1045                        let s = std::str::from_utf8(&buffer_data[offset..offset + size]).unwrap();
1046                        builder.append_value(s);
1047                    }
1048                }
1049            }
1050            Ok(Arc::new(builder.finish()))
1051        }
1052        DataType::BinaryView => {
1053            let views = json_col.views.ok_or_else(|| {
1054                ArrowError::JsonError("BinaryView requires VIEWS field".to_string())
1055            })?;
1056            let variadic_buffers = json_col.variadic_data_buffers.unwrap_or_default();
1057            let validity = json_col.validity.as_ref();
1058
1059            let mut builder = BinaryViewBuilder::new();
1060            for (i, view) in views.iter().enumerate() {
1061                let is_valid = validity.map_or(1, |v| v[i]);
1062                if is_valid == 0 {
1063                    builder.append_null();
1064                } else {
1065                    let view_obj = view.as_object().unwrap();
1066                    let size = view_obj["SIZE"].as_u64().unwrap() as usize;
1067                    // Check for INLINED key presence - inlined if SIZE <= 12
1068                    if let Some(inlined) = view_obj.get("INLINED") {
1069                        let data = hex::decode(inlined.as_str().unwrap()).unwrap();
1070                        builder.append_value(&data);
1071                    } else {
1072                        // Reference to variadic buffer
1073                        let buffer_index = view_obj["BUFFER_INDEX"].as_u64().unwrap() as usize;
1074                        let offset = view_obj["OFFSET"].as_u64().unwrap() as usize;
1075                        let buffer_data = hex::decode(&variadic_buffers[buffer_index]).unwrap();
1076                        builder.append_value(&buffer_data[offset..offset + size]);
1077                    }
1078                }
1079            }
1080            Ok(Arc::new(builder.finish()))
1081        }
1082        DataType::RunEndEncoded(run_ends_field, values_field) => {
1083            let children = json_col.children.clone().unwrap();
1084            if children.len() != 2 {
1085                return Err(ArrowError::JsonError(
1086                    "RunEndEncoded requires exactly 2 children".to_string(),
1087                ));
1088            }
1089            let run_ends_array =
1090                array_from_json(run_ends_field, children[0].clone(), dictionaries)?;
1091            let values_array = array_from_json(values_field, children[1].clone(), dictionaries)?;
1092
1093            let run_array_data = ArrayData::builder(field.data_type().clone())
1094                .len(json_col.count)
1095                .add_child_data(run_ends_array.into_data())
1096                .add_child_data(values_array.into_data())
1097                .build()
1098                .unwrap();
1099
1100            Ok(make_array(run_array_data))
1101        }
1102        t => Err(ArrowError::JsonError(format!(
1103            "data type {t} not supported"
1104        ))),
1105    }
1106}
1107
1108/// Construct a [`DictionaryArray`] from a partially typed JSON column
1109pub fn dictionary_array_from_json(
1110    field: &Field,
1111    json_col: ArrowJsonColumn,
1112    dict_key: &DataType,
1113    dict_value: &DataType,
1114    dictionary: &ArrowJsonDictionaryBatch,
1115    dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
1116) -> Result<ArrayRef> {
1117    match dict_key {
1118        DataType::Int8
1119        | DataType::Int16
1120        | DataType::Int32
1121        | DataType::Int64
1122        | DataType::UInt8
1123        | DataType::UInt16
1124        | DataType::UInt32
1125        | DataType::UInt64 => {
1126            let null_buf = create_null_buf(&json_col);
1127
1128            // build the key data into a buffer, then construct values separately
1129            #[expect(deprecated)]
1130            let key_field = Field::new_dict(
1131                "key",
1132                dict_key.clone(),
1133                field.is_nullable(),
1134                #[expect(deprecated)]
1135                field
1136                    .dict_id()
1137                    .expect("Dictionary fields must have a dict_id value"),
1138                field
1139                    .dict_is_ordered()
1140                    .expect("Dictionary fields must have a dict_is_ordered value"),
1141            );
1142            let keys = array_from_json(&key_field, json_col, None)?;
1143            // note: not enough info on nullability of dictionary
1144            let value_field = Field::new("value", dict_value.clone(), true);
1145            let values = array_from_json(
1146                &value_field,
1147                dictionary.data.columns[0].clone(),
1148                dictionaries,
1149            )?;
1150
1151            // convert key and value to dictionary data
1152            let dict_data = ArrayData::builder(field.data_type().clone())
1153                .len(keys.len())
1154                .add_buffer(keys.to_data().buffers()[0].clone())
1155                .null_bit_buffer(Some(null_buf))
1156                .add_child_data(values.into_data())
1157                .build()
1158                .unwrap();
1159
1160            let array = match dict_key {
1161                DataType::Int8 => Arc::new(Int8DictionaryArray::from(dict_data)) as ArrayRef,
1162                DataType::Int16 => Arc::new(Int16DictionaryArray::from(dict_data)),
1163                DataType::Int32 => Arc::new(Int32DictionaryArray::from(dict_data)),
1164                DataType::Int64 => Arc::new(Int64DictionaryArray::from(dict_data)),
1165                DataType::UInt8 => Arc::new(UInt8DictionaryArray::from(dict_data)),
1166                DataType::UInt16 => Arc::new(UInt16DictionaryArray::from(dict_data)),
1167                DataType::UInt32 => Arc::new(UInt32DictionaryArray::from(dict_data)),
1168                DataType::UInt64 => Arc::new(UInt64DictionaryArray::from(dict_data)),
1169                _ => unreachable!(),
1170            };
1171            Ok(array)
1172        }
1173        _ => Err(ArrowError::JsonError(format!(
1174            "Dictionary key type {dict_key:?} not supported"
1175        ))),
1176    }
1177}
1178
1179/// A helper to create a null buffer from a `Vec<bool>`
1180fn create_null_buf(json_col: &ArrowJsonColumn) -> Buffer {
1181    let num_bytes = bit_util::ceil(json_col.count, 8);
1182    let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, false);
1183    json_col
1184        .validity
1185        .clone()
1186        .unwrap()
1187        .iter()
1188        .enumerate()
1189        .for_each(|(i, v)| {
1190            let null_slice = null_buf.as_slice_mut();
1191            if *v != 0 {
1192                bit_util::set_bit(null_slice, i);
1193            }
1194        });
1195    null_buf.into()
1196}
1197
1198impl ArrowJsonBatch {
1199    /// Convert a [`RecordBatch`] to an [`ArrowJsonBatch`]
1200    ///
1201    /// <div class="warning">
1202    ///
1203    /// This function is **deliberately incomplete**! As noted in the crate-level documentation,
1204    /// this crate is only intended for use within the Arrow project itself.
1205    ///
1206    /// Right now, this function only supports `DataType::Int8` columns. Other data types will lead
1207    /// to an empty `ArrowJsonColumn`.
1208    ///
1209    /// </div>
1210    pub fn from_batch(batch: &RecordBatch) -> ArrowJsonBatch {
1211        let mut json_batch = ArrowJsonBatch {
1212            count: batch.num_rows(),
1213            columns: Vec::with_capacity(batch.num_columns()),
1214        };
1215
1216        for (col, field) in batch.columns().iter().zip(batch.schema().fields.iter()) {
1217            let json_col = match field.data_type() {
1218                DataType::Int8 => {
1219                    let col = col.as_any().downcast_ref::<Int8Array>().unwrap();
1220
1221                    let mut validity: Vec<u8> = Vec::with_capacity(col.len());
1222                    let mut data: Vec<Value> = Vec::with_capacity(col.len());
1223
1224                    for i in 0..col.len() {
1225                        if col.is_null(i) {
1226                            validity.push(1);
1227                            data.push(0i8.into());
1228                        } else {
1229                            validity.push(0);
1230                            data.push(col.value(i).into());
1231                        }
1232                    }
1233
1234                    ArrowJsonColumn {
1235                        name: field.name().clone(),
1236                        count: col.len(),
1237                        validity: Some(validity),
1238                        data: Some(data),
1239                        offset: None,
1240                        type_id: None,
1241                        size: None,
1242                        views: None,
1243                        variadic_data_buffers: None,
1244                        children: None,
1245                    }
1246                }
1247                _ => ArrowJsonColumn {
1248                    name: field.name().clone(),
1249                    count: col.len(),
1250                    validity: None,
1251                    data: None,
1252                    offset: None,
1253                    type_id: None,
1254                    size: None,
1255                    views: None,
1256                    variadic_data_buffers: None,
1257                    children: None,
1258                },
1259            };
1260
1261            json_batch.columns.push(json_col);
1262        }
1263
1264        json_batch
1265    }
1266}
1267
1268#[cfg(test)]
1269mod tests {
1270    use super::*;
1271
1272    #[test]
1273    fn test_schema_equality() {
1274        let json = r#"
1275        {
1276            "fields": [
1277                {
1278                    "name": "c1",
1279                    "type": {"name": "int", "isSigned": true, "bitWidth": 32},
1280                    "nullable": true,
1281                    "children": []
1282                },
1283                {
1284                    "name": "c2",
1285                    "type": {"name": "floatingpoint", "precision": "DOUBLE"},
1286                    "nullable": true,
1287                    "children": []
1288                },
1289                {
1290                    "name": "c3",
1291                    "type": {"name": "utf8"},
1292                    "nullable": true,
1293                    "children": []
1294                },
1295                {
1296                    "name": "c4",
1297                    "type": {
1298                        "name": "list"
1299                    },
1300                    "nullable": true,
1301                    "children": [
1302                        {
1303                            "name": "custom_item",
1304                            "type": {
1305                                "name": "int",
1306                                "isSigned": true,
1307                                "bitWidth": 32
1308                            },
1309                            "nullable": false,
1310                            "children": []
1311                        }
1312                    ]
1313                }
1314            ]
1315        }"#;
1316        let json_schema: ArrowJsonSchema = serde_json::from_str(json).unwrap();
1317        let schema = Schema::new(vec![
1318            Field::new("c1", DataType::Int32, true),
1319            Field::new("c2", DataType::Float64, true),
1320            Field::new("c3", DataType::Utf8, true),
1321            Field::new(
1322                "c4",
1323                DataType::List(Arc::new(Field::new("custom_item", DataType::Int32, false))),
1324                true,
1325            ),
1326        ]);
1327        assert!(json_schema.equals_schema(&schema));
1328    }
1329
1330    #[test]
1331    fn test_arrow_data_equality() {
1332        let secs_tz = Some("Europe/Budapest".into());
1333        let millis_tz = Some("America/New_York".into());
1334        let micros_tz = Some("UTC".into());
1335        let nanos_tz = Some("Africa/Johannesburg".into());
1336
1337        let schema = Schema::new(vec![
1338            Field::new("bools-with-metadata-map", DataType::Boolean, true)
1339                .with_metadata([("k", "v")]),
1340            Field::new("bools-with-metadata-vec", DataType::Boolean, true)
1341                .with_metadata([("k2", "v2")]),
1342            Field::new("bools", DataType::Boolean, true),
1343            Field::new("int8s", DataType::Int8, true),
1344            Field::new("int16s", DataType::Int16, true),
1345            Field::new("int32s", DataType::Int32, true),
1346            Field::new("int64s", DataType::Int64, true),
1347            Field::new("uint8s", DataType::UInt8, true),
1348            Field::new("uint16s", DataType::UInt16, true),
1349            Field::new("uint32s", DataType::UInt32, true),
1350            Field::new("uint64s", DataType::UInt64, true),
1351            Field::new("float32s", DataType::Float32, true),
1352            Field::new("float64s", DataType::Float64, true),
1353            Field::new("date_days", DataType::Date32, true),
1354            Field::new("date_millis", DataType::Date64, true),
1355            Field::new("time_secs", DataType::Time32(TimeUnit::Second), true),
1356            Field::new("time_millis", DataType::Time32(TimeUnit::Millisecond), true),
1357            Field::new("time_micros", DataType::Time64(TimeUnit::Microsecond), true),
1358            Field::new("time_nanos", DataType::Time64(TimeUnit::Nanosecond), true),
1359            Field::new("ts_secs", DataType::Timestamp(TimeUnit::Second, None), true),
1360            Field::new(
1361                "ts_millis",
1362                DataType::Timestamp(TimeUnit::Millisecond, None),
1363                true,
1364            ),
1365            Field::new(
1366                "ts_micros",
1367                DataType::Timestamp(TimeUnit::Microsecond, None),
1368                true,
1369            ),
1370            Field::new(
1371                "ts_nanos",
1372                DataType::Timestamp(TimeUnit::Nanosecond, None),
1373                true,
1374            ),
1375            Field::new(
1376                "ts_secs_tz",
1377                DataType::Timestamp(TimeUnit::Second, secs_tz.clone()),
1378                true,
1379            ),
1380            Field::new(
1381                "ts_millis_tz",
1382                DataType::Timestamp(TimeUnit::Millisecond, millis_tz.clone()),
1383                true,
1384            ),
1385            Field::new(
1386                "ts_micros_tz",
1387                DataType::Timestamp(TimeUnit::Microsecond, micros_tz.clone()),
1388                true,
1389            ),
1390            Field::new(
1391                "ts_nanos_tz",
1392                DataType::Timestamp(TimeUnit::Nanosecond, nanos_tz.clone()),
1393                true,
1394            ),
1395            Field::new("utf8s", DataType::Utf8, true),
1396            Field::new(
1397                "lists",
1398                DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
1399                true,
1400            ),
1401            Field::new(
1402                "structs",
1403                DataType::Struct(Fields::from(vec![
1404                    Field::new("int32s", DataType::Int32, true),
1405                    Field::new("utf8s", DataType::Utf8, true),
1406                ])),
1407                true,
1408            ),
1409            Field::new("utf8views", DataType::Utf8View, true),
1410            Field::new("binaryviews", DataType::BinaryView, true),
1411            Field::new(
1412                "listviews",
1413                DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
1414                true,
1415            ),
1416            Field::new(
1417                "largelistviews",
1418                DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
1419                true,
1420            ),
1421            Field::new(
1422                "runendencoded",
1423                DataType::RunEndEncoded(
1424                    Arc::new(Field::new("run_ends", DataType::Int16, false)),
1425                    Arc::new(Field::new("values", DataType::Int32, true)),
1426                ),
1427                true,
1428            ),
1429        ]);
1430
1431        let bools_with_metadata_map = BooleanArray::from(vec![Some(true), None, Some(false)]);
1432        let bools_with_metadata_vec = BooleanArray::from(vec![Some(true), None, Some(false)]);
1433        let bools = BooleanArray::from(vec![Some(true), None, Some(false)]);
1434        let int8s = Int8Array::from(vec![Some(1), None, Some(3)]);
1435        let int16s = Int16Array::from(vec![Some(1), None, Some(3)]);
1436        let int32s = Int32Array::from(vec![Some(1), None, Some(3)]);
1437        let int64s = Int64Array::from(vec![Some(1), None, Some(3)]);
1438        let uint8s = UInt8Array::from(vec![Some(1), None, Some(3)]);
1439        let uint16s = UInt16Array::from(vec![Some(1), None, Some(3)]);
1440        let uint32s = UInt32Array::from(vec![Some(1), None, Some(3)]);
1441        let uint64s = UInt64Array::from(vec![Some(1), None, Some(3)]);
1442        let float32s = Float32Array::from(vec![Some(1.0), None, Some(3.0)]);
1443        let float64s = Float64Array::from(vec![Some(1.0), None, Some(3.0)]);
1444        let date_days = Date32Array::from(vec![Some(1196848), None, None]);
1445        let date_millis = Date64Array::from(vec![
1446            Some(167903550396207),
1447            Some(29923997007884),
1448            Some(30612271819236),
1449        ]);
1450        let time_secs = Time32SecondArray::from(vec![Some(27974), Some(78592), Some(43207)]);
1451        let time_millis =
1452            Time32MillisecondArray::from(vec![Some(6613125), Some(74667230), Some(52260079)]);
1453        let time_micros = Time64MicrosecondArray::from(vec![Some(62522958593), None, None]);
1454        let time_nanos =
1455            Time64NanosecondArray::from(vec![Some(73380123595985), None, Some(16584393546415)]);
1456        let ts_secs = TimestampSecondArray::from(vec![None, Some(193438817552), None]);
1457        let ts_millis =
1458            TimestampMillisecondArray::from(vec![None, Some(38606916383008), Some(58113709376587)]);
1459        let ts_micros = TimestampMicrosecondArray::from(vec![None, None, None]);
1460        let ts_nanos = TimestampNanosecondArray::from(vec![None, None, Some(-6473623571954960143)]);
1461        let ts_secs_tz = TimestampSecondArray::from(vec![None, Some(193438817552), None])
1462            .with_timezone_opt(secs_tz);
1463        let ts_millis_tz =
1464            TimestampMillisecondArray::from(vec![None, Some(38606916383008), Some(58113709376587)])
1465                .with_timezone_opt(millis_tz);
1466        let ts_micros_tz =
1467            TimestampMicrosecondArray::from(vec![None, None, None]).with_timezone_opt(micros_tz);
1468        let ts_nanos_tz =
1469            TimestampNanosecondArray::from(vec![None, None, Some(-6473623571954960143)])
1470                .with_timezone_opt(nanos_tz);
1471        let utf8s = StringArray::from(vec![Some("aa"), None, Some("bbb")]);
1472
1473        let value_data = Int32Array::from(vec![None, Some(2), None, None]);
1474        let value_offsets = Buffer::from_slice_ref([0, 3, 4, 4]);
1475        let list_data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
1476        let list_data = ArrayData::builder(list_data_type)
1477            .len(3)
1478            .add_buffer(value_offsets)
1479            .add_child_data(value_data.into_data())
1480            .null_bit_buffer(Some(Buffer::from([0b00000011])))
1481            .build()
1482            .unwrap();
1483        let lists = ListArray::from(list_data);
1484
1485        let structs_int32s = Int32Array::from(vec![None, Some(-2), None]);
1486        let structs_utf8s = StringArray::from(vec![None, None, Some("aaaaaa")]);
1487        let struct_data_type = DataType::Struct(Fields::from(vec![
1488            Field::new("int32s", DataType::Int32, true),
1489            Field::new("utf8s", DataType::Utf8, true),
1490        ]));
1491        let struct_data = ArrayData::builder(struct_data_type)
1492            .len(3)
1493            .add_child_data(structs_int32s.into_data())
1494            .add_child_data(structs_utf8s.into_data())
1495            .null_bit_buffer(Some(Buffer::from([0b00000011])))
1496            .build()
1497            .unwrap();
1498        let structs = StructArray::from(struct_data);
1499
1500        let utf8views =
1501            StringViewArray::from(vec![Some("hello"), None, Some("this is not inlined")]);
1502        let binaryviews = BinaryViewArray::from_iter(vec![
1503            Some(b"\xf3\x4d".as_slice()),
1504            Some(b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f".as_slice()),
1505            None,
1506        ]);
1507
1508        let listview_value_data = Int32Array::from(vec![Some(1), Some(2), Some(3), None, Some(5)]);
1509        let listview_offsets = Buffer::from_slice_ref([0i32, 2, 2]);
1510        let listview_sizes = Buffer::from_slice_ref([2i32, 0, 3]);
1511        let listview_data_type =
1512            DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true)));
1513        let listview_data = ArrayData::builder(listview_data_type)
1514            .len(3)
1515            .add_buffer(listview_offsets)
1516            .add_buffer(listview_sizes)
1517            .add_child_data(listview_value_data.into_data())
1518            .null_bit_buffer(Some(Buffer::from([0b00000101])))
1519            .build()
1520            .unwrap();
1521        let listviews = ListViewArray::from(listview_data);
1522
1523        let largelistview_value_data = Int32Array::from(vec![Some(10), None, Some(30)]);
1524        let largelistview_offsets = Buffer::from_slice_ref([0i64, 2, 3]);
1525        let largelistview_sizes = Buffer::from_slice_ref([2i64, 1, 0]);
1526        let largelistview_data_type =
1527            DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true)));
1528        let largelistview_data = ArrayData::builder(largelistview_data_type)
1529            .len(3)
1530            .add_buffer(largelistview_offsets)
1531            .add_buffer(largelistview_sizes)
1532            .add_child_data(largelistview_value_data.into_data())
1533            .null_bit_buffer(Some(Buffer::from([0b00000011])))
1534            .build()
1535            .unwrap();
1536        let largelistviews = LargeListViewArray::from(largelistview_data);
1537
1538        let ree_run_ends = Int16Array::from(vec![2, 3]);
1539        let ree_values = Int32Array::from(vec![Some(100), None]);
1540        let ree_data_type = DataType::RunEndEncoded(
1541            Arc::new(Field::new("run_ends", DataType::Int16, false)),
1542            Arc::new(Field::new("values", DataType::Int32, true)),
1543        );
1544        let ree_data = ArrayData::builder(ree_data_type)
1545            .len(3)
1546            .add_child_data(ree_run_ends.into_data())
1547            .add_child_data(ree_values.into_data())
1548            .build()
1549            .unwrap();
1550        let runendencoded = RunArray::<Int16Type>::from(ree_data);
1551
1552        let record_batch = RecordBatch::try_new(
1553            Arc::new(schema.clone()),
1554            vec![
1555                Arc::new(bools_with_metadata_map),
1556                Arc::new(bools_with_metadata_vec),
1557                Arc::new(bools),
1558                Arc::new(int8s),
1559                Arc::new(int16s),
1560                Arc::new(int32s),
1561                Arc::new(int64s),
1562                Arc::new(uint8s),
1563                Arc::new(uint16s),
1564                Arc::new(uint32s),
1565                Arc::new(uint64s),
1566                Arc::new(float32s),
1567                Arc::new(float64s),
1568                Arc::new(date_days),
1569                Arc::new(date_millis),
1570                Arc::new(time_secs),
1571                Arc::new(time_millis),
1572                Arc::new(time_micros),
1573                Arc::new(time_nanos),
1574                Arc::new(ts_secs),
1575                Arc::new(ts_millis),
1576                Arc::new(ts_micros),
1577                Arc::new(ts_nanos),
1578                Arc::new(ts_secs_tz),
1579                Arc::new(ts_millis_tz),
1580                Arc::new(ts_micros_tz),
1581                Arc::new(ts_nanos_tz),
1582                Arc::new(utf8s),
1583                Arc::new(lists),
1584                Arc::new(structs),
1585                Arc::new(utf8views),
1586                Arc::new(binaryviews),
1587                Arc::new(listviews),
1588                Arc::new(largelistviews),
1589                Arc::new(runendencoded),
1590            ],
1591        )
1592        .unwrap();
1593        let json = std::fs::read_to_string("data/integration.json").unwrap();
1594        let arrow_json: ArrowJson = serde_json::from_str(&json).unwrap();
1595        // test schemas
1596        assert!(arrow_json.schema.equals_schema(&schema));
1597        // test record batch
1598        assert_eq!(arrow_json.get_record_batches().unwrap()[0], record_batch);
1599    }
1600}
arrow_integration_test/lib.rs

arrow_integration_test/
lib.rs
