use arrow2::array::{Array, Int32Array, Utf8Array};
use arrow2::chunk::Chunk;
use arrow2::datatypes::{DataType, Field};
use arrow2::error::Result;
use arrow2::io::odbc::api;
use arrow2::io::odbc::api::Cursor;
use arrow2::io::odbc::read;
use arrow2::io::odbc::write;
fn main() -> Result<()> {
let connector = "Driver={SQLite3};Database=sqlite-test.db";
let env = api::Environment::new()?;
let connection = env.connect_with_connection_string(connector)?;
connection.execute("DROP TABLE IF EXISTS example;", ())?;
connection.execute("CREATE TABLE example (c1 INT, c2 TEXT);", ())?;
let query = "INSERT INTO example (c1, c2) VALUES (?, ?)";
let prepared = connection.prepare(query).unwrap();
let fields = vec![
Field::new("unused", DataType::Int32, true),
Field::new("unused", DataType::LargeUtf8, true),
];
let mut writer = write::Writer::try_new(prepared, fields)?;
let chunk = Chunk::new(vec![
Box::new(Int32Array::from_slice([1, 2, 3])) as Box<dyn Array>,
Box::new(Utf8Array::<i64>::from([Some("Hello"), None, Some("World")])),
]);
writer.write(&chunk)?;
let chunks = read(&connection, "SELECT c1 FROM example")?;
assert_eq!(chunks[0].columns()[0], chunk.columns()[0]);
Ok(())
}
pub fn read(connection: &api::Connection<'_>, query: &str) -> Result<Vec<Chunk<Box<dyn Array>>>> {
let mut a = connection.prepare(query)?;
let fields = read::infer_schema(&a)?;
let max_batch_size = 100;
let buffer = read::buffer_from_metadata(&a, max_batch_size)?;
let cursor = a.execute(())?.unwrap();
let mut cursor = cursor.bind_buffer(buffer)?;
let mut chunks = vec![];
while let Some(batch) = cursor.fetch()? {
let arrays = (0..batch.num_cols())
.zip(fields.iter())
.map(|(index, field)| {
let column_view = batch.column(index);
read::deserialize(column_view, field.data_type.clone())
})
.collect::<Vec<_>>();
chunks.push(Chunk::new(arrays));
}
Ok(chunks)
}