deltalake-rb 0.1.0 → 0.1.2

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

@@ -26,22 +26,51 @@ 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)
+        @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,10 +83,35 @@ 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,
-      enforce_retention_duration: true
+      enforce_retention_duration: true,
+      post_commithook_properties: nil,
+      commit_properties: nil
     )
       if retention_hours
         if retention_hours < 0
@@ -68,10 +122,72 @@ module DeltaLake
       @table.vacuum(
         dry_run,
         retention_hours,
-        enforce_retention_duration
+        enforce_retention_duration,
+        commit_properties,
+        post_commithook_properties
       )
     end
 
+    def optimize
+      TableOptimizer.new(self)
+    end
+
+    def alter
+      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,
+      commit_properties: nil
+    )
+      if target.is_a?(Time)
+        metrics =
+          @table.restore(
+            target.utc.iso8601(9),
+            ignore_missing_files,
+            protocol_downgrade_allowed,
+            commit_properties
+          )
+      else
+        metrics =
+          @table.restore(
+            target,
+            ignore_missing_files,
+            protocol_downgrade_allowed,
+            commit_properties
+          )
+      end
+      JSON.parse(metrics)
+    end
+
     def to_polars(eager: true)
       require "polars-df"
 
@@ -80,7 +196,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
@@ -90,14 +212,52 @@ 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,
+      post_commithook_properties: nil,
+      commit_properties: nil
+    )
+      metrics =
+        @table.repair(
+          dry_run,
+          commit_properties,
+          post_commithook_properties
+        )
+      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,58 @@
+module DeltaLake
+  class TableAlterer
+    def initialize(table)
+      @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,
+          "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
+
+    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

@@ -0,0 +1,67 @@
+module DeltaLake
+  class TableOptimizer
+    def initialize(table)
+      @table = table
+    end
+
+    def compact(
+      partition_filters: nil,
+      target_size: nil,
+      max_concurrent_tasks: 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,
+          writer_properties,
+          post_commithook_properties,
+          commit_properties
+        )
+      @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,
+      partition_filters: nil,
+      target_size: nil,
+      max_concurrent_tasks:  nil,
+      max_spill_size: 20 * 1024 * 1024 * 1024,
+      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,
+          writer_properties,
+          post_commithook_properties,
+          commit_properties
+        )
+      @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/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.0"
+  VERSION = "0.1.2"
 end