feat: Add DuckDB checks for unsupported column types

Author: peasee

src/duckdb.rs

@@ -1,14 +1,3 @@
-use crate::sql::db_connection_pool::{
-    self,
-    dbconnection::{
-        duckdbconn::{
-            flatten_table_function_name, is_table_function, DuckDBParameter, DuckDbConnection,
-        },
-        get_schema, DbConnection,
-    },
-    duckdbpool::DuckDbConnectionPool,
-    DbConnectionPool, DbInstanceKey, Mode,
-};
 use crate::sql::sql_provider_datafusion;
 use crate::util::{
     self,
@@ -17,6 +6,20 @@ use crate::util::{
     indexes::IndexType,
     on_conflict::{self, OnConflict},
 };
+use crate::{
+    sql::db_connection_pool::{
+        self,
+        dbconnection::{
+            duckdbconn::{
+                flatten_table_function_name, is_table_function, DuckDBParameter, DuckDbConnection,
+            },
+            get_schema, DbConnection,
+        },
+        duckdbpool::DuckDbConnectionPool,
+        DbConnectionPool, DbInstanceKey, Mode,
+    },
+    InvalidTypeAction,
+};
 use arrow::{array::RecordBatch, datatypes::SchemaRef};
 use async_trait::async_trait;
 use datafusion::{
@@ -120,6 +123,7 @@ type Result<T, E = Error> = std::result::Result<T, E>;
 pub struct DuckDBTableProviderFactory {
     access_mode: AccessMode,
     instances: Arc<Mutex<HashMap<DbInstanceKey, DuckDbConnectionPool>>>,
+    invalid_type_action: InvalidTypeAction,
 }
 
 const DUCKDB_DB_PATH_PARAM: &str = "open";
@@ -132,9 +136,16 @@ impl DuckDBTableProviderFactory {
         Self {
             access_mode,
             instances: Arc::new(Mutex::new(HashMap::new())),
+            invalid_type_action: InvalidTypeAction::Error,
         }
     }
 
+    #[must_use]
+    pub fn with_invalid_type_action(mut self, invalid_type_action: InvalidTypeAction) -> Self {
+        self.invalid_type_action = invalid_type_action;
+        self
+    }
+
     #[must_use]
     pub fn attach_databases(&self, options: &HashMap<String, String>) -> Vec<Arc<str>> {
         options
@@ -181,7 +192,9 @@ impl DuckDBTableProviderFactory {
             return Ok(instance.clone());
         }
 
-        let pool = DuckDbConnectionPool::new_memory().context(DbConnectionPoolSnafu)?;
+        let pool = DuckDbConnectionPool::new_memory()
+            .context(DbConnectionPoolSnafu)?
+            .with_invalid_type_action(self.invalid_type_action.clone());
 
         instances.insert(key, pool.clone());
 
@@ -201,7 +214,8 @@ impl DuckDBTableProviderFactory {
         }
 
         let pool = DuckDbConnectionPool::new_file(&db_path, &self.access_mode)
-            .context(DbConnectionPoolSnafu)?;
+            .context(DbConnectionPoolSnafu)?
+            .with_invalid_type_action(self.invalid_type_action.clone());
 
         instances.insert(key, pool.clone());
 

src/lib.rs

@@ -1,3 +1,4 @@
+use serde::{Deserialize, Serialize};
 use snafu::prelude::*;
 
 pub mod sql;
@@ -23,3 +24,11 @@ pub enum Error {
     #[snafu(display("Error reading file: {source}"))]
     FileReadError { source: std::io::Error },
 }
+
+#[derive(PartialEq, Eq, Clone, Debug, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum InvalidTypeAction {
+    Error,
+    Warn,
+    Ignore,
+}

src/sql/db_connection_pool/dbconnection.rs

@@ -1,5 +1,6 @@
 use std::{any::Any, sync::Arc};
 
+use arrow_schema::DataType;
 use datafusion::{
     arrow::datatypes::SchemaRef, execution::SendableRecordBatchStream, sql::TableReference,
 };
@@ -25,6 +26,12 @@ pub enum Error {
     #[snafu(display("Unable to get schema: {source}"))]
     UnableToGetSchema { source: GenericError },
 
+    #[snafu(display("The field '{field_name}' has an unsupported data type: {data_type}"))]
+    UnsupportedDataType {
+        data_type: DataType,
+        field_name: String,
+    },
+
     #[snafu(display("Unable to query arrow: {source}"))]
     UnableToQueryArrow { source: GenericError },