Merge pull request #50 from Embucket/sql_extention

Support new datatypes and UDFs

Author: osipovartem

crates/catalog/src/models.rs

@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::fmt::Display;
 use uuid::Uuid;
 
 pub use iceberg::{Namespace, NamespaceIdent, TableCreation, TableRequirement, TableUpdate};

crates/catalog/src/repository.rs

@@ -1,10 +1,8 @@
 use async_trait::async_trait;
 use std::fmt::{Display, Formatter};
 use std::sync::Arc;
-use uuid::Uuid;
 
 use utils::Db;
-use utils::{Entity, Repository};
 
 use crate::error::Result;
 use crate::models::{Database, DatabaseIdent, Table, TableIdent, WarehouseIdent};
@@ -244,6 +242,7 @@ mod tests {
             },
             metadata_location: "s3://bucket/path".to_string(),
             metadata: create_table_metadata(),
+            properties: Default::default(),
         };
 
         repo.put(&table).await.expect("failed to create table");

crates/catalog/src/service.rs

@@ -9,7 +9,6 @@ use async_trait::async_trait;
 use bytes::Bytes;
 use chrono::Utc;
 use control_plane::models::{StorageProfile, Warehouse};
-use control_plane::service::ControlService;
 use iceberg::spec::{FormatVersion, TableMetadata};
 use iceberg::{spec::TableMetadataBuilder, TableCreation};
 use object_store::path::Path;
@@ -326,7 +325,7 @@ impl Catalog for CatalogImpl {
         &self,
         namespace: &DatabaseIdent,
         storage_profile: &StorageProfile,
-        warehouse: &Warehouse,
+        _warehouse: &Warehouse,
         table_name: String,
         metadata_location: String,
         properties: Option<HashMap<String, String>>,
@@ -871,7 +870,6 @@ mod tests {
             ident: res.unwrap().ident,
             requirements: vec![],
             updates: vec![update],
-            properties: None,
         };
 
         let res = service.update_table(commit).await;

crates/control_plane/Cargo.toml

@@ -13,9 +13,9 @@ utils = { path = "../utils" }
 futures = { workspace = true }
 serde = { workspace = true }
 datafusion = { workspace = true }
-iceberg-rust = { git = "https://github.com/JanKaul/iceberg-rust.git"}
-iceberg-rest-catalog = { git = "https://github.com/JanKaul/iceberg-rust.git"}
-datafusion_iceberg = { git = "https://github.com/JanKaul/iceberg-rust.git"}
+iceberg-rust = { git = "https://github.com/JanKaul/iceberg-rust.git" }
+iceberg-rest-catalog = { git = "https://github.com/JanKaul/iceberg-rust.git" }
+datafusion_iceberg = { git = "https://github.com/JanKaul/iceberg-rust.git" }
 arrow = { version = "53" }
 arrow-json = { version = "53" }
 datafusion-functions-json = { version = "0.43.0" }
@@ -28,6 +28,9 @@ quick-xml = { version = "0.36.2" }
 icelake = { git = "https://github.com/Embucket/icelake.git" }
 bytes = { version = "1.8.0" }
 url = { version = "2.5.2" }
+serde_json = "1.0.128"
+once_cell = "1.19.0"
+regex = "1.11.0"
 
 [dev-dependencies]
 slatedb = { workspace = true }

crates/control_plane/src/service.rs

@@ -1,4 +1,3 @@
-use datafusion_iceberg::planner::IcebergQueryPlanner;
 use crate::error::{extract_error_message, Error, Result};
 use crate::models::{Credentials, StorageProfile, StorageProfileCreateRequest};
 use crate::models::{Warehouse, WarehouseCreateRequest};
@@ -10,7 +9,6 @@ use async_trait::async_trait;
 use bytes::Bytes;
 use datafusion::execution::context::SessionContext;
 use datafusion::prelude::CsvReadOptions;
-use datafusion::execution::SessionStateBuilder;
 use datafusion_iceberg::catalog::catalog::IcebergCatalog;
 use iceberg_rest_catalog::apis::configuration::Configuration;
 use iceberg_rest_catalog::catalog::RestCatalog;
@@ -20,7 +18,6 @@ use object_store::{ObjectStore, PutPayload};
 use rusoto_core::{HttpClient, Region};
 use rusoto_credential::StaticProvider;
 use rusoto_s3::{GetBucketAclRequest, S3Client, S3};
-use std::any::Any;
 use std::collections::HashMap;
 use std::sync::Arc;
 use url::Url;
@@ -174,7 +171,7 @@ impl ControlService for ControlServiceImpl {
         &self,
         warehouse_id: &Uuid,
         database_name: &String,
-        table_name: &String,
+        _table_name: &String,
         query: &String,
     ) -> Result<String> {
         let warehouse = self.get_warehouse(*warehouse_id).await?;
@@ -266,12 +263,9 @@ impl ControlService for ControlServiceImpl {
         };
         let endpoint_url = Url::parse(storage_profile.endpoint.clone().unwrap().as_str())
             .map_err(|e| Error::DataFusionError(format!("Invalid endpoint URL: {}", e)))?;
-
         ctx.register_object_store(&endpoint_url, Arc::from(object_store));
-        let df = ctx.read_csv(path_string, CsvReadOptions::new()).await?;
-        let data = df.collect().await?;
 
-        println!("{:?}", data);
+        // println!("{:?}", data);
         // Commented code is writing with iceberg-rust-jankaul
         // Let it sit here just in case
         //////////////////////////////////////
@@ -337,22 +331,22 @@ impl ControlService for ControlServiceImpl {
                 // ),
             ])
         };
-        let catalog = icelake::catalog::load_catalog(&config).await.unwrap();
-
-        let table_ident = TableIdentifier::new(vec![database_name, table_name]).unwrap();
-        let mut table = catalog.load_table(&table_ident).await.unwrap();
+        let catalog = icelake::catalog::load_catalog(&config).await?;
+        let table_ident = TableIdentifier::new(vec![database_name, table_name])?;
+        let mut table = catalog.load_table(&table_ident).await?;
+        let table_schema = table.current_arrow_schema()?;
         println!("{:?}", table.table_name());
+
+        let df = ctx.read_csv(path_string, CsvReadOptions::new().schema(&*table_schema)).await?;
+        let data = df.collect().await?;
+
         let builder = table
-            .writer_builder()
-            .unwrap()
-            .rolling_writer_builder(None)
-            .unwrap();
+            .writer_builder()?
+            .rolling_writer_builder(None)?;
         let mut writer = table
-            .writer_builder()
-            .unwrap()
+            .writer_builder()?
             .build_append_only_writer(builder)
-            .await
-            .unwrap();
+            .await?;
 
         for r in data {
             writer.write(&r).await?;
@@ -361,7 +355,7 @@ impl ControlService for ControlServiceImpl {
         let res: Vec<icelake::types::DataFile> = writer.close().await?;
         let mut txn = icelake::transaction::Transaction::new(&mut table);
         txn.append_data_file(res);
-        txn.commit().await.unwrap();
+        txn.commit().await?;
 
         Ok(())
     }

crates/control_plane/src/sql/functions/common.rs

@@ -17,7 +17,6 @@ pub fn first_non_empty_type(union_array: &UnionArray) -> Option<(DataType, Array
     None
 }
 
-
 pub fn convert_record_batches(records: Vec<RecordBatch>) -> DataFusionResult<Vec<RecordBatch>> {
     let mut converted_batches = Vec::new();