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//! 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#![warn(missing_docs)]
25use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, ScalarBuffer};
26use hex::decode;
27use num::BigInt;
28use num::Signed;
29use serde::{Deserialize, Serialize};
30use serde_json::{Map as SJMap, Value};
31use std::collections::HashMap;
32use std::sync::Arc;
33
34use arrow::array::*;
35use arrow::buffer::{Buffer, MutableBuffer};
36use arrow::datatypes::*;
37use arrow::error::{ArrowError, Result};
38use arrow::util::bit_util;
39
40mod datatype;
41mod field;
42mod schema;
43
44pub use datatype::*;
45pub use field::*;
46pub use schema::*;
47
48/// A struct that represents an Arrow file with a schema and record batches
49///
50/// See <https://github.com/apache/arrow/blob/master/docs/source/format/Integration.rst#json-test-data-format>
51#[derive(Deserialize, Serialize, Debug)]
52pub struct ArrowJson {
53    /// The Arrow schema for JSON file
54    pub schema: ArrowJsonSchema,
55    /// The `RecordBatch`es in the JSON file
56    pub batches: Vec<ArrowJsonBatch>,
57    /// The dictionaries in the JSON file
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub dictionaries: Option<Vec<ArrowJsonDictionaryBatch>>,
60}
61
62/// A struct that partially reads the Arrow JSON schema.
63///
64/// Fields are left as JSON `Value` as they vary by `DataType`
65#[derive(Deserialize, Serialize, Debug)]
66pub struct ArrowJsonSchema {
67    /// An array of JSON fields
68    pub fields: Vec<ArrowJsonField>,
69    /// An array of metadata key-value pairs
70    #[serde(skip_serializing_if = "Option::is_none")]
71    pub metadata: Option<Vec<HashMap<String, String>>>,
72}
73
74/// Fields are left as JSON `Value` as they vary by `DataType`
75#[derive(Deserialize, Serialize, Debug)]
76pub struct ArrowJsonField {
77    /// The name of the field
78    pub name: String,
79    /// The data type of the field,
80    /// can be any valid JSON value
81    #[serde(rename = "type")]
82    pub field_type: Value,
83    /// Whether the field is nullable
84    pub nullable: bool,
85    /// The children fields
86    pub children: Vec<ArrowJsonField>,
87    /// The dictionary for the field
88    #[serde(skip_serializing_if = "Option::is_none")]
89    pub dictionary: Option<ArrowJsonFieldDictionary>,
90    /// The metadata for the field, if any
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub metadata: Option<Value>,
93}
94
95impl From<&FieldRef> for ArrowJsonField {
96    fn from(value: &FieldRef) -> Self {
97        Self::from(value.as_ref())
98    }
99}
100
101impl From<&Field> for ArrowJsonField {
102    fn from(field: &Field) -> Self {
103        let metadata_value = match field.metadata().is_empty() {
104            false => {
105                let mut array = Vec::new();
106                for (k, v) in field.metadata() {
107                    let mut kv_map = SJMap::new();
108                    kv_map.insert(k.clone(), Value::String(v.clone()));
109                    array.push(Value::Object(kv_map));
110                }
111                if !array.is_empty() {
112                    Some(Value::Array(array))
113                } else {
114                    None
115                }
116            }
117            _ => None,
118        };
119
120        Self {
121            name: field.name().to_string(),
122            field_type: data_type_to_json(field.data_type()),
123            nullable: field.is_nullable(),
124            children: vec![],
125            dictionary: None, // TODO: not enough info
126            metadata: metadata_value,
127        }
128    }
129}
130
131/// Represents a dictionary-encoded field in the Arrow JSON format
132#[derive(Deserialize, Serialize, Debug)]
133pub struct ArrowJsonFieldDictionary {
134    /// A unique identifier for the dictionary
135    pub id: i64,
136    /// The type of the dictionary index
137    #[serde(rename = "indexType")]
138    pub index_type: DictionaryIndexType,
139    /// Whether the dictionary is ordered
140    #[serde(rename = "isOrdered")]
141    pub is_ordered: bool,
142}
143
144/// Type of an index for a dictionary-encoded field in the Arrow JSON format
145#[derive(Deserialize, Serialize, Debug)]
146pub struct DictionaryIndexType {
147    /// The name of the dictionary index type
148    pub name: String,
149    /// Whether the dictionary index type is signed
150    #[serde(rename = "isSigned")]
151    pub is_signed: bool,
152    /// The bit width of the dictionary index type
153    #[serde(rename = "bitWidth")]
154    pub bit_width: i64,
155}
156
157/// A struct that partially reads the Arrow JSON record batch
158#[derive(Deserialize, Serialize, Debug, Clone)]
159pub struct ArrowJsonBatch {
160    count: usize,
161    /// The columns in the record batch
162    pub columns: Vec<ArrowJsonColumn>,
163}
164
165/// A struct that partially reads the Arrow JSON dictionary batch
166#[derive(Deserialize, Serialize, Debug, Clone)]
167#[allow(non_snake_case)]
168pub struct ArrowJsonDictionaryBatch {
169    /// The unique identifier for the dictionary
170    pub id: i64,
171    /// The data for the dictionary
172    pub data: ArrowJsonBatch,
173}
174
175/// A struct that partially reads the Arrow JSON column/array
176#[derive(Deserialize, Serialize, Clone, Debug)]
177pub struct ArrowJsonColumn {
178    name: String,
179    /// The number of elements in the column
180    pub count: usize,
181    /// The validity bitmap to determine null values
182    #[serde(rename = "VALIDITY")]
183    pub validity: Option<Vec<u8>>,
184    /// The data values in the column
185    #[serde(rename = "DATA")]
186    pub data: Option<Vec<Value>>,
187    /// The offsets for variable-sized data types
188    #[serde(rename = "OFFSET")]
189    pub offset: Option<Vec<Value>>, // leaving as Value as 64-bit offsets are strings
190    /// The type id for union types
191    #[serde(rename = "TYPE_ID")]
192    pub type_id: Option<Vec<i8>>,
193    /// The children columns for nested types
194    pub children: Option<Vec<ArrowJsonColumn>>,
195}
196
197impl ArrowJson {
198    /// Compare the Arrow JSON with a record batch reader
199    pub fn equals_reader(&self, reader: &mut dyn RecordBatchReader) -> Result<bool> {
200        if !self.schema.equals_schema(&reader.schema()) {
201            return Ok(false);
202        }
203
204        for json_batch in self.get_record_batches()?.into_iter() {
205            let batch = reader.next();
206            match batch {
207                Some(Ok(batch)) => {
208                    if json_batch != batch {
209                        println!("json: {json_batch:?}");
210                        println!("batch: {batch:?}");
211                        return Ok(false);
212                    }
213                }
214                Some(Err(e)) => return Err(e),
215                None => return Ok(false),
216            }
217        }
218
219        Ok(true)
220    }
221
222    /// Convert the stored dictionaries to `Vec[RecordBatch]`
223    pub fn get_record_batches(&self) -> Result<Vec<RecordBatch>> {
224        let schema = self.schema.to_arrow_schema()?;
225
226        let mut dictionaries = HashMap::new();
227        self.dictionaries.iter().for_each(|dict_batches| {
228            dict_batches.iter().for_each(|d| {
229                dictionaries.insert(d.id, d.clone());
230            });
231        });
232
233        let batches: Result<Vec<_>> = self
234            .batches
235            .iter()
236            .map(|col| record_batch_from_json(&schema, col.clone(), Some(&dictionaries)))
237            .collect();
238
239        batches
240    }
241}
242
243impl ArrowJsonSchema {
244    /// Compare the Arrow JSON schema with the Arrow `Schema`
245    fn equals_schema(&self, schema: &Schema) -> bool {
246        let field_len = self.fields.len();
247        if field_len != schema.fields().len() {
248            return false;
249        }
250        for i in 0..field_len {
251            let json_field = &self.fields[i];
252            let field = schema.field(i);
253            if !json_field.equals_field(field) {
254                return false;
255            }
256        }
257        true
258    }
259
260    fn to_arrow_schema(&self) -> Result<Schema> {
261        let arrow_fields: Result<Vec<_>> = self
262            .fields
263            .iter()
264            .map(|field| field.to_arrow_field())
265            .collect();
266
267        if let Some(metadatas) = &self.metadata {
268            let mut metadata: HashMap<String, String> = HashMap::new();
269
270            metadatas.iter().for_each(|pair| {
271                let key = pair.get("key").unwrap();
272                let value = pair.get("value").unwrap();
273                metadata.insert(key.clone(), value.clone());
274            });
275
276            Ok(Schema::new_with_metadata(arrow_fields?, metadata))
277        } else {
278            Ok(Schema::new(arrow_fields?))
279        }
280    }
281}
282
283impl ArrowJsonField {
284    /// Compare the Arrow JSON field with the Arrow `Field`
285    fn equals_field(&self, field: &Field) -> bool {
286        // convert to a field
287        match self.to_arrow_field() {
288            Ok(self_field) => {
289                assert_eq!(&self_field, field, "Arrow fields not the same");
290                true
291            }
292            Err(e) => {
293                eprintln!("Encountered error while converting JSON field to Arrow field: {e:?}");
294                false
295            }
296        }
297    }
298
299    /// Convert to an Arrow Field
300    /// TODO: convert to use an Into
301    fn to_arrow_field(&self) -> Result<Field> {
302        // a bit regressive, but we have to convert the field to JSON in order to convert it
303        let field =
304            serde_json::to_value(self).map_err(|error| ArrowError::JsonError(error.to_string()))?;
305        field_from_json(&field)
306    }
307}
308
309/// Generates a [`RecordBatch`] from an Arrow JSON batch, given a schema
310pub fn record_batch_from_json(
311    schema: &Schema,
312    json_batch: ArrowJsonBatch,
313    json_dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
314) -> Result<RecordBatch> {
315    let mut columns = vec![];
316
317    for (field, json_col) in schema.fields().iter().zip(json_batch.columns) {
318        let col = array_from_json(field, json_col, json_dictionaries)?;
319        columns.push(col);
320    }
321
322    RecordBatch::try_new(Arc::new(schema.clone()), columns)
323}
324
325/// Construct an Arrow array from a partially typed JSON column
326pub fn array_from_json(
327    field: &Field,
328    json_col: ArrowJsonColumn,
329    dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
330) -> Result<ArrayRef> {
331    match field.data_type() {
332        DataType::Null => Ok(Arc::new(NullArray::new(json_col.count))),
333        DataType::Boolean => {
334            let mut b = BooleanBuilder::with_capacity(json_col.count);
335            for (is_valid, value) in json_col
336                .validity
337                .as_ref()
338                .unwrap()
339                .iter()
340                .zip(json_col.data.unwrap())
341            {
342                match is_valid {
343                    1 => b.append_value(value.as_bool().unwrap()),
344                    _ => b.append_null(),
345                };
346            }
347            Ok(Arc::new(b.finish()))
348        }
349        DataType::Int8 => {
350            let mut b = Int8Builder::with_capacity(json_col.count);
351            for (is_valid, value) in json_col
352                .validity
353                .as_ref()
354                .unwrap()
355                .iter()
356                .zip(json_col.data.unwrap())
357            {
358                match is_valid {
359                    1 => b.append_value(value.as_i64().ok_or_else(|| {
360                        ArrowError::JsonError(format!("Unable to get {value:?} as int64"))
361                    })? as i8),
362                    _ => b.append_null(),
363                };
364            }
365            Ok(Arc::new(b.finish()))
366        }
367        DataType::Int16 => {
368            let mut b = Int16Builder::with_capacity(json_col.count);
369            for (is_valid, value) in json_col
370                .validity
371                .as_ref()
372                .unwrap()
373                .iter()
374                .zip(json_col.data.unwrap())
375            {
376                match is_valid {
377                    1 => b.append_value(value.as_i64().unwrap() as i16),
378                    _ => b.append_null(),
379                };
380            }
381            Ok(Arc::new(b.finish()))
382        }
383        DataType::Int32 | DataType::Date32 | DataType::Time32(_) => {
384            let mut b = Int32Builder::with_capacity(json_col.count);
385            for (is_valid, value) in json_col
386                .validity
387                .as_ref()
388                .unwrap()
389                .iter()
390                .zip(json_col.data.unwrap())
391            {
392                match is_valid {
393                    1 => b.append_value(value.as_i64().unwrap() as i32),
394                    _ => b.append_null(),
395                };
396            }
397            let array = Arc::new(b.finish()) as ArrayRef;
398            arrow::compute::cast(&array, field.data_type())
399        }
400        DataType::Interval(IntervalUnit::YearMonth) => {
401            let mut b = IntervalYearMonthBuilder::with_capacity(json_col.count);
402            for (is_valid, value) in json_col
403                .validity
404                .as_ref()
405                .unwrap()
406                .iter()
407                .zip(json_col.data.unwrap())
408            {
409                match is_valid {
410                    1 => b.append_value(value.as_i64().unwrap() as i32),
411                    _ => b.append_null(),
412                };
413            }
414            Ok(Arc::new(b.finish()))
415        }
416        DataType::Int64
417        | DataType::Date64
418        | DataType::Time64(_)
419        | DataType::Timestamp(_, _)
420        | DataType::Duration(_) => {
421            let mut b = Int64Builder::with_capacity(json_col.count);
422            for (is_valid, value) in json_col
423                .validity
424                .as_ref()
425                .unwrap()
426                .iter()
427                .zip(json_col.data.unwrap())
428            {
429                match is_valid {
430                    1 => b.append_value(match value {
431                        Value::Number(n) => n.as_i64().unwrap(),
432                        Value::String(s) => s.parse().expect("Unable to parse string as i64"),
433                        _ => panic!("Unable to parse {value:?} as number"),
434                    }),
435                    _ => b.append_null(),
436                };
437            }
438            let array = Arc::new(b.finish()) as ArrayRef;
439            arrow::compute::cast(&array, field.data_type())
440        }
441        DataType::Interval(IntervalUnit::DayTime) => {
442            let mut b = IntervalDayTimeBuilder::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::Object(ref map)
453                            if map.contains_key("days") && map.contains_key("milliseconds") =>
454                        {
455                            match field.data_type() {
456                                DataType::Interval(IntervalUnit::DayTime) => {
457                                    let days = map.get("days").unwrap();
458                                    let milliseconds = map.get("milliseconds").unwrap();
459
460                                    match (days, milliseconds) {
461                                        (Value::Number(d), Value::Number(m)) => {
462                                            let days = d.as_i64().unwrap() as _;
463                                            let millis = m.as_i64().unwrap() as _;
464                                            IntervalDayTime::new(days, millis)
465                                        }
466                                        _ => {
467                                            panic!("Unable to parse {value:?} as interval daytime")
468                                        }
469                                    }
470                                }
471                                _ => panic!("Unable to parse {value:?} as interval daytime"),
472                            }
473                        }
474                        _ => panic!("Unable to parse {value:?} as number"),
475                    }),
476                    _ => b.append_null(),
477                };
478            }
479            Ok(Arc::new(b.finish()))
480        }
481        DataType::UInt8 => {
482            let mut b = UInt8Builder::with_capacity(json_col.count);
483            for (is_valid, value) in json_col
484                .validity
485                .as_ref()
486                .unwrap()
487                .iter()
488                .zip(json_col.data.unwrap())
489            {
490                match is_valid {
491                    1 => b.append_value(value.as_u64().unwrap() as u8),
492                    _ => b.append_null(),
493                };
494            }
495            Ok(Arc::new(b.finish()))
496        }
497        DataType::UInt16 => {
498            let mut b = UInt16Builder::with_capacity(json_col.count);
499            for (is_valid, value) in json_col
500                .validity
501                .as_ref()
502                .unwrap()
503                .iter()
504                .zip(json_col.data.unwrap())
505            {
506                match is_valid {
507                    1 => b.append_value(value.as_u64().unwrap() as u16),
508                    _ => b.append_null(),
509                };
510            }
511            Ok(Arc::new(b.finish()))
512        }
513        DataType::UInt32 => {
514            let mut b = UInt32Builder::with_capacity(json_col.count);
515            for (is_valid, value) in json_col
516                .validity
517                .as_ref()
518                .unwrap()
519                .iter()
520                .zip(json_col.data.unwrap())
521            {
522                match is_valid {
523                    1 => b.append_value(value.as_u64().unwrap() as u32),
524                    _ => b.append_null(),
525                };
526            }
527            Ok(Arc::new(b.finish()))
528        }
529        DataType::UInt64 => {
530            let mut b = UInt64Builder::with_capacity(json_col.count);
531            for (is_valid, value) in json_col
532                .validity
533                .as_ref()
534                .unwrap()
535                .iter()
536                .zip(json_col.data.unwrap())
537            {
538                match is_valid {
539                    1 => {
540                        if value.is_string() {
541                            b.append_value(
542                                value
543                                    .as_str()
544                                    .unwrap()
545                                    .parse()
546                                    .expect("Unable to parse string as u64"),
547                            )
548                        } else if value.is_number() {
549                            b.append_value(value.as_u64().expect("Unable to read number as u64"))
550                        } else {
551                            panic!("Unable to parse value {value:?} as u64")
552                        }
553                    }
554                    _ => b.append_null(),
555                };
556            }
557            Ok(Arc::new(b.finish()))
558        }
559        DataType::Interval(IntervalUnit::MonthDayNano) => {
560            let mut b = IntervalMonthDayNanoBuilder::with_capacity(json_col.count);
561            for (is_valid, value) in json_col
562                .validity
563                .as_ref()
564                .unwrap()
565                .iter()
566                .zip(json_col.data.unwrap())
567            {
568                match is_valid {
569                    1 => b.append_value(match value {
570                        Value::Object(v) => {
571                            let months = v.get("months").unwrap();
572                            let days = v.get("days").unwrap();
573                            let nanoseconds = v.get("nanoseconds").unwrap();
574                            match (months, days, nanoseconds) {
575                                (
576                                    Value::Number(months),
577                                    Value::Number(days),
578                                    Value::Number(nanoseconds),
579                                ) => {
580                                    let months = months.as_i64().unwrap() as i32;
581                                    let days = days.as_i64().unwrap() as i32;
582                                    let nanoseconds = nanoseconds.as_i64().unwrap();
583                                    IntervalMonthDayNano::new(months, days, nanoseconds)
584                                }
585                                (_, _, _) => {
586                                    panic!("Unable to parse {v:?} as MonthDayNano")
587                                }
588                            }
589                        }
590                        _ => panic!("Unable to parse {value:?} as MonthDayNano"),
591                    }),
592                    _ => b.append_null(),
593                };
594            }
595            Ok(Arc::new(b.finish()))
596        }
597        DataType::Float32 => {
598            let mut b = Float32Builder::with_capacity(json_col.count);
599            for (is_valid, value) in json_col
600                .validity
601                .as_ref()
602                .unwrap()
603                .iter()
604                .zip(json_col.data.unwrap())
605            {
606                match is_valid {
607                    1 => b.append_value(value.as_f64().unwrap() as f32),
608                    _ => b.append_null(),
609                };
610            }
611            Ok(Arc::new(b.finish()))
612        }
613        DataType::Float64 => {
614            let mut b = Float64Builder::with_capacity(json_col.count);
615            for (is_valid, value) in json_col
616                .validity
617                .as_ref()
618                .unwrap()
619                .iter()
620                .zip(json_col.data.unwrap())
621            {
622                match is_valid {
623                    1 => b.append_value(value.as_f64().unwrap()),
624                    _ => b.append_null(),
625                };
626            }
627            Ok(Arc::new(b.finish()))
628        }
629        DataType::Binary => {
630            let mut b = BinaryBuilder::with_capacity(json_col.count, 1024);
631            for (is_valid, value) in json_col
632                .validity
633                .as_ref()
634                .unwrap()
635                .iter()
636                .zip(json_col.data.unwrap())
637            {
638                match is_valid {
639                    1 => {
640                        let v = decode(value.as_str().unwrap()).unwrap();
641                        b.append_value(&v)
642                    }
643                    _ => b.append_null(),
644                };
645            }
646            Ok(Arc::new(b.finish()))
647        }
648        DataType::LargeBinary => {
649            let mut b = LargeBinaryBuilder::with_capacity(json_col.count, 1024);
650            for (is_valid, value) in json_col
651                .validity
652                .as_ref()
653                .unwrap()
654                .iter()
655                .zip(json_col.data.unwrap())
656            {
657                match is_valid {
658                    1 => {
659                        let v = decode(value.as_str().unwrap()).unwrap();
660                        b.append_value(&v)
661                    }
662                    _ => b.append_null(),
663                };
664            }
665            Ok(Arc::new(b.finish()))
666        }
667        DataType::Utf8 => {
668            let mut b = StringBuilder::with_capacity(json_col.count, 1024);
669            for (is_valid, value) in json_col
670                .validity
671                .as_ref()
672                .unwrap()
673                .iter()
674                .zip(json_col.data.unwrap())
675            {
676                match is_valid {
677                    1 => b.append_value(value.as_str().unwrap()),
678                    _ => b.append_null(),
679                };
680            }
681            Ok(Arc::new(b.finish()))
682        }
683        DataType::LargeUtf8 => {
684            let mut b = LargeStringBuilder::with_capacity(json_col.count, 1024);
685            for (is_valid, value) in json_col
686                .validity
687                .as_ref()
688                .unwrap()
689                .iter()
690                .zip(json_col.data.unwrap())
691            {
692                match is_valid {
693                    1 => b.append_value(value.as_str().unwrap()),
694                    _ => b.append_null(),
695                };
696            }
697            Ok(Arc::new(b.finish()))
698        }
699        DataType::FixedSizeBinary(len) => {
700            let mut b = FixedSizeBinaryBuilder::with_capacity(json_col.count, *len);
701            for (is_valid, value) in json_col
702                .validity
703                .as_ref()
704                .unwrap()
705                .iter()
706                .zip(json_col.data.unwrap())
707            {
708                match is_valid {
709                    1 => {
710                        let v = hex::decode(value.as_str().unwrap()).unwrap();
711                        b.append_value(&v)?
712                    }
713                    _ => b.append_null(),
714                };
715            }
716            Ok(Arc::new(b.finish()))
717        }
718        DataType::List(child_field) => {
719            let null_buf = create_null_buf(&json_col);
720            let children = json_col.children.clone().unwrap();
721            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
722            let offsets: Vec<i32> = json_col
723                .offset
724                .unwrap()
725                .iter()
726                .map(|v| v.as_i64().unwrap() as i32)
727                .collect();
728            let list_data = ArrayData::builder(field.data_type().clone())
729                .len(json_col.count)
730                .offset(0)
731                .add_buffer(Buffer::from(offsets.to_byte_slice()))
732                .add_child_data(child_array.into_data())
733                .null_bit_buffer(Some(null_buf))
734                .build()
735                .unwrap();
736            Ok(Arc::new(ListArray::from(list_data)))
737        }
738        DataType::LargeList(child_field) => {
739            let null_buf = create_null_buf(&json_col);
740            let children = json_col.children.clone().unwrap();
741            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
742            let offsets: Vec<i64> = json_col
743                .offset
744                .unwrap()
745                .iter()
746                .map(|v| match v {
747                    Value::Number(n) => n.as_i64().unwrap(),
748                    Value::String(s) => s.parse::<i64>().unwrap(),
749                    _ => panic!("64-bit offset must be either string or number"),
750                })
751                .collect();
752            let list_data = ArrayData::builder(field.data_type().clone())
753                .len(json_col.count)
754                .offset(0)
755                .add_buffer(Buffer::from(offsets.to_byte_slice()))
756                .add_child_data(child_array.into_data())
757                .null_bit_buffer(Some(null_buf))
758                .build()
759                .unwrap();
760            Ok(Arc::new(LargeListArray::from(list_data)))
761        }
762        DataType::FixedSizeList(child_field, _) => {
763            let children = json_col.children.clone().unwrap();
764            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
765            let null_buf = create_null_buf(&json_col);
766            let list_data = ArrayData::builder(field.data_type().clone())
767                .len(json_col.count)
768                .add_child_data(child_array.into_data())
769                .null_bit_buffer(Some(null_buf))
770                .build()
771                .unwrap();
772            Ok(Arc::new(FixedSizeListArray::from(list_data)))
773        }
774        DataType::Struct(fields) => {
775            // construct struct with null data
776            let null_buf = create_null_buf(&json_col);
777            let mut array_data = ArrayData::builder(field.data_type().clone())
778                .len(json_col.count)
779                .null_bit_buffer(Some(null_buf));
780
781            for (field, col) in fields.iter().zip(json_col.children.unwrap()) {
782                let array = array_from_json(field, col, dictionaries)?;
783                array_data = array_data.add_child_data(array.into_data());
784            }
785
786            let array = StructArray::from(array_data.build().unwrap());
787            Ok(Arc::new(array))
788        }
789        DataType::Dictionary(key_type, value_type) => {
790            #[allow(deprecated)]
791            let dict_id = field.dict_id().ok_or_else(|| {
792                ArrowError::JsonError(format!("Unable to find dict_id for field {field:?}"))
793            })?;
794            // find dictionary
795            let dictionary = dictionaries
796                .ok_or_else(|| {
797                    ArrowError::JsonError(format!(
798                        "Unable to find any dictionaries for field {field:?}"
799                    ))
800                })?
801                .get(&dict_id);
802            match dictionary {
803                Some(dictionary) => dictionary_array_from_json(
804                    field,
805                    json_col,
806                    key_type,
807                    value_type,
808                    dictionary,
809                    dictionaries,
810                ),
811                None => Err(ArrowError::JsonError(format!(
812                    "Unable to find dictionary for field {field:?}"
813                ))),
814            }
815        }
816        DataType::Decimal128(precision, scale) => {
817            let mut b = Decimal128Builder::with_capacity(json_col.count);
818            for (is_valid, value) in json_col
819                .validity
820                .as_ref()
821                .unwrap()
822                .iter()
823                .zip(json_col.data.unwrap())
824            {
825                match is_valid {
826                    1 => b.append_value(value.as_str().unwrap().parse::<i128>().unwrap()),
827                    _ => b.append_null(),
828                };
829            }
830            Ok(Arc::new(
831                b.finish().with_precision_and_scale(*precision, *scale)?,
832            ))
833        }
834        DataType::Decimal256(precision, scale) => {
835            let mut b = Decimal256Builder::with_capacity(json_col.count);
836            for (is_valid, value) in json_col
837                .validity
838                .as_ref()
839                .unwrap()
840                .iter()
841                .zip(json_col.data.unwrap())
842            {
843                match is_valid {
844                    1 => {
845                        let str = value.as_str().unwrap();
846                        let integer = BigInt::parse_bytes(str.as_bytes(), 10).unwrap();
847                        let integer_bytes = integer.to_signed_bytes_le();
848                        let mut bytes = if integer.is_positive() {
849                            [0_u8; 32]
850                        } else {
851                            [255_u8; 32]
852                        };
853                        bytes[0..integer_bytes.len()].copy_from_slice(integer_bytes.as_slice());
854                        b.append_value(i256::from_le_bytes(bytes));
855                    }
856                    _ => b.append_null(),
857                }
858            }
859            Ok(Arc::new(
860                b.finish().with_precision_and_scale(*precision, *scale)?,
861            ))
862        }
863        DataType::Map(child_field, _) => {
864            let null_buf = create_null_buf(&json_col);
865            let children = json_col.children.clone().unwrap();
866            let child_array = array_from_json(child_field, children[0].clone(), dictionaries)?;
867            let offsets: Vec<i32> = json_col
868                .offset
869                .unwrap()
870                .iter()
871                .map(|v| v.as_i64().unwrap() as i32)
872                .collect();
873            let array_data = ArrayData::builder(field.data_type().clone())
874                .len(json_col.count)
875                .add_buffer(Buffer::from(offsets.to_byte_slice()))
876                .add_child_data(child_array.into_data())
877                .null_bit_buffer(Some(null_buf))
878                .build()
879                .unwrap();
880
881            let array = MapArray::from(array_data);
882            Ok(Arc::new(array))
883        }
884        DataType::Union(fields, _) => {
885            let type_ids = if let Some(type_id) = json_col.type_id {
886                type_id
887            } else {
888                return Err(ArrowError::JsonError(
889                    "Cannot find expected type_id in json column".to_string(),
890                ));
891            };
892
893            let offset: Option<ScalarBuffer<i32>> = json_col
894                .offset
895                .map(|offsets| offsets.iter().map(|v| v.as_i64().unwrap() as i32).collect());
896
897            let mut children = Vec::with_capacity(fields.len());
898            for ((_, field), col) in fields.iter().zip(json_col.children.unwrap()) {
899                let array = array_from_json(field, col, dictionaries)?;
900                children.push(array);
901            }
902
903            let array =
904                UnionArray::try_new(fields.clone(), type_ids.into(), offset, children).unwrap();
905            Ok(Arc::new(array))
906        }
907        t => Err(ArrowError::JsonError(format!(
908            "data type {t:?} not supported"
909        ))),
910    }
911}
912
913/// Construct a [`DictionaryArray`] from a partially typed JSON column
914pub fn dictionary_array_from_json(
915    field: &Field,
916    json_col: ArrowJsonColumn,
917    dict_key: &DataType,
918    dict_value: &DataType,
919    dictionary: &ArrowJsonDictionaryBatch,
920    dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
921) -> Result<ArrayRef> {
922    match dict_key {
923        DataType::Int8
924        | DataType::Int16
925        | DataType::Int32
926        | DataType::Int64
927        | DataType::UInt8
928        | DataType::UInt16
929        | DataType::UInt32
930        | DataType::UInt64 => {
931            let null_buf = create_null_buf(&json_col);
932
933            // build the key data into a buffer, then construct values separately
934            #[allow(deprecated)]
935            let key_field = Field::new_dict(
936                "key",
937                dict_key.clone(),
938                field.is_nullable(),
939                #[allow(deprecated)]
940                field
941                    .dict_id()
942                    .expect("Dictionary fields must have a dict_id value"),
943                field
944                    .dict_is_ordered()
945                    .expect("Dictionary fields must have a dict_is_ordered value"),
946            );
947            let keys = array_from_json(&key_field, json_col, None)?;
948            // note: not enough info on nullability of dictionary
949            let value_field = Field::new("value", dict_value.clone(), true);
950            let values = array_from_json(
951                &value_field,
952                dictionary.data.columns[0].clone(),
953                dictionaries,
954            )?;
955
956            // convert key and value to dictionary data
957            let dict_data = ArrayData::builder(field.data_type().clone())
958                .len(keys.len())
959                .add_buffer(keys.to_data().buffers()[0].clone())
960                .null_bit_buffer(Some(null_buf))
961                .add_child_data(values.into_data())
962                .build()
963                .unwrap();
964
965            let array = match dict_key {
966                DataType::Int8 => Arc::new(Int8DictionaryArray::from(dict_data)) as ArrayRef,
967                DataType::Int16 => Arc::new(Int16DictionaryArray::from(dict_data)),
968                DataType::Int32 => Arc::new(Int32DictionaryArray::from(dict_data)),
969                DataType::Int64 => Arc::new(Int64DictionaryArray::from(dict_data)),
970                DataType::UInt8 => Arc::new(UInt8DictionaryArray::from(dict_data)),
971                DataType::UInt16 => Arc::new(UInt16DictionaryArray::from(dict_data)),
972                DataType::UInt32 => Arc::new(UInt32DictionaryArray::from(dict_data)),
973                DataType::UInt64 => Arc::new(UInt64DictionaryArray::from(dict_data)),
974                _ => unreachable!(),
975            };
976            Ok(array)
977        }
978        _ => Err(ArrowError::JsonError(format!(
979            "Dictionary key type {dict_key:?} not supported"
980        ))),
981    }
982}
983
984/// A helper to create a null buffer from a `Vec<bool>`
985fn create_null_buf(json_col: &ArrowJsonColumn) -> Buffer {
986    let num_bytes = bit_util::ceil(json_col.count, 8);
987    let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, false);
988    json_col
989        .validity
990        .clone()
991        .unwrap()
992        .iter()
993        .enumerate()
994        .for_each(|(i, v)| {
995            let null_slice = null_buf.as_slice_mut();
996            if *v != 0 {
997                bit_util::set_bit(null_slice, i);
998            }
999        });
1000    null_buf.into()
1001}
1002
1003impl ArrowJsonBatch {
1004    /// Convert a [`RecordBatch`] to an [`ArrowJsonBatch`]
1005    pub fn from_batch(batch: &RecordBatch) -> ArrowJsonBatch {
1006        let mut json_batch = ArrowJsonBatch {
1007            count: batch.num_rows(),
1008            columns: Vec::with_capacity(batch.num_columns()),
1009        };
1010
1011        for (col, field) in batch.columns().iter().zip(batch.schema().fields.iter()) {
1012            let json_col = match field.data_type() {
1013                DataType::Int8 => {
1014                    let col = col.as_any().downcast_ref::<Int8Array>().unwrap();
1015
1016                    let mut validity: Vec<u8> = Vec::with_capacity(col.len());
1017                    let mut data: Vec<Value> = Vec::with_capacity(col.len());
1018
1019                    for i in 0..col.len() {
1020                        if col.is_null(i) {
1021                            validity.push(1);
1022                            data.push(0i8.into());
1023                        } else {
1024                            validity.push(0);
1025                            data.push(col.value(i).into());
1026                        }
1027                    }
1028
1029                    ArrowJsonColumn {
1030                        name: field.name().clone(),
1031                        count: col.len(),
1032                        validity: Some(validity),
1033                        data: Some(data),
1034                        offset: None,
1035                        type_id: None,
1036                        children: None,
1037                    }
1038                }
1039                _ => ArrowJsonColumn {
1040                    name: field.name().clone(),
1041                    count: col.len(),
1042                    validity: None,
1043                    data: None,
1044                    offset: None,
1045                    type_id: None,
1046                    children: None,
1047                },
1048            };
1049
1050            json_batch.columns.push(json_col);
1051        }
1052
1053        json_batch
1054    }
1055}
1056
1057#[cfg(test)]
1058mod tests {
1059    use super::*;
1060
1061    use std::fs::File;
1062    use std::io::Read;
1063
1064    #[test]
1065    fn test_schema_equality() {
1066        let json = r#"
1067        {
1068            "fields": [
1069                {
1070                    "name": "c1",
1071                    "type": {"name": "int", "isSigned": true, "bitWidth": 32},
1072                    "nullable": true,
1073                    "children": []
1074                },
1075                {
1076                    "name": "c2",
1077                    "type": {"name": "floatingpoint", "precision": "DOUBLE"},
1078                    "nullable": true,
1079                    "children": []
1080                },
1081                {
1082                    "name": "c3",
1083                    "type": {"name": "utf8"},
1084                    "nullable": true,
1085                    "children": []
1086                },
1087                {
1088                    "name": "c4",
1089                    "type": {
1090                        "name": "list"
1091                    },
1092                    "nullable": true,
1093                    "children": [
1094                        {
1095                            "name": "custom_item",
1096                            "type": {
1097                                "name": "int",
1098                                "isSigned": true,
1099                                "bitWidth": 32
1100                            },
1101                            "nullable": false,
1102                            "children": []
1103                        }
1104                    ]
1105                }
1106            ]
1107        }"#;
1108        let json_schema: ArrowJsonSchema = serde_json::from_str(json).unwrap();
1109        let schema = Schema::new(vec![
1110            Field::new("c1", DataType::Int32, true),
1111            Field::new("c2", DataType::Float64, true),
1112            Field::new("c3", DataType::Utf8, true),
1113            Field::new(
1114                "c4",
1115                DataType::List(Arc::new(Field::new("custom_item", DataType::Int32, false))),
1116                true,
1117            ),
1118        ]);
1119        assert!(json_schema.equals_schema(&schema));
1120    }
1121
1122    #[test]
1123    fn test_arrow_data_equality() {
1124        let secs_tz = Some("Europe/Budapest".into());
1125        let millis_tz = Some("America/New_York".into());
1126        let micros_tz = Some("UTC".into());
1127        let nanos_tz = Some("Africa/Johannesburg".into());
1128
1129        let schema = Schema::new(vec![
1130            Field::new("bools-with-metadata-map", DataType::Boolean, true).with_metadata(
1131                [("k".to_string(), "v".to_string())]
1132                    .iter()
1133                    .cloned()
1134                    .collect(),
1135            ),
1136            Field::new("bools-with-metadata-vec", DataType::Boolean, true).with_metadata(
1137                [("k2".to_string(), "v2".to_string())]
1138                    .iter()
1139                    .cloned()
1140                    .collect(),
1141            ),
1142            Field::new("bools", DataType::Boolean, true),
1143            Field::new("int8s", DataType::Int8, true),
1144            Field::new("int16s", DataType::Int16, true),
1145            Field::new("int32s", DataType::Int32, true),
1146            Field::new("int64s", DataType::Int64, true),
1147            Field::new("uint8s", DataType::UInt8, true),
1148            Field::new("uint16s", DataType::UInt16, true),
1149            Field::new("uint32s", DataType::UInt32, true),
1150            Field::new("uint64s", DataType::UInt64, true),
1151            Field::new("float32s", DataType::Float32, true),
1152            Field::new("float64s", DataType::Float64, true),
1153            Field::new("date_days", DataType::Date32, true),
1154            Field::new("date_millis", DataType::Date64, true),
1155            Field::new("time_secs", DataType::Time32(TimeUnit::Second), true),
1156            Field::new("time_millis", DataType::Time32(TimeUnit::Millisecond), true),
1157            Field::new("time_micros", DataType::Time64(TimeUnit::Microsecond), true),
1158            Field::new("time_nanos", DataType::Time64(TimeUnit::Nanosecond), true),
1159            Field::new("ts_secs", DataType::Timestamp(TimeUnit::Second, None), true),
1160            Field::new(
1161                "ts_millis",
1162                DataType::Timestamp(TimeUnit::Millisecond, None),
1163                true,
1164            ),
1165            Field::new(
1166                "ts_micros",
1167                DataType::Timestamp(TimeUnit::Microsecond, None),
1168                true,
1169            ),
1170            Field::new(
1171                "ts_nanos",
1172                DataType::Timestamp(TimeUnit::Nanosecond, None),
1173                true,
1174            ),
1175            Field::new(
1176                "ts_secs_tz",
1177                DataType::Timestamp(TimeUnit::Second, secs_tz.clone()),
1178                true,
1179            ),
1180            Field::new(
1181                "ts_millis_tz",
1182                DataType::Timestamp(TimeUnit::Millisecond, millis_tz.clone()),
1183                true,
1184            ),
1185            Field::new(
1186                "ts_micros_tz",
1187                DataType::Timestamp(TimeUnit::Microsecond, micros_tz.clone()),
1188                true,
1189            ),
1190            Field::new(
1191                "ts_nanos_tz",
1192                DataType::Timestamp(TimeUnit::Nanosecond, nanos_tz.clone()),
1193                true,
1194            ),
1195            Field::new("utf8s", DataType::Utf8, true),
1196            Field::new(
1197                "lists",
1198                DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
1199                true,
1200            ),
1201            Field::new(
1202                "structs",
1203                DataType::Struct(Fields::from(vec![
1204                    Field::new("int32s", DataType::Int32, true),
1205                    Field::new("utf8s", DataType::Utf8, true),
1206                ])),
1207                true,
1208            ),
1209        ]);
1210
1211        let bools_with_metadata_map = BooleanArray::from(vec![Some(true), None, Some(false)]);
1212        let bools_with_metadata_vec = BooleanArray::from(vec![Some(true), None, Some(false)]);
1213        let bools = BooleanArray::from(vec![Some(true), None, Some(false)]);
1214        let int8s = Int8Array::from(vec![Some(1), None, Some(3)]);
1215        let int16s = Int16Array::from(vec![Some(1), None, Some(3)]);
1216        let int32s = Int32Array::from(vec![Some(1), None, Some(3)]);
1217        let int64s = Int64Array::from(vec![Some(1), None, Some(3)]);
1218        let uint8s = UInt8Array::from(vec![Some(1), None, Some(3)]);
1219        let uint16s = UInt16Array::from(vec![Some(1), None, Some(3)]);
1220        let uint32s = UInt32Array::from(vec![Some(1), None, Some(3)]);
1221        let uint64s = UInt64Array::from(vec![Some(1), None, Some(3)]);
1222        let float32s = Float32Array::from(vec![Some(1.0), None, Some(3.0)]);
1223        let float64s = Float64Array::from(vec![Some(1.0), None, Some(3.0)]);
1224        let date_days = Date32Array::from(vec![Some(1196848), None, None]);
1225        let date_millis = Date64Array::from(vec![
1226            Some(167903550396207),
1227            Some(29923997007884),
1228            Some(30612271819236),
1229        ]);
1230        let time_secs = Time32SecondArray::from(vec![Some(27974), Some(78592), Some(43207)]);
1231        let time_millis =
1232            Time32MillisecondArray::from(vec![Some(6613125), Some(74667230), Some(52260079)]);
1233        let time_micros = Time64MicrosecondArray::from(vec![Some(62522958593), None, None]);
1234        let time_nanos =
1235            Time64NanosecondArray::from(vec![Some(73380123595985), None, Some(16584393546415)]);
1236        let ts_secs = TimestampSecondArray::from(vec![None, Some(193438817552), None]);
1237        let ts_millis =
1238            TimestampMillisecondArray::from(vec![None, Some(38606916383008), Some(58113709376587)]);
1239        let ts_micros = TimestampMicrosecondArray::from(vec![None, None, None]);
1240        let ts_nanos = TimestampNanosecondArray::from(vec![None, None, Some(-6473623571954960143)]);
1241        let ts_secs_tz = TimestampSecondArray::from(vec![None, Some(193438817552), None])
1242            .with_timezone_opt(secs_tz);
1243        let ts_millis_tz =
1244            TimestampMillisecondArray::from(vec![None, Some(38606916383008), Some(58113709376587)])
1245                .with_timezone_opt(millis_tz);
1246        let ts_micros_tz =
1247            TimestampMicrosecondArray::from(vec![None, None, None]).with_timezone_opt(micros_tz);
1248        let ts_nanos_tz =
1249            TimestampNanosecondArray::from(vec![None, None, Some(-6473623571954960143)])
1250                .with_timezone_opt(nanos_tz);
1251        let utf8s = StringArray::from(vec![Some("aa"), None, Some("bbb")]);
1252
1253        let value_data = Int32Array::from(vec![None, Some(2), None, None]);
1254        let value_offsets = Buffer::from_slice_ref([0, 3, 4, 4]);
1255        let list_data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
1256        let list_data = ArrayData::builder(list_data_type)
1257            .len(3)
1258            .add_buffer(value_offsets)
1259            .add_child_data(value_data.into_data())
1260            .null_bit_buffer(Some(Buffer::from([0b00000011])))
1261            .build()
1262            .unwrap();
1263        let lists = ListArray::from(list_data);
1264
1265        let structs_int32s = Int32Array::from(vec![None, Some(-2), None]);
1266        let structs_utf8s = StringArray::from(vec![None, None, Some("aaaaaa")]);
1267        let struct_data_type = DataType::Struct(Fields::from(vec![
1268            Field::new("int32s", DataType::Int32, true),
1269            Field::new("utf8s", DataType::Utf8, true),
1270        ]));
1271        let struct_data = ArrayData::builder(struct_data_type)
1272            .len(3)
1273            .add_child_data(structs_int32s.into_data())
1274            .add_child_data(structs_utf8s.into_data())
1275            .null_bit_buffer(Some(Buffer::from([0b00000011])))
1276            .build()
1277            .unwrap();
1278        let structs = StructArray::from(struct_data);
1279
1280        let record_batch = RecordBatch::try_new(
1281            Arc::new(schema.clone()),
1282            vec![
1283                Arc::new(bools_with_metadata_map),
1284                Arc::new(bools_with_metadata_vec),
1285                Arc::new(bools),
1286                Arc::new(int8s),
1287                Arc::new(int16s),
1288                Arc::new(int32s),
1289                Arc::new(int64s),
1290                Arc::new(uint8s),
1291                Arc::new(uint16s),
1292                Arc::new(uint32s),
1293                Arc::new(uint64s),
1294                Arc::new(float32s),
1295                Arc::new(float64s),
1296                Arc::new(date_days),
1297                Arc::new(date_millis),
1298                Arc::new(time_secs),
1299                Arc::new(time_millis),
1300                Arc::new(time_micros),
1301                Arc::new(time_nanos),
1302                Arc::new(ts_secs),
1303                Arc::new(ts_millis),
1304                Arc::new(ts_micros),
1305                Arc::new(ts_nanos),
1306                Arc::new(ts_secs_tz),
1307                Arc::new(ts_millis_tz),
1308                Arc::new(ts_micros_tz),
1309                Arc::new(ts_nanos_tz),
1310                Arc::new(utf8s),
1311                Arc::new(lists),
1312                Arc::new(structs),
1313            ],
1314        )
1315        .unwrap();
1316        let mut file = File::open("data/integration.json").unwrap();
1317        let mut json = String::new();
1318        file.read_to_string(&mut json).unwrap();
1319        let arrow_json: ArrowJson = serde_json::from_str(&json).unwrap();
1320        // test schemas
1321        assert!(arrow_json.schema.equals_schema(&schema));
1322        // test record batch
1323        assert_eq!(arrow_json.get_record_batches().unwrap()[0], record_batch);
1324    }
1325}
arrow_integration_test/lib.rs

arrow_integration_test/
lib.rs
