iceberg 0.11.2 → 0.12.1

data/ext/iceberg/src/catalog.rs

@@ -1,9 +1,15 @@
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+
 #[cfg(feature = "datafusion")]
-use datafusion::execution::context::SessionContext;
-use iceberg::io::LocalFsStorageFactory;
+use datafusion::common::ScalarValue;
+#[cfg(feature = "datafusion")]
+use datafusion::execution::context::{SessionConfig, SessionContext};
 use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
-use iceberg::spec::Schema;
-use iceberg::{Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, TableCreation, TableIdent};
+use iceberg::spec::FormatVersion;
+use iceberg::{
+    Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Runtime, TableCreation, TableIdent,
+};
 #[cfg(feature = "glue")]
 use iceberg_catalog_glue::{GLUE_CATALOG_PROP_WAREHOUSE, GlueCatalog, GlueCatalogBuilder};
 #[cfg(feature = "rest")]
@@ -21,13 +27,23 @@ use iceberg_catalog_sql::{
 };
 #[cfg(feature = "datafusion")]
 use iceberg_datafusion::IcebergCatalogProvider;
-use std::collections::HashMap;
-use std::sync::{Arc, RwLock};
+use iceberg_storage_opendal::OpenDalStorageFactory;
+use magnus::{Error as RbErr, Ruby};
+#[cfg(feature = "datafusion")]
+use magnus::{
+    Float, Integer, RArray, RString, Value, value::Qfalse, value::Qtrue, value::ReprValue,
+};
 
 use crate::error::to_rb_err;
-use crate::runtime::runtime;
+#[cfg(feature = "datafusion")]
+use crate::error::{datafusion_error, todo_error};
+#[cfg(feature = "datafusion")]
+use crate::result::collect_batches;
+use crate::runtime::{runtime, tokio_runtime};
 use crate::utils::Wrap;
-use crate::{RbResult, RbTable};
+#[cfg(feature = "datafusion")]
+use crate::utils::date_to_i32;
+use crate::{RbPartitionSpec, RbResult, RbSchema, RbSortOrder, RbTable};
 
 pub enum RbCatalogType {
     #[cfg(feature = "glue")]
@@ -82,7 +98,11 @@ impl RbCatalog {
     pub fn new_glue(warehouse: String) -> RbResult<Self> {
         let props = HashMap::from([(GLUE_CATALOG_PROP_WAREHOUSE.to_string(), warehouse)]);
         let catalog = runtime()
-            .block_on(GlueCatalogBuilder::default().load("glue", props))
+            .block_on(
+                GlueCatalogBuilder::default()
+                    .with_runtime(Runtime::new(tokio_runtime()))
+                    .load("glue", props),
+            )
             .map_err(to_rb_err)?;
         Ok(Self {
             catalog: RbCatalogType::Glue(catalog.into()).into(),
@@ -91,13 +111,14 @@ impl RbCatalog {
 
     pub fn new_memory(warehouse: Option<String>) -> RbResult<Self> {
         let mut props = HashMap::new();
-        if let Some(v) = warehouse {
-            props.insert(MEMORY_CATALOG_WAREHOUSE.to_string(), v);
+        if let Some(ref v) = warehouse {
+            props.insert(MEMORY_CATALOG_WAREHOUSE.to_string(), v.clone());
         }
         let catalog = runtime()
             .block_on(
                 MemoryCatalogBuilder::default()
-                    .with_storage_factory(Arc::new(LocalFsStorageFactory))
+                    .with_runtime(Runtime::new(tokio_runtime()))
+                    .with_storage_factory(Arc::new(storage_factory(warehouse)))
                     .load("memory", props),
             )
             .map_err(to_rb_err)?;
@@ -114,13 +135,14 @@ impl RbCatalog {
     ) -> RbResult<Self> {
         let mut props = props;
         props.insert(REST_CATALOG_PROP_URI.to_string(), uri);
-        if let Some(v) = warehouse {
-            props.insert(REST_CATALOG_PROP_WAREHOUSE.to_string(), v);
+        if let Some(ref v) = warehouse {
+            props.insert(REST_CATALOG_PROP_WAREHOUSE.to_string(), v.clone());
         }
         let catalog = runtime()
             .block_on(
                 RestCatalogBuilder::default()
-                    .with_storage_factory(Arc::new(LocalFsStorageFactory))
+                    .with_runtime(Runtime::new(tokio_runtime()))
+                    .with_storage_factory(Arc::new(storage_factory(warehouse)))
                     .load("rest", props),
             )
             .map_err(to_rb_err)?;
@@ -134,7 +156,11 @@ impl RbCatalog {
         let mut props = HashMap::new();
         props.insert(S3TABLES_CATALOG_PROP_TABLE_BUCKET_ARN.to_string(), arn);
         let catalog = runtime()
-            .block_on(S3TablesCatalogBuilder::default().load("s3tables", props))
+            .block_on(
+                S3TablesCatalogBuilder::default()
+                    .with_runtime(Runtime::new(tokio_runtime()))
+                    .load("s3tables", props),
+            )
             .map_err(to_rb_err)?;
         Ok(Self {
             catalog: RbCatalogType::S3Tables(catalog.into()).into(),
@@ -150,7 +176,7 @@ impl RbCatalog {
     ) -> RbResult<Self> {
         let mut props = props;
         props.insert(SQL_CATALOG_PROP_URI.to_string(), uri);
-        props.insert(SQL_CATALOG_PROP_WAREHOUSE.to_string(), warehouse);
+        props.insert(SQL_CATALOG_PROP_WAREHOUSE.to_string(), warehouse.clone());
         props.insert(
             SQL_CATALOG_PROP_BIND_STYLE.to_string(),
             SqlBindStyle::DollarNumeric.to_string(),
@@ -158,7 +184,8 @@ impl RbCatalog {
         let catalog = runtime()
             .block_on(
                 SqlCatalogBuilder::default()
-                    .with_storage_factory(Arc::new(LocalFsStorageFactory))
+                    .with_runtime(Runtime::new(tokio_runtime()))
+                    .with_storage_factory(Arc::new(storage_factory(Some(warehouse))))
                     .load(name, props),
             )
             .map_err(to_rb_err)?;
@@ -279,20 +306,42 @@ impl RbCatalog {
             .collect())
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn create_table(
-        &self,
+        ruby: &Ruby,
+        rb_self: &Self,
         name: Wrap<TableIdent>,
-        schema: Wrap<Schema>,
+        schema: &RbSchema,
         location: Option<String>,
+        partition_spec: Option<&RbPartitionSpec>,
+        sort_order: Option<&RbSortOrder>,
+        properties: HashMap<String, String>,
+        format_version: u8,
     ) -> RbResult<RbTable> {
+        let format_version = match format_version {
+            1 => FormatVersion::V1,
+            2 => FormatVersion::V2,
+            3 => FormatVersion::V3,
+            _ => {
+                return Err(RbErr::new(
+                    ruby.exception_arg_error(),
+                    "Unsupported format version",
+                ));
+            }
+        };
         let creation = TableCreation::builder()
             .name(name.0.name)
-            .schema(schema.0)
+            .schema(schema.schema.clone())
             .location_opt(location)
+            .partition_spec_opt(partition_spec.map(|v| v.spec.clone()))
+            .sort_order_opt(sort_order.map(|v| v.order.clone()))
+            .properties(properties)
+            .format_version(format_version)
             .build();
         let table = runtime()
             .block_on(
-                self.catalog
+                rb_self
+                    .catalog
                     .read()
                     .unwrap()
                     .as_catalog()
@@ -332,6 +381,19 @@ impl RbCatalog {
         Ok(())
     }
 
+    pub fn purge_table(&self, name: Wrap<TableIdent>) -> RbResult<()> {
+        runtime()
+            .block_on(
+                self.catalog
+                    .read()
+                    .unwrap()
+                    .as_catalog()
+                    .purge_table(&name.0),
+            )
+            .map_err(to_rb_err)?;
+        Ok(())
+    }
+
     pub fn table_exists(&self, name: Wrap<TableIdent>) -> RbResult<bool> {
         let exists = runtime()
             .block_on(
@@ -376,22 +438,84 @@ impl RbCatalog {
     }
 
     #[cfg(feature = "datafusion")]
-    pub fn sql(&self, sql: String) -> RbResult<()> {
-        let runtime = runtime();
-
-        // TODO only create context once
+    pub fn session_context(&self, default_namespace: Option<String>) -> RbResult<RbSessionContext> {
         let catalog = self.catalog.read().unwrap().as_arc();
-        let provider = runtime
+        let provider = runtime()
             .block_on(IcebergCatalogProvider::try_new(catalog))
-            .unwrap();
-        let ctx = SessionContext::new();
+            .map_err(to_rb_err)?;
+        let mut config = SessionConfig::new();
+        if let Some(namespace) = default_namespace {
+            config = config.with_default_catalog_and_schema("datafusion", namespace);
+        }
+        let ctx = SessionContext::new_with_config(config);
         ctx.register_catalog("datafusion", Arc::new(provider));
+        Ok(RbSessionContext { ctx })
+    }
+}
 
-        let df = runtime.block_on(ctx.sql(&sql)).unwrap();
-        let _results = runtime.block_on(df.collect()).unwrap();
+#[cfg(feature = "datafusion")]
+#[magnus::wrap(class = "Iceberg::SessionContext")]
+pub struct RbSessionContext {
+    pub(crate) ctx: SessionContext,
+}
 
-        // println!("{:?}", _results);
+#[cfg(feature = "datafusion")]
+impl RbSessionContext {
+    pub fn sql(ruby: &Ruby, rb_self: &Self, sql: String, rb_params: RArray) -> RbResult<Value> {
+        let mut params = Vec::new();
+        for param in rb_params.into_iter() {
+            if param.is_nil() {
+                params.push(ScalarValue::Null);
+            } else if Qtrue::from_value(param).is_some() {
+                params.push(ScalarValue::from(true));
+            } else if Qfalse::from_value(param).is_some() {
+                params.push(ScalarValue::from(false));
+            } else if let Some(v) = Integer::from_value(param) {
+                params.push(ScalarValue::from(v.to_i64()?));
+            } else if let Some(v) = Float::from_value(param) {
+                params.push(ScalarValue::from(v.to_f64()));
+            } else if let Some(v) = RString::from_value(param) {
+                // TODO support binary strings
+                params.push(ScalarValue::from(v.to_string()?));
+            } else if unsafe { param.classname() } == "Date" {
+                params.push(ScalarValue::Date32(Some(date_to_i32(param)?)));
+            } else if unsafe { param.classname() } == "Time" {
+                let sec: i64 = param.funcall("to_i", ())?;
+                let nsec: i64 = param.funcall("nsec", ())?;
+                params.push(ScalarValue::TimestampNanosecond(
+                    Some(sec * 1_000_000_000 + nsec),
+                    None,
+                ));
+            } else {
+                return Err(todo_error(param));
+            }
+        }
 
-        Ok(())
+        let runtime = runtime();
+        let stream = runtime
+            .block_on(rb_self.ctx.sql(&sql))
+            .map_err(datafusion_error)?
+            .with_param_values(params)
+            .map_err(datafusion_error)?;
+        let batches = runtime
+            .block_on(stream.collect())
+            .map_err(datafusion_error)?;
+        collect_batches(ruby, batches)
+    }
+}
+
+fn storage_factory(warehouse: Option<String>) -> OpenDalStorageFactory {
+    match warehouse {
+        Some(ref v) => {
+            if v.starts_with("s3://") {
+                OpenDalStorageFactory::S3 {
+                    customized_credential_load: None,
+                }
+            } else {
+                OpenDalStorageFactory::Fs
+            }
+        }
+        // TODO consider OpenDalStorageFactory::Memory
+        None => OpenDalStorageFactory::Fs,
     }
 }

data/ext/iceberg/src/encryption.rs

@@ -0,0 +1,32 @@
+use std::collections::HashMap;
+
+use iceberg::spec::EncryptedKey;
+use magnus::{IntoValue, Ruby, value::ReprValue};
+
+#[magnus::wrap(class = "Iceberg::EncryptedKey")]
+pub struct RbEncryptedKey {
+    pub(crate) key: EncryptedKey,
+}
+
+impl RbEncryptedKey {
+    pub fn key_id(&self) -> &str {
+        self.key.key_id()
+    }
+
+    pub fn encrypted_by_id(&self) -> Option<&str> {
+        self.key.encrypted_by_id()
+    }
+
+    pub fn properties(&self) -> HashMap<String, String> {
+        self.key.properties().clone()
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::EncryptedKey key_id={}, encrypted_by_id={}, properties={}>",
+            rb_self.key_id().into_value_with(ruby).inspect(),
+            rb_self.encrypted_by_id().into_value_with(ruby).inspect(),
+            rb_self.properties().into_value_with(ruby).inspect(),
+        )
+    }
+}

data/ext/iceberg/src/error.rs

@@ -1,17 +1,49 @@
+use std::error::Error as _;
+
+use iceberg::{Error, ErrorKind};
 use magnus::{Error as RbErr, RModule, Ruby, prelude::*};
 
-pub fn to_rb_err(err: iceberg::Error) -> RbErr {
-    let class_name = match err.kind() {
-        iceberg::ErrorKind::NamespaceAlreadyExists => "NamespaceAlreadyExistsError",
-        iceberg::ErrorKind::NamespaceNotFound => "NamespaceNotFoundError",
-        iceberg::ErrorKind::TableAlreadyExists => "TableAlreadyExistsError",
-        iceberg::ErrorKind::TableNotFound => "TableNotFoundError",
-        iceberg::ErrorKind::FeatureUnsupported => "UnsupportedFeatureError",
-        iceberg::ErrorKind::DataInvalid => "InvalidDataError",
+pub fn to_rb_err(err: Error) -> RbErr {
+    let mut class_name = match err.kind() {
+        ErrorKind::NamespaceAlreadyExists => "NamespaceAlreadyExistsError",
+        ErrorKind::NamespaceNotFound => "NoSuchNamespaceError",
+        ErrorKind::TableAlreadyExists => "TableAlreadyExistsError",
+        ErrorKind::TableNotFound => "NoSuchTableError",
+        ErrorKind::FeatureUnsupported => "UnsupportedFeatureError",
+        ErrorKind::DataInvalid => "InvalidDataError",
         _ => "Error",
     };
 
-    let mut class = Ruby::get()
+    // no way to get context separately
+    // https://github.com/apache/iceberg-rust/issues/1071
+    let mut message = err.message().to_string();
+
+    let mut source = err.source();
+    while let Some(s) = source {
+        let m = s.to_string();
+        if message.contains(&m) {
+            break;
+        }
+        message.push_str(": ");
+        message.push_str(&m);
+        source = s.source();
+    }
+
+    if class_name == "Error" {
+        let s = err.to_string();
+        // for Glue
+        if s.contains("EntityNotFoundException") {
+            class_name = "NoSuchTableError";
+
+        // for Glue and S3 Tables
+        } else if s.contains("AlreadyExistsException")
+            || s.contains("A table with an identical name already exists in the namespace.")
+        {
+            class_name = "TableAlreadyExistsError";
+        }
+    }
+
+    let class = Ruby::get()
         .unwrap()
         .class_object()
         .const_get::<_, RModule>("Iceberg")
@@ -19,20 +51,28 @@ pub fn to_rb_err(err: iceberg::Error) -> RbErr {
         .const_get(class_name)
         .unwrap();
 
-    // no way to get context separately
-    // https://github.com/apache/iceberg-rust/issues/1071
-    let message = err.to_string();
-
-    // TODO remove in 0.12.0
-    if message.contains("target schema is not superset of current schema") {
-        class = Ruby::get().unwrap().exception_arg_error();
-    }
+    RbErr::new(class, message)
+}
 
-    let message = message
-        // TODO improve
-        .strip_prefix("Unexpected => ")
-        .map(|v| v.to_string())
-        .unwrap_or(message);
+pub fn todo_error<T: std::fmt::Debug>(message: T) -> RbErr {
+    let class = Ruby::get()
+        .unwrap()
+        .class_object()
+        .const_get::<_, RModule>("Iceberg")
+        .unwrap()
+        .const_get("Todo")
+        .unwrap();
+    RbErr::new(class, format!("not implemented yet: {:?}", message))
+}
 
-    RbErr::new(class, message)
+#[cfg(feature = "datafusion")]
+pub fn datafusion_error(err: datafusion::common::DataFusionError) -> RbErr {
+    let class = Ruby::get()
+        .unwrap()
+        .class_object()
+        .const_get::<_, RModule>("Iceberg")
+        .unwrap()
+        .const_get("Error")
+        .unwrap();
+    RbErr::new(class, err.to_string())
 }