DuckDB: don't use ArrowVTab while creating the table (more types support) (#319)

Author: sgrebnov

core/src/duckdb.rs

@@ -156,6 +156,9 @@ pub enum Error {
     #[snafu(display("Failed to register Arrow scan view for DuckDB ingestion: {source}"))]
     UnableToRegisterArrowScanView { source: duckdb::Error },
 
+    #[snafu(display("Failed to register Arrow scan view to build table creation statement: {source}"))]
+    UnableToRegisterArrowScanViewForTableCreation { source: duckdb::Error },
+
     #[snafu(display("Failed to drop Arrow scan view for DuckDB ingestion: {source}"))]
     UnableToDropArrowScanView { source: duckdb::Error },
 }

core/src/duckdb/creator.rs

@@ -2,11 +2,15 @@ use crate::sql::arrow_sql_gen::statement::IndexBuilder;
 use crate::sql::db_connection_pool::dbconnection::duckdbconn::DuckDbConnection;
 use crate::sql::db_connection_pool::duckdbpool::DuckDbConnectionPool;
 use crate::util::on_conflict::OnConflict;
-use arrow::{array::RecordBatch, datatypes::SchemaRef};
+use arrow::{
+    array::{RecordBatch, RecordBatchIterator, RecordBatchReader},
+    datatypes::SchemaRef,
+    ffi_stream::FFI_ArrowArrayStream,
+};
 use datafusion::common::utils::quote_identifier;
 use datafusion::common::Constraints;
 use datafusion::sql::TableReference;
-use duckdb::{vtab::arrow_recordbatch_to_query_params, ToSql, Transaction};
+use duckdb::Transaction;
 use itertools::Itertools;
 use snafu::prelude::*;
 use std::collections::HashSet;
@@ -392,19 +396,25 @@ impl TableManager {
             .transaction()
             .context(super::UnableToBeginTransactionSnafu)?;
         let table_name = self.table_name();
-        let empty_record = RecordBatch::new_empty(Arc::clone(&self.table_definition.schema));
+        let record_batch_reader =
+            create_empty_record_batch_reader(Arc::clone(&self.table_definition.schema));
+        let stream = FFI_ArrowArrayStream::new(Box::new(record_batch_reader));
 
-        let arrow_params = arrow_recordbatch_to_query_params(empty_record);
-        let arrow_params_vec: Vec<&dyn ToSql> = arrow_params
-            .iter()
-            .map(|p| p as &dyn ToSql)
-            .collect::<Vec<_>>();
-        let arrow_params_ref: &[&dyn ToSql] = &arrow_params_vec;
-        let sql =
-            format!(r#"CREATE TABLE IF NOT EXISTS "{table_name}" AS SELECT * FROM arrow(?, ?)"#,);
+        let current_ts = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .context(super::UnableToGetSystemTimeSnafu)?
+            .as_millis();
+
+        let view_name = format!("__scan_{}_{current_ts}", table_name);
+        tx.register_arrow_scan_view(&view_name, &stream)
+            .context(super::UnableToRegisterArrowScanViewForTableCreationSnafu)?;
+
+        let sql = format!(
+            r#"CREATE TABLE IF NOT EXISTS "{table_name}" AS SELECT * FROM "{view_name}""#,
+        );
         tracing::debug!("{sql}");
 
-        tx.execute(&sql, arrow_params_ref)
+        tx.execute(&sql, [])
             .context(super::UnableToCreateDuckDBTableSnafu)?;
 
         let create_stmt = tx
@@ -670,6 +680,12 @@ impl TableManager {
     }
 }
 
+fn create_empty_record_batch_reader(schema: SchemaRef) -> impl RecordBatchReader {
+    let empty_batch = RecordBatch::new_empty(Arc::clone(&schema));
+    let batches = vec![empty_batch];
+    RecordBatchIterator::new(batches.into_iter().map(Ok), schema)
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct ViewCreator {
     name: RelationName,

core/tests/arrow_record_batch_gen/mod.rs

@@ -766,6 +766,60 @@ pub(crate) fn get_arrow_dictionary_array_record_batch() -> (RecordBatch, SchemaR
     (record_batch, schema)
 }
 
+pub(crate) fn get_arrow_map_record_batch() -> (RecordBatch, SchemaRef) {
+    let keys = vec!["a", "b", "c", "d", "e", "f", "g", "h"];
+    let values_data = UInt32Array::from(vec![
+        Some(0u32),
+        None,
+        Some(20),
+        Some(30),
+        None,
+        Some(50),
+        Some(60),
+        Some(70),
+    ]);
+    // Construct a buffer for value offsets, for the nested array:
+    //  [[a, b, c], [d, e, f], [g, h]]
+    let entry_offsets = [0, 3, 6, 8];
+    let map_array =
+        MapArray::new_from_strings(keys.clone().into_iter(), &values_data, &entry_offsets)
+            .expect("Failed to create MapArray");
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "map_array",
+        map_array.data_type().clone(),
+        true,
+    )]));
+    let rb = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(map_array)])
+        .expect("Failed to created arrow Map array record batch");
+    (rb, schema)
+}
+
+// Custom Test Case for Sqlite <-> Arrow Decimal Roundtrip
+// SQLite supports up to 16 precision for decimal numbers through REAL type, conforming to IEEE 754 Binary-64 format - https://www.sqlite.org/floatingpoint.html
+pub(crate) fn get_sqlite_arrow_decimal_record_batch() -> (RecordBatch, SchemaRef) {
+    let decimal128_array =
+        Decimal128Array::from(vec![i128::from(123), i128::from(222), i128::from(321)])
+            .with_precision_and_scale(16, 10)
+            .expect("Fail to create Decimal128 array");
+    let decimal256_array =
+        Decimal256Array::from(vec![i256::from(-123), i256::from(222), i256::from(0)])
+            .with_precision_and_scale(16, 10)
+            .expect("Fail to create Decimal256 array");
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("decimal128", DataType::Decimal128(16, 10), false),
+        Field::new("decimal256", DataType::Decimal256(16, 10), false),
+    ]));
+
+    let record_batch = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![Arc::new(decimal128_array), Arc::new(decimal256_array)],
+    )
+    .expect("Failed to created arrow decimal record batch");
+
+    (record_batch, schema)
+}
+
 fn parse_json_to_batch(json_data: &str, schema: SchemaRef) -> RecordBatch {
     let reader = arrow_json::ReaderBuilder::new(schema)
         .build(std::io::Cursor::new(json_data))

core/tests/duckdb/mod.rs

@@ -83,8 +83,6 @@ async fn arrow_duckdb_round_trip(
 }
 
 #[rstest]
-#[ignore]
-// Binder Error: Unsupported data type: FixedSizeBinary(2), please file an issue https://github.com/wangfenjin/duckdb-rs"
 #[case::binary(get_arrow_binary_record_batch(), "binary")]
 #[case::int(get_arrow_int_record_batch(), "int")]
 #[case::float(get_arrow_float_record_batch(), "float")]
@@ -93,11 +91,11 @@ async fn arrow_duckdb_round_trip(
 #[case::timestamp(get_arrow_timestamp_record_batch(), "timestamp")]
 #[case::date(get_arrow_date_record_batch(), "date")]
 #[case::struct_type(get_arrow_struct_record_batch(), "struct")]
-#[ignore] // Decimal256(76,10) is not yet supported for ArrowVTab
+#[ignore] // DuckDB does not support Decimal256 / duckdb_arrow_scan failed to register view
 #[case::decimal(get_arrow_decimal_record_batch(), "decimal")]
-#[ignore] // Interval(DayTime) is not yet supported for ArrowVTab
+#[ignore] // Interval(DayTime) is not supported: / "Conversion Error: Could not convert Interval to Microsecond"
 #[case::interval(get_arrow_interval_record_batch(), "interval")]
-#[ignore] // Duration(NanoSecond) is not yet supported for ArrowVTab
+#[ignore] // TimeUnit::Nanosecond is not correctly supported; written values are zeros
 #[case::duration(get_arrow_duration_record_batch(), "duration")]
 #[case::list(get_arrow_list_record_batch(), "list")]
 #[case::null(get_arrow_null_record_batch(), "null")]
@@ -107,6 +105,8 @@ async fn arrow_duckdb_round_trip(
     "list_of_fixed_size_lists"
 )]
 #[case::list_of_lists(get_arrow_list_of_lists_record_batch(), "list_of_lists")]
+#[case::map(get_arrow_map_record_batch(), "map")]
+#[case::dictionary(get_arrow_dictionary_array_record_batch(), "dictionary")]
 #[test_log::test(tokio::test)]
 async fn test_arrow_duckdb_roundtrip(
     #[case] arrow_result: (RecordBatch, SchemaRef),