iceberg 0.11.1 → 0.12.0

data/ext/iceberg/src/ruby.rs

@@ -0,0 +1,51 @@
+use std::ffi::c_void;
+use std::ptr::null_mut;
+
+use magnus::Ruby;
+use rb_sys::rb_thread_call_without_gvl;
+
+pub trait GvlExt {
+    fn detach<T, F>(&self, func: F) -> T
+    where
+        F: Send + FnOnce() -> T,
+        T: Send;
+}
+
+impl GvlExt for Ruby {
+    fn detach<T, F>(&self, func: F) -> T
+    where
+        F: Send + FnOnce() -> T,
+        T: Send,
+    {
+        let mut data = CallbackData {
+            func: Some(func),
+            result: None,
+        };
+
+        unsafe {
+            rb_thread_call_without_gvl(
+                Some(call_without_gvl::<F, T>),
+                &mut data as *mut _ as *mut c_void,
+                None,
+                null_mut(),
+            );
+        }
+
+        data.result.unwrap()
+    }
+}
+
+struct CallbackData<F, T> {
+    func: Option<F>,
+    result: Option<T>,
+}
+
+extern "C" fn call_without_gvl<F, T>(data: *mut c_void) -> *mut c_void
+where
+    F: FnOnce() -> T,
+{
+    let data = unsafe { &mut *(data as *mut CallbackData<F, T>) };
+    let func = data.func.take().unwrap();
+    data.result = Some(func());
+    null_mut()
+}

data/ext/iceberg/src/runtime.rs

@@ -26,8 +26,12 @@ pub fn runtime() -> Handle {
     match Handle::try_current() {
         Ok(h) => h.clone(),
         _ => {
-            let rt = RUNTIME.get_or_init(|| Runtime::new().unwrap());
+            let rt = tokio_runtime();
             rt.handle().clone()
         }
     }
 }
+
+pub fn tokio_runtime() -> &'static Runtime {
+    RUNTIME.get_or_init(|| Runtime::new().unwrap())
+}

data/ext/iceberg/src/scan.rs

@@ -1,55 +1,138 @@
+use std::sync::RwLock;
+
+use arrow_array::RecordBatchIterator;
+use arrow_array::ffi_stream::FFI_ArrowArrayStream;
 use futures::TryStreamExt;
-use iceberg::scan::TableScan;
-use magnus::{RArray, Ruby, Value};
-use std::cell::RefCell;
+use iceberg::scan::{FileScanTask, TableScan};
+use magnus::{IntoValue, RArray, Ruby, Value, value::ReprValue};
 
 use crate::RbResult;
+use crate::capsule::RbCapsule;
 use crate::error::to_rb_err;
+use crate::result::collect_batches;
+use crate::ruby::GvlExt;
 use crate::runtime::runtime;
-use crate::utils::rb_snapshot;
+use crate::snapshot::RbSnapshot;
 
 #[magnus::wrap(class = "Iceberg::RbTableScan")]
 pub struct RbTableScan {
-    pub scan: RefCell<TableScan>,
+    pub scan: RwLock<TableScan>,
+}
+
+#[magnus::wrap(class = "Iceberg::FileScanTask")]
+pub struct RbFileScanTask {
+    pub scan: FileScanTask,
+}
+
+#[magnus::wrap(class = "Iceberg::DataFile")]
+pub struct RbDataFile {
+    pub file_path: String,
+    pub record_count: Option<u64>,
+    pub file_size_in_bytes: u64,
+    pub equality_ids: Option<Vec<i32>>,
 }
 
 impl RbTableScan {
     pub fn plan_files(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
-        let scan = rb_self.scan.borrow();
-
-        let runtime = runtime();
-        let plan_files = runtime.block_on(scan.plan_files()).map_err(to_rb_err)?;
-        let plan_files: Vec<_> = runtime
-            .block_on(plan_files.try_collect())
+        let plan_files: Vec<_> = ruby
+            .detach(|| {
+                let scan = rb_self.scan.read().unwrap();
+                let runtime = runtime();
+                let plan_files = runtime.block_on(scan.plan_files())?;
+                runtime.block_on(plan_files.try_collect())
+            })
             .map_err(to_rb_err)?;
+
         let files = ruby.ary_new();
         for v in plan_files {
-            let file = ruby.hash_new();
-            file.aset(ruby.to_symbol("start"), v.start)?;
-            file.aset(ruby.to_symbol("length"), v.length)?;
-            file.aset(ruby.to_symbol("record_count"), v.record_count)?;
-            file.aset(ruby.to_symbol("data_file_path"), v.data_file_path)?;
-            file.aset(ruby.to_symbol("project_field_ids"), v.project_field_ids)?;
-
-            let deletes = ruby.ary_new();
-            for d in v.deletes {
-                let delete = ruby.hash_new();
-                delete.aset(ruby.to_symbol("file_path"), d.file_path)?;
-                delete.aset(ruby.to_symbol("partition_spec_id"), d.partition_spec_id)?;
-                delete.aset(ruby.to_symbol("equality_ids"), d.equality_ids)?;
-                deletes.push(delete)?;
-            }
-            file.aset(ruby.to_symbol("deletes"), deletes)?;
-
-            files.push(file)?;
+            files.push(RbFileScanTask { scan: v })?;
         }
         Ok(files)
     }
 
-    pub fn snapshot(ruby: &Ruby, rb_self: &Self) -> RbResult<Option<Value>> {
-        match rb_self.scan.borrow().snapshot() {
-            Some(s) => Ok(Some(rb_snapshot(ruby, s)?)),
-            None => Ok(None),
+    pub fn snapshot(&self) -> Option<RbSnapshot> {
+        self.scan.read().unwrap().snapshot().map(|v| v.into())
+    }
+
+    pub fn collect(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
+        let runtime = runtime();
+        let scan = rb_self.scan.read().unwrap();
+        let stream = runtime.block_on(scan.to_arrow()).map_err(to_rb_err)?;
+        let batches: Vec<_> = runtime.block_on(stream.try_collect()).map_err(to_rb_err)?;
+        collect_batches(ruby, batches)
+    }
+
+    pub fn arrow_c_stream(&self) -> RbResult<RbCapsule> {
+        let runtime = runtime();
+        let scan = self.scan.read().unwrap();
+        let stream = runtime.block_on(scan.to_arrow()).map_err(to_rb_err)?;
+        let batches: Vec<_> = runtime.block_on(stream.try_collect()).map_err(to_rb_err)?;
+        let stream = if batches.is_empty() {
+            // TODO fix schema
+            FFI_ArrowArrayStream::empty()
+        } else {
+            let schema = batches[0].schema();
+            let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
+            FFI_ArrowArrayStream::new(Box::new(reader))
+        };
+        Ok(RbCapsule::new(stream, Some("arrow_array_stream".into())))
+    }
+}
+
+impl RbFileScanTask {
+    pub fn file(&self) -> RbDataFile {
+        RbDataFile {
+            file_path: self.scan.data_file_path.clone(),
+            record_count: self.scan.record_count,
+            file_size_in_bytes: self.scan.file_size_in_bytes,
+            equality_ids: None,
         }
     }
+
+    pub fn delete_files(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(rb_self.scan.deletes.iter().map(|v| RbDataFile {
+            file_path: v.file_path.clone(),
+            record_count: None,
+            file_size_in_bytes: v.file_size_in_bytes,
+            equality_ids: v.equality_ids.clone(),
+        }))
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::FileScanTask file={}, delete_files={}>",
+            rb_self.file().into_value_with(ruby).inspect(),
+            Self::delete_files(ruby, rb_self)
+                .into_value_with(ruby)
+                .inspect(),
+        )
+    }
+}
+
+impl RbDataFile {
+    pub fn file_path(&self) -> &str {
+        &self.file_path
+    }
+
+    pub fn record_count(&self) -> Option<u64> {
+        self.record_count
+    }
+
+    pub fn file_size_in_bytes(&self) -> u64 {
+        self.file_size_in_bytes
+    }
+
+    pub fn equality_ids(&self) -> Option<Vec<i32>> {
+        self.equality_ids.clone()
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::DataFile file_path={}, record_count={}, file_size_in_bytes={}, equality_ids={}>",
+            rb_self.file_path().into_value_with(ruby).inspect(),
+            rb_self.record_count().into_value_with(ruby).inspect(),
+            rb_self.file_size_in_bytes().into_value_with(ruby).inspect(),
+            rb_self.equality_ids().into_value_with(ruby).inspect(),
+        )
+    }
 }

data/ext/iceberg/src/schema.rs

@@ -0,0 +1,305 @@
+use std::sync::Arc;
+
+use arrow_schema::Schema as ArrowSchema;
+use arrow_schema::ffi::FFI_ArrowSchema;
+use iceberg::arrow::{arrow_schema_to_schema_auto_assign_ids, schema_to_arrow_schema};
+use iceberg::spec::{ListType, MapType, NestedField, PrimitiveType, Schema, StructType, Type};
+use magnus::{
+    Error as RbErr, IntoValue, RArray, RClass, RHash, RModule, Ruby, TryConvert, Value, prelude::*,
+};
+
+use crate::RbResult;
+use crate::capsule::RbCapsule;
+use crate::error::to_rb_err;
+use crate::utils::{Wrap, default_value, rb_literal};
+
+#[magnus::wrap(class = "Iceberg::Schema")]
+pub struct RbSchema {
+    pub(crate) schema: Schema,
+}
+
+#[magnus::wrap(class = "Iceberg::NestedField")]
+pub struct RbNestedField {
+    pub(crate) field: Arc<NestedField>,
+}
+
+impl RbSchema {
+    pub fn new(args: &[Value]) -> RbResult<Self> {
+        let schema = if args.len() == 1
+            && let Ok(arrow_schema) =
+                args[0].funcall::<_, _, Wrap<ArrowSchema>>("arrow_c_schema", ())
+        {
+            arrow_schema_to_schema_auto_assign_ids(&arrow_schema.0).map_err(to_rb_err)?
+        } else {
+            let mut fields = Vec::new();
+            for rb_field in args {
+                fields.push(<&RbNestedField>::try_convert(*rb_field)?.field.clone());
+            }
+            Schema::builder()
+                .with_fields(fields)
+                .build()
+                .map_err(to_rb_err)?
+        };
+        Ok(Self { schema })
+    }
+
+    pub fn fields(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(
+            rb_self
+                .schema
+                .as_struct()
+                .fields()
+                .iter()
+                .map(RbNestedField::from),
+        )
+    }
+
+    pub fn schema_id(&self) -> i32 {
+        self.schema.schema_id()
+    }
+
+    pub fn identifier_field_ids(&self) -> Vec<i32> {
+        self.schema.identifier_field_ids().collect()
+    }
+
+    pub fn highest_field_id(&self) -> i32 {
+        self.schema.highest_field_id()
+    }
+
+    pub fn as_struct(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
+        let iceberg = ruby.class_object().const_get::<_, RModule>("Iceberg")?;
+        let fields = Self::fields(ruby, rb_self);
+        iceberg
+            .const_get::<_, RClass>("StructType")?
+            .new_instance(unsafe { fields.as_slice() })
+    }
+
+    pub fn arrow_c_schema(&self) -> RbResult<RbCapsule> {
+        let schema = schema_to_arrow_schema(&self.schema).map_err(to_rb_err)?;
+        let schema = FFI_ArrowSchema::try_from(&schema).unwrap();
+        Ok(RbCapsule::new(schema, Some("arrow_schema".into())))
+    }
+
+    pub fn eq(&self, other: &Self) -> bool {
+        self.schema == other.schema
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::Schema fields={}, schema_id={}, identifier_field_ids={}>",
+            Self::fields(ruby, rb_self).inspect(),
+            rb_self.schema_id().into_value_with(ruby).inspect(),
+            rb_self
+                .identifier_field_ids()
+                .into_value_with(ruby)
+                .inspect(),
+        )
+    }
+}
+
+impl RbNestedField {
+    pub fn new(ruby: &Ruby, rb_field: RHash) -> RbResult<Self> {
+        let rb_type: Value = rb_field.aref(ruby.to_symbol("field_type"))?;
+        let field_type = match &*unsafe { rb_type.classname() } {
+            "Iceberg::BooleanType" => Type::Primitive(PrimitiveType::Boolean),
+            "Iceberg::IntType" => Type::Primitive(PrimitiveType::Int),
+            "Iceberg::LongType" => Type::Primitive(PrimitiveType::Long),
+            "Iceberg::FloatType" => Type::Primitive(PrimitiveType::Float),
+            "Iceberg::DoubleType" => Type::Primitive(PrimitiveType::Double),
+            "Iceberg::DecimalType" => {
+                let precision: u32 = rb_type.funcall("precision", ())?;
+                let scale: u32 = rb_type.funcall("scale", ())?;
+                Type::Primitive(PrimitiveType::Decimal { precision, scale })
+            }
+            "Iceberg::DateType" => Type::Primitive(PrimitiveType::Date),
+            "Iceberg::TimeType" => Type::Primitive(PrimitiveType::Time),
+            "Iceberg::TimestampType" => Type::Primitive(PrimitiveType::Timestamp),
+            "Iceberg::TimestamptzType" => Type::Primitive(PrimitiveType::Timestamptz),
+            "Iceberg::TimestampNanoType" => Type::Primitive(PrimitiveType::TimestampNs),
+            "Iceberg::TimestamptzNanoType" => Type::Primitive(PrimitiveType::TimestamptzNs),
+            "Iceberg::StringType" => Type::Primitive(PrimitiveType::String),
+            "Iceberg::UUIDType" => Type::Primitive(PrimitiveType::Uuid),
+            "Iceberg::FixedType" => {
+                let length: u64 = rb_type.funcall("length", ())?;
+                Type::Primitive(PrimitiveType::Fixed(length))
+            }
+            "Iceberg::BinaryType" => Type::Primitive(PrimitiveType::Binary),
+            "Iceberg::StructType" => {
+                let fields = rb_type
+                    .funcall::<_, _, RArray>("fields", ())?
+                    .typecheck::<&RbNestedField>()?
+                    .into_iter()
+                    .map(|v| v.field.clone())
+                    .collect::<Vec<_>>();
+                Type::Struct(StructType::new(fields))
+            }
+            "Iceberg::ListType" => {
+                let element_field = rb_type
+                    .funcall::<_, _, &RbNestedField>("element_field", ())?
+                    .field
+                    .clone();
+                Type::List(ListType::new(element_field))
+            }
+            "Iceberg::MapType" => {
+                let key_field = rb_type
+                    .funcall::<_, _, &RbNestedField>("key_field", ())?
+                    .field
+                    .clone();
+                let value_field = rb_type
+                    .funcall::<_, _, &RbNestedField>("value_field", ())?
+                    .field
+                    .clone();
+                Type::Map(MapType::new(key_field, value_field))
+            }
+            _ => {
+                return Err(RbErr::new(
+                    ruby.exception_arg_error(),
+                    format!("Type not supported: {}", rb_type),
+                ));
+            }
+        };
+
+        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)?;
+
+        let field = NestedField {
+            id: rb_field.aref(ruby.to_symbol("field_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,
+        };
+
+        Ok(Self {
+            field: field.into(),
+        })
+    }
+
+    pub fn field_id(&self) -> i32 {
+        self.field.id
+    }
+
+    pub fn name(&self) -> &str {
+        &self.field.name
+    }
+
+    pub fn field_type(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
+        let iceberg = ruby.class_object().const_get::<_, RModule>("Iceberg")?;
+        let field_type = &*rb_self.field.field_type;
+        let v = match field_type {
+            Type::Primitive(ty) => match ty {
+                PrimitiveType::Boolean => iceberg
+                    .const_get::<_, RClass>("BooleanType")?
+                    .new_instance(())?,
+                PrimitiveType::Int => iceberg
+                    .const_get::<_, RClass>("IntType")?
+                    .new_instance(())?,
+                PrimitiveType::Long => iceberg
+                    .const_get::<_, RClass>("LongType")?
+                    .new_instance(())?,
+                PrimitiveType::Float => iceberg
+                    .const_get::<_, RClass>("FloatType")?
+                    .new_instance(())?,
+                PrimitiveType::Double => iceberg
+                    .const_get::<_, RClass>("DoubleType")?
+                    .new_instance(())?,
+                PrimitiveType::Decimal { precision, scale } => iceberg
+                    .const_get::<_, RClass>("DecimalType")?
+                    .new_instance((*precision, *scale))?,
+                PrimitiveType::Date => iceberg
+                    .const_get::<_, RClass>("DateType")?
+                    .new_instance(())?,
+                PrimitiveType::Time => iceberg
+                    .const_get::<_, RClass>("TimeType")?
+                    .new_instance(())?,
+                PrimitiveType::Timestamp => iceberg
+                    .const_get::<_, RClass>("TimestampType")?
+                    .new_instance(())?,
+                PrimitiveType::Timestamptz => iceberg
+                    .const_get::<_, RClass>("TimestamptzType")?
+                    .new_instance(())?,
+                PrimitiveType::TimestampNs => iceberg
+                    .const_get::<_, RClass>("TimestampNanoType")?
+                    .new_instance(())?,
+                PrimitiveType::TimestamptzNs => iceberg
+                    .const_get::<_, RClass>("TimestamptzNanoType")?
+                    .new_instance(())?,
+                PrimitiveType::String => iceberg
+                    .const_get::<_, RClass>("StringType")?
+                    .new_instance(())?,
+                PrimitiveType::Uuid => iceberg
+                    .const_get::<_, RClass>("UUIDType")?
+                    .new_instance(())?,
+                PrimitiveType::Fixed(length) => iceberg
+                    .const_get::<_, RClass>("FixedType")?
+                    .new_instance((*length,))?,
+                PrimitiveType::Binary => iceberg
+                    .const_get::<_, RClass>("BinaryType")?
+                    .new_instance(())?,
+            },
+            Type::Struct(ty) => iceberg
+                .const_get::<_, RClass>("StructType")?
+                .new_instance((ruby.ary_from_iter(ty.fields().iter().map(RbNestedField::from)),))?,
+            Type::List(ty) => iceberg
+                .const_get::<_, RClass>("ListType")?
+                .new_instance((RbNestedField::from(&ty.element_field),))?,
+            Type::Map(ty) => iceberg.const_get::<_, RClass>("MapType")?.new_instance((
+                RbNestedField::from(&ty.key_field),
+                RbNestedField::from(&ty.value_field),
+            ))?,
+        };
+        Ok(v)
+    }
+
+    pub fn required(&self) -> bool {
+        self.field.required
+    }
+
+    pub fn doc(&self) -> Option<&str> {
+        self.field.doc.as_deref()
+    }
+
+    pub fn initial_default(ruby: &Ruby, rb_self: &Self) -> RbResult<Option<Value>> {
+        rb_self
+            .field
+            .initial_default
+            .as_ref()
+            .map(|v| rb_literal(ruby, v))
+            .transpose()
+    }
+
+    pub fn write_default(ruby: &Ruby, rb_self: &Self) -> RbResult<Option<Value>> {
+        rb_self
+            .field
+            .write_default
+            .as_ref()
+            .map(|v| rb_literal(ruby, v))
+            .transpose()
+    }
+
+    pub fn eq(&self, other: &Self) -> bool {
+        self.field == other.field
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> RbResult<String> {
+        Ok(format!(
+            "#<Iceberg::NestedField field_id={}, name={}, field_type={}, required={}, doc={}, initial_default={}, write_default={}>",
+            rb_self.field_id().into_value_with(ruby).inspect(),
+            rb_self.name().into_value_with(ruby).inspect(),
+            Self::field_type(ruby, rb_self)?.inspect(),
+            rb_self.required().into_value_with(ruby).inspect(),
+            rb_self.doc().into_value_with(ruby).inspect(),
+            Self::initial_default(ruby, rb_self)?
+                .into_value_with(ruby)
+                .inspect(),
+            Self::write_default(ruby, rb_self)?
+                .into_value_with(ruby)
+                .inspect(),
+        ))
+    }
+}

data/ext/iceberg/src/snapshot.rs

@@ -0,0 +1,85 @@
+use iceberg::spec::{MetadataLog, Snapshot, SnapshotLog};
+use magnus::{IntoValue, Ruby, value::ReprValue};
+
+#[magnus::wrap(class = "Iceberg::Snapshot")]
+pub struct RbSnapshot {
+    pub(crate) snapshot: Snapshot,
+}
+
+#[magnus::wrap(class = "Iceberg::SnapshotLogEntry")]
+pub struct RbSnapshotLogEntry {
+    pub(crate) log: SnapshotLog,
+}
+
+#[magnus::wrap(class = "Iceberg::MetadataLogEntry")]
+pub struct RbMetadataLogEntry {
+    pub(crate) log: MetadataLog,
+}
+
+impl RbSnapshot {
+    pub fn snapshot_id(&self) -> i64 {
+        self.snapshot.snapshot_id()
+    }
+
+    pub fn parent_snapshot_id(&self) -> Option<i64> {
+        self.snapshot.parent_snapshot_id()
+    }
+
+    pub fn sequence_number(&self) -> i64 {
+        self.snapshot.sequence_number()
+    }
+
+    pub fn manifest_list(&self) -> &str {
+        self.snapshot.manifest_list()
+    }
+
+    pub fn schema_id(&self) -> Option<i32> {
+        self.snapshot.schema_id()
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::Snapshot snapshot_id={}, parent_snapshot_id={}, sequence_number={}, schema_id={}>",
+            rb_self.snapshot_id().into_value_with(ruby).inspect(),
+            rb_self.parent_snapshot_id().into_value_with(ruby).inspect(),
+            rb_self.sequence_number().into_value_with(ruby).inspect(),
+            rb_self.schema_id().into_value_with(ruby).inspect(),
+        )
+    }
+}
+
+impl RbSnapshotLogEntry {
+    pub fn snapshot_id(&self) -> i64 {
+        self.log.snapshot_id
+    }
+
+    pub fn timestamp_ms(&self) -> i64 {
+        self.log.timestamp_ms
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::SnapshotLogEntry snapshot_id={}, timestamp_ms={}>",
+            rb_self.snapshot_id().into_value_with(ruby).inspect(),
+            rb_self.timestamp_ms().into_value_with(ruby).inspect(),
+        )
+    }
+}
+
+impl RbMetadataLogEntry {
+    pub fn metadata_file(&self) -> &str {
+        &self.log.metadata_file
+    }
+
+    pub fn timestamp_ms(&self) -> i64 {
+        self.log.timestamp_ms
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::MetadataLogEntry metadata_file={}, timestamp_ms={}>",
+            rb_self.metadata_file().into_value_with(ruby).inspect(),
+            rb_self.timestamp_ms().into_value_with(ruby).inspect(),
+        )
+    }
+}