Merge pull request #37 from wangfenjin/arrow-schema

add get_schema api

Author: wangfenjin

src/arrow_batch.rs

@@ -1,5 +1,5 @@
 use super::Statement;
-use arrow::record_batch::RecordBatch;
+use arrow::{datatypes::SchemaRef, record_batch::RecordBatch};
 
 /// An handle for the resulting RecordBatch of a query.
 #[must_use = "Arrow is lazy and will do nothing unless consumed"]
@@ -12,6 +12,12 @@ impl<'stmt> Arrow<'stmt> {
     pub(crate) fn new(stmt: &'stmt Statement<'stmt>) -> Arrow<'stmt> {
         Arrow { stmt: Some(stmt) }
     }
+
+    /// return arrow schema
+    #[inline]
+    pub fn get_schema(&self) -> SchemaRef {
+        self.stmt.unwrap().stmt.schema()
+    }
 }
 
 impl<'stmt> Iterator for Arrow<'stmt> {

src/column.rs

@@ -29,8 +29,13 @@ impl Statement<'_> {
     /// If associated DB schema can be altered concurrently, you should make
     /// sure that current statement has already been stepped once before
     /// calling this method.
-    pub fn column_names(&self) -> Vec<&String> {
-        self.stmt.schema().fields().iter().map(|f| f.name()).collect()
+    pub fn column_names(&self) -> Vec<String> {
+        self.stmt
+            .schema()
+            .fields()
+            .iter()
+            .map(|f| f.name().to_owned())
+            .collect()
     }
 
     /// Return the number of columns in the result set returned by the prepared

src/lib.rs

@@ -528,6 +528,7 @@ mod test {
     use std::fmt;
 
     use arrow::array::Int32Array;
+    use arrow::datatypes::DataType;
     use arrow::record_batch::RecordBatch;
     use fallible_iterator::FallibleIterator;
 
@@ -635,7 +636,7 @@ mod test {
         let filename = "no_such_file.db";
         let result =
             Connection::open_with_flags(filename, Config::default().access_mode(config::AccessMode::ReadOnly)?);
-        assert!(!result.is_ok());
+        assert!(result.is_err());
         let err = result.err().unwrap();
         if let Error::DuckDBFailure(_e, Some(msg)) = err {
             // TODO: update error code
@@ -1180,6 +1181,11 @@ mod test {
         let mut stmt = db.prepare("select t from test order by t desc")?;
         let mut arr = stmt.query_arrow([])?;
 
+        let schema = arr.get_schema();
+        assert_eq!(schema.fields().len(), 1);
+        assert_eq!(schema.field(0).name(), "t");
+        assert_eq!(schema.field(0).data_type(), &DataType::Int32);
+
         let rb = arr.next().unwrap();
         let column = rb.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
         assert_eq!(column.len(), 5);

src/raw_statement.rs

@@ -18,7 +18,6 @@ use arrow::ffi::{ArrowArray, FFI_ArrowArray, FFI_ArrowSchema};
 pub struct RawStatement {
     ptr: ffi::duckdb_prepared_statement,
     result: Option<ffi::duckdb_arrow>,
-    c_schema: Option<*const FFI_ArrowSchema>,
     schema: Option<SchemaRef>,
 }
 
@@ -28,7 +27,6 @@ impl RawStatement {
         RawStatement {
             ptr: stmt,
             result: None,
-            c_schema: None,
             schema: None,
         }
     }
@@ -60,11 +58,6 @@ impl RawStatement {
             let (mut arrays, mut schema) = ArrowArray::into_raw(ArrowArray::empty());
             let schema = &mut schema;
             let arrays = &mut arrays;
-            // TODO: Can we reuse schema?
-            // destroy schema as we don't need it...
-            // Arc::from_raw(schema);
-            // TODO: use this after https://github.com/apache/arrow-rs/pull/612
-            // let mut arrays = Arc::into_raw(Arc::new(FFI_ArrowArray::empty()));
             if ffi::duckdb_query_arrow_array(self.result_unwrap(), arrays as *mut _ as *mut *mut c_void)
                 != ffi::DuckDBSuccess
             {
@@ -104,8 +97,8 @@ impl RawStatement {
     }
 
     #[inline]
-    pub fn schema(&self) -> &SchemaRef {
-        self.schema.as_ref().unwrap()
+    pub fn schema(&self) -> SchemaRef {
+        self.schema.clone().unwrap()
     }
 
     #[inline]
@@ -161,14 +154,13 @@ impl RawStatement {
 
             let rows_changed = ffi::duckdb_arrow_rows_changed(out);
             let mut c_schema = Arc::into_raw(Arc::new(FFI_ArrowSchema::empty()));
-            let schema = &mut c_schema;
-            let rc = ffi::duckdb_query_arrow_schema(out, schema as *mut _ as *mut *mut c_void);
+            let rc = ffi::duckdb_query_arrow_schema(out, &mut c_schema as *mut _ as *mut *mut c_void);
             if rc != ffi::DuckDBSuccess {
                 Arc::from_raw(c_schema);
                 result_from_duckdb_arrow(rc, out)?;
             }
             self.schema = Some(Arc::new(Schema::try_from(&*c_schema).unwrap()));
-            self.c_schema = Some(c_schema);
+            Arc::from_raw(c_schema);
 
             self.result = Some(out);
             Ok(rows_changed as usize)
@@ -178,12 +170,6 @@ impl RawStatement {
     #[inline]
     pub fn reset_result(&mut self) {
         self.schema = None;
-        if self.c_schema.is_some() {
-            unsafe {
-                Arc::from_raw(self.c_schema.unwrap());
-            }
-            self.c_schema = None;
-        }
         if self.result.is_some() {
             unsafe {
                 ffi::duckdb_destroy_arrow(&mut self.result_unwrap());