iceberg 0.7.0 → 0.10.1

data/ext/iceberg/src/table.rs

@@ -0,0 +1,369 @@
+use arrow_array::ffi_stream::ArrowArrayStreamReader;
+use iceberg::TableIdent;
+use iceberg::io::FileIO;
+use iceberg::spec::FormatVersion;
+use iceberg::table::{StaticTable, Table};
+use iceberg::transaction::{ApplyTransactionAction, Transaction};
+use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
+use iceberg::writer::file_writer::ParquetWriterBuilder;
+use iceberg::writer::file_writer::location_generator::{
+    DefaultFileNameGenerator, DefaultLocationGenerator,
+};
+use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
+use magnus::{Error as RbErr, RArray, Ruby, Value};
+use parquet::file::properties::WriterProperties;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::sync::Arc;
+use uuid::Uuid;
+
+use crate::RbResult;
+use crate::arrow::RbArrowType;
+use crate::catalog::RbCatalog;
+use crate::error::to_rb_err;
+use crate::runtime::runtime;
+use crate::scan::RbTableScan;
+use crate::utils::*;
+
+#[magnus::wrap(class = "Iceberg::RbTable")]
+pub struct RbTable {
+    pub table: RefCell<Table>,
+}
+
+impl RbTable {
+    pub fn scan(&self, snapshot_id: Option<i64>) -> RbResult<RbTableScan> {
+        let table = self.table.borrow();
+        let mut builder = table.scan();
+        if let Some(si) = snapshot_id {
+            builder = builder.snapshot_id(si);
+        }
+        let scan = builder.build().map_err(to_rb_err)?;
+        Ok(RbTableScan { scan: scan.into() })
+    }
+
+    pub fn append(
+        ruby: &Ruby,
+        rb_self: &Self,
+        data: RbArrowType<ArrowArrayStreamReader>,
+        catalog: &RbCatalog,
+    ) -> RbResult<RbTable> {
+        let runtime = runtime();
+        let table = rb_self.table.borrow();
+        let catalog = catalog.catalog.borrow();
+
+        let table_schema: Arc<arrow_schema::Schema> = Arc::new(
+            table
+                .metadata()
+                .current_schema()
+                .as_ref()
+                .try_into()
+                .unwrap(),
+        );
+
+        let location_generator =
+            DefaultLocationGenerator::new(table.metadata().clone()).map_err(to_rb_err)?;
+        let file_name_generator = DefaultFileNameGenerator::new(
+            // TODO move task id to suffix to match Python and Java
+            "0".to_string(),
+            Some(Uuid::new_v4().to_string()),
+            iceberg::spec::DataFileFormat::Parquet,
+        );
+
+        let parquet_writer_builder = ParquetWriterBuilder::new(
+            WriterProperties::default(),
+            table.metadata().current_schema().clone(),
+            table.file_io().clone(),
+            location_generator.clone(),
+            file_name_generator.clone(),
+        );
+        let data_file_writer_builder = DataFileWriterBuilder::new(parquet_writer_builder, None, 0);
+        let mut data_file_writer = runtime
+            .block_on(data_file_writer_builder.build())
+            .map_err(to_rb_err)?;
+
+        for batch in data.0 {
+            let batch = batch
+                .unwrap()
+                .with_schema(table_schema.clone())
+                .map_err(|e| RbErr::new(ruby.exception_arg_error(), e.to_string()))?;
+            runtime
+                .block_on(data_file_writer.write(batch))
+                .map_err(to_rb_err)?;
+        }
+
+        let data_files = runtime
+            .block_on(data_file_writer.close())
+            .map_err(to_rb_err)?;
+
+        let tx = Transaction::new(&table);
+        let append_action = tx.fast_append().add_data_files(data_files.clone());
+        let tx = append_action.apply(tx).map_err(to_rb_err)?;
+        let table = runtime
+            .block_on(tx.commit(catalog.as_catalog()))
+            .map_err(to_rb_err)?;
+
+        Ok(RbTable {
+            table: table.into(),
+        })
+    }
+
+    pub fn format_version(&self) -> i32 {
+        match self.table.borrow().metadata().format_version() {
+            FormatVersion::V1 => 1,
+            FormatVersion::V2 => 2,
+        }
+    }
+
+    pub fn uuid(&self) -> String {
+        self.table.borrow().metadata().uuid().to_string()
+    }
+
+    pub fn location(&self) -> String {
+        self.table.borrow().metadata().location().to_string()
+    }
+
+    pub fn last_sequence_number(&self) -> i64 {
+        self.table.borrow().metadata().last_sequence_number()
+    }
+
+    pub fn next_sequence_number(&self) -> i64 {
+        self.table.borrow().metadata().next_sequence_number()
+    }
+
+    pub fn last_column_id(&self) -> i32 {
+        self.table.borrow().metadata().last_column_id()
+    }
+
+    pub fn last_partition_id(&self) -> i32 {
+        self.table.borrow().metadata().last_partition_id()
+    }
+
+    pub fn last_updated_ms(&self) -> i64 {
+        self.table.borrow().metadata().last_updated_ms()
+    }
+
+    pub fn schemas(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let schemas = ruby.ary_new();
+        for s in rb_self.table.borrow().metadata().schemas_iter() {
+            schemas.push(rb_schema(s)?)?;
+        }
+        Ok(schemas)
+    }
+
+    pub fn schema_by_id(&self, schema_id: i32) -> RbResult<Option<Value>> {
+        let schema = match self.table.borrow().metadata().schema_by_id(schema_id) {
+            Some(s) => Some(rb_schema(s)?),
+            None => None,
+        };
+        Ok(schema)
+    }
+
+    pub fn current_schema(&self) -> RbResult<Value> {
+        rb_schema(self.table.borrow().metadata().current_schema())
+    }
+
+    pub fn current_schema_id(&self) -> i32 {
+        self.table.borrow().metadata().current_schema_id()
+    }
+
+    pub fn partition_specs(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let partition_specs = ruby.ary_new();
+        for s in rb_self.table.borrow().metadata().partition_specs_iter() {
+            partition_specs.push(rb_partition_spec(s)?)?;
+        }
+        Ok(partition_specs)
+    }
+
+    pub fn partition_spec_by_id(&self, partition_spec_id: i32) -> RbResult<Option<Value>> {
+        let partition_spec = match self
+            .table
+            .borrow()
+            .metadata()
+            .partition_spec_by_id(partition_spec_id)
+        {
+            Some(s) => Some(rb_partition_spec(s)?),
+            None => None,
+        };
+        Ok(partition_spec)
+    }
+
+    pub fn default_partition_spec(&self) -> RbResult<Value> {
+        rb_partition_spec(self.table.borrow().metadata().default_partition_spec())
+    }
+
+    pub fn default_partition_spec_id(&self) -> i32 {
+        self.table.borrow().metadata().default_partition_spec_id()
+    }
+
+    pub fn snapshots(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let snapshots = ruby.ary_new();
+        for s in rb_self.table.borrow().metadata().snapshots() {
+            snapshots.push(rb_snapshot(s)?)?;
+        }
+        Ok(snapshots)
+    }
+
+    pub fn snapshot_by_id(&self, snapshot_id: i64) -> RbResult<Option<Value>> {
+        let snapshot = match self.table.borrow().metadata().snapshot_by_id(snapshot_id) {
+            Some(s) => Some(rb_snapshot(s)?),
+            None => None,
+        };
+        Ok(snapshot)
+    }
+
+    pub fn history(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let history = ruby.ary_new();
+        for s in rb_self.table.borrow().metadata().history() {
+            let snapshot_log = ruby.hash_new();
+            snapshot_log.aset(ruby.to_symbol("snapshot_id"), s.snapshot_id)?;
+            // TODO timestamp
+            history.push(snapshot_log)?;
+        }
+        Ok(history)
+    }
+
+    pub fn metadata_log(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let metadata_logs = ruby.ary_new();
+        for s in rb_self.table.borrow().metadata().metadata_log() {
+            let metadata_log = ruby.hash_new();
+            metadata_log.aset(
+                ruby.to_symbol("metadata_file"),
+                ruby.str_new(&s.metadata_file),
+            )?;
+            // TODO timestamp
+            metadata_logs.push(metadata_log)?;
+        }
+        Ok(metadata_logs)
+    }
+
+    pub fn current_snapshot(&self) -> RbResult<Option<Value>> {
+        let snapshot = match self.table.borrow().metadata().current_snapshot() {
+            Some(s) => Some(rb_snapshot(s)?),
+            None => None,
+        };
+        Ok(snapshot)
+    }
+
+    pub fn current_snapshot_id(&self) -> Option<i64> {
+        self.table.borrow().metadata().current_snapshot_id()
+    }
+
+    pub fn snapshot_for_ref(&self, ref_name: String) -> RbResult<Option<Value>> {
+        let snapshot = match self.table.borrow().metadata().snapshot_for_ref(&ref_name) {
+            Some(s) => Some(rb_snapshot(s)?),
+            None => None,
+        };
+        Ok(snapshot)
+    }
+
+    pub fn sort_orders(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let sort_orders = ruby.ary_new();
+        for s in rb_self.table.borrow().metadata().sort_orders_iter() {
+            sort_orders.push(rb_sort_order(s)?)?;
+        }
+        Ok(sort_orders)
+    }
+
+    pub fn sort_order_by_id(&self, sort_order_id: i64) -> RbResult<Option<Value>> {
+        let sort_order = match self
+            .table
+            .borrow()
+            .metadata()
+            .sort_order_by_id(sort_order_id)
+        {
+            Some(s) => Some(rb_sort_order(s)?),
+            None => None,
+        };
+        Ok(sort_order)
+    }
+
+    pub fn default_sort_order(&self) -> RbResult<Value> {
+        rb_sort_order(self.table.borrow().metadata().default_sort_order())
+    }
+
+    pub fn default_sort_order_id(&self) -> i64 {
+        self.table.borrow().metadata().default_sort_order_id()
+    }
+
+    pub fn properties(&self) -> HashMap<String, String> {
+        self.table.borrow().metadata().properties().clone()
+    }
+
+    pub fn statistics(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let statistics = ruby.ary_new();
+        for s in rb_self.table.borrow().metadata().statistics_iter() {
+            statistics.push(rb_statistics_file(s)?)?;
+        }
+        Ok(statistics)
+    }
+
+    pub fn partition_statistics(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let statistics = ruby.ary_new();
+        for s in rb_self
+            .table
+            .borrow()
+            .metadata()
+            .partition_statistics_iter()
+        {
+            statistics.push(rb_partition_statistics_file(s)?)?;
+        }
+        Ok(statistics)
+    }
+
+    pub fn statistics_for_snapshot(&self, snapshot_id: i64) -> RbResult<Option<Value>> {
+        let statistics = match self
+            .table
+            .borrow()
+            .metadata()
+            .statistics_for_snapshot(snapshot_id)
+        {
+            Some(s) => Some(rb_statistics_file(s)?),
+            None => None,
+        };
+        Ok(statistics)
+    }
+
+    pub fn partition_statistics_for_snapshot(&self, snapshot_id: i64) -> RbResult<Option<Value>> {
+        let statistics = match self
+            .table
+            .borrow()
+            .metadata()
+            .partition_statistics_for_snapshot(snapshot_id)
+        {
+            Some(s) => Some(rb_partition_statistics_file(s)?),
+            None => None,
+        };
+        Ok(statistics)
+    }
+
+    pub fn encryption_keys(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        let encryption_keys = ruby.ary_new();
+        for (k, v) in rb_self.table.borrow().metadata().encryption_keys_iter() {
+            encryption_keys.push((ruby.str_new(k), ruby.str_new(v)))?;
+        }
+        Ok(encryption_keys)
+    }
+
+    pub fn encryption_key(&self, key_id: String) -> Option<String> {
+        self.table
+            .borrow()
+            .metadata()
+            .encryption_key(&key_id)
+            .cloned()
+    }
+
+    pub fn from_metadata_file(location: String) -> RbResult<Self> {
+        let file_io = FileIO::from_path(&location).unwrap().build().unwrap();
+        let table_ident = TableIdent::from_strs(["static-table", &location]).unwrap();
+        let static_table = runtime()
+            .block_on(StaticTable::from_metadata_file(
+                &location,
+                table_ident,
+                file_io,
+            ))
+            .map_err(to_rb_err)?;
+        Ok(Self {
+            table: static_table.into_table().into(),
+        })
+    }
+}

data/ext/iceberg/src/utils.rs

@@ -0,0 +1,245 @@
+use iceberg::spec::{
+    Literal, NestedField, PartitionSpec, PartitionStatisticsFile, PrimitiveLiteral, PrimitiveType,
+    Schema, Snapshot, SortOrder, StatisticsFile, Type,
+};
+use iceberg::{NamespaceIdent, TableIdent};
+use magnus::{
+    Error as RbErr, IntoValue, RArray, RClass, RHash, RModule, Ruby, TryConvert, Value, kwargs,
+    prelude::*,
+};
+
+use crate::RbResult;
+use crate::error::to_rb_err;
+
+pub struct Wrap<T>(pub T);
+
+impl TryConvert for Wrap<NamespaceIdent> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let ident = if let Ok(vec) = Vec::<String>::try_convert(ob) {
+            // prevent dots due to buggy behavior with iceberg-catalog-{rest,sql} crates
+            if vec.iter().any(|v| v.contains(".")) {
+                return Err(RbErr::new(
+                    Ruby::get_with(ob).exception_arg_error(),
+                    "Unsupported namespace",
+                ));
+            }
+            NamespaceIdent::from_vec(vec).map_err(to_rb_err)?
+        } else {
+            NamespaceIdent::from_strs(String::try_convert(ob)?.split(".")).map_err(to_rb_err)?
+        };
+        Ok(Wrap(ident))
+    }
+}
+
+impl TryConvert for Wrap<TableIdent> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let ident = if let Ok(vec) = Vec::<String>::try_convert(ob) {
+            TableIdent::from_strs(vec.iter()).map_err(to_rb_err)?
+        } else {
+            TableIdent::from_strs(String::try_convert(ob)?.split(".")).map_err(to_rb_err)?
+        };
+        Ok(Wrap(ident))
+    }
+}
+
+impl TryConvert for Wrap<Schema> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let ruby = Ruby::get_with(ob);
+        let mut fields = Vec::new();
+        let rb_fields: RArray = ob.funcall("fields", ())?;
+        for rb_field in rb_fields {
+            let rb_field = RHash::try_convert(rb_field)?;
+            let rb_type: Value = rb_field.aref(ruby.to_symbol("type"))?;
+            let field_type = if let Ok(s) = String::try_convert(rb_type) {
+                match s.as_str() {
+                    "boolean" => Type::Primitive(PrimitiveType::Boolean),
+                    "int" => Type::Primitive(PrimitiveType::Int),
+                    "long" => Type::Primitive(PrimitiveType::Long),
+                    "float" => Type::Primitive(PrimitiveType::Float),
+                    "double" => Type::Primitive(PrimitiveType::Double),
+                    // TODO PrimitiveType::Decimal
+                    "date" => Type::Primitive(PrimitiveType::Date),
+                    "time" => Type::Primitive(PrimitiveType::Time),
+                    "timestamp" => Type::Primitive(PrimitiveType::Timestamp),
+                    "timestamptz" => Type::Primitive(PrimitiveType::Timestamptz),
+                    "timestamp_ns" => Type::Primitive(PrimitiveType::TimestampNs),
+                    "timestamptz_ns" => Type::Primitive(PrimitiveType::TimestamptzNs),
+                    "string" => Type::Primitive(PrimitiveType::String),
+                    "uuid" => Type::Primitive(PrimitiveType::Uuid),
+                    // TODO PrimitiveType::Fixed
+                    "binary" => Type::Primitive(PrimitiveType::Binary),
+                    _ => {
+                        return Err(RbErr::new(
+                            ruby.exception_arg_error(),
+                            format!("Type not supported: {}", s),
+                        ));
+                    }
+                }
+            } else {
+                let class_name = unsafe { rb_type.classname() }.to_string();
+                match class_name.as_str() {
+                    "Polars::Boolean" => Type::Primitive(PrimitiveType::Boolean),
+                    "Polars::Int32" => Type::Primitive(PrimitiveType::Int),
+                    "Polars::Int64" => Type::Primitive(PrimitiveType::Long),
+                    "Polars::Float32" => Type::Primitive(PrimitiveType::Float),
+                    "Polars::Float64" => Type::Primitive(PrimitiveType::Double),
+                    "Polars::Date" => Type::Primitive(PrimitiveType::Date),
+                    "Polars::Time" => Type::Primitive(PrimitiveType::Time),
+                    "Polars::String" => Type::Primitive(PrimitiveType::String),
+                    "Polars::Binary" => Type::Primitive(PrimitiveType::Binary),
+                    _ => {
+                        return Err(RbErr::new(
+                            ruby.exception_arg_error(),
+                            format!("Type not supported: {}", class_name),
+                        ));
+                    }
+                }
+            };
+
+            let initial_default = rb_field.aref(ruby.to_symbol("initial_default"))?;
+            let write_default = rb_field.aref(ruby.to_symbol("write_default"))?;
+
+            let initial_default = default_value(initial_default, &field_type)?;
+            let write_default = default_value(write_default, &field_type)?;
+
+            fields.push(
+                NestedField {
+                    id: rb_field.aref(ruby.to_symbol("id"))?,
+                    name: rb_field.aref(ruby.to_symbol("name"))?,
+                    required: rb_field.aref(ruby.to_symbol("required"))?,
+                    field_type: field_type.into(),
+                    doc: rb_field.aref(ruby.to_symbol("doc"))?,
+                    initial_default,
+                    write_default,
+                }
+                .into(),
+            );
+        }
+        let schema = Schema::builder()
+            .with_fields(fields)
+            .build()
+            .map_err(to_rb_err)?;
+        Ok(Wrap(schema))
+    }
+}
+
+fn default_value(ob: Value, field_type: &Type) -> RbResult<Option<Literal>> {
+    if ob.is_nil() {
+        return Ok(None);
+    }
+
+    let lit = match field_type {
+        Type::Primitive(ty) => {
+            let pl = match ty {
+                PrimitiveType::Boolean => PrimitiveLiteral::Boolean(bool::try_convert(ob)?),
+                PrimitiveType::Int => PrimitiveLiteral::Int(i32::try_convert(ob)?),
+                PrimitiveType::Long => PrimitiveLiteral::Long(i64::try_convert(ob)?),
+                PrimitiveType::Float => PrimitiveLiteral::Float(f32::try_convert(ob)?.into()),
+                PrimitiveType::Double => PrimitiveLiteral::Double(f64::try_convert(ob)?.into()),
+                PrimitiveType::String => PrimitiveLiteral::String(String::try_convert(ob)?),
+                _ => todo!(),
+            };
+            Literal::Primitive(pl)
+        }
+        _ => todo!(),
+    };
+    Ok(Some(lit))
+}
+
+pub fn rb_schema(schema: &Schema) -> RbResult<Value> {
+    let ruby = Ruby::get().unwrap();
+    let fields = ruby.ary_new();
+    for f in schema.as_struct().fields() {
+        let field = ruby.hash_new();
+        field.aset(ruby.to_symbol("id"), f.id)?;
+        field.aset(ruby.to_symbol("name"), ruby.str_new(&f.name))?;
+
+        let field_type = match &*f.field_type {
+            Type::Primitive(ty) => match ty {
+                PrimitiveType::Boolean => "boolean",
+                PrimitiveType::Int => "int",
+                PrimitiveType::Long => "long",
+                PrimitiveType::Float => "float",
+                PrimitiveType::Double => "double",
+                PrimitiveType::Decimal {
+                    precision: _,
+                    scale: _,
+                } => todo!(),
+                PrimitiveType::Date => "date",
+                PrimitiveType::Time => "time",
+                PrimitiveType::Timestamp => "timestamp",
+                PrimitiveType::Timestamptz => "timestamptz",
+                PrimitiveType::TimestampNs => "timestamp_ns",
+                PrimitiveType::TimestamptzNs => "timestamptz_ns",
+                PrimitiveType::String => "string",
+                PrimitiveType::Uuid => "uuid",
+                PrimitiveType::Binary => "binary",
+                PrimitiveType::Fixed(_) => todo!(),
+            },
+            _ => todo!(),
+        };
+        field.aset(ruby.to_symbol("type"), field_type)?;
+
+        field.aset(ruby.to_symbol("required"), f.required)?;
+
+        let initial_default = f.initial_default.as_ref().map(rb_literal);
+        field.aset(ruby.to_symbol("initial_default"), initial_default)?;
+
+        let write_default = f.write_default.as_ref().map(rb_literal);
+        field.aset(ruby.to_symbol("write_default"), write_default)?;
+
+        field.aset(
+            ruby.to_symbol("doc"),
+            f.doc.as_ref().map(|v| ruby.str_new(v)),
+        )?;
+
+        fields.push(field)?;
+    }
+    let schema_id = schema.schema_id();
+
+    ruby.class_object()
+        .const_get::<_, RModule>("Iceberg")
+        .unwrap()
+        .const_get::<_, RClass>("Schema")
+        .unwrap()
+        .funcall("new", (fields, kwargs!("schema_id" => schema_id)))
+}
+
+pub fn rb_snapshot(_snapshot: &Snapshot) -> RbResult<Value> {
+    todo!();
+}
+
+pub fn rb_partition_spec(_partition_spec: &PartitionSpec) -> RbResult<Value> {
+    todo!();
+}
+
+pub fn rb_sort_order(_sort_order: &SortOrder) -> RbResult<Value> {
+    todo!();
+}
+
+pub fn rb_statistics_file(_statistics_file: &StatisticsFile) -> RbResult<Value> {
+    todo!();
+}
+
+pub fn rb_partition_statistics_file(
+    _partition_statistics_file: &PartitionStatisticsFile,
+) -> RbResult<Value> {
+    todo!();
+}
+
+pub fn rb_literal(literal: &Literal) -> Value {
+    let ruby = Ruby::get().unwrap();
+    match literal {
+        Literal::Primitive(pl) => match pl {
+            PrimitiveLiteral::Boolean(v) => v.into_value_with(&ruby),
+            PrimitiveLiteral::Int(v) => v.into_value_with(&ruby),
+            PrimitiveLiteral::Long(v) => v.into_value_with(&ruby),
+            PrimitiveLiteral::Float(v) => v.into_value_with(&ruby),
+            PrimitiveLiteral::Double(v) => v.into_value_with(&ruby),
+            PrimitiveLiteral::String(v) => ruby.str_new(v).as_value(),
+            PrimitiveLiteral::Binary(v) => ruby.str_from_slice(v).as_value(),
+            _ => todo!(),
+        },
+        _ => todo!(),
+    }
+}