arrow_json_integration_test/
arrow-json-integration-test.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//! Converts between the Arrow IPC file format and the integration test JSON format.
19
20// The unused_crate_dependencies lint does not work well for crates defining additional examples/bin targets
21#![allow(unused_crate_dependencies)]
22
23use arrow::error::{ArrowError, Result};
24use arrow::ipc::reader::FileReader;
25use arrow::ipc::writer::FileWriter;
26use arrow_integration_test::*;
27use arrow_integration_testing::{canonicalize_schema, open_json_file};
28use clap::Parser;
29use std::fs::File;
30
31#[derive(clap::ValueEnum, Debug, Clone)]
32#[clap(rename_all = "SCREAMING_SNAKE_CASE")]
33enum Mode {
34    ArrowToJson,
35    JsonToArrow,
36    Validate,
37}
38
39#[derive(Debug, Parser)]
40#[clap(author, version, about("rust arrow-json-integration-test"), long_about = None)]
41struct Args {
42    #[clap(short, long)]
43    integration: bool,
44    #[clap(short, long, help("Path to ARROW file"))]
45    arrow: String,
46    #[clap(short, long, help("Path to JSON file"))]
47    json: String,
48    #[clap(
49        value_enum,
50        short,
51        long,
52        default_value = "VALIDATE",
53        help = "Mode of integration testing tool"
54    )]
55    mode: Mode,
56    #[clap(short, long)]
57    verbose: bool,
58}
59
60fn main() -> Result<()> {
61    let args = Args::parse();
62    let arrow_file = args.arrow;
63    let json_file = args.json;
64    let verbose = args.verbose;
65    match args.mode {
66        Mode::JsonToArrow => json_to_arrow(&json_file, &arrow_file, verbose),
67        Mode::ArrowToJson => arrow_to_json(&arrow_file, &json_file, verbose),
68        Mode::Validate => validate(&arrow_file, &json_file, verbose),
69    }
70}
71
72fn json_to_arrow(json_name: &str, arrow_name: &str, verbose: bool) -> Result<()> {
73    if verbose {
74        eprintln!("Converting {json_name} to {arrow_name}");
75    }
76
77    let json_file = open_json_file(json_name)?;
78
79    let arrow_file = File::create(arrow_name)?;
80    let mut writer = FileWriter::try_new(arrow_file, &json_file.schema)?;
81
82    for b in json_file.read_batches()? {
83        writer.write(&b)?;
84    }
85
86    writer.finish()?;
87
88    Ok(())
89}
90
91fn arrow_to_json(arrow_name: &str, json_name: &str, verbose: bool) -> Result<()> {
92    if verbose {
93        eprintln!("Converting {arrow_name} to {json_name}");
94    }
95
96    let arrow_file = File::open(arrow_name)?;
97    let reader = FileReader::try_new(arrow_file, None)?;
98
99    let mut fields: Vec<ArrowJsonField> = vec![];
100    for f in reader.schema().fields() {
101        fields.push(ArrowJsonField::from(f));
102    }
103    let schema = ArrowJsonSchema {
104        fields,
105        metadata: None,
106    };
107
108    let batches = reader
109        .map(|batch| Ok(ArrowJsonBatch::from_batch(&batch?)))
110        .collect::<Result<Vec<_>>>()?;
111
112    let arrow_json = ArrowJson {
113        schema,
114        batches,
115        dictionaries: None,
116    };
117
118    let json_file = File::create(json_name)?;
119    serde_json::to_writer(&json_file, &arrow_json).unwrap();
120
121    Ok(())
122}
123
124fn validate(arrow_name: &str, json_name: &str, verbose: bool) -> Result<()> {
125    if verbose {
126        eprintln!("Validating {arrow_name} and {json_name}");
127    }
128
129    // open JSON file
130    let json_file = open_json_file(json_name)?;
131
132    // open Arrow file
133    let arrow_file = File::open(arrow_name)?;
134    let mut arrow_reader = FileReader::try_new(arrow_file, None)?;
135    let arrow_schema = arrow_reader.schema().as_ref().to_owned();
136
137    // compare schemas
138    if canonicalize_schema(&json_file.schema) != canonicalize_schema(&arrow_schema) {
139        return Err(ArrowError::ComputeError(format!(
140            "Schemas do not match. JSON: {:?}. Arrow: {:?}",
141            json_file.schema, arrow_schema
142        )));
143    }
144
145    let json_batches = json_file.read_batches()?;
146
147    // compare number of batches
148    assert!(
149        json_batches.len() == arrow_reader.num_batches(),
150        "JSON batches and Arrow batches are unequal"
151    );
152
153    if verbose {
154        eprintln!(
155            "Schemas match. JSON file has {} batches.",
156            json_batches.len()
157        );
158    }
159
160    for json_batch in json_batches {
161        if let Some(Ok(arrow_batch)) = arrow_reader.next() {
162            // compare batches
163            let num_columns = arrow_batch.num_columns();
164            assert!(num_columns == json_batch.num_columns());
165            assert!(arrow_batch.num_rows() == json_batch.num_rows());
166
167            for i in 0..num_columns {
168                assert_eq!(
169                    arrow_batch.column(i).as_ref(),
170                    json_batch.column(i).as_ref(),
171                    "Arrow and JSON batch columns not the same"
172                );
173            }
174        } else {
175            return Err(ArrowError::ComputeError(
176                "no more arrow batches left".to_owned(),
177            ));
178        }
179    }
180
181    if arrow_reader.next().is_some() {
182        return Err(ArrowError::ComputeError(
183            "no more json batches left".to_owned(),
184        ));
185    }
186
187    Ok(())
188}
arrow_json_integration_test/arrow-json-integration-test.rs

arrow_json_integration_test/
arrow-json-integration-test.rs
