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

arrow_integration_testing/
lib.rs
