arrow_integration_testing/
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//! Common code used in the integration test binaries
19
20// The unused_crate_dependencies lint does not work well for crates defining additional examples/bin targets
21#![allow(unused_crate_dependencies)]
22#![warn(missing_docs)]
23use serde_json::Value;
24
25use arrow::array::{Array, StructArray};
26use arrow::datatypes::{DataType, Field, Fields, Schema};
27use arrow::error::{ArrowError, Result};
28use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema, from_ffi_and_data_type};
29use arrow::record_batch::RecordBatch;
30use arrow::util::test_util::arrow_test_data;
31use arrow_integration_test::*;
32use std::collections::HashMap;
33use std::ffi::{CStr, CString, c_char, c_int};
34use std::fs::File;
35use std::io::BufReader;
36use std::iter::zip;
37use std::ptr;
38use std::sync::Arc;
39
40/// The expected username for the basic auth integration test.
41pub const AUTH_USERNAME: &str = "arrow";
42/// The expected password for the basic auth integration test.
43pub const AUTH_PASSWORD: &str = "flight";
44
45pub mod flight_client_scenarios;
46pub mod flight_server_scenarios;
47
48/// An Arrow file in JSON format
49pub struct ArrowFile {
50    /// The schema of the file
51    pub schema: Schema,
52    // we can evolve this into a concrete Arrow type
53    // this is temporarily not being read from
54    dictionaries: HashMap<i64, ArrowJsonDictionaryBatch>,
55    arrow_json: Value,
56}
57
58impl ArrowFile {
59    /// Read a single [RecordBatch] from the file
60    pub fn read_batch(&self, batch_num: usize) -> Result<RecordBatch> {
61        let b = self.arrow_json["batches"].get(batch_num).ok_or_else(|| {
62            ArrowError::ParseError(format!("Arrow JSON has no batch {batch_num}"))
63        })?;
64        self.batch_from_json(b)
65    }
66
67    /// Read all [RecordBatch]es from the file
68    pub fn read_batches(&self) -> Result<Vec<RecordBatch>> {
69        self.arrow_json["batches"]
70            .as_array()
71            .ok_or_else(|| ArrowError::ParseError("Arrow JSON has no 'batches' array".to_string()))?
72            .iter()
73            .map(|b| self.batch_from_json(b))
74            .collect()
75    }
76
77    fn batch_from_json(&self, batch: &Value) -> Result<RecordBatch> {
78        let json_batch: ArrowJsonBatch = serde_json::from_value(batch.clone())
79            .map_err(|err| ArrowError::ParseError(format!("Invalid Arrow JSON batch: {err}")))?;
80        record_batch_from_json(&self.schema, json_batch, Some(&self.dictionaries))
81    }
82}
83
84/// Canonicalize the names of map fields in a schema
85pub fn canonicalize_schema(schema: &Schema) -> Schema {
86    let fields = schema
87        .fields()
88        .iter()
89        .map(|field| match field.data_type() {
90            DataType::Map(child_field, sorted) => match child_field.data_type() {
91                DataType::Struct(fields) if fields.len() == 2 => {
92                    let first_field = &fields[0];
93                    let key_field = Arc::new(Field::new(
94                        Field::MAP_KEY_FIELD_DEFAULT_NAME,
95                        first_field.data_type().clone(),
96                        false,
97                    ));
98                    let second_field = &fields[1];
99                    let value_field = Arc::new(Field::new(
100                        Field::MAP_VALUE_FIELD_DEFAULT_NAME,
101                        second_field.data_type().clone(),
102                        second_field.is_nullable(),
103                    ));
104
105                    let fields = Fields::from([key_field, value_field]);
106                    let struct_type = DataType::Struct(fields);
107                    let child_field =
108                        Field::new(Field::MAP_ENTRIES_FIELD_DEFAULT_NAME, struct_type, false);
109
110                    Arc::new(Field::new(
111                        field.name().as_str(),
112                        DataType::Map(Arc::new(child_field), *sorted),
113                        field.is_nullable(),
114                    ))
115                }
116                _ => panic!("The child field of Map type should be Struct type with 2 fields."),
117            },
118            _ => field.clone(),
119        })
120        .collect::<Fields>();
121
122    Schema::new(fields).with_metadata(schema.metadata().clone())
123}
124
125/// Read an Arrow file in JSON format
126pub fn open_json_file(json_name: &str) -> Result<ArrowFile> {
127    let json_file = File::open(json_name)?;
128    let reader = BufReader::new(json_file);
129    let arrow_json: Value = serde_json::from_reader(reader)
130        .map_err(|err| ArrowError::ParseError(format!("Invalid Arrow JSON: {err}")))?;
131    let schema = schema_from_json(&arrow_json["schema"])?;
132    // read dictionaries
133    let mut dictionaries = HashMap::new();
134    if let Some(dicts) = arrow_json.get("dictionaries") {
135        let dicts = dicts.as_array().ok_or_else(|| {
136            ArrowError::ParseError("Arrow JSON 'dictionaries' is not an array".to_string())
137        })?;
138        for d in dicts {
139            let json_dict: ArrowJsonDictionaryBatch =
140                serde_json::from_value(d.clone()).map_err(|err| {
141                    ArrowError::ParseError(format!("Invalid Arrow JSON dictionary: {err}"))
142                })?;
143            // TODO: convert to a concrete Arrow type
144            dictionaries.insert(json_dict.id, json_dict);
145        }
146    }
147    Ok(ArrowFile {
148        schema,
149        dictionaries,
150        arrow_json,
151    })
152}
153
154/// Read gzipped JSON test file
155///
156/// For example given the input:
157/// version = `0.17.1`
158/// path = `generated_union`
159///
160/// Returns the contents of
161/// `arrow-ipc-stream/integration/0.17.1/generated_union.json.gz`
162pub fn read_gzip_json(version: &str, path: &str) -> ArrowJson {
163    use flate2::read::GzDecoder;
164    use std::io::Read;
165
166    let testdata = arrow_test_data();
167    let file = File::open(format!(
168        "{testdata}/arrow-ipc-stream/integration/{version}/{path}.json.gz"
169    ))
170    .unwrap();
171    let mut gz = GzDecoder::new(&file);
172    let mut s = String::new();
173    gz.read_to_string(&mut s).unwrap();
174    // convert to Arrow JSON
175    let arrow_json: ArrowJson = serde_json::from_str(&s).unwrap();
176    arrow_json
177}
178
179/// C Data Integration entrypoint to export the schema from a JSON file
180fn cdata_integration_export_schema_from_json(
181    c_json_name: *const c_char,
182    out: *mut FFI_ArrowSchema,
183) -> Result<()> {
184    let json_name = unsafe { CStr::from_ptr(c_json_name) };
185    let f = open_json_file(json_name.to_str()?)?;
186    let c_schema = FFI_ArrowSchema::try_from(&f.schema)?;
187    // Move exported schema into output struct
188    unsafe { ptr::write(out, c_schema) };
189    Ok(())
190}
191
192/// C Data Integration entrypoint to export a batch from a JSON file
193fn cdata_integration_export_batch_from_json(
194    c_json_name: *const c_char,
195    batch_num: c_int,
196    out: *mut FFI_ArrowArray,
197) -> Result<()> {
198    let json_name = unsafe { CStr::from_ptr(c_json_name) };
199    let b = open_json_file(json_name.to_str()?)?.read_batch(batch_num.try_into().unwrap())?;
200    let a = StructArray::from(b).into_data();
201    let c_array = FFI_ArrowArray::new(&a);
202    // Move exported array into output struct
203    unsafe { ptr::write(out, c_array) };
204    Ok(())
205}
206
207fn cdata_integration_import_schema_and_compare_to_json(
208    c_json_name: *const c_char,
209    c_schema: *mut FFI_ArrowSchema,
210) -> Result<()> {
211    let json_name = unsafe { CStr::from_ptr(c_json_name) };
212    let json_schema = open_json_file(json_name.to_str()?)?.schema;
213
214    // The source ArrowSchema will be released when this is dropped
215    let imported_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema) };
216    let imported_schema = Schema::try_from(&imported_schema)?;
217
218    // compare schemas
219    if canonicalize_schema(&json_schema) != canonicalize_schema(&imported_schema) {
220        return Err(ArrowError::ComputeError(format!(
221            "Schemas do not match.\n- JSON: {json_schema:?}\n- Imported: {imported_schema:?}",
222        )));
223    }
224    Ok(())
225}
226
227fn compare_batches(a: &RecordBatch, b: &RecordBatch) -> Result<()> {
228    if a.num_columns() != b.num_columns() {
229        return Err(ArrowError::InvalidArgumentError(
230            "batches do not have the same number of columns".to_string(),
231        ));
232    }
233    for (a_column, b_column) in zip(a.columns(), b.columns()) {
234        if a_column != b_column {
235            return Err(ArrowError::InvalidArgumentError(
236                "batch columns are not the same".to_string(),
237            ));
238        }
239    }
240    Ok(())
241}
242
243fn cdata_integration_import_batch_and_compare_to_json(
244    c_json_name: *const c_char,
245    batch_num: c_int,
246    c_array: *mut FFI_ArrowArray,
247) -> Result<()> {
248    let json_name = unsafe { CStr::from_ptr(c_json_name) };
249    let json_batch =
250        open_json_file(json_name.to_str()?)?.read_batch(batch_num.try_into().unwrap())?;
251    let schema = json_batch.schema();
252
253    let data_type_for_import = DataType::Struct(schema.fields.clone());
254    let imported_array = unsafe { FFI_ArrowArray::from_raw(c_array) };
255    let imported_array = unsafe { from_ffi_and_data_type(imported_array, data_type_for_import) }?;
256    imported_array.validate_full()?;
257    let imported_batch = RecordBatch::from(StructArray::from(imported_array));
258
259    compare_batches(&json_batch, &imported_batch)
260}
261
262// If Result is an error, then export a const char* to its string display, otherwise NULL
263fn result_to_c_error<T, E: std::fmt::Display>(result: &std::result::Result<T, E>) -> *mut c_char {
264    match result {
265        Ok(_) => ptr::null_mut(),
266        Err(e) => CString::new(format!("{e}")).unwrap().into_raw(),
267    }
268}
269
270/// Release a const char* exported by result_to_c_error()
271///
272/// # Safety
273///
274/// The pointer is assumed to have been obtained using CString::into_raw.
275#[unsafe(no_mangle)]
276pub unsafe extern "C" fn arrow_rs_free_error(c_error: *mut c_char) {
277    if !c_error.is_null() {
278        drop(unsafe { CString::from_raw(c_error) });
279    }
280}
281
282/// A C-ABI for exporting an Arrow schema from a JSON file
283#[unsafe(no_mangle)]
284pub extern "C" fn arrow_rs_cdata_integration_export_schema_from_json(
285    c_json_name: *const c_char,
286    out: *mut FFI_ArrowSchema,
287) -> *mut c_char {
288    let r = cdata_integration_export_schema_from_json(c_json_name, out);
289    result_to_c_error(&r)
290}
291
292/// A C-ABI to compare an Arrow schema against a JSON file
293#[unsafe(no_mangle)]
294pub extern "C" fn arrow_rs_cdata_integration_import_schema_and_compare_to_json(
295    c_json_name: *const c_char,
296    c_schema: *mut FFI_ArrowSchema,
297) -> *mut c_char {
298    let r = cdata_integration_import_schema_and_compare_to_json(c_json_name, c_schema);
299    result_to_c_error(&r)
300}
301
302/// A C-ABI for exporting a RecordBatch from a JSON file
303#[unsafe(no_mangle)]
304pub extern "C" fn arrow_rs_cdata_integration_export_batch_from_json(
305    c_json_name: *const c_char,
306    batch_num: c_int,
307    out: *mut FFI_ArrowArray,
308) -> *mut c_char {
309    let r = cdata_integration_export_batch_from_json(c_json_name, batch_num, out);
310    result_to_c_error(&r)
311}
312
313/// A C-ABI to compare a RecordBatch against a JSON file
314#[unsafe(no_mangle)]
315pub extern "C" fn arrow_rs_cdata_integration_import_batch_and_compare_to_json(
316    c_json_name: *const c_char,
317    batch_num: c_int,
318    c_array: *mut FFI_ArrowArray,
319) -> *mut c_char {
320    let r = cdata_integration_import_batch_and_compare_to_json(c_json_name, batch_num, c_array);
321    result_to_c_error(&r)
322}
arrow_integration_testing/lib.rs

arrow_integration_testing/
lib.rs
