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

arrow_integration_test/
lib.rs
