iceberg 0.11.2 → 0.12.0

data/ext/iceberg/src/statistics.rs

@@ -0,0 +1,71 @@
+use iceberg::spec::{PartitionStatisticsFile, StatisticsFile};
+use magnus::{IntoValue, Ruby, value::ReprValue};
+
+#[magnus::wrap(class = "Iceberg::StatisticsFile")]
+pub struct RbStatisticsFile {
+    pub(crate) file: StatisticsFile,
+}
+
+#[magnus::wrap(class = "Iceberg::PartitionStatisticsFile")]
+pub struct RbPartitionStatisticsFile {
+    pub(crate) file: PartitionStatisticsFile,
+}
+
+impl RbStatisticsFile {
+    pub fn snapshot_id(&self) -> i64 {
+        self.file.snapshot_id
+    }
+
+    pub fn statistics_path(&self) -> &str {
+        &self.file.statistics_path
+    }
+
+    pub fn file_size_in_bytes(&self) -> i64 {
+        self.file.file_size_in_bytes
+    }
+
+    pub fn file_footer_size_in_bytes(&self) -> i64 {
+        self.file.file_footer_size_in_bytes
+    }
+
+    pub fn key_metadata(&self) -> Option<&str> {
+        self.file.key_metadata.as_deref()
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::StatisticsFile snapshot_id={}, statistics_path={}, file_size_in_bytes={}, file_footer_size_in_bytes={}, key_metadata={}>",
+            rb_self.snapshot_id().into_value_with(ruby).inspect(),
+            rb_self.statistics_path().into_value_with(ruby).inspect(),
+            rb_self.file_size_in_bytes().into_value_with(ruby).inspect(),
+            rb_self
+                .file_footer_size_in_bytes()
+                .into_value_with(ruby)
+                .inspect(),
+            rb_self.key_metadata().into_value_with(ruby).inspect(),
+        )
+    }
+}
+
+impl RbPartitionStatisticsFile {
+    pub fn snapshot_id(&self) -> i64 {
+        self.file.snapshot_id
+    }
+
+    pub fn statistics_path(&self) -> &str {
+        &self.file.statistics_path
+    }
+
+    pub fn file_size_in_bytes(&self) -> i64 {
+        self.file.file_size_in_bytes
+    }
+
+    pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
+        format!(
+            "#<Iceberg::PartitionStatisticsFile snapshot_id={}, statistics_path={}, file_size_in_bytes={}>",
+            rb_self.snapshot_id().into_value_with(ruby).inspect(),
+            rb_self.statistics_path().into_value_with(ruby).inspect(),
+            rb_self.file_size_in_bytes().into_value_with(ruby).inspect(),
+        )
+    }
+}

data/ext/iceberg/src/table.rs

@@ -1,7 +1,13 @@
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+
+use arrow_array::RecordBatch;
 use arrow_array::ffi_stream::ArrowArrayStreamReader;
+use arrow_cast::cast;
+use arrow_schema::{ArrowError, DataType, Field, Schema};
 use iceberg::TableIdent;
 use iceberg::io::FileIO;
-use iceberg::spec::FormatVersion;
+use iceberg::spec::{FormatVersion, TableMetadata};
 use iceberg::table::{StaticTable, Table};
 use iceberg::transaction::{ApplyTransactionAction, Transaction};
 use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
@@ -11,20 +17,23 @@ use iceberg::writer::file_writer::location_generator::{
 };
 use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
 use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
-use magnus::{RArray, Ruby, Value};
+use magnus::{RArray, Ruby};
 use parquet::file::properties::WriterProperties;
-use std::collections::HashMap;
-use std::sync::{Arc, RwLock};
 use uuid::Uuid;
 
 use crate::RbResult;
-use crate::arrow::RbArrowType;
 use crate::catalog::RbCatalog;
+use crate::encryption::RbEncryptedKey;
 use crate::error::to_rb_err;
+use crate::partitioning::RbPartitionSpec;
 use crate::ruby::GvlExt;
 use crate::runtime::runtime;
 use crate::scan::RbTableScan;
-use crate::utils::*;
+use crate::schema::RbSchema;
+use crate::snapshot::{RbMetadataLogEntry, RbSnapshot, RbSnapshotLogEntry};
+use crate::sorting::RbSortOrder;
+use crate::statistics::{RbPartitionStatisticsFile, RbStatisticsFile};
+use crate::utils::Wrap;
 
 #[magnus::wrap(class = "Iceberg::RbTable")]
 pub struct RbTable {
@@ -32,6 +41,14 @@ pub struct RbTable {
 }
 
 impl RbTable {
+    pub fn identifier(&self) -> Vec<String> {
+        let table = self.table.read().unwrap();
+        let ident = table.identifier();
+        let mut vec = ident.namespace.clone().inner();
+        vec.push(ident.name.clone());
+        vec
+    }
+
     pub fn scan(&self, snapshot_id: Option<i64>) -> RbResult<RbTableScan> {
         let table = self.table.read().unwrap();
         let mut builder = table.scan();
@@ -45,7 +62,7 @@ impl RbTable {
     pub fn append(
         ruby: &Ruby,
         rb_self: &Self,
-        data: RbArrowType<ArrowArrayStreamReader>,
+        data: Wrap<ArrowArrayStreamReader>,
         catalog: &RbCatalog,
     ) -> RbResult<RbTable> {
         let table = ruby
@@ -63,7 +80,7 @@ impl RbTable {
                         .unwrap(),
                 );
 
-                let location_generator = DefaultLocationGenerator::new(table.metadata().clone())?;
+                let location_generator = DefaultLocationGenerator::new(table.metadata())?;
                 let file_name_generator = DefaultFileNameGenerator::new(
                     // TODO move task id to suffix to match Python and Java
                     "0".to_string(),
@@ -88,7 +105,7 @@ impl RbTable {
                     runtime.block_on(data_file_writer_builder.build(None))?;
 
                 for batch in data.0 {
-                    let batch = batch.unwrap().with_schema(table_schema.clone())?;
+                    let batch = cast_batch(batch.unwrap(), table_schema.clone())?;
                     runtime.block_on(data_file_writer.write(batch))?;
                 }
 
@@ -107,8 +124,36 @@ impl RbTable {
         })
     }
 
+    pub fn metadata(&self) -> RbTableMetadata {
+        RbTableMetadata {
+            metadata: self.table.read().unwrap().metadata().clone(),
+        }
+    }
+
+    pub fn from_metadata_file(location: String) -> RbResult<Self> {
+        let file_io = FileIO::new_with_fs();
+        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(),
+        })
+    }
+}
+
+#[magnus::wrap(class = "Iceberg::TableMetadata")]
+pub struct RbTableMetadata {
+    pub metadata: TableMetadata,
+}
+
+impl RbTableMetadata {
     pub fn format_version(&self) -> i32 {
-        match self.table.read().unwrap().metadata().format_version() {
+        match self.metadata.format_version() {
             FormatVersion::V1 => 1,
             FormatVersion::V2 => 2,
             FormatVersion::V3 => 3,
@@ -116,321 +161,224 @@ impl RbTable {
     }
 
     pub fn uuid(&self) -> String {
-        self.table.read().unwrap().metadata().uuid().to_string()
+        self.metadata.uuid().to_string()
     }
 
     pub fn location(&self) -> String {
-        self.table.read().unwrap().metadata().location().to_string()
+        self.metadata.location().to_string()
     }
 
     pub fn last_sequence_number(&self) -> i64 {
-        self.table.read().unwrap().metadata().last_sequence_number()
+        self.metadata.last_sequence_number()
     }
 
     pub fn next_sequence_number(&self) -> i64 {
-        self.table.read().unwrap().metadata().next_sequence_number()
+        self.metadata.next_sequence_number()
     }
 
     pub fn last_column_id(&self) -> i32 {
-        self.table.read().unwrap().metadata().last_column_id()
+        self.metadata.last_column_id()
     }
 
     pub fn last_partition_id(&self) -> i32 {
-        self.table.read().unwrap().metadata().last_partition_id()
+        self.metadata.last_partition_id()
     }
 
     pub fn last_updated_ms(&self) -> i64 {
-        self.table.read().unwrap().metadata().last_updated_ms()
+        self.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.read().unwrap().metadata().schemas_iter() {
-            schemas.push(rb_schema(ruby, s)?)?;
-        }
-        Ok(schemas)
-    }
-
-    pub fn schema_by_id(ruby: &Ruby, rb_self: &Self, schema_id: i32) -> RbResult<Option<Value>> {
-        let schema = match rb_self
-            .table
-            .read()
-            .unwrap()
-            .metadata()
-            .schema_by_id(schema_id)
-        {
-            Some(s) => Some(rb_schema(ruby, s)?),
-            None => None,
-        };
-        Ok(schema)
-    }
-
-    pub fn current_schema(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
-        rb_schema(
-            ruby,
-            rb_self.table.read().unwrap().metadata().current_schema(),
-        )
+    pub fn schemas(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(rb_self.metadata.schemas_iter().map(RbSchema::from))
+    }
+
+    pub fn schema_by_id(&self, schema_id: i32) -> Option<RbSchema> {
+        self.metadata.schema_by_id(schema_id).map(|v| v.into())
+    }
+
+    pub fn current_schema(&self) -> RbSchema {
+        self.metadata.current_schema().into()
     }
 
     pub fn current_schema_id(&self) -> i32 {
-        self.table.read().unwrap().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
-            .read()
-            .unwrap()
-            .metadata()
-            .partition_specs_iter()
-        {
-            partition_specs.push(rb_partition_spec(s)?)?;
-        }
-        Ok(partition_specs)
+        self.metadata.current_schema_id()
     }
 
-    pub fn partition_spec_by_id(&self, partition_spec_id: i32) -> RbResult<Option<Value>> {
-        let partition_spec = match self
-            .table
-            .read()
-            .unwrap()
-            .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
-                .read()
-                .unwrap()
-                .metadata()
-                .default_partition_spec(),
+    pub fn partition_specs(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(
+            rb_self
+                .metadata
+                .partition_specs_iter()
+                .map(RbPartitionSpec::from),
         )
     }
 
+    pub fn partition_spec_by_id(&self, partition_spec_id: i32) -> Option<RbPartitionSpec> {
+        self.metadata
+            .partition_spec_by_id(partition_spec_id)
+            .map(|v| v.into())
+    }
+
+    pub fn default_partition_spec(&self) -> RbPartitionSpec {
+        self.metadata.default_partition_spec().into()
+    }
+
     pub fn default_partition_spec_id(&self) -> i32 {
-        self.table
-            .read()
-            .unwrap()
-            .metadata()
-            .default_partition_spec_id()
+        self.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.read().unwrap().metadata().snapshots() {
-            snapshots.push(rb_snapshot(ruby, s)?)?;
-        }
-        Ok(snapshots)
+    pub fn snapshots(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(rb_self.metadata.snapshots().map(RbSnapshot::from))
     }
 
-    pub fn snapshot_by_id(
-        ruby: &Ruby,
-        rb_self: &Self,
-        snapshot_id: i64,
-    ) -> RbResult<Option<Value>> {
-        let snapshot = match rb_self
-            .table
-            .read()
-            .unwrap()
-            .metadata()
-            .snapshot_by_id(snapshot_id)
-        {
-            Some(s) => Some(rb_snapshot(ruby, 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.read().unwrap().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.read().unwrap().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 snapshot_by_id(&self, snapshot_id: i64) -> Option<RbSnapshot> {
+        self.metadata.snapshot_by_id(snapshot_id).map(|v| v.into())
     }
 
-    pub fn current_snapshot(ruby: &Ruby, rb_self: &Self) -> RbResult<Option<Value>> {
-        let snapshot = match rb_self.table.read().unwrap().metadata().current_snapshot() {
-            Some(s) => Some(rb_snapshot(ruby, s)?),
-            None => None,
-        };
-        Ok(snapshot)
+    pub fn history(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(
+            rb_self
+                .metadata
+                .history()
+                .iter()
+                .map(RbSnapshotLogEntry::from),
+        )
+    }
+
+    pub fn metadata_log(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(
+            rb_self
+                .metadata
+                .metadata_log()
+                .iter()
+                .map(RbMetadataLogEntry::from),
+        )
+    }
+
+    pub fn current_snapshot(&self) -> Option<RbSnapshot> {
+        self.metadata.current_snapshot().map(|v| v.into())
     }
 
     pub fn current_snapshot_id(&self) -> Option<i64> {
-        self.table.read().unwrap().metadata().current_snapshot_id()
+        self.metadata.current_snapshot_id()
     }
 
-    pub fn snapshot_for_ref(
-        ruby: &Ruby,
-        rb_self: &Self,
-        ref_name: String,
-    ) -> RbResult<Option<Value>> {
-        let snapshot = match rb_self
-            .table
-            .read()
-            .unwrap()
-            .metadata()
-            .snapshot_for_ref(&ref_name)
-        {
-            Some(s) => Some(rb_snapshot(ruby, 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.read().unwrap().metadata().sort_orders_iter() {
-            sort_orders.push(rb_sort_order(s)?)?;
-        }
-        Ok(sort_orders)
+    pub fn snapshot_for_ref(&self, ref_name: String) -> Option<RbSnapshot> {
+        self.metadata.snapshot_for_ref(&ref_name).map(|v| v.into())
+    }
+
+    pub fn sort_orders(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(rb_self.metadata.sort_orders_iter().map(RbSortOrder::from))
     }
 
-    pub fn sort_order_by_id(&self, sort_order_id: i64) -> RbResult<Option<Value>> {
-        let sort_order = match self
-            .table
-            .read()
-            .unwrap()
-            .metadata()
+    pub fn sort_order_by_id(&self, sort_order_id: i64) -> Option<RbSortOrder> {
+        self.metadata
             .sort_order_by_id(sort_order_id)
-        {
-            Some(s) => Some(rb_sort_order(s)?),
-            None => None,
-        };
-        Ok(sort_order)
+            .map(|v| v.into())
     }
 
-    pub fn default_sort_order(&self) -> RbResult<Value> {
-        rb_sort_order(self.table.read().unwrap().metadata().default_sort_order())
+    pub fn default_sort_order(&self) -> RbSortOrder {
+        self.metadata.default_sort_order().into()
     }
 
     pub fn default_sort_order_id(&self) -> i64 {
-        self.table
-            .read()
-            .unwrap()
-            .metadata()
-            .default_sort_order_id()
+        self.metadata.default_sort_order_id()
     }
 
     pub fn properties(&self) -> HashMap<String, String> {
-        self.table.read().unwrap().metadata().properties().clone()
+        self.metadata.properties().clone()
     }
 
-    pub fn statistics(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
-        let statistics = ruby.ary_new();
-        for s in rb_self.table.read().unwrap().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
-            .read()
-            .unwrap()
-            .metadata()
-            .partition_statistics_iter()
-        {
-            statistics.push(rb_partition_statistics_file(s)?)?;
-        }
-        Ok(statistics)
+    pub fn statistics(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(
+            rb_self
+                .metadata
+                .statistics_iter()
+                .map(RbStatisticsFile::from),
+        )
     }
 
-    pub fn statistics_for_snapshot(&self, snapshot_id: i64) -> RbResult<Option<Value>> {
-        let statistics = match self
-            .table
-            .read()
-            .unwrap()
-            .metadata()
+    pub fn partition_statistics(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(
+            rb_self
+                .metadata
+                .partition_statistics_iter()
+                .map(RbPartitionStatisticsFile::from),
+        )
+    }
+
+    pub fn statistics_for_snapshot(&self, snapshot_id: i64) -> Option<RbStatisticsFile> {
+        self.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
-            .read()
-            .unwrap()
-            .metadata()
+            .map(|v| v.into())
+    }
+
+    pub fn partition_statistics_for_snapshot(
+        &self,
+        snapshot_id: i64,
+    ) -> Option<RbPartitionStatisticsFile> {
+        self.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 in rb_self
-            .table
-            .read()
-            .unwrap()
-            .metadata()
-            .encryption_keys_iter()
-        {
-            encryption_keys.push(rb_encrypted_key(k)?)?;
-        }
-        Ok(encryption_keys)
+            .map(|v| v.into())
     }
 
-    pub fn encryption_key(&self, key_id: String) -> RbResult<Option<Value>> {
-        let key = match self
-            .table
-            .read()
-            .unwrap()
-            .metadata()
-            .encryption_key(&key_id)
-        {
-            Some(k) => Some(rb_encrypted_key(k)?),
-            None => None,
-        };
-        Ok(key)
+    pub fn encryption_keys(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(
+            rb_self
+                .metadata
+                .encryption_keys_iter()
+                .map(RbEncryptedKey::from),
+        )
     }
 
-    pub fn from_metadata_file(location: String) -> RbResult<Self> {
-        let file_io = FileIO::new_with_fs();
-        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(),
-        })
+    pub fn encryption_key(&self, key_id: String) -> Option<RbEncryptedKey> {
+        self.metadata.encryption_key(&key_id).map(|v| v.into())
+    }
+
+    pub fn next_row_id(&self) -> u64 {
+        self.metadata.next_row_id()
+    }
+}
+
+fn cast_batch(batch: RecordBatch, table_schema: Arc<Schema>) -> Result<RecordBatch, ArrowError> {
+    let mut fields = Vec::new();
+    let mut columns = Vec::new();
+    for (field, column) in batch.schema().fields.iter().zip(batch.columns()) {
+        match field.data_type() {
+            DataType::Utf8View => {
+                fields.push(Arc::new(Field::new(
+                    field.name(),
+                    DataType::Utf8,
+                    field.is_nullable(),
+                )));
+                columns.push(cast(column, &DataType::Utf8)?);
+            }
+            DataType::BinaryView => {
+                // TODO convert to FixedSizeBinary if needed
+                fields.push(Arc::new(Field::new(
+                    field.name(),
+                    DataType::LargeBinary,
+                    field.is_nullable(),
+                )));
+                columns.push(cast(column, &DataType::LargeBinary)?);
+            }
+            DataType::Timestamp(time_unit, Some(_)) => {
+                fields.push(Arc::new(Field::new(
+                    field.name(),
+                    DataType::Timestamp(*time_unit, Some("+00:00".into())),
+                    field.is_nullable(),
+                )));
+                columns.push(cast(
+                    column,
+                    &DataType::Timestamp(*time_unit, Some("+00:00".into())),
+                )?);
+            }
+            _ => {
+                // cloning Arc is cheap
+                fields.push(field.clone());
+                columns.push(column.clone());
+            }
+        }
     }
+    RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?.with_schema(table_schema)
 }