deltalake-rb 0.1.0 → 0.1.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 55b36ac54ad1d38070efaa4d555cbe368dcb47c98c88a56fd27d90cc3633e563
-  data.tar.gz: 6a4464866554b770725d2f7e9e3ad39983933349d3ac57df97aa6204f90111d4
+  metadata.gz: 1ad1a2f352a83da63ccbde0126430b052d44801be630f3ad5e8326832205dc52
+  data.tar.gz: 1ce59b16589b891390d4ab1e81284d57389b6e8ff950faf11d1b3cf4736ce235
 SHA512:
-  metadata.gz: 9a545da9d049c2519dfe075f9ccd0b186e0c5369401bbb258548205a72e0dfe3b57a1167675641a7df235d717159ec1fa9547858ba7db2f78ba5503fa31a3bd2
-  data.tar.gz: 976cf69691f1e13e02ea374232be12a0afbef4cb35e1e7889d618c04dca10b1bd05e6b359765fe752a36b4109920bb88c1a95cffbeedcb36fca7040db5afe3aa
+  metadata.gz: 28c7d8f93e8dc78d9e81490e62d27be6d431efea6e1242c14ecb931b82b766926cb65a35b9440d372d27513e95c1f5c90e042fa42fb97a6407f1c3d30abb15ca
+  data.tar.gz: 547aa9019ac83f8ae9955f6f2cafe062a14223bdf18a29471b98e68d86d9655c69d47e84f520f9e1b5331e10f00e807141ed18d3786ae9995b120230c7217855

data/CHANGELOG.md

@@ -1,3 +1,13 @@
+## 0.1.1 (2024-11-22)
+
+- Added support for constraints
+- Added support for small file compaction
+- Added support for Z Ordering
+- Added `history`, `partitions`, `protocol`, `repair`, and `restore` methods to `Table`
+- Added experimental `load_cdf` method to `Table`
+- Fixed handling of unsigned integers
+- Fixed error with timestamps
+
 ## 0.1.0 (2024-11-20)
 
 - First release

data/Cargo.lock

@@ -1488,13 +1488,15 @@ dependencies = [
 
 [[package]]
 name = "deltalake"
-version = "0.1.0"
+version = "0.1.1"
 dependencies = [
  "arrow",
  "arrow-schema",
  "chrono",
  "deltalake 0.21.0",
+ "futures",
  "magnus",
+ "num_cpus",
  "serde",
  "serde_json",
  "tokio",

data/README.md

@@ -14,7 +14,7 @@ Add this line to your application’s Gemfile:
 gem "deltalake-rb"
 ```
 
-It can take a few minutes to compile the gem.
+It can take 5-10 minutes to compile the gem.
 
 ## Getting Started
 
@@ -50,6 +50,18 @@ Overwrite a table
 DeltaLake.write("./data/delta", df, mode: "overwrite")
 ```
 
+Add a constraint
+
+```ruby
+dt.alter.add_constraint({"a_gt_0" => "a > 0"})
+```
+
+Drop a constraint
+
+```ruby
+dt.alter.drop_constraint("a_gt_0")
+```
+
 Delete rows
 
 ```ruby
@@ -62,6 +74,18 @@ Vacuum
 dt.vacuum(dry_run: false)
 ```
 
+Perform small file compaction
+
+```ruby
+dt.optimize.compact
+```
+
+Colocate similar data in the same files
+
+```ruby
+dt.optimize.z_order(["a"])
+```
+
 Load a previous version of a table
 
 ```ruby
@@ -70,16 +94,22 @@ dt = DeltaLake::Table.new("./data/delta", version: 1)
 dt.load_as_version(1)
 ```
 
+Get the schema
+
+```ruby
+dt.schema
+```
+
 Get metadata
 
 ```ruby
 dt.metadata
 ```
 
-Get the schema
+Get history
 
 ```ruby
-dt.schema
+dt.history
 ```
 
 ## API

data/ext/deltalake/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "deltalake"
-version = "0.1.0"
+version = "0.1.1"
 license = "Apache-2.0"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
@@ -15,7 +15,9 @@ arrow = { version = "52", features = ["ffi"] }
 arrow-schema = { version = "52", features = ["serde"] }
 chrono = "0.4"
 deltalake = { version = "=0.21.0", features = ["datafusion", "s3"] }
+futures = "0.3"
 magnus = "0.7"
+num_cpus = "1"
 serde = "1"
 serde_json = "1"
 tokio = { version = "1", features = ["rt-multi-thread"] }

data/ext/deltalake/src/lib.rs

@@ -3,20 +3,39 @@ mod schema;
 mod utils;
 
 use std::cell::RefCell;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::future::IntoFuture;
+use std::str::FromStr;
+use std::time;
 
-use chrono::Duration;
+use chrono::{DateTime, Duration, FixedOffset, Utc};
 use deltalake::arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
-use deltalake::kernel::StructType;
+use deltalake::arrow::record_batch::RecordBatchIterator;
+use deltalake::datafusion::physical_plan::ExecutionPlan;
+use deltalake::datafusion::prelude::SessionContext;
+use deltalake::errors::DeltaTableError;
+use deltalake::kernel::{scalars::ScalarExt, StructType, Transaction};
+use deltalake::operations::collect_sendable_stream;
+use deltalake::operations::constraints::ConstraintBuilder;
 use deltalake::operations::delete::DeleteBuilder;
+use deltalake::operations::drop_constraints::DropConstraintBuilder;
+use deltalake::operations::filesystem_check::FileSystemCheckBuilder;
+use deltalake::operations::load_cdf::CdfLoadBuilder;
+use deltalake::operations::optimize::{OptimizeBuilder, OptimizeType};
+use deltalake::operations::restore::RestoreBuilder;
+use deltalake::operations::transaction::TableReference;
 use deltalake::operations::vacuum::VacuumBuilder;
+use deltalake::partitions::PartitionFilter;
 use deltalake::storage::IORuntime;
 use deltalake::DeltaOps;
 use error::DeltaError;
+use futures::future::join_all;
 
-use magnus::{function, method, prelude::*, Error, Module, Ruby, Value};
+use magnus::{
+    exception, function, method, prelude::*, Error, Integer, Module, RArray, RHash, Ruby, Value,
+};
 
+use crate::error::DeltaProtocolError;
 use crate::error::RubyError;
 use crate::schema::{schema_to_rbobject, Field};
 use crate::utils::rt;
@@ -38,6 +57,19 @@ struct RawDeltaTableMetaData {
     configuration: HashMap<String, Option<String>>,
 }
 
+#[magnus::wrap(class = "DeltaLake::ArrowArrayStream")]
+pub struct ArrowArrayStream {
+    stream: FFI_ArrowArrayStream,
+}
+
+impl ArrowArrayStream {
+    pub fn to_i(&self) -> usize {
+        (&self.stream as *const _) as usize
+    }
+}
+
+type StringVec = Vec<String>;
+
 impl RawDeltaTable {
     pub fn new(
         table_uri: String,
@@ -113,37 +145,138 @@ impl RawDeltaTable {
         })
     }
 
+    pub fn protocol_versions(&self) -> RbResult<(i32, i32, Option<StringVec>, Option<StringVec>)> {
+        let binding = self._table.borrow();
+        let table_protocol = binding.protocol().map_err(RubyError::from)?;
+        Ok((
+            table_protocol.min_reader_version,
+            table_protocol.min_writer_version,
+            table_protocol
+                .writer_features
+                .as_ref()
+                .and_then(|features| {
+                    let empty_set = !features.is_empty();
+                    empty_set.then(|| {
+                        features
+                            .iter()
+                            .map(|v| v.to_string())
+                            .collect::<Vec<String>>()
+                    })
+                }),
+            table_protocol
+                .reader_features
+                .as_ref()
+                .and_then(|features| {
+                    let empty_set = !features.is_empty();
+                    empty_set.then(|| {
+                        features
+                            .iter()
+                            .map(|v| v.to_string())
+                            .collect::<Vec<String>>()
+                    })
+                }),
+        ))
+    }
+
     pub fn load_version(&self, version: i64) -> RbResult<()> {
         Ok(rt()
             .block_on(self._table.borrow_mut().load_version(version))
             .map_err(RubyError::from)?)
     }
 
-    pub fn files(&self) -> RbResult<Vec<String>> {
-        if !self.has_files()? {
-            return Err(DeltaError::new_err("Table is instantiated without files."));
-        }
+    pub fn get_latest_version(&self) -> RbResult<i64> {
+        Ok(rt()
+            .block_on(self._table.borrow().get_latest_version())
+            .map_err(RubyError::from)?)
+    }
 
+    pub fn get_earliest_version(&self) -> RbResult<i64> {
+        Ok(rt()
+            .block_on(self._table.borrow().get_earliest_version())
+            .map_err(RubyError::from)?)
+    }
+
+    pub fn get_num_index_cols(&self) -> RbResult<i32> {
         Ok(self
             ._table
             .borrow()
-            .get_files_iter()
+            .snapshot()
             .map_err(RubyError::from)?
-            .map(|f| f.to_string())
-            .collect())
+            .config()
+            .num_indexed_cols())
     }
 
-    pub fn file_uris(&self) -> RbResult<Vec<String>> {
-        if !self._table.borrow().config.require_files {
-            return Err(DeltaError::new_err("Table is initiated without files."));
-        }
-
+    pub fn get_stats_columns(&self) -> RbResult<Option<Vec<String>>> {
         Ok(self
             ._table
             .borrow()
-            .get_file_uris()
+            .snapshot()
             .map_err(RubyError::from)?
-            .collect())
+            .config()
+            .stats_columns()
+            .map(|v| v.iter().map(|v| v.to_string()).collect::<Vec<String>>()))
+    }
+
+    pub fn load_with_datetime(&self, ds: String) -> RbResult<()> {
+        let datetime = DateTime::<Utc>::from(
+            DateTime::<FixedOffset>::parse_from_rfc3339(&ds).map_err(|err| {
+                Error::new(
+                    exception::arg_error(),
+                    format!("Failed to parse datetime string: {err}"),
+                )
+            })?,
+        );
+        Ok(rt()
+            .block_on(self._table.borrow_mut().load_with_datetime(datetime))
+            .map_err(RubyError::from)?)
+    }
+
+    pub fn files(&self, partition_filters: Option<Value>) -> RbResult<Vec<String>> {
+        if !self.has_files()? {
+            return Err(DeltaError::new_err("Table is instantiated without files."));
+        }
+
+        if let Some(filters) = partition_filters {
+            let filters = convert_partition_filters(filters).map_err(RubyError::from)?;
+            Ok(self
+                ._table
+                .borrow()
+                .get_files_by_partitions(&filters)
+                .map_err(RubyError::from)?
+                .into_iter()
+                .map(|p| p.to_string())
+                .collect())
+        } else {
+            Ok(self
+                ._table
+                .borrow()
+                .get_files_iter()
+                .map_err(RubyError::from)?
+                .map(|f| f.to_string())
+                .collect())
+        }
+    }
+
+    pub fn file_uris(&self, partition_filters: Option<Value>) -> RbResult<Vec<String>> {
+        if !self._table.borrow().config.require_files {
+            return Err(DeltaError::new_err("Table is initiated without files."));
+        }
+
+        if let Some(filters) = partition_filters {
+            let filters = convert_partition_filters(filters).map_err(RubyError::from)?;
+            Ok(self
+                ._table
+                .borrow()
+                .get_file_uris_by_partitions(&filters)
+                .map_err(RubyError::from)?)
+        } else {
+            Ok(self
+                ._table
+                .borrow()
+                .get_file_uris()
+                .map_err(RubyError::from)?
+                .collect())
+        }
     }
 
     pub fn schema(&self) -> RbResult<Value> {
@@ -177,6 +310,214 @@ impl RawDeltaTable {
         Ok(metrics.files_deleted)
     }
 
+    pub fn compact_optimize(
+        &self,
+        target_size: Option<i64>,
+        max_concurrent_tasks: Option<usize>,
+        min_commit_interval: Option<u64>,
+    ) -> RbResult<String> {
+        let mut cmd = OptimizeBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        )
+        .with_max_concurrent_tasks(max_concurrent_tasks.unwrap_or_else(num_cpus::get));
+        if let Some(size) = target_size {
+            cmd = cmd.with_target_size(size);
+        }
+        if let Some(commit_interval) = min_commit_interval {
+            cmd = cmd.with_min_commit_interval(time::Duration::from_secs(commit_interval));
+        }
+
+        let (table, metrics) = rt().block_on(cmd.into_future()).map_err(RubyError::from)?;
+        self._table.borrow_mut().state = table.state;
+        Ok(serde_json::to_string(&metrics).unwrap())
+    }
+
+    pub fn z_order_optimize(
+        &self,
+        z_order_columns: Vec<String>,
+        target_size: Option<i64>,
+        max_concurrent_tasks: Option<usize>,
+        max_spill_size: usize,
+        min_commit_interval: Option<u64>,
+    ) -> RbResult<String> {
+        let mut cmd = OptimizeBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        )
+        .with_max_concurrent_tasks(max_concurrent_tasks.unwrap_or_else(num_cpus::get))
+        .with_max_spill_size(max_spill_size)
+        .with_type(OptimizeType::ZOrder(z_order_columns));
+        if let Some(size) = target_size {
+            cmd = cmd.with_target_size(size);
+        }
+        if let Some(commit_interval) = min_commit_interval {
+            cmd = cmd.with_min_commit_interval(time::Duration::from_secs(commit_interval));
+        }
+
+        let (table, metrics) = rt().block_on(cmd.into_future()).map_err(RubyError::from)?;
+        self._table.borrow_mut().state = table.state;
+        Ok(serde_json::to_string(&metrics).unwrap())
+    }
+
+    pub fn add_constraints(&self, constraints: HashMap<String, String>) -> RbResult<()> {
+        let mut cmd = ConstraintBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        );
+
+        for (col_name, expression) in constraints {
+            cmd = cmd.with_constraint(col_name.clone(), expression.clone());
+        }
+
+        let table = rt().block_on(cmd.into_future()).map_err(RubyError::from)?;
+        self._table.borrow_mut().state = table.state;
+        Ok(())
+    }
+
+    pub fn drop_constraints(&self, name: String, raise_if_not_exists: bool) -> RbResult<()> {
+        let cmd = DropConstraintBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        )
+        .with_constraint(name)
+        .with_raise_if_not_exists(raise_if_not_exists);
+
+        let table = rt().block_on(cmd.into_future()).map_err(RubyError::from)?;
+        self._table.borrow_mut().state = table.state;
+        Ok(())
+    }
+
+    pub fn load_cdf(
+        &self,
+        starting_version: i64,
+        ending_version: Option<i64>,
+        starting_timestamp: Option<String>,
+        ending_timestamp: Option<String>,
+        columns: Option<Vec<String>>,
+    ) -> RbResult<ArrowArrayStream> {
+        let ctx = SessionContext::new();
+        let mut cdf_read = CdfLoadBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        )
+        .with_starting_version(starting_version);
+
+        if let Some(ev) = ending_version {
+            cdf_read = cdf_read.with_ending_version(ev);
+        }
+        if let Some(st) = starting_timestamp {
+            let starting_ts: DateTime<Utc> = DateTime::<Utc>::from_str(&st)
+                .map_err(|pe| Error::new(exception::arg_error(), pe.to_string()))?
+                .to_utc();
+            cdf_read = cdf_read.with_starting_timestamp(starting_ts);
+        }
+        if let Some(et) = ending_timestamp {
+            let ending_ts = DateTime::<Utc>::from_str(&et)
+                .map_err(|pe| Error::new(exception::arg_error(), pe.to_string()))?
+                .to_utc();
+            cdf_read = cdf_read.with_starting_timestamp(ending_ts);
+        }
+
+        if let Some(columns) = columns {
+            cdf_read = cdf_read.with_columns(columns);
+        }
+
+        cdf_read = cdf_read.with_session_ctx(ctx.clone());
+
+        let plan = rt().block_on(cdf_read.build()).map_err(RubyError::from)?;
+
+        let mut tasks = vec![];
+        for p in 0..plan.properties().output_partitioning().partition_count() {
+            let inner_plan = plan.clone();
+            let partition_batch = inner_plan.execute(p, ctx.task_ctx()).unwrap();
+            let handle = rt().spawn(collect_sendable_stream(partition_batch));
+            tasks.push(handle);
+        }
+
+        // This is unfortunate.
+        let batches = rt()
+            .block_on(join_all(tasks))
+            .into_iter()
+            .flatten()
+            .collect::<Result<Vec<Vec<_>>, _>>()
+            .unwrap()
+            .into_iter()
+            .flatten()
+            .map(Ok);
+        let batch_iter = RecordBatchIterator::new(batches, plan.schema());
+        let ffi_stream = FFI_ArrowArrayStream::new(Box::new(batch_iter));
+        Ok(ArrowArrayStream { stream: ffi_stream })
+    }
+
+    pub fn restore(
+        &self,
+        target: Option<Value>,
+        ignore_missing_files: bool,
+        protocol_downgrade_allowed: bool,
+    ) -> RbResult<String> {
+        let mut cmd = RestoreBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        );
+        if let Some(val) = target {
+            if let Some(version) = Integer::from_value(val) {
+                cmd = cmd.with_version_to_restore(version.to_i64()?)
+            }
+            if let Ok(ds) = String::try_convert(val) {
+                let datetime = DateTime::<Utc>::from(
+                    DateTime::<FixedOffset>::parse_from_rfc3339(ds.as_ref()).map_err(|err| {
+                        Error::new(
+                            exception::arg_error(),
+                            format!("Failed to parse datetime string: {err}"),
+                        )
+                    })?,
+                );
+                cmd = cmd.with_datetime_to_restore(datetime)
+            }
+        }
+        cmd = cmd.with_ignore_missing_files(ignore_missing_files);
+        cmd = cmd.with_protocol_downgrade_allowed(protocol_downgrade_allowed);
+
+        let (table, metrics) = rt().block_on(cmd.into_future()).map_err(RubyError::from)?;
+        self._table.borrow_mut().state = table.state;
+        Ok(serde_json::to_string(&metrics).unwrap())
+    }
+
+    pub fn history(&self, limit: Option<usize>) -> RbResult<Vec<String>> {
+        let history = rt()
+            .block_on(self._table.borrow().history(limit))
+            .map_err(RubyError::from)?;
+        Ok(history
+            .iter()
+            .map(|c| serde_json::to_string(c).unwrap())
+            .collect())
+    }
+
     pub fn update_incremental(&self) -> RbResult<()> {
         #[allow(deprecated)]
         Ok(rt()
@@ -184,6 +525,56 @@ impl RawDeltaTable {
             .map_err(RubyError::from)?)
     }
 
+    fn get_active_partitions(&self) -> RbResult<RArray> {
+        let binding = self._table.borrow();
+        let _column_names: HashSet<&str> = binding
+            .get_schema()
+            .map_err(|_| DeltaProtocolError::new_err("table does not yet have a schema"))?
+            .fields()
+            .map(|field| field.name().as_str())
+            .collect();
+        let partition_columns: HashSet<&str> = binding
+            .metadata()
+            .map_err(RubyError::from)?
+            .partition_columns
+            .iter()
+            .map(|col| col.as_str())
+            .collect();
+
+        let converted_filters = Vec::new();
+
+        let partition_columns: Vec<&str> = partition_columns.into_iter().collect();
+
+        let adds = binding
+            .snapshot()
+            .map_err(RubyError::from)?
+            .get_active_add_actions_by_partitions(&converted_filters)
+            .map_err(RubyError::from)?
+            .collect::<Result<Vec<_>, _>>()
+            .map_err(RubyError::from)?;
+        let active_partitions: HashSet<Vec<(&str, Option<String>)>> = adds
+            .iter()
+            .flat_map(|add| {
+                Ok::<_, RubyError>(
+                    partition_columns
+                        .iter()
+                        .flat_map(|col| {
+                            Ok::<_, RubyError>((
+                                *col,
+                                add.partition_values()
+                                    .map_err(RubyError::from)?
+                                    .get(*col)
+                                    .map(|v| v.serialize()),
+                            ))
+                        })
+                        .collect(),
+                )
+            })
+            .collect();
+
+        Ok(RArray::from_iter(active_partitions))
+    }
+
     pub fn delete(&self, predicate: Option<String>) -> RbResult<String> {
         let mut cmd = DeleteBuilder::new(
             self._table.borrow().log_store(),
@@ -201,6 +592,38 @@ impl RawDeltaTable {
         self._table.borrow_mut().state = table.state;
         Ok(serde_json::to_string(&metrics).unwrap())
     }
+
+    pub fn repair(&self, dry_run: bool) -> RbResult<String> {
+        let cmd = FileSystemCheckBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        )
+        .with_dry_run(dry_run);
+
+        let (table, metrics) = rt().block_on(cmd.into_future()).map_err(RubyError::from)?;
+        self._table.borrow_mut().state = table.state;
+        Ok(serde_json::to_string(&metrics).unwrap())
+    }
+
+    pub fn transaction_versions(&self) -> RHash {
+        RHash::from_iter(
+            self._table
+                .borrow()
+                .get_app_transaction_version()
+                .into_iter()
+                .map(|(app_id, transaction)| (app_id, RbTransaction::from(transaction))),
+        )
+    }
+}
+
+fn convert_partition_filters(
+    _partitions_filters: Value,
+) -> Result<Vec<PartitionFilter>, DeltaTableError> {
+    todo!()
 }
 
 impl RawDeltaTableMetaData {
@@ -229,6 +652,23 @@ impl RawDeltaTableMetaData {
     }
 }
 
+#[magnus::wrap(class = "DeltaLake::Transaction")]
+pub struct RbTransaction {
+    pub app_id: String,
+    pub version: i64,
+    pub last_updated: Option<i64>,
+}
+
+impl From<Transaction> for RbTransaction {
+    fn from(value: Transaction) -> Self {
+        RbTransaction {
+            app_id: value.app_id,
+            version: value.version,
+            last_updated: value.last_updated,
+        }
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 fn write_to_deltalake(
     table_uri: String,
@@ -313,16 +753,68 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("version", method!(RawDeltaTable::version, 0))?;
     class.define_method("has_files", method!(RawDeltaTable::has_files, 0))?;
     class.define_method("metadata", method!(RawDeltaTable::metadata, 0))?;
+    class.define_method(
+        "protocol_versions",
+        method!(RawDeltaTable::protocol_versions, 0),
+    )?;
     class.define_method("load_version", method!(RawDeltaTable::load_version, 1))?;
-    class.define_method("files", method!(RawDeltaTable::files, 0))?;
-    class.define_method("file_uris", method!(RawDeltaTable::file_uris, 0))?;
+    class.define_method(
+        "get_latest_version",
+        method!(RawDeltaTable::get_latest_version, 0),
+    )?;
+    class.define_method(
+        "get_earliest_version",
+        method!(RawDeltaTable::get_earliest_version, 0),
+    )?;
+    class.define_method(
+        "get_num_index_cols",
+        method!(RawDeltaTable::get_num_index_cols, 0),
+    )?;
+    class.define_method(
+        "get_stats_columns",
+        method!(RawDeltaTable::get_stats_columns, 0),
+    )?;
+    class.define_method(
+        "load_with_datetime",
+        method!(RawDeltaTable::load_with_datetime, 1),
+    )?;
+    class.define_method("files", method!(RawDeltaTable::files, 1))?;
+    class.define_method("file_uris", method!(RawDeltaTable::file_uris, 1))?;
     class.define_method("schema", method!(RawDeltaTable::schema, 0))?;
     class.define_method("vacuum", method!(RawDeltaTable::vacuum, 3))?;
+    class.define_method(
+        "compact_optimize",
+        method!(RawDeltaTable::compact_optimize, 3),
+    )?;
+    class.define_method(
+        "z_order_optimize",
+        method!(RawDeltaTable::z_order_optimize, 5),
+    )?;
+    class.define_method(
+        "add_constraints",
+        method!(RawDeltaTable::add_constraints, 1),
+    )?;
+    class.define_method(
+        "drop_constraints",
+        method!(RawDeltaTable::drop_constraints, 2),
+    )?;
+    class.define_method("load_cdf", method!(RawDeltaTable::load_cdf, 5))?;
+    class.define_method("restore", method!(RawDeltaTable::restore, 3))?;
+    class.define_method("history", method!(RawDeltaTable::history, 1))?;
     class.define_method(
         "update_incremental",
         method!(RawDeltaTable::update_incremental, 0),
     )?;
+    class.define_method(
+        "get_active_partitions",
+        method!(RawDeltaTable::get_active_partitions, 0),
+    )?;
     class.define_method("delete", method!(RawDeltaTable::delete, 1))?;
+    class.define_method("repair", method!(RawDeltaTable::repair, 1))?;
+    class.define_method(
+        "transaction_versions",
+        method!(RawDeltaTable::transaction_versions, 0),
+    )?;
 
     let class = module.define_class("RawDeltaTableMetaData", ruby.class_object())?;
     class.define_method("id", method!(RawDeltaTableMetaData::id, 0))?;
@@ -344,6 +836,9 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         method!(RawDeltaTableMetaData::configuration, 0),
     )?;
 
+    let class = module.define_class("ArrowArrayStream", ruby.class_object())?;
+    class.define_method("to_i", method!(ArrowArrayStream::to_i, 0))?;
+
     let class = module.define_class("Field", ruby.class_object())?;
     class.define_method("name", method!(Field::name, 0))?;
     class.define_method("type", method!(Field::get_type, 0))?;

data/lib/deltalake/table.rb

@@ -26,22 +26,54 @@ module DeltaLake
       @table.version
     end
 
-    def files
-      @table.files
+    def partitions
+      partitions = []
+      @table.get_active_partitions.each do |partition|
+        next unless partition
+        partitions << partition.to_h
+      end
+      partitions
+    end
+
+    def files(partition_filters: nil)
+      @table.files(_stringify_partition_values(partition_filters))
     end
 
-    def file_uris
-      @table.file_uris
+    def file_uris(partition_filters: nil)
+      @table.file_uris(_stringify_partition_values(partition_filters))
     end
 
     def load_as_version(version)
       if version.is_a?(Integer)
         @table.load_version(version)
+      elsif version.is_a?(Time)
+        # needed for iso8601
+        require "time"
+
+        @table.load_with_datetime(version.utc.iso8601(9))
+      elsif version.is_a?(String)
+        @table.load_with_datetime(version)
       else
-        raise TypeError, "Invalid datatype provided for version, only Integer is accepted."
+        raise TypeError, "Invalid datatype provided for version, only Integer, String, and Time are accepted."
       end
     end
 
+    def load_cdf(
+      starting_version: 0,
+      ending_version: nil,
+      starting_timestamp: nil,
+      ending_timestamp: nil,
+      columns: nil
+    )
+      @table.load_cdf(
+        starting_version,
+        ending_version,
+        starting_timestamp,
+        ending_timestamp,
+        columns
+      )
+    end
+
     def table_uri
       @table.table_uri
     end
@@ -54,6 +86,29 @@ module DeltaLake
       Metadata.new(@table)
     end
 
+    def protocol
+      ProtocolVersions.new(*@table.protocol_versions)
+    end
+
+    def history(limit: nil)
+      backwards_enumerate = lambda do |iterable, start_end, &block|
+        n = start_end
+        iterable.each do |elem|
+          block.call(n, elem)
+          n -= 1
+        end
+      end
+
+      commits = @table.history(limit)
+      history = []
+      backwards_enumerate.(commits, @table.get_latest_version) do |version, commit_info_raw|
+        commit = JSON.parse(commit_info_raw)
+        commit["version"] = version
+        history << commit
+      end
+      history
+    end
+
     def vacuum(
       retention_hours: nil,
       dry_run: true,
@@ -72,6 +127,40 @@ module DeltaLake
       )
     end
 
+    def optimize
+      TableOptimizer.new(self)
+    end
+
+    def alter
+      TableAlterer.new(self)
+    end
+
+    def restore(
+      target,
+      ignore_missing_files: false,
+      protocol_downgrade_allowed: false
+    )
+      if target.is_a?(Time)
+        # needed for iso8601
+        require "time"
+
+        metrics =
+          @table.restore(
+            target.utc.iso8601(9),
+            ignore_missing_files,
+            protocol_downgrade_allowed
+          )
+      else
+        metrics =
+          @table.restore(
+            target,
+            ignore_missing_files,
+            protocol_downgrade_allowed
+          )
+      end
+      JSON.parse(metrics)
+    end
+
     def to_polars(eager: true)
       require "polars-df"
 
@@ -80,7 +169,13 @@ module DeltaLake
         if sources.empty?
           Polars::LazyFrame.new
         else
-          storage_options = @storage_options&.except("AWS_S3_ALLOW_UNSAFE_RENAME")
+          delta_keys = [
+            "AWS_S3_ALLOW_UNSAFE_RENAME",
+            "AWS_S3_LOCKING_PROVIDER",
+            "CONDITIONAL_PUT",
+            "DELTA_DYNAMO_TABLE_NAME"
+          ]
+          storage_options = @storage_options&.reject { |k, _| delta_keys.include?(k.to_s.upcase) }
           Polars.scan_parquet(sources, storage_options: storage_options)
         end
       eager ? lf.collect : lf
@@ -95,9 +190,30 @@ module DeltaLake
       JSON.parse(metrics).transform_keys(&:to_sym)
     end
 
+    def repair(dry_run: false)
+      metrics =
+        @table.repair(
+          dry_run
+        )
+      JSON.parse(metrics).transform_keys(&:to_sym)
+    end
+
+    def transaction_versions
+      @table.transaction_versions
+    end
+
     # private
     def _table
       @table
     end
+
+    # private
+    def _stringify_partition_values(partition_filters)
+      if partition_filters.nil?
+        return partition_filters
+      end
+
+      raise Todo
+    end
   end
 end

data/lib/deltalake/table_alterer.rb

@@ -0,0 +1,25 @@
+module DeltaLake
+  class TableAlterer
+    def initialize(table)
+      @table = table
+    end
+
+    def add_constraint(constraints)
+      if constraints.length > 1
+        raise ArgumentError,
+          "add_constraints is limited to a single constraint addition at once for now."
+      end
+
+      @table._table.add_constraints(
+        constraints
+      )
+    end
+
+    def drop_constraint(name, raise_if_not_exists: true)
+      @table._table.drop_constraints(
+        name,
+        raise_if_not_exists
+      )
+    end
+  end
+end

data/lib/deltalake/table_optimizer.rb

@@ -0,0 +1,51 @@
+module DeltaLake
+  class TableOptimizer
+    def initialize(table)
+      @table = table
+    end
+
+    def compact(
+      target_size: nil,
+      max_concurrent_tasks: nil,
+      min_commit_interval: nil
+    )
+      metrics =
+        @table._table.compact_optimize(
+          target_size,
+          max_concurrent_tasks,
+          min_commit_interval
+        )
+      @table.update_incremental
+      result = JSON.parse(metrics)
+      ["filesAdded", "filesRemoved"].each do |key|
+        result[key] = JSON.parse(result[key]) if result[key].is_a?(String)
+      end
+      # TODO return underscore symbols like delete
+      result
+    end
+
+    def z_order(
+      columns,
+      target_size: nil,
+      max_concurrent_tasks:  nil,
+      max_spill_size: 20 * 1024 * 1024 * 1024,
+      min_commit_interval: nil
+    )
+      metrics =
+        @table._table.z_order_optimize(
+          Array(columns),
+          target_size,
+          max_concurrent_tasks,
+          max_spill_size,
+          min_commit_interval
+        )
+      @table.update_incremental
+      result = JSON.parse(metrics)
+      ["filesAdded", "filesRemoved"].each do |key|
+        result[key] = JSON.parse(result[key]) if result[key].is_a?(String)
+      end
+      # TODO return underscore symbols like delete
+      result
+    end
+  end
+end

data/lib/deltalake/version.rb

@@ -1,3 +1,3 @@
 module DeltaLake
-  VERSION = "0.1.0"
+  VERSION = "0.1.1"
 end

data/lib/deltalake.rb

@@ -13,6 +13,8 @@ require_relative "deltalake/field"
 require_relative "deltalake/metadata"
 require_relative "deltalake/schema"
 require_relative "deltalake/table"
+require_relative "deltalake/table_alterer"
+require_relative "deltalake/table_optimizer"
 require_relative "deltalake/version"
 
 module DeltaLake
@@ -22,8 +24,21 @@ module DeltaLake
   class CommitFailedError < Error; end
   class SchemaMismatchError < Error; end
 
-  class << self
+  class Todo < Error
+    def message
+      "not implemented yet"
+    end
+  end
+
+  ProtocolVersions =
+    Struct.new(
+      :min_reader_version,
+      :min_writer_version,
+      :writer_features,
+      :reader_features
+    )
 
+  class << self
     def write(
       table_or_uri,
       data,
@@ -95,10 +110,58 @@ module DeltaLake
 
     def convert_data(data)
       if data.respond_to?(:arrow_c_stream)
+        # TODO convert other object types
+        # should probably move logic to Rust
+        if defined?(Polars::DataFrame) && data.is_a?(Polars::DataFrame)
+          data = convert_polars_data(data)
+        end
+
         data.arrow_c_stream
       else
         raise TypeError, "Only objects implementing the Arrow C stream interface are valid inputs for source."
       end
     end
+
+    # unsigned integers are not part of the protocol
+    # https://github.com/delta-io/delta/blob/master/PROTOCOL.md#primitive-types
+    def convert_polars_data(data)
+      new_schema = {}
+      data.schema.each do |k, v|
+        new_type = convert_polars_type(v)
+        new_schema[k] = new_type if new_type
+      end
+
+      if new_schema.any?
+        data.cast(new_schema)
+      else
+        data
+      end
+    end
+
+    def convert_polars_type(t)
+      case t
+      when Polars::UInt8
+        Polars::Int8
+      when Polars::UInt16
+        Polars::Int16
+      when Polars::UInt32
+        Polars::Int32
+      when Polars::UInt64
+        Polars::Int64
+      when Polars::Datetime
+        Polars::Datetime.new("us", t.time_zone) if t.time_unit != "us"
+      when Polars::List
+        inner = convert_polars_type(t.inner)
+        Polars::List.new(inner) if inner
+      when Polars::Array
+        inner = convert_polars_type(t.inner)
+        Polars::Array.new(t.inner, t.width) if inner
+      when Polars::Struct
+        if t.fields.any? { |f| convert_polars_type(f.dtype) }
+          fields = t.fields.map { |f| Polars::Field.new(f.name, convert_polars_type(f.dtype) || f.dtype) }
+          Polars::Struct.new(fields)
+        end
+      end
+    end
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: deltalake-rb
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.1.1
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-11-21 00:00:00.000000000 Z
+date: 2024-11-23 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys
@@ -48,6 +48,8 @@ files:
 - lib/deltalake/metadata.rb
 - lib/deltalake/schema.rb
 - lib/deltalake/table.rb
+- lib/deltalake/table_alterer.rb
+- lib/deltalake/table_optimizer.rb
 - lib/deltalake/version.rb
 homepage: https://github.com/ankane/delta-ruby
 licenses: