Fixing merge issues + use spiceai_duckdb_fork

Author: phillipleblanc

Cargo.lock

@@ -48,4 +48,4 @@ datafusion-proto = { version = "46" }
 datafusion-physical-expr = { version = "46" }
 datafusion-physical-plan = { version = "46" }
 datafusion-table-providers = { path = "core" }
-duckdb = { version = "=1.2.1" }
+duckdb = { version = "=1.2.1", package = "spiceai_duckdb_fork" } # Forked to add support for duckdb_scan_arrow, pending: https://github.com/duckdb/duckdb-rs/pull/488

Cargo.toml

@@ -109,6 +109,7 @@ rstest = "0.25.0"
 test-log = { version = "0.2", features = ["trace"] }
 tokio-stream = { version = "0.1", features = ["net"] }
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+tempfile = "3.19.1"
 
 [features]
 duckdb = [

core/Cargo.toml

@@ -1,5 +1,4 @@
 use crate::sql::sql_provider_datafusion;
-use crate::util::constraints::get_primary_keys_from_constraints;
 use crate::util::{
     self,
     column_reference::{self, ColumnReference},
@@ -36,10 +35,7 @@ use datafusion::{
 use duckdb::{AccessMode, DuckdbConnectionManager};
 use itertools::Itertools;
 use snafu::prelude::*;
-use std::{
-    collections::{HashMap, HashSet},
-    sync::Arc,
-};
+use std::{collections::HashMap, sync::Arc};
 use tokio::sync::Mutex;
 use write::DuckDBTableWriterBuilder;
 
@@ -404,29 +400,6 @@ impl TableProviderFactory for DuckDBTableProviderFactory {
             .with_pool(pool)
             .set_on_conflict(on_conflict);
 
-        // If the table is already created, we don't create it again and don't apply primary keys and remove previosly created indexes (if any).
-        // Thus we verify that primary keys and indexes for the table created match the configuration.
-        let mut table_schema_matches = true;
-        table_schema_matches &= duckdb
-            .verify_primary_keys_match()
-            .await
-            .map_err(to_datafusion_error)?;
-
-        table_schema_matches &= duckdb
-            .verify_indexes_match(&indexes)
-            .await
-            .map_err(to_datafusion_error)?;
-
-        if !table_schema_matches {
-            tracing::warn!(
-                "Schema mismatch detected for table '{table_name}' in database '{db_path}'.\n\
-         The local table definition does not match the expected schema.\n\
-         To resolve this issue, drop the existing table. A new table with the correct schema will be created automatically on the next access.",
-                db_path = duckdb.pool.db_path(),
-                table_name = duckdb.table_name
-            );
-        }
-
         let dyn_pool: Arc<DynDuckDbConnectionPool> = Arc::new(read_pool);
 
         if let Some(memory_limit) = options.get(DUCKDB_SETTING_MEMORY_LIMIT) {

core/src/duckdb.rs

@@ -840,6 +840,7 @@ pub(crate) mod tests {
                 Arc::clone(&table_definition),
                 *overwrite,
                 None,
+                table_definition.schema(),
             );
             let data_sink: Arc<dyn DataSink> = Arc::new(duckdb_sink);
             let rows_written = data_sink
@@ -953,6 +954,7 @@ pub(crate) mod tests {
                 Arc::clone(&table_definition),
                 *overwrite,
                 None,
+                table_definition.schema(),
             );
             let data_sink: Arc<dyn DataSink> = Arc::new(duckdb_sink);
             let rows_written = data_sink

core/src/duckdb/creator.rs

@@ -13,7 +13,6 @@ use arrow::ffi_stream::FFI_ArrowArrayStream;
 use arrow::{array::RecordBatch, datatypes::SchemaRef};
 use arrow_schema::ArrowError;
 use async_trait::async_trait;
-use datafusion::arrow::{array::RecordBatch, datatypes::SchemaRef};
 use datafusion::catalog::Session;
 use datafusion::common::{Constraints, SchemaExt};
 use datafusion::logical_expr::dml::InsertOp;
@@ -178,7 +177,7 @@ impl TableProvider for DuckDBTableWriter {
             Arc::new(DuckDBDataSink::new(
                 Arc::clone(&self.pool),
                 Arc::clone(&self.table_definition),
-                overwrite,
+                op,
                 self.on_conflict.clone(),
                 self.schema(),
             )),
@@ -649,7 +648,7 @@ impl RecordBatchReader for RecordBatchReaderFromStream {
 #[cfg(test)]
 mod test {
     use arrow::array::{Int64Array, StringArray};
-    use datafusion_physical_plan::memory::MemoryStream;
+    use datafusion::physical_plan::memory::MemoryStream;
 
     use super::*;
     use crate::{
@@ -672,6 +671,7 @@ mod test {
             Arc::clone(&table_definition),
             InsertOp::Overwrite,
             None,
+            table_definition.schema(),
         );
         let data_sink: Arc<dyn DataSink> = Arc::new(duckdb_sink);
 
@@ -771,6 +771,7 @@ mod test {
             Arc::clone(&table_definition),
             InsertOp::Overwrite,
             None,
+            table_definition.schema(),
         );
         let data_sink: Arc<dyn DataSink> = Arc::new(duckdb_sink);
 
@@ -876,6 +877,7 @@ mod test {
             Arc::clone(&table_definition),
             InsertOp::Overwrite,
             None,
+            table_definition.schema(),
         );
         let data_sink: Arc<dyn DataSink> = Arc::new(duckdb_sink);
 
@@ -975,6 +977,7 @@ mod test {
             Arc::clone(&table_definition),
             InsertOp::Append,
             None,
+            table_definition.schema(),
         );
         let data_sink: Arc<dyn DataSink> = Arc::new(duckdb_sink);
 
@@ -1075,6 +1078,7 @@ mod test {
             Arc::clone(&table_definition),
             InsertOp::Append,
             None,
+            table_definition.schema(),
         );
         let data_sink: Arc<dyn DataSink> = Arc::new(duckdb_sink);
 

core/src/duckdb/write.rs

@@ -69,10 +69,8 @@ impl MySQLTable {
         Ok(Arc::new(MySQLSQLExec::new(
             projections,
             schema,
-            &self.base_table.table_reference,
             Arc::clone(&self.pool),
-            filters,
-            limit,
+            sql,
         )?))
     }
 }
@@ -123,10 +121,8 @@ impl MySQLSQLExec {
     fn new(
         projections: Option<&Vec<usize>>,
         schema: &SchemaRef,
-        table_reference: &TableReference,
         pool: Arc<MySQLConnectionPool>,
-        filters: &[Expr],
-        limit: Option<usize>,
+        sql: String,
     ) -> DataFusionResult<Self> {
         let base_exec = SqlExec::new(projections, schema, pool, sql)?;
 

core/src/mysql/sql_table.rs

@@ -1,9 +1,9 @@
-use datafusion::execution::context::SessionContext;
+use datafusion::{datasource::memory::MemorySourceConfig, execution::context::SessionContext};
 use datafusion_table_providers::sql::{
     db_connection_pool::DbConnectionPool, sql_provider_datafusion::SqlTable,
 };
 use mysql_async::prelude::ToValue;
-use rstest::rstest;
+use rstest::{fixture, rstest};
 use std::sync::Arc;
 
 use arrow::{
@@ -13,6 +13,7 @@ use arrow::{
 
 use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection;
 
+use crate::arrow_record_batch_gen::*;
 use crate::docker::RunningContainer;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::catalog::TableProviderFactory;