iceberg 0.11.2 → 0.12.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,17 +1,36 @@
1
1
  mod arrow;
2
+ mod batch;
3
+ mod capsule;
2
4
  mod catalog;
5
+ mod encryption;
3
6
  mod error;
7
+ mod partitioning;
8
+ mod result;
4
9
  mod ruby;
5
10
  mod runtime;
6
11
  mod scan;
12
+ mod schema;
13
+ mod snapshot;
14
+ mod sorting;
15
+ mod statistics;
7
16
  mod table;
8
17
  mod utils;
9
18
 
10
19
  use magnus::{Error as RbErr, Ruby, function, method, prelude::*};
11
20
 
21
+ use crate::batch::RbArrowRecordBatch;
22
+ use crate::capsule::RbCapsule;
12
23
  use crate::catalog::RbCatalog;
13
- use crate::scan::RbTableScan;
14
- use crate::table::RbTable;
24
+ #[cfg(feature = "datafusion")]
25
+ use crate::catalog::RbSessionContext;
26
+ use crate::encryption::RbEncryptedKey;
27
+ use crate::partitioning::{RbPartitionField, RbPartitionSpec};
28
+ use crate::scan::{RbDataFile, RbFileScanTask, RbTableScan};
29
+ use crate::schema::{RbNestedField, RbSchema};
30
+ use crate::snapshot::{RbMetadataLogEntry, RbSnapshot, RbSnapshotLogEntry};
31
+ use crate::sorting::{RbSortField, RbSortOrder};
32
+ use crate::statistics::{RbPartitionStatisticsFile, RbStatisticsFile};
33
+ use crate::table::{RbTable, RbTableMetadata};
15
34
 
16
35
  type RbResult<T> = Result<T, RbErr>;
17
36
 
@@ -39,93 +58,297 @@ fn init(ruby: &Ruby) -> RbResult<()> {
39
58
  class.define_method("update_namespace", method!(RbCatalog::update_namespace, 2))?;
40
59
  class.define_method("drop_namespace", method!(RbCatalog::drop_namespace, 1))?;
41
60
  class.define_method("list_tables", method!(RbCatalog::list_tables, 1))?;
42
- class.define_method("create_table", method!(RbCatalog::create_table, 3))?;
61
+ class.define_method("create_table", method!(RbCatalog::create_table, 6))?;
43
62
  class.define_method("load_table", method!(RbCatalog::load_table, 1))?;
44
63
  class.define_method("drop_table", method!(RbCatalog::drop_table, 1))?;
64
+ class.define_method("purge_table", method!(RbCatalog::purge_table, 1))?;
45
65
  class.define_method("table_exists?", method!(RbCatalog::table_exists, 1))?;
46
66
  class.define_method("rename_table", method!(RbCatalog::rename_table, 2))?;
47
67
  class.define_method("register_table", method!(RbCatalog::register_table, 2))?;
48
68
  #[cfg(feature = "datafusion")]
49
- class.define_method("sql", method!(RbCatalog::sql, 1))?;
69
+ class.define_method("session_context", method!(RbCatalog::session_context, 1))?;
70
+
71
+ #[cfg(feature = "datafusion")]
72
+ {
73
+ let class = module.define_class("SessionContext", ruby.class_object())?;
74
+ class.define_method("sql", method!(RbSessionContext::sql, 2))?;
75
+ }
50
76
 
51
77
  let class = module.define_class("RbTable", ruby.class_object())?;
78
+ class.define_method("identifier", method!(RbTable::identifier, 0))?;
52
79
  class.define_method("scan", method!(RbTable::scan, 1))?;
53
80
  class.define_method("append", method!(RbTable::append, 2))?;
54
- class.define_method("format_version", method!(RbTable::format_version, 0))?;
55
- class.define_method("uuid", method!(RbTable::uuid, 0))?;
56
- class.define_method("location", method!(RbTable::location, 0))?;
81
+ class.define_method("metadata", method!(RbTable::metadata, 0))?;
82
+ class.define_singleton_method(
83
+ "from_metadata_file",
84
+ function!(RbTable::from_metadata_file, 1),
85
+ )?;
86
+
87
+ let class = module.define_class("TableMetadata", ruby.class_object())?;
88
+ class.define_method(
89
+ "format_version",
90
+ method!(RbTableMetadata::format_version, 0),
91
+ )?;
92
+ class.define_method("uuid", method!(RbTableMetadata::uuid, 0))?;
93
+ class.define_method("location", method!(RbTableMetadata::location, 0))?;
57
94
  class.define_method(
58
95
  "last_sequence_number",
59
- method!(RbTable::last_sequence_number, 0),
96
+ method!(RbTableMetadata::last_sequence_number, 0),
60
97
  )?;
61
98
  class.define_method(
62
99
  "next_sequence_number",
63
- method!(RbTable::next_sequence_number, 0),
100
+ method!(RbTableMetadata::next_sequence_number, 0),
101
+ )?;
102
+ class.define_method(
103
+ "last_column_id",
104
+ method!(RbTableMetadata::last_column_id, 0),
105
+ )?;
106
+ class.define_method(
107
+ "last_partition_id",
108
+ method!(RbTableMetadata::last_partition_id, 0),
109
+ )?;
110
+ class.define_method(
111
+ "last_updated_ms",
112
+ method!(RbTableMetadata::last_updated_ms, 0),
113
+ )?;
114
+ class.define_method("schemas", method!(RbTableMetadata::schemas, 0))?;
115
+ class.define_method("schema_by_id", method!(RbTableMetadata::schema_by_id, 1))?;
116
+ class.define_method(
117
+ "current_schema",
118
+ method!(RbTableMetadata::current_schema, 0),
119
+ )?;
120
+ class.define_method(
121
+ "current_schema_id",
122
+ method!(RbTableMetadata::current_schema_id, 0),
123
+ )?;
124
+ class.define_method(
125
+ "partition_specs",
126
+ method!(RbTableMetadata::partition_specs, 0),
64
127
  )?;
65
- class.define_method("last_column_id", method!(RbTable::last_column_id, 0))?;
66
- class.define_method("last_partition_id", method!(RbTable::last_partition_id, 0))?;
67
- class.define_method("last_updated_ms", method!(RbTable::last_updated_ms, 0))?;
68
- class.define_method("schemas", method!(RbTable::schemas, 0))?;
69
- class.define_method("schema_by_id", method!(RbTable::schema_by_id, 1))?;
70
- class.define_method("current_schema", method!(RbTable::current_schema, 0))?;
71
- class.define_method("current_schema_id", method!(RbTable::current_schema_id, 0))?;
72
- class.define_method("partition_specs", method!(RbTable::partition_specs, 0))?;
73
128
  class.define_method(
74
129
  "partition_spec_by_id",
75
- method!(RbTable::partition_spec_by_id, 1),
130
+ method!(RbTableMetadata::partition_spec_by_id, 1),
76
131
  )?;
77
132
  class.define_method(
78
133
  "default_partition_spec",
79
- method!(RbTable::default_partition_spec, 0),
134
+ method!(RbTableMetadata::default_partition_spec, 0),
80
135
  )?;
81
136
  class.define_method(
82
137
  "default_partition_spec_id",
83
- method!(RbTable::default_partition_spec_id, 0),
138
+ method!(RbTableMetadata::default_partition_spec_id, 0),
139
+ )?;
140
+ class.define_method("snapshots", method!(RbTableMetadata::snapshots, 0))?;
141
+ class.define_method(
142
+ "snapshot_by_id",
143
+ method!(RbTableMetadata::snapshot_by_id, 1),
144
+ )?;
145
+ class.define_method("history", method!(RbTableMetadata::history, 0))?;
146
+ class.define_method("metadata_log", method!(RbTableMetadata::metadata_log, 0))?;
147
+ class.define_method(
148
+ "current_snapshot",
149
+ method!(RbTableMetadata::current_snapshot, 0),
84
150
  )?;
85
- class.define_method("snapshots", method!(RbTable::snapshots, 0))?;
86
- class.define_method("snapshot_by_id", method!(RbTable::snapshot_by_id, 1))?;
87
- class.define_method("history", method!(RbTable::history, 0))?;
88
- class.define_method("metadata_log", method!(RbTable::metadata_log, 0))?;
89
- class.define_method("current_snapshot", method!(RbTable::current_snapshot, 0))?;
90
151
  class.define_method(
91
152
  "current_snapshot_id",
92
- method!(RbTable::current_snapshot_id, 0),
153
+ method!(RbTableMetadata::current_snapshot_id, 0),
154
+ )?;
155
+ class.define_method(
156
+ "snapshot_for_ref",
157
+ method!(RbTableMetadata::snapshot_for_ref, 1),
158
+ )?;
159
+ class.define_method("sort_orders", method!(RbTableMetadata::sort_orders, 0))?;
160
+ class.define_method(
161
+ "sort_order_by_id",
162
+ method!(RbTableMetadata::sort_order_by_id, 1),
93
163
  )?;
94
- class.define_method("snapshot_for_ref", method!(RbTable::snapshot_for_ref, 1))?;
95
- class.define_method("sort_orders", method!(RbTable::sort_orders, 0))?;
96
- class.define_method("sort_order_by_id", method!(RbTable::sort_order_by_id, 1))?;
97
164
  class.define_method(
98
165
  "default_sort_order",
99
- method!(RbTable::default_sort_order, 0),
166
+ method!(RbTableMetadata::default_sort_order, 0),
100
167
  )?;
101
168
  class.define_method(
102
169
  "default_sort_order_id",
103
- method!(RbTable::default_sort_order_id, 0),
170
+ method!(RbTableMetadata::default_sort_order_id, 0),
104
171
  )?;
105
- class.define_method("properties", method!(RbTable::properties, 0))?;
106
- class.define_method("statistics", method!(RbTable::statistics, 0))?;
172
+ class.define_method("properties", method!(RbTableMetadata::properties, 0))?;
173
+ class.define_method("statistics", method!(RbTableMetadata::statistics, 0))?;
107
174
  class.define_method(
108
175
  "partition_statistics",
109
- method!(RbTable::partition_statistics, 0),
176
+ method!(RbTableMetadata::partition_statistics, 0),
110
177
  )?;
111
178
  class.define_method(
112
179
  "statistics_for_snapshot",
113
- method!(RbTable::statistics_for_snapshot, 1),
180
+ method!(RbTableMetadata::statistics_for_snapshot, 1),
114
181
  )?;
115
182
  class.define_method(
116
183
  "partition_statistics_for_snapshot",
117
- method!(RbTable::partition_statistics_for_snapshot, 1),
184
+ method!(RbTableMetadata::partition_statistics_for_snapshot, 1),
118
185
  )?;
119
- class.define_method("encryption_keys", method!(RbTable::encryption_keys, 0))?;
120
- class.define_method("encryption_key", method!(RbTable::encryption_key, 1))?;
121
- class.define_singleton_method(
122
- "from_metadata_file",
123
- function!(RbTable::from_metadata_file, 1),
186
+ class.define_method(
187
+ "encryption_keys",
188
+ method!(RbTableMetadata::encryption_keys, 0),
189
+ )?;
190
+ class.define_method(
191
+ "encryption_key",
192
+ method!(RbTableMetadata::encryption_key, 1),
124
193
  )?;
194
+ class.define_method("next_row_id", method!(RbTableMetadata::next_row_id, 0))?;
125
195
 
126
196
  let class = module.define_class("RbTableScan", ruby.class_object())?;
127
197
  class.define_method("plan_files", method!(RbTableScan::plan_files, 0))?;
128
198
  class.define_method("snapshot", method!(RbTableScan::snapshot, 0))?;
199
+ class.define_method("collect", method!(RbTableScan::collect, 0))?;
200
+ class.define_method("arrow_c_stream", method!(RbTableScan::arrow_c_stream, 0))?;
201
+
202
+ let class = module.define_class("FileScanTask", ruby.class_object())?;
203
+ class.define_method("file", method!(RbFileScanTask::file, 0))?;
204
+ class.define_method("delete_files", method!(RbFileScanTask::delete_files, 0))?;
205
+ class.define_method("inspect", method!(RbFileScanTask::inspect, 0))?;
206
+
207
+ let class = module.define_class("DataFile", ruby.class_object())?;
208
+ class.define_method("file_path", method!(RbDataFile::file_path, 0))?;
209
+ class.define_method("record_count", method!(RbDataFile::record_count, 0))?;
210
+ class.define_method(
211
+ "file_size_in_bytes",
212
+ method!(RbDataFile::file_size_in_bytes, 0),
213
+ )?;
214
+ class.define_method("equality_ids", method!(RbDataFile::equality_ids, 0))?;
215
+ class.define_method("inspect", method!(RbDataFile::inspect, 0))?;
216
+
217
+ let class = module.define_class("Schema", ruby.class_object())?;
218
+ class.define_singleton_method("new", function!(RbSchema::new, -1))?;
219
+ class.define_method("fields", method!(RbSchema::fields, 0))?;
220
+ class.define_method("schema_id", method!(RbSchema::schema_id, 0))?;
221
+ class.define_method(
222
+ "identifier_field_ids",
223
+ method!(RbSchema::identifier_field_ids, 0),
224
+ )?;
225
+ class.define_method("highest_field_id", method!(RbSchema::highest_field_id, 0))?;
226
+ class.define_method("as_struct", method!(RbSchema::as_struct, 0))?;
227
+ class.define_method("arrow_c_schema", method!(RbSchema::arrow_c_schema, 0))?;
228
+ class.define_method("==", method!(RbSchema::eq, 1))?;
229
+ class.define_method("inspect", method!(RbSchema::inspect, 0))?;
230
+
231
+ let class = module.define_class("NestedField", ruby.class_object())?;
232
+ class.define_singleton_method("new", function!(RbNestedField::new, 1))?;
233
+ class.define_method("field_id", method!(RbNestedField::field_id, 0))?;
234
+ class.define_method("name", method!(RbNestedField::name, 0))?;
235
+ class.define_method("field_type", method!(RbNestedField::field_type, 0))?;
236
+ class.define_method("required", method!(RbNestedField::required, 0))?;
237
+ class.define_method("doc", method!(RbNestedField::doc, 0))?;
238
+ class.define_method(
239
+ "initial_default",
240
+ method!(RbNestedField::initial_default, 0),
241
+ )?;
242
+ class.define_method("write_default", method!(RbNestedField::write_default, 0))?;
243
+ class.define_method("==", method!(RbNestedField::eq, 1))?;
244
+ class.define_method("inspect", method!(RbNestedField::inspect, 0))?;
245
+
246
+ let class = module.define_class("Capsule", ruby.class_object())?;
247
+ class.define_method("to_i", method!(RbCapsule::to_i, 0))?;
248
+ class.define_method("name", method!(RbCapsule::name, 0))?;
249
+
250
+ let class = module.define_class("ArrowRecordBatch", ruby.class_object())?;
251
+ class.define_singleton_method("new", function!(RbArrowRecordBatch::new, 2))?;
252
+ class.define_method(
253
+ "arrow_c_stream",
254
+ method!(RbArrowRecordBatch::arrow_c_stream, 0),
255
+ )?;
256
+
257
+ let class = module.define_class("PartitionSpec", ruby.class_object())?;
258
+ class.define_singleton_method("new", function!(RbPartitionSpec::new, -1))?;
259
+ class.define_method("spec_id", method!(RbPartitionSpec::spec_id, 0))?;
260
+ class.define_method("==", method!(RbPartitionSpec::eq, 1))?;
261
+ class.define_method("inspect", method!(RbPartitionSpec::inspect, 0))?;
262
+
263
+ let class = module.define_class("PartitionField", ruby.class_object())?;
264
+ class.define_singleton_method("new", function!(RbPartitionField::new, 1))?;
265
+ class.define_method("source_id", method!(RbPartitionField::source_id, 0))?;
266
+ class.define_method("field_id", method!(RbPartitionField::field_id, 0))?;
267
+ class.define_method("name", method!(RbPartitionField::name, 0))?;
268
+ class.define_method("transform", method!(RbPartitionField::transform, 0))?;
269
+ class.define_method("==", method!(RbPartitionField::eq, 1))?;
270
+ class.define_method("inspect", method!(RbPartitionField::inspect, 0))?;
271
+
272
+ let class = module.define_class("SortOrder", ruby.class_object())?;
273
+ class.define_singleton_method("new", function!(RbSortOrder::new, -1))?;
274
+ class.define_method("order_id", method!(RbSortOrder::order_id, 0))?;
275
+ class.define_method("fields", method!(RbSortOrder::fields, 0))?;
276
+ class.define_method("==", method!(RbSortOrder::eq, 1))?;
277
+ class.define_method("inspect", method!(RbSortOrder::inspect, 0))?;
278
+
279
+ let class = module.define_class("SortField", ruby.class_object())?;
280
+ class.define_singleton_method("new", function!(RbSortField::new, 1))?;
281
+ class.define_method("source_id", method!(RbSortField::source_id, 0))?;
282
+ class.define_method("transform", method!(RbSortField::transform, 0))?;
283
+ class.define_method("direction", method!(RbSortField::direction, 0))?;
284
+ class.define_method("null_order", method!(RbSortField::null_order, 0))?;
285
+ class.define_method("==", method!(RbSortField::eq, 1))?;
286
+ class.define_method("inspect", method!(RbSortField::inspect, 0))?;
287
+
288
+ let class = module.define_class("Snapshot", ruby.class_object())?;
289
+ class.define_method("snapshot_id", method!(RbSnapshot::snapshot_id, 0))?;
290
+ class.define_method(
291
+ "parent_snapshot_id",
292
+ method!(RbSnapshot::parent_snapshot_id, 0),
293
+ )?;
294
+ class.define_method("sequence_number", method!(RbSnapshot::sequence_number, 0))?;
295
+ class.define_method("manifest_list", method!(RbSnapshot::manifest_list, 0))?;
296
+ class.define_method("schema_id", method!(RbSnapshot::schema_id, 0))?;
297
+ class.define_method("inspect", method!(RbSnapshot::inspect, 0))?;
298
+
299
+ let class = module.define_class("SnapshotLogEntry", ruby.class_object())?;
300
+ class.define_method("snapshot_id", method!(RbSnapshotLogEntry::snapshot_id, 0))?;
301
+ class.define_method("timestamp_ms", method!(RbSnapshotLogEntry::timestamp_ms, 0))?;
302
+ class.define_method("inspect", method!(RbSnapshotLogEntry::inspect, 0))?;
303
+
304
+ let class = module.define_class("MetadataLogEntry", ruby.class_object())?;
305
+ class.define_method(
306
+ "metadata_file",
307
+ method!(RbMetadataLogEntry::metadata_file, 0),
308
+ )?;
309
+ class.define_method("timestamp_ms", method!(RbMetadataLogEntry::timestamp_ms, 0))?;
310
+ class.define_method("inspect", method!(RbMetadataLogEntry::inspect, 0))?;
311
+
312
+ let class = module.define_class("EncryptedKey", ruby.class_object())?;
313
+ class.define_method("key_id", method!(RbEncryptedKey::key_id, 0))?;
314
+ class.define_method(
315
+ "encrypted_by_id",
316
+ method!(RbEncryptedKey::encrypted_by_id, 0),
317
+ )?;
318
+ class.define_method("properties", method!(RbEncryptedKey::properties, 0))?;
319
+ class.define_method("inspect", method!(RbEncryptedKey::inspect, 0))?;
320
+
321
+ let class = module.define_class("StatisticsFile", ruby.class_object())?;
322
+ class.define_method("snapshot_id", method!(RbStatisticsFile::snapshot_id, 0))?;
323
+ class.define_method(
324
+ "statistics_path",
325
+ method!(RbStatisticsFile::statistics_path, 0),
326
+ )?;
327
+ class.define_method(
328
+ "file_size_in_bytes",
329
+ method!(RbStatisticsFile::file_size_in_bytes, 0),
330
+ )?;
331
+ class.define_method(
332
+ "file_footer_size_in_bytes",
333
+ method!(RbStatisticsFile::file_footer_size_in_bytes, 0),
334
+ )?;
335
+ class.define_method("key_metadata", method!(RbStatisticsFile::key_metadata, 0))?;
336
+ class.define_method("inspect", method!(RbStatisticsFile::inspect, 0))?;
337
+
338
+ let class = module.define_class("PartitionStatisticsFile", ruby.class_object())?;
339
+ class.define_method(
340
+ "snapshot_id",
341
+ method!(RbPartitionStatisticsFile::snapshot_id, 0),
342
+ )?;
343
+ class.define_method(
344
+ "statistics_path",
345
+ method!(RbPartitionStatisticsFile::statistics_path, 0),
346
+ )?;
347
+ class.define_method(
348
+ "file_size_in_bytes",
349
+ method!(RbPartitionStatisticsFile::file_size_in_bytes, 0),
350
+ )?;
351
+ class.define_method("inspect", method!(RbPartitionStatisticsFile::inspect, 0))?;
129
352
 
130
353
  Ok(())
131
354
  }
@@ -0,0 +1,102 @@
1
+ use iceberg::spec::{Transform, UnboundPartitionField, UnboundPartitionSpec};
2
+ use magnus::{IntoValue, RArray, RHash, Ruby, TryConvert, Value, value::ReprValue};
3
+
4
+ use crate::RbResult;
5
+ use crate::error::to_rb_err;
6
+ use crate::utils::{Wrap, rb_transform};
7
+
8
+ #[magnus::wrap(class = "Iceberg::PartitionSpec")]
9
+ pub struct RbPartitionSpec {
10
+ pub(crate) spec: UnboundPartitionSpec,
11
+ }
12
+
13
+ #[magnus::wrap(class = "Iceberg::PartitionField")]
14
+ pub struct RbPartitionField {
15
+ pub(crate) field: UnboundPartitionField,
16
+ }
17
+
18
+ impl RbPartitionSpec {
19
+ pub fn new(args: &[Value]) -> RbResult<Self> {
20
+ let mut fields = Vec::new();
21
+ for v in args {
22
+ fields.push(<&RbPartitionField>::try_convert(*v)?.field.clone());
23
+ }
24
+ let spec = UnboundPartitionSpec::builder()
25
+ .with_spec_id(0)
26
+ .add_partition_fields(fields)
27
+ .map_err(to_rb_err)?
28
+ .build();
29
+ Ok(Self { spec })
30
+ }
31
+
32
+ pub fn spec_id(&self) -> Option<i32> {
33
+ self.spec.spec_id()
34
+ }
35
+
36
+ pub fn fields(ruby: &Ruby, rb_self: &Self) -> RArray {
37
+ ruby.ary_from_iter(
38
+ rb_self
39
+ .spec
40
+ .fields()
41
+ .iter()
42
+ .map(|v| RbPartitionField { field: v.clone() }),
43
+ )
44
+ }
45
+
46
+ pub fn eq(&self, other: &Self) -> bool {
47
+ self.spec == other.spec
48
+ }
49
+
50
+ pub fn inspect(ruby: &Ruby, rb_self: &Self) -> String {
51
+ format!(
52
+ "#<Iceberg::PartitionSpec spec_id={}, fields={}>",
53
+ rb_self.spec_id().into_value_with(ruby).inspect(),
54
+ Self::fields(ruby, rb_self).inspect(),
55
+ )
56
+ }
57
+ }
58
+
59
+ impl RbPartitionField {
60
+ pub fn new(ruby: &Ruby, ob: RHash) -> RbResult<Self> {
61
+ let field = UnboundPartitionField::builder()
62
+ .source_id(ob.aref(ruby.to_symbol("source_id"))?)
63
+ .field_id(ob.aref(ruby.to_symbol("field_id"))?)
64
+ .name(ob.aref(ruby.to_symbol("name"))?)
65
+ .transform(
66
+ ob.aref::<_, Wrap<Transform>>(ruby.to_symbol("transform"))?
67
+ .0,
68
+ )
69
+ .build();
70
+ Ok(Self { field })
71
+ }
72
+
73
+ pub fn source_id(&self) -> i32 {
74
+ self.field.source_id
75
+ }
76
+
77
+ pub fn field_id(&self) -> Option<i32> {
78
+ self.field.field_id
79
+ }
80
+
81
+ pub fn name(&self) -> &str {
82
+ &self.field.name
83
+ }
84
+
85
+ pub fn transform(&self) -> RbResult<Value> {
86
+ rb_transform(&self.field.transform)
87
+ }
88
+
89
+ pub fn eq(&self, other: &Self) -> bool {
90
+ self.field == other.field
91
+ }
92
+
93
+ pub fn inspect(ruby: &Ruby, rb_self: &Self) -> RbResult<String> {
94
+ Ok(format!(
95
+ "#<Iceberg::PartitionField source_id={}, field_id={}, name={}, transform={}>",
96
+ rb_self.source_id().into_value_with(ruby).inspect(),
97
+ rb_self.field_id().into_value_with(ruby).inspect(),
98
+ rb_self.name().into_value_with(ruby).inspect(),
99
+ rb_self.transform()?.into_value_with(ruby).inspect(),
100
+ ))
101
+ }
102
+ }
@@ -0,0 +1,201 @@
1
+ use std::sync::Arc;
2
+
3
+ use arrow::array::{
4
+ Array, BooleanArray, Date32Array, Decimal128Array, Float32Array, Float64Array, Int8Array,
5
+ Int16Array, Int32Array, Int64Array, LargeBinaryArray, StringArray, TimestampMicrosecondArray,
6
+ TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array,
7
+ };
8
+ use arrow::datatypes::{DataType as ArrowDataType, Decimal128Type};
9
+ use arrow_array::RecordBatch;
10
+ use arrow_array::types::DecimalType;
11
+ use arrow_schema::TimeUnit;
12
+ use magnus::{Class, IntoValue, Module, RArray, RClass, RModule, Ruby, Value, value::ReprValue};
13
+
14
+ use crate::RbResult;
15
+ use crate::error::todo_error;
16
+ use crate::utils::EPOCH;
17
+
18
+ pub fn collect_batches(ruby: &Ruby, batches: Vec<RecordBatch>) -> RbResult<Value> {
19
+ let columns = ruby.ary_new();
20
+ let rows = ruby.ary_new();
21
+
22
+ for batch in batches {
23
+ if columns.is_empty() {
24
+ for field in &batch.schema().fields {
25
+ columns.push(ruby.str_new(field.name()))?;
26
+ }
27
+ }
28
+
29
+ for _ in 0..batch.num_rows() {
30
+ rows.push(ruby.ary_new())?;
31
+ }
32
+
33
+ for column in batch.columns() {
34
+ match column.data_type() {
35
+ ArrowDataType::Null => collect_column_null(ruby, column, rows)?,
36
+ ArrowDataType::Boolean => collect_column_boolean(ruby, column, rows)?,
37
+ ArrowDataType::Int8 => collect_column_int8(ruby, column, rows)?,
38
+ ArrowDataType::Int16 => collect_column_int16(ruby, column, rows)?,
39
+ ArrowDataType::Int32 => collect_column_int32(ruby, column, rows)?,
40
+ ArrowDataType::Int64 => collect_column_int64(ruby, column, rows)?,
41
+ ArrowDataType::UInt8 => collect_column_uint8(ruby, column, rows)?,
42
+ ArrowDataType::UInt16 => collect_column_uint16(ruby, column, rows)?,
43
+ ArrowDataType::UInt32 => collect_column_uint32(ruby, column, rows)?,
44
+ ArrowDataType::UInt64 => collect_column_uint64(ruby, column, rows)?,
45
+ ArrowDataType::Float32 => collect_column_float32(ruby, column, rows)?,
46
+ ArrowDataType::Float64 => collect_column_float64(ruby, column, rows)?,
47
+ ArrowDataType::Decimal128(precision, scale) => {
48
+ collect_column_decimal128(ruby, column, rows, *precision, *scale)?
49
+ }
50
+ ArrowDataType::Date32 => collect_column_date32(ruby, column, rows)?,
51
+ ArrowDataType::Timestamp(TimeUnit::Microsecond, None) => {
52
+ collect_column_timestamp_us(ruby, column, rows)?
53
+ }
54
+ ArrowDataType::Timestamp(TimeUnit::Nanosecond, None) => {
55
+ collect_column_timestamp_ns(ruby, column, rows)?
56
+ }
57
+ ArrowDataType::Utf8 => collect_column_utf8(ruby, column, rows)?,
58
+ ArrowDataType::LargeBinary => collect_column_large_binary(ruby, column, rows)?,
59
+ _ => return Err(todo_error(column.data_type())),
60
+ }
61
+ }
62
+ }
63
+
64
+ ruby.class_object()
65
+ .const_get::<_, RModule>("Iceberg")?
66
+ .const_get::<_, RClass>("Result")?
67
+ .new_instance((columns, rows))
68
+ }
69
+
70
+ pub fn collect_column_null(ruby: &Ruby, column: &Arc<dyn Array>, rows: RArray) -> RbResult<()> {
71
+ for i in 0..column.len() {
72
+ rows.entry::<RArray>(i.try_into().unwrap())?
73
+ .push(ruby.qnil())?;
74
+ }
75
+ Ok(())
76
+ }
77
+
78
+ macro_rules! collect_column {
79
+ ($name:ident, $type:ty) => {
80
+ pub fn $name(ruby: &Ruby, column: &Arc<dyn Array>, rows: RArray) -> RbResult<()> {
81
+ let array = column.as_any().downcast_ref::<$type>().unwrap();
82
+ for (i, value) in array.iter().enumerate() {
83
+ rows.entry::<RArray>(i.try_into().unwrap())?
84
+ .push(value.map(|v| v.into_value_with(ruby)))?;
85
+ }
86
+ Ok(())
87
+ }
88
+ };
89
+ }
90
+
91
+ collect_column!(collect_column_boolean, BooleanArray);
92
+ collect_column!(collect_column_int8, Int8Array);
93
+ collect_column!(collect_column_int16, Int16Array);
94
+ collect_column!(collect_column_int32, Int32Array);
95
+ collect_column!(collect_column_int64, Int64Array);
96
+ collect_column!(collect_column_uint8, UInt8Array);
97
+ collect_column!(collect_column_uint16, UInt16Array);
98
+ collect_column!(collect_column_uint32, UInt32Array);
99
+ collect_column!(collect_column_uint64, UInt64Array);
100
+ collect_column!(collect_column_float32, Float32Array);
101
+ collect_column!(collect_column_float64, Float64Array);
102
+ collect_column!(collect_column_utf8, StringArray);
103
+
104
+ pub fn collect_column_decimal128(
105
+ ruby: &Ruby,
106
+ column: &Arc<dyn Array>,
107
+ rows: RArray,
108
+ precision: u8,
109
+ scale: i8,
110
+ ) -> RbResult<()> {
111
+ ruby.require("bigdecimal")?;
112
+ let array = column.as_any().downcast_ref::<Decimal128Array>().unwrap();
113
+ for (i, value) in array.iter().enumerate() {
114
+ let value: Option<Value> = value
115
+ .map(|v| {
116
+ ruby.class_object().funcall(
117
+ "BigDecimal",
118
+ (Decimal128Type::format_decimal(v, precision, scale).into_value_with(ruby),),
119
+ )
120
+ })
121
+ .transpose()?;
122
+ rows.entry::<RArray>(i.try_into().unwrap())?.push(value)?;
123
+ }
124
+ Ok(())
125
+ }
126
+
127
+ pub fn collect_column_date32(ruby: &Ruby, column: &Arc<dyn Array>, rows: RArray) -> RbResult<()> {
128
+ let epoch = ruby.get_inner(&EPOCH);
129
+ let array = column.as_any().downcast_ref::<Date32Array>().unwrap();
130
+ for (i, value) in array.iter().enumerate() {
131
+ rows.entry::<RArray>(i.try_into().unwrap())?.push(
132
+ value
133
+ .map(|v| epoch.funcall::<_, _, Value>("+", (v,)))
134
+ .transpose()?,
135
+ )?;
136
+ }
137
+ Ok(())
138
+ }
139
+
140
+ pub fn collect_column_timestamp_us(
141
+ ruby: &Ruby,
142
+ column: &Arc<dyn Array>,
143
+ rows: RArray,
144
+ ) -> RbResult<()> {
145
+ let time_class = ruby.class_object().const_get::<_, Value>("Time")?;
146
+ let time_unit = ruby.to_symbol("usec");
147
+ let array = column
148
+ .as_any()
149
+ .downcast_ref::<TimestampMicrosecondArray>()
150
+ .unwrap();
151
+ for (i, value) in array.iter().enumerate() {
152
+ let value: Option<Value> = match value {
153
+ Some(v) => {
154
+ let sec = v / 1_000_000;
155
+ let usec = v % 1_000_000;
156
+ Some(time_class.funcall("at", (sec, usec, time_unit))?)
157
+ }
158
+ None => None,
159
+ };
160
+ rows.entry::<RArray>(i.try_into().unwrap())?.push(value)?;
161
+ }
162
+ Ok(())
163
+ }
164
+
165
+ pub fn collect_column_timestamp_ns(
166
+ ruby: &Ruby,
167
+ column: &Arc<dyn Array>,
168
+ rows: RArray,
169
+ ) -> RbResult<()> {
170
+ let time_class = ruby.class_object().const_get::<_, Value>("Time")?;
171
+ let time_unit = ruby.to_symbol("nsec");
172
+ let array = column
173
+ .as_any()
174
+ .downcast_ref::<TimestampNanosecondArray>()
175
+ .unwrap();
176
+ for (i, value) in array.iter().enumerate() {
177
+ let value: Option<Value> = match value {
178
+ Some(v) => {
179
+ let sec = v / 1_000_000_000;
180
+ let nsec = v % 1_000_000_000;
181
+ Some(time_class.funcall("at", (sec, nsec, time_unit))?)
182
+ }
183
+ None => None,
184
+ };
185
+ rows.entry::<RArray>(i.try_into().unwrap())?.push(value)?;
186
+ }
187
+ Ok(())
188
+ }
189
+
190
+ pub fn collect_column_large_binary(
191
+ ruby: &Ruby,
192
+ column: &Arc<dyn Array>,
193
+ rows: RArray,
194
+ ) -> RbResult<()> {
195
+ let array = column.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
196
+ for (i, value) in array.iter().enumerate() {
197
+ rows.entry::<RArray>(i.try_into().unwrap())?
198
+ .push(value.map(|v| ruby.str_from_slice(v)))?;
199
+ }
200
+ Ok(())
201
+ }