deltalake-rb 0.1.1 → 0.1.3

data/ext/deltalake/src/merge.rs

@@ -0,0 +1,205 @@
+use deltalake::arrow::array::RecordBatchReader;
+use deltalake::arrow::datatypes::Schema as ArrowSchema;
+use deltalake::arrow::ffi_stream::ArrowArrayStreamReader;
+use deltalake::datafusion::catalog::TableProvider;
+use deltalake::datafusion::datasource::MemTable;
+use deltalake::datafusion::prelude::SessionContext;
+use deltalake::logstore::LogStoreRef;
+use deltalake::operations::merge::MergeBuilder;
+use deltalake::table::state::DeltaTableState;
+use deltalake::{DeltaResult, DeltaTable};
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::future::IntoFuture;
+use std::sync::Arc;
+
+use crate::error::RubyError;
+use crate::utils::rt;
+use crate::RbResult;
+use crate::{
+    maybe_create_commit_properties, set_writer_properties, RbCommitProperties,
+    RbPostCommitHookProperties, RbWriterProperties,
+};
+
+#[magnus::wrap(class = "DeltaLake::RbMergeBuilder")]
+pub(crate) struct RbMergeBuilder {
+    _builder: RefCell<Option<MergeBuilder>>,
+    source_alias: Option<String>,
+    target_alias: Option<String>,
+    #[allow(dead_code)]
+    arrow_schema: Arc<ArrowSchema>,
+}
+
+// getters
+impl RbMergeBuilder {
+    pub fn source_alias(&self) -> Option<String> {
+        self.source_alias.clone()
+    }
+
+    pub fn target_alias(&self) -> Option<String> {
+        self.target_alias.clone()
+    }
+}
+
+impl RbMergeBuilder {
+    pub fn new(
+        log_store: LogStoreRef,
+        snapshot: DeltaTableState,
+        source: ArrowArrayStreamReader,
+        predicate: String,
+        source_alias: Option<String>,
+        target_alias: Option<String>,
+        safe_cast: bool,
+        writer_properties: Option<RbWriterProperties>,
+        post_commithook_properties: Option<RbPostCommitHookProperties>,
+        commit_properties: Option<RbCommitProperties>,
+    ) -> DeltaResult<Self> {
+        let ctx = SessionContext::new();
+        let schema = source.schema();
+        let batches = vec![source.map(|batch| batch.unwrap()).collect::<Vec<_>>()];
+        let table_provider: Arc<dyn TableProvider> =
+            Arc::new(MemTable::try_new(schema.clone(), batches).unwrap());
+        let source_df = ctx.read_table(table_provider).unwrap();
+
+        let mut cmd =
+            MergeBuilder::new(log_store, snapshot, predicate, source_df).with_safe_cast(safe_cast);
+
+        if let Some(src_alias) = &source_alias {
+            cmd = cmd.with_source_alias(src_alias);
+        }
+
+        if let Some(trgt_alias) = &target_alias {
+            cmd = cmd.with_target_alias(trgt_alias);
+        }
+
+        if let Some(writer_props) = writer_properties {
+            cmd = cmd.with_writer_properties(set_writer_properties(writer_props)?);
+        }
+
+        if let Some(commit_properties) =
+            maybe_create_commit_properties(commit_properties, post_commithook_properties)
+        {
+            cmd = cmd.with_commit_properties(commit_properties);
+        }
+
+        Ok(Self {
+            _builder: RefCell::new(Some(cmd)),
+            source_alias,
+            target_alias,
+            arrow_schema: schema,
+        })
+    }
+
+    pub fn execute(&self) -> DeltaResult<(DeltaTable, String)> {
+        let (table, metrics) = rt().block_on(self._builder.take().unwrap().into_future())?;
+        Ok((table, serde_json::to_string(&metrics).unwrap()))
+    }
+}
+
+impl RbMergeBuilder {
+    pub fn when_matched_update(
+        &self,
+        updates: HashMap<String, String>,
+        predicate: Option<String>,
+    ) -> RbResult<()> {
+        let mut binding = self._builder.borrow_mut();
+        *binding = match binding.take() {
+            Some(cmd) => Some(
+                cmd.when_matched_update(|mut update| {
+                    for (column, expression) in updates {
+                        update = update.update(column, expression)
+                    }
+                    if let Some(predicate) = predicate {
+                        update = update.predicate(predicate)
+                    };
+                    update
+                })
+                .map_err(RubyError::from)?,
+            ),
+            None => unreachable!(),
+        };
+        Ok(())
+    }
+
+    pub fn when_matched_delete(&self, predicate: Option<String>) -> RbResult<()> {
+        let mut binding = self._builder.borrow_mut();
+        *binding = match binding.take() {
+            Some(cmd) => Some(
+                cmd.when_matched_delete(|mut delete| {
+                    if let Some(predicate) = predicate {
+                        delete = delete.predicate(predicate)
+                    };
+                    delete
+                })
+                .map_err(RubyError::from)?,
+            ),
+            None => unreachable!(),
+        };
+        Ok(())
+    }
+
+    pub fn when_not_matched_insert(
+        &self,
+        updates: HashMap<String, String>,
+        predicate: Option<String>,
+    ) -> RbResult<()> {
+        let mut binding = self._builder.borrow_mut();
+        *binding = match binding.take() {
+            Some(cmd) => Some(
+                cmd.when_not_matched_insert(|mut insert| {
+                    for (column, expression) in updates {
+                        insert = insert.set(column, expression)
+                    }
+                    if let Some(predicate) = predicate {
+                        insert = insert.predicate(predicate)
+                    };
+                    insert
+                })
+                .map_err(RubyError::from)?,
+            ),
+            None => unreachable!(),
+        };
+        Ok(())
+    }
+
+    pub fn when_not_matched_by_source_update(
+        &self,
+        updates: HashMap<String, String>,
+        predicate: Option<String>,
+    ) -> RbResult<()> {
+        let mut binding = self._builder.borrow_mut();
+        *binding = match binding.take() {
+            Some(cmd) => Some(
+                cmd.when_not_matched_by_source_update(|mut update| {
+                    for (column, expression) in updates {
+                        update = update.update(column, expression)
+                    }
+                    if let Some(predicate) = predicate {
+                        update = update.predicate(predicate)
+                    };
+                    update
+                })
+                .map_err(RubyError::from)?,
+            ),
+            None => unreachable!(),
+        };
+        Ok(())
+    }
+
+    pub fn when_not_matched_by_source_delete(&self, predicate: Option<String>) -> RbResult<()> {
+        let mut binding = self._builder.borrow_mut();
+        *binding = match binding.take() {
+            Some(cmd) => Some(
+                cmd.when_not_matched_by_source_delete(|mut delete| {
+                    if let Some(predicate) = predicate {
+                        delete = delete.predicate(predicate)
+                    };
+                    delete
+                })
+                .map_err(RubyError::from)?,
+            ),
+            None => unreachable!(),
+        };
+        Ok(())
+    }
+}

data/lib/deltalake/table.rb

@@ -47,9 +47,6 @@ module DeltaLake
       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)
@@ -112,7 +109,9 @@ module DeltaLake
     def vacuum(
       retention_hours: nil,
       dry_run: true,
-      enforce_retention_duration: true
+      enforce_retention_duration: true,
+      post_commithook_properties: nil,
+      commit_properties: nil
     )
       if retention_hours
         if retention_hours < 0
@@ -123,7 +122,9 @@ module DeltaLake
       @table.vacuum(
         dry_run,
         retention_hours,
-        enforce_retention_duration
+        enforce_retention_duration,
+        commit_properties,
+        post_commithook_properties
       )
     end
 
@@ -135,49 +136,80 @@ module DeltaLake
       TableAlterer.new(self)
     end
 
+    def merge(
+      source,
+      predicate,
+      source_alias: nil,
+      target_alias: nil,
+      error_on_type_mismatch: true,
+      writer_properties: nil,
+      post_commithook_properties: nil,
+      commit_properties: nil
+    )
+      source = Utils.convert_data(source)
+
+      rb_merge_builder =
+        @table.create_merge_builder(
+          source,
+          predicate,
+          source_alias,
+          target_alias,
+          !error_on_type_mismatch,
+          writer_properties,
+          post_commithook_properties,
+          commit_properties
+        )
+      TableMerger.new(rb_merge_builder, @table)
+    end
+
     def restore(
       target,
       ignore_missing_files: false,
-      protocol_downgrade_allowed: false
+      protocol_downgrade_allowed: false,
+      commit_properties: nil
     )
       if target.is_a?(Time)
-        # needed for iso8601
-        require "time"
-
         metrics =
           @table.restore(
             target.utc.iso8601(9),
             ignore_missing_files,
-            protocol_downgrade_allowed
+            protocol_downgrade_allowed,
+            commit_properties
           )
       else
         metrics =
           @table.restore(
             target,
             ignore_missing_files,
-            protocol_downgrade_allowed
+            protocol_downgrade_allowed,
+            commit_properties
           )
       end
       JSON.parse(metrics)
     end
 
-    def to_polars(eager: true)
+    def to_polars(eager: true, rechunk: false, columns: nil)
       require "polars-df"
 
       sources = file_uris
-      lf =
-        if sources.empty?
-          Polars::LazyFrame.new
-        else
-          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)
+      if sources.empty?
+        lf = Polars::LazyFrame.new
+      else
+        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) }
+        lf = Polars.scan_parquet(sources, storage_options: storage_options, rechunk: rechunk)
+
+        if columns
+          # by_name requires polars-df > 0.15.0
+          lf = lf.select(Polars.cs.by_name(*columns))
         end
+      end
+
       eager ? lf.collect : lf
     end
 
@@ -185,15 +217,32 @@ module DeltaLake
       @table.update_incremental
     end
 
-    def delete(predicate = nil)
-      metrics = @table.delete(predicate)
+    def delete(
+      predicate = nil,
+      writer_properties: nil,
+      post_commithook_properties: nil,
+      commit_properties: nil
+    )
+      metrics =
+        @table.delete(
+          predicate,
+          writer_properties,
+          post_commithook_properties,
+          commit_properties
+        )
       JSON.parse(metrics).transform_keys(&:to_sym)
     end
 
-    def repair(dry_run: false)
+    def repair(
+      dry_run: false,
+      post_commithook_properties: nil,
+      commit_properties: nil
+    )
       metrics =
         @table.repair(
-          dry_run
+          dry_run,
+          commit_properties,
+          post_commithook_properties
         )
       JSON.parse(metrics).transform_keys(&:to_sym)
     end

data/lib/deltalake/table_alterer.rb

@@ -4,6 +4,29 @@ module DeltaLake
       @table = table
     end
 
+    def add_feature(
+      feature,
+      allow_protocol_versions_increase: false
+    )
+      if !feature.is_a?(Array)
+        feature = [feature]
+      end
+      @table._table.add_feature(
+        feature,
+        allow_protocol_versions_increase
+      )
+    end
+
+    def add_columns(fields)
+      if fields.is_a?(DeltaLake::Field)
+        fields = [fields]
+      end
+
+      @table._table.add_columns(
+        fields
+      )
+    end
+
     def add_constraint(constraints)
       if constraints.length > 1
         raise ArgumentError,
@@ -21,5 +44,15 @@ module DeltaLake
         raise_if_not_exists
       )
     end
+
+    def set_table_properties(
+      properties,
+      raise_if_not_exists: true
+    )
+      @table._table.set_table_properties(
+        properties,
+        raise_if_not_exists
+      )
+    end
   end
 end

data/lib/deltalake/table_merger.rb

@@ -0,0 +1,38 @@
+module DeltaLake
+  class TableMerger
+    def initialize(builder, table)
+      @builder = builder
+      @table = table
+    end
+
+    def when_matched_update(updates, predicate: nil)
+      @builder.when_matched_update(updates, predicate)
+      self
+    end
+
+    def when_not_matched_insert(updates, predicate: nil)
+      @builder.when_not_matched_insert(updates, predicate)
+      self
+    end
+
+    def when_matched_delete(predicate: nil)
+      @builder.when_matched_delete(predicate)
+      self
+    end
+
+    def when_not_matched_by_source_update(updates, predicate: nil)
+      @builder.when_not_matched_by_source_update(updates, predicate)
+      self
+    end
+
+    def when_not_matched_by_source_delete(predicate: nil)
+      @builder.when_not_matched_by_source_delete(predicate)
+      self
+    end
+
+    def execute
+      metrics = @table.merge_execute(@builder)
+      JSON.parse(metrics).transform_keys(&:to_sym)
+    end
+  end
+end

data/lib/deltalake/table_optimizer.rb

@@ -5,15 +5,23 @@ module DeltaLake
     end
 
     def compact(
+      partition_filters: nil,
       target_size: nil,
       max_concurrent_tasks: nil,
-      min_commit_interval: nil
+      min_commit_interval: nil,
+      writer_properties: nil,
+      post_commithook_properties: nil,
+      commit_properties: nil
     )
       metrics =
         @table._table.compact_optimize(
+          @table._stringify_partition_values(partition_filters),
           target_size,
           max_concurrent_tasks,
-          min_commit_interval
+          min_commit_interval,
+          writer_properties,
+          post_commithook_properties,
+          commit_properties
         )
       @table.update_incremental
       result = JSON.parse(metrics)
@@ -26,18 +34,26 @@ module DeltaLake
 
     def z_order(
       columns,
+      partition_filters: nil,
       target_size: nil,
       max_concurrent_tasks:  nil,
       max_spill_size: 20 * 1024 * 1024 * 1024,
-      min_commit_interval: nil
+      min_commit_interval: nil,
+      writer_properties: nil,
+      post_commithook_properties: nil,
+      commit_properties: nil
     )
       metrics =
         @table._table.z_order_optimize(
           Array(columns),
+          @table._stringify_partition_values(partition_filters),
           target_size,
           max_concurrent_tasks,
           max_spill_size,
-          min_commit_interval
+          min_commit_interval,
+          writer_properties,
+          post_commithook_properties,
+          commit_properties
         )
       @table.update_incremental
       result = JSON.parse(metrics)

data/lib/deltalake/utils.rb

@@ -0,0 +1,59 @@
+module DeltaLake
+  module Utils
+    def self.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 self.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 self.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

data/lib/deltalake/version.rb

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