From 1251e0952b3849e7d0c3c42ce192d6248cd18b5d Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Fri, 18 Feb 2022 05:12:56 +0000 Subject: [PATCH] Moved functions --- arrow-odbc-integration-testing/src/lib.rs | 140 +-------------------- arrow-odbc-integration-testing/src/read.rs | 139 ++++++++++++++++++++ 2 files changed, 143 insertions(+), 136 deletions(-) create mode 100644 arrow-odbc-integration-testing/src/read.rs diff --git a/arrow-odbc-integration-testing/src/lib.rs b/arrow-odbc-integration-testing/src/lib.rs index 373211850c6..124c57c3620 100644 --- a/arrow-odbc-integration-testing/src/lib.rs +++ b/arrow-odbc-integration-testing/src/lib.rs @@ -1,12 +1,9 @@ #![cfg(test)] -use arrow2::array::{Array, BinaryArray, BooleanArray, Int32Array, Utf8Array}; -use arrow2::chunk::Chunk; -use arrow2::datatypes::Field; -use arrow2::error::Result; -use arrow2::io::odbc::api::{Connection, Cursor, Environment, Error as OdbcError}; -use arrow2::io::odbc::{buffer_from_metadata, deserialize, infer_schema}; + +mod read; + +use arrow2::io::odbc::api::{Connection, Environment, Error as OdbcError}; use lazy_static::lazy_static; -use stdext::function_name; lazy_static! { /// This is an example for using doc comment attributes @@ -17,135 +14,6 @@ lazy_static! { const MSSQL: &str = "Driver={ODBC Driver 17 for SQL Server};Server=localhost;UID=SA;PWD=My@Test@Password1;"; -#[test] -fn int() -> Result<()> { - let table_name = function_name!().rsplit_once(':').unwrap().1; - let expected = vec![Chunk::new(vec![Box::new(Int32Array::from_slice([1])) as _])]; - - test(expected, "INT", "(1)", table_name) -} - -#[test] -fn int_nullable() -> Result<()> { - let table_name = function_name!().rsplit_once(':').unwrap().1; - let expected = vec![Chunk::new(vec![ - Box::new(Int32Array::from([Some(1), None])) as _, - ])]; - - test(expected, "INT", "(1),(NULL)", table_name) -} - -#[test] -fn bool() -> Result<()> { - let table_name = function_name!().rsplit_once(':').unwrap().1; - let expected = vec![Chunk::new(vec![ - Box::new(BooleanArray::from_slice([true])) as _ - ])]; - - test(expected, "BIT", "(1)", table_name) -} - -#[test] -fn bool_nullable() -> Result<()> { - let table_name = function_name!().rsplit_once(':').unwrap().1; - let expected = vec![Chunk::new(vec![ - Box::new(BooleanArray::from([Some(true), None])) as _, - ])]; - - test(expected, "BIT", "(1),(NULL)", table_name) -} - -#[test] -fn binary() -> Result<()> { - let table_name = function_name!().rsplit_once(':').unwrap().1; - let expected = vec![Chunk::new(vec![ - Box::new(BinaryArray::::from([Some(b"ab")])) as _, - ])]; - - test( - expected, - "VARBINARY(2)", - "(CAST('ab' AS VARBINARY(2)))", - table_name, - ) -} - -#[test] -fn binary_nullable() -> Result<()> { - let table_name = function_name!().rsplit_once(':').unwrap().1; - let expected = - vec![Chunk::new(vec![ - Box::new(BinaryArray::::from([Some(b"ab"), None, Some(b"ac")])) as _, - ])]; - - test( - expected, - "VARBINARY(2)", - "(CAST('ab' AS VARBINARY(2))),(NULL),(CAST('ac' AS VARBINARY(2)))", - table_name, - ) -} - -#[test] -fn utf8_nullable() -> Result<()> { - let table_name = function_name!().rsplit_once(':').unwrap().1; - let expected = - vec![Chunk::new(vec![ - Box::new(Utf8Array::::from([Some("ab"), None, Some("ac")])) as _, - ])]; - - test(expected, "VARCHAR(2)", "('ab'),(NULL),('ac')", table_name) -} - -fn test( - expected: Vec>>, - type_: &str, - insert: &str, - table_name: &str, -) -> Result<()> { - let connection = ENV.connect_with_connection_string(MSSQL).unwrap(); - setup_empty_table(&connection, table_name, &[type_]).unwrap(); - connection - .execute(&format!("INSERT INTO {table_name} (a) VALUES {insert}"), ()) - .unwrap(); - - // When - let query = format!("SELECT a FROM {table_name} ORDER BY id"); - - let chunks = read(&connection, &query)?.1; - - assert_eq!(chunks, expected); - Ok(()) -} - -fn read( - connection: &Connection<'_>, - query: &str, -) -> Result<(Vec, Vec>>)> { - let mut a = connection.prepare(query).unwrap(); - let fields = infer_schema(&a)?; - - let max_batch_size = 100; - let buffer = buffer_from_metadata(&a, max_batch_size).unwrap(); - - let cursor = a.execute(()).unwrap().unwrap(); - let mut cursor = cursor.bind_buffer(buffer).unwrap(); - - let mut chunks = vec![]; - while let Some(batch) = cursor.fetch().unwrap() { - let arrays = (0..batch.num_cols()) - .zip(fields.iter()) - .map(|(index, field)| { - let column_view = batch.column(index); - deserialize(column_view, field.data_type.clone()) - }) - .collect::>(); - chunks.push(Chunk::new(arrays)); - } - - Ok((fields, chunks)) -} - /// Creates the table and assures it is empty. Columns are named a,b,c, etc. pub fn setup_empty_table( conn: &Connection<'_>, diff --git a/arrow-odbc-integration-testing/src/read.rs b/arrow-odbc-integration-testing/src/read.rs new file mode 100644 index 00000000000..2f68b29d483 --- /dev/null +++ b/arrow-odbc-integration-testing/src/read.rs @@ -0,0 +1,139 @@ +use stdext::function_name; + +use arrow2::array::{Array, BinaryArray, BooleanArray, Int32Array, Utf8Array}; +use arrow2::chunk::Chunk; +use arrow2::datatypes::Field; +use arrow2::error::Result; +use arrow2::io::odbc::api::{Connection, Cursor}; +use arrow2::io::odbc::{buffer_from_metadata, deserialize, infer_schema}; + +use super::{setup_empty_table, ENV, MSSQL}; + +#[test] +fn int() -> Result<()> { + let table_name = function_name!().rsplit_once(':').unwrap().1; + let expected = vec![Chunk::new(vec![Box::new(Int32Array::from_slice([1])) as _])]; + + test(expected, "INT", "(1)", table_name) +} + +#[test] +fn int_nullable() -> Result<()> { + let table_name = function_name!().rsplit_once(':').unwrap().1; + let expected = vec![Chunk::new(vec![ + Box::new(Int32Array::from([Some(1), None])) as _, + ])]; + + test(expected, "INT", "(1),(NULL)", table_name) +} + +#[test] +fn bool() -> Result<()> { + let table_name = function_name!().rsplit_once(':').unwrap().1; + let expected = vec![Chunk::new(vec![ + Box::new(BooleanArray::from_slice([true])) as _ + ])]; + + test(expected, "BIT", "(1)", table_name) +} + +#[test] +fn bool_nullable() -> Result<()> { + let table_name = function_name!().rsplit_once(':').unwrap().1; + let expected = vec![Chunk::new(vec![ + Box::new(BooleanArray::from([Some(true), None])) as _, + ])]; + + test(expected, "BIT", "(1),(NULL)", table_name) +} + +#[test] +fn binary() -> Result<()> { + let table_name = function_name!().rsplit_once(':').unwrap().1; + let expected = vec![Chunk::new(vec![ + Box::new(BinaryArray::::from([Some(b"ab")])) as _, + ])]; + + test( + expected, + "VARBINARY(2)", + "(CAST('ab' AS VARBINARY(2)))", + table_name, + ) +} + +#[test] +fn binary_nullable() -> Result<()> { + let table_name = function_name!().rsplit_once(':').unwrap().1; + let expected = + vec![Chunk::new(vec![ + Box::new(BinaryArray::::from([Some(b"ab"), None, Some(b"ac")])) as _, + ])]; + + test( + expected, + "VARBINARY(2)", + "(CAST('ab' AS VARBINARY(2))),(NULL),(CAST('ac' AS VARBINARY(2)))", + table_name, + ) +} + +#[test] +fn utf8_nullable() -> Result<()> { + let table_name = function_name!().rsplit_once(':').unwrap().1; + let expected = + vec![Chunk::new(vec![ + Box::new(Utf8Array::::from([Some("ab"), None, Some("ac")])) as _, + ])]; + + test(expected, "VARCHAR(2)", "('ab'),(NULL),('ac')", table_name) +} + +fn test( + expected: Vec>>, + type_: &str, + insert: &str, + table_name: &str, +) -> Result<()> { + let connection = ENV.connect_with_connection_string(MSSQL).unwrap(); + setup_empty_table(&connection, table_name, &[type_]).unwrap(); + connection + .execute(&format!("INSERT INTO {table_name} (a) VALUES {insert}"), ()) + .unwrap(); + + // When + let query = format!("SELECT a FROM {table_name} ORDER BY id"); + + let chunks = read(&connection, &query)?.1; + + assert_eq!(chunks, expected); + Ok(()) +} + +fn read( + connection: &Connection<'_>, + query: &str, +) -> Result<(Vec, Vec>>)> { + let mut a = connection.prepare(query).unwrap(); + let fields = infer_schema(&a)?; + + let max_batch_size = 100; + let buffer = buffer_from_metadata(&a, max_batch_size).unwrap(); + + let cursor = a.execute(()).unwrap().unwrap(); + let mut cursor = cursor.bind_buffer(buffer).unwrap(); + + let mut chunks = vec![]; + while let Some(batch) = cursor.fetch().unwrap() { + let arrays = (0..batch.num_cols()) + .zip(fields.iter()) + .map(|(index, field)| { + let column_view = batch.column(index); + deserialize(column_view, field.data_type.clone()) + }) + .collect::>(); + chunks.push(Chunk::new(arrays)); + } + + Ok((fields, chunks)) +}