deltalake-rb 0.1.0

data/ext/deltalake/src/lib.rs

@@ -0,0 +1,353 @@
+mod error;
+mod schema;
+mod utils;
+
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::future::IntoFuture;
+
+use chrono::Duration;
+use deltalake::arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
+use deltalake::kernel::StructType;
+use deltalake::operations::delete::DeleteBuilder;
+use deltalake::operations::vacuum::VacuumBuilder;
+use deltalake::storage::IORuntime;
+use deltalake::DeltaOps;
+use error::DeltaError;
+
+use magnus::{function, method, prelude::*, Error, Module, Ruby, Value};
+
+use crate::error::RubyError;
+use crate::schema::{schema_to_rbobject, Field};
+use crate::utils::rt;
+
+type RbResult<T> = Result<T, Error>;
+
+#[magnus::wrap(class = "DeltaLake::RawDeltaTable")]
+struct RawDeltaTable {
+    _table: RefCell<deltalake::DeltaTable>,
+}
+
+#[magnus::wrap(class = "DeltaLake::RawDeltaTableMetaData")]
+struct RawDeltaTableMetaData {
+    id: String,
+    name: Option<String>,
+    description: Option<String>,
+    partition_columns: Vec<String>,
+    created_time: Option<i64>,
+    configuration: HashMap<String, Option<String>>,
+}
+
+impl RawDeltaTable {
+    pub fn new(
+        table_uri: String,
+        version: Option<i64>,
+        storage_options: Option<HashMap<String, String>>,
+        without_files: bool,
+        log_buffer_size: Option<usize>,
+    ) -> RbResult<Self> {
+        let mut builder = deltalake::DeltaTableBuilder::from_uri(&table_uri)
+            .with_io_runtime(IORuntime::default());
+
+        if let Some(storage_options) = storage_options {
+            builder = builder.with_storage_options(storage_options)
+        }
+        if let Some(version) = version {
+            builder = builder.with_version(version)
+        }
+        if without_files {
+            builder = builder.without_files()
+        }
+        if let Some(buf_size) = log_buffer_size {
+            builder = builder
+                .with_log_buffer_size(buf_size)
+                .map_err(RubyError::from)?;
+        }
+
+        let table = rt().block_on(builder.load()).map_err(RubyError::from)?;
+        Ok(RawDeltaTable {
+            _table: RefCell::new(table),
+        })
+    }
+
+    pub fn is_deltatable(
+        table_uri: String,
+        storage_options: Option<HashMap<String, String>>,
+    ) -> RbResult<bool> {
+        let mut builder = deltalake::DeltaTableBuilder::from_uri(&table_uri);
+        if let Some(storage_options) = storage_options {
+            builder = builder.with_storage_options(storage_options)
+        }
+        Ok(rt()
+            .block_on(async {
+                match builder.build() {
+                    Ok(table) => table.verify_deltatable_existence().await,
+                    Err(err) => Err(err),
+                }
+            })
+            .map_err(RubyError::from)?)
+    }
+
+    pub fn table_uri(&self) -> RbResult<String> {
+        Ok(self._table.borrow().table_uri())
+    }
+
+    pub fn version(&self) -> RbResult<i64> {
+        Ok(self._table.borrow().version())
+    }
+
+    pub fn has_files(&self) -> RbResult<bool> {
+        Ok(self._table.borrow().config.require_files)
+    }
+
+    pub fn metadata(&self) -> RbResult<RawDeltaTableMetaData> {
+        let binding = self._table.borrow();
+        let metadata = binding.metadata().map_err(RubyError::from)?;
+        Ok(RawDeltaTableMetaData {
+            id: metadata.id.clone(),
+            name: metadata.name.clone(),
+            description: metadata.description.clone(),
+            partition_columns: metadata.partition_columns.clone(),
+            created_time: metadata.created_time,
+            configuration: metadata.configuration.clone(),
+        })
+    }
+
+    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."));
+        }
+
+        Ok(self
+            ._table
+            .borrow()
+            .get_files_iter()
+            .map_err(RubyError::from)?
+            .map(|f| f.to_string())
+            .collect())
+    }
+
+    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."));
+        }
+
+        Ok(self
+            ._table
+            .borrow()
+            .get_file_uris()
+            .map_err(RubyError::from)?
+            .collect())
+    }
+
+    pub fn schema(&self) -> RbResult<Value> {
+        let binding = self._table.borrow();
+        let schema: &StructType = binding.get_schema().map_err(RubyError::from)?;
+        schema_to_rbobject(schema.to_owned())
+    }
+
+    pub fn vacuum(
+        &self,
+        dry_run: bool,
+        retention_hours: Option<u64>,
+        enforce_retention_duration: bool,
+    ) -> RbResult<Vec<String>> {
+        let mut cmd = VacuumBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        )
+        .with_enforce_retention_duration(enforce_retention_duration)
+        .with_dry_run(dry_run);
+        if let Some(retention_period) = retention_hours {
+            cmd = cmd.with_retention_period(Duration::hours(retention_period as i64));
+        }
+
+        let (table, metrics) = rt().block_on(cmd.into_future()).map_err(RubyError::from)?;
+        self._table.borrow_mut().state = table.state;
+        Ok(metrics.files_deleted)
+    }
+
+    pub fn update_incremental(&self) -> RbResult<()> {
+        #[allow(deprecated)]
+        Ok(rt()
+            .block_on(self._table.borrow_mut().update_incremental(None))
+            .map_err(RubyError::from)?)
+    }
+
+    pub fn delete(&self, predicate: Option<String>) -> RbResult<String> {
+        let mut cmd = DeleteBuilder::new(
+            self._table.borrow().log_store(),
+            self._table
+                .borrow()
+                .snapshot()
+                .map_err(RubyError::from)?
+                .clone(),
+        );
+        if let Some(predicate) = predicate {
+            cmd = cmd.with_predicate(predicate);
+        }
+
+        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())
+    }
+}
+
+impl RawDeltaTableMetaData {
+    fn id(&self) -> String {
+        self.id.clone()
+    }
+
+    fn name(&self) -> Option<String> {
+        self.name.clone()
+    }
+
+    fn description(&self) -> Option<String> {
+        self.description.clone()
+    }
+
+    fn partition_columns(&self) -> Vec<String> {
+        self.partition_columns.clone()
+    }
+
+    fn created_time(&self) -> Option<i64> {
+        self.created_time
+    }
+
+    fn configuration(&self) -> HashMap<String, Option<String>> {
+        self.configuration.clone()
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn write_to_deltalake(
+    table_uri: String,
+    data: Value,
+    mode: String,
+    table: Option<&RawDeltaTable>,
+    schema_mode: Option<String>,
+    partition_by: Option<Vec<String>>,
+    predicate: Option<String>,
+    target_file_size: Option<usize>,
+    name: Option<String>,
+    description: Option<String>,
+    configuration: Option<HashMap<String, Option<String>>>,
+    storage_options: Option<HashMap<String, String>>,
+) -> RbResult<()> {
+    let capsule_pointer: usize = data.funcall("to_i", ())?;
+
+    // use similar approach as Polars to avoid copy
+    let stream_ptr =
+        Box::new(unsafe { std::ptr::replace(capsule_pointer as _, FFI_ArrowArrayStream::empty()) });
+    let stream = ArrowArrayStreamReader::try_new(*stream_ptr)
+        .map_err(|err| DeltaError::new_err(err.to_string()))?;
+
+    let batches = stream.map(|batch| batch.unwrap()).collect::<Vec<_>>();
+    let save_mode = mode.parse().map_err(RubyError::from)?;
+
+    let options = storage_options.clone().unwrap_or_default();
+    let table = if let Some(table) = table {
+        DeltaOps(table._table.borrow().clone())
+    } else {
+        rt().block_on(DeltaOps::try_from_uri_with_storage_options(
+            &table_uri, options,
+        ))
+        .map_err(RubyError::from)?
+    };
+
+    let mut builder = table.write(batches).with_save_mode(save_mode);
+    if let Some(schema_mode) = schema_mode {
+        builder = builder.with_schema_mode(schema_mode.parse().map_err(RubyError::from)?);
+    }
+    if let Some(partition_columns) = partition_by {
+        builder = builder.with_partition_columns(partition_columns);
+    }
+
+    if let Some(name) = &name {
+        builder = builder.with_table_name(name);
+    };
+
+    if let Some(description) = &description {
+        builder = builder.with_description(description);
+    };
+
+    if let Some(predicate) = predicate {
+        builder = builder.with_replace_where(predicate);
+    };
+
+    if let Some(target_file_size) = target_file_size {
+        builder = builder.with_target_file_size(target_file_size)
+    };
+
+    if let Some(config) = configuration {
+        builder = builder.with_configuration(config);
+    };
+
+    rt().block_on(builder.into_future())
+        .map_err(RubyError::from)?;
+
+    Ok(())
+}
+
+#[magnus::init]
+fn init(ruby: &Ruby) -> RbResult<()> {
+    deltalake::aws::register_handlers(None);
+
+    let module = ruby.define_module("DeltaLake")?;
+    module.define_singleton_method("write_deltalake_rust", function!(write_to_deltalake, 12))?;
+
+    let class = module.define_class("RawDeltaTable", ruby.class_object())?;
+    class.define_singleton_method("new", function!(RawDeltaTable::new, 5))?;
+    class.define_singleton_method("is_deltatable", function!(RawDeltaTable::is_deltatable, 2))?;
+    class.define_method("table_uri", method!(RawDeltaTable::table_uri, 0))?;
+    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("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("schema", method!(RawDeltaTable::schema, 0))?;
+    class.define_method("vacuum", method!(RawDeltaTable::vacuum, 3))?;
+    class.define_method(
+        "update_incremental",
+        method!(RawDeltaTable::update_incremental, 0),
+    )?;
+    class.define_method("delete", method!(RawDeltaTable::delete, 1))?;
+
+    let class = module.define_class("RawDeltaTableMetaData", ruby.class_object())?;
+    class.define_method("id", method!(RawDeltaTableMetaData::id, 0))?;
+    class.define_method("name", method!(RawDeltaTableMetaData::name, 0))?;
+    class.define_method(
+        "description",
+        method!(RawDeltaTableMetaData::description, 0),
+    )?;
+    class.define_method(
+        "partition_columns",
+        method!(RawDeltaTableMetaData::partition_columns, 0),
+    )?;
+    class.define_method(
+        "created_time",
+        method!(RawDeltaTableMetaData::created_time, 0),
+    )?;
+    class.define_method(
+        "configuration",
+        method!(RawDeltaTableMetaData::configuration, 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))?;
+    class.define_method("nullable", method!(Field::nullable, 0))?;
+
+    Ok(())
+}

data/ext/deltalake/src/schema.rs

@@ -0,0 +1,37 @@
+use deltalake::kernel::{StructField, StructType as DeltaStructType};
+use magnus::{value::ReprValue, Module, RArray, RModule, Ruby, Value};
+
+use crate::RbResult;
+
+pub fn schema_to_rbobject(schema: DeltaStructType) -> RbResult<Value> {
+    let fields = schema.fields().map(|field| Field {
+        inner: field.clone(),
+    });
+
+    let rb_schema: Value = Ruby::get()
+        .unwrap()
+        .class_object()
+        .const_get::<_, RModule>("DeltaLake")?
+        .const_get("Schema")?;
+
+    rb_schema.funcall("new", (RArray::from_iter(fields),))
+}
+
+#[magnus::wrap(class = "DeltaLake::Field")]
+pub struct Field {
+    pub inner: StructField,
+}
+
+impl Field {
+    pub fn name(&self) -> String {
+        self.inner.name().to_string()
+    }
+
+    pub fn get_type(&self) -> String {
+        self.inner.data_type().to_string()
+    }
+
+    pub fn nullable(&self) -> bool {
+        self.inner.is_nullable()
+    }
+}

data/ext/deltalake/src/utils.rs

@@ -0,0 +1,21 @@
+use std::sync::OnceLock;
+
+use tokio::runtime::Runtime;
+
+#[inline]
+pub fn rt() -> &'static Runtime {
+    static TOKIO_RT: OnceLock<Runtime> = OnceLock::new();
+    static PID: OnceLock<u32> = OnceLock::new();
+    let pid = std::process::id();
+    let runtime_pid = *PID.get_or_init(|| pid);
+    if pid != runtime_pid {
+        panic!(
+            "Forked process detected - current PID is {} but the tokio runtime was created by {}. The tokio \
+            runtime does not support forked processes https://github.com/tokio-rs/tokio/issues/4301. If you are \
+            seeing this message while using Ruby multithreading make sure to use the `spawn` or `forkserver` \
+            mode.",
+            pid, runtime_pid
+        );
+    }
+    TOKIO_RT.get_or_init(|| Runtime::new().expect("Failed to create a tokio runtime."))
+}

data/lib/deltalake/field.rb

@@ -0,0 +1,12 @@
+module DeltaLake
+  class Field
+    def inspect
+      attributes = {
+        name: name,
+        type: type,
+        nullable: nullable
+      }
+      "<#{self.class.name} #{attributes.map { |k, v| "#{k}=#{v.inspect}" }.join(", ")}>"
+    end
+  end
+end

data/lib/deltalake/metadata.rb

@@ -0,0 +1,43 @@
+module DeltaLake
+  class Metadata
+    def initialize(table)
+      @metadata = table.metadata
+    end
+
+    def id
+      @metadata.id
+    end
+
+    def name
+      @metadata.name
+    end
+
+    def description
+      @metadata.description
+    end
+
+    def partition_columns
+      @metadata.partition_columns
+    end
+
+    def created_time
+      @metadata.created_time
+    end
+
+    def configuration
+      @metadata.configuration
+    end
+
+    def inspect
+      attributes = {
+        id: id,
+        name: name,
+        description: description,
+        partition_columns: partition_columns,
+        created_time: created_time,
+        configuration: configuration
+      }
+      "<#{self.class.name} #{attributes.map { |k, v| "#{k}=#{v.inspect}" }.join(", ")}>"
+    end
+  end
+end

data/lib/deltalake/schema.rb

@@ -0,0 +1,9 @@
+module DeltaLake
+  class Schema
+    attr_reader :fields
+
+    def initialize(fields)
+      @fields = fields
+    end
+  end
+end

data/lib/deltalake/table.rb

@@ -0,0 +1,103 @@
+module DeltaLake
+  class Table
+    def initialize(
+      table_uri,
+      version: nil,
+      storage_options: nil,
+      without_files: false,
+      log_buffer_size: nil
+    )
+      @storage_options = storage_options
+      @table =
+        RawDeltaTable.new(
+          table_uri,
+          version,
+          storage_options,
+          without_files,
+          log_buffer_size
+        )
+    end
+
+    def self.exists?(table_uri, storage_options: nil)
+      RawDeltaTable.is_deltatable(table_uri, storage_options)
+    end
+
+    def version
+      @table.version
+    end
+
+    def files
+      @table.files
+    end
+
+    def file_uris
+      @table.file_uris
+    end
+
+    def load_as_version(version)
+      if version.is_a?(Integer)
+        @table.load_version(version)
+      else
+        raise TypeError, "Invalid datatype provided for version, only Integer is accepted."
+      end
+    end
+
+    def table_uri
+      @table.table_uri
+    end
+
+    def schema
+      @table.schema
+    end
+
+    def metadata
+      Metadata.new(@table)
+    end
+
+    def vacuum(
+      retention_hours: nil,
+      dry_run: true,
+      enforce_retention_duration: true
+    )
+      if retention_hours
+        if retention_hours < 0
+          raise ArgumentError, "The retention periods should be positive."
+        end
+      end
+
+      @table.vacuum(
+        dry_run,
+        retention_hours,
+        enforce_retention_duration
+      )
+    end
+
+    def to_polars(eager: true)
+      require "polars-df"
+
+      sources = file_uris
+      lf =
+        if sources.empty?
+          Polars::LazyFrame.new
+        else
+          storage_options = @storage_options&.except("AWS_S3_ALLOW_UNSAFE_RENAME")
+          Polars.scan_parquet(sources, storage_options: storage_options)
+        end
+      eager ? lf.collect : lf
+    end
+
+    def update_incremental
+      @table.update_incremental
+    end
+
+    def delete(predicate = nil)
+      metrics = @table.delete(predicate)
+      JSON.parse(metrics).transform_keys(&:to_sym)
+    end
+
+    # private
+    def _table
+      @table
+    end
+  end
+end

data/lib/deltalake/version.rb

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

data/lib/deltalake-rb.rb

@@ -0,0 +1 @@
+require_relative "deltalake"

data/lib/deltalake.rb

@@ -0,0 +1,104 @@
+# ext
+begin
+  require "deltalake/#{RUBY_VERSION.to_f}/deltalake"
+rescue LoadError
+  require "deltalake/deltalake"
+end
+
+# stdlib
+require "json"
+
+# modules
+require_relative "deltalake/field"
+require_relative "deltalake/metadata"
+require_relative "deltalake/schema"
+require_relative "deltalake/table"
+require_relative "deltalake/version"
+
+module DeltaLake
+  class Error < StandardError; end
+  class TableNotFoundError < Error; end
+  class DeltaProtocolError < Error; end
+  class CommitFailedError < Error; end
+  class SchemaMismatchError < Error; end
+
+  class << self
+
+    def write(
+      table_or_uri,
+      data,
+      partition_by: nil,
+      mode: "error",
+      name: nil,
+      description: nil,
+      configuration: nil,
+      schema_mode: nil,
+      storage_options: nil,
+      predicate: nil,
+      target_file_size: nil
+    )
+      table, table_uri = try_get_table_and_table_uri(table_or_uri, storage_options)
+
+      if partition_by.is_a?(String)
+        partition_by = [partition_by]
+      end
+
+      if !table.nil? && mode == "ignore"
+        return
+      end
+
+      data = convert_data(data)
+
+      write_deltalake_rust(
+        table_uri,
+        data,
+        mode,
+        table&._table,
+        schema_mode,
+        partition_by,
+        predicate,
+        target_file_size,
+        name,
+        description,
+        configuration,
+        storage_options
+      )
+
+      if table
+        table.update_incremental
+      end
+    end
+
+    private
+
+    def try_get_table_and_table_uri(table_or_uri, storage_options)
+      if !table_or_uri.is_a?(String) && !table_or_uri.is_a?(Table)
+        raise ArgumentError, "table_or_uri must be a String or Table"
+      end
+
+      if table_or_uri.is_a?(String)
+        table = try_get_deltatable(table_or_uri, storage_options)
+        table_uri = table_or_uri.to_s
+      else
+        table = table_or_uri
+        table_uri = table._table.table_uri
+      end
+
+      [table, table_uri]
+    end
+
+    def try_get_deltatable(table_uri, storage_options)
+      Table.new(table_uri, storage_options: storage_options)
+    rescue TableNotFoundError
+      nil
+    end
+
+    def convert_data(data)
+      if data.respond_to?(:arrow_c_stream)
+        data.arrow_c_stream
+      else
+        raise TypeError, "Only objects implementing the Arrow C stream interface are valid inputs for source."
+      end
+    end
+  end
+end