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

arrow_json_integration_test/
arrow-json-integration-test.rs
