parquet 0.0.4 → 0.2.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b88d6751418f21c4ec032d05b6d0a6e9dbd37304983ed80e1a290508c787d118
-  data.tar.gz: 948702f38cad3c4d4e76efccbd9d7d8ad4c81366c4dcba2c71058cc4d013c237
+  metadata.gz: 90e876ca198a0e1871f692a382f09ceaeec670d162da26f2c102ea4eca4244bf
+  data.tar.gz: 96743e260cbd2fb55f6cdeaf256fbb1e915c57651fdc3f20fdd58b6a34596544
 SHA512:
-  metadata.gz: 30f90ee2f597aa6e2d5a84b8ab9780af3d71fa41d3a1152f47d7a12b34bc203b8ff06b04c3f929f689c93be9e962186a3e6c305f61724b36ad4e6ad551c11f49
-  data.tar.gz: 5a83b007e0c4789c6cfde1f8037228b0b00f2f0ef7ea0f932d7eaafefb91669db422450bbfd923f4388e2bfc644cae57f514828a2e4a2868ee6a20b492af428e
+  metadata.gz: 1609a37c5a9bd9f1d57bb31dd02b2fdb5b608a7c044686e6ef2513c95e53e830bd7bf7048a36904465a32a5915425c7b6bf581c5b35a4fb19f950cbca20913b2
+  data.tar.gz: 96ec18377fc5944556760329c126f440de61d3b378bfa976a66437db03f0a51220c880afd14098a5b1968daa968d2e836c50f83bef21507789ba4df314c48148

data/Cargo.lock

@@ -387,6 +387,22 @@ version = "1.13.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"
 
+[[package]]
+name = "errno"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d"
+dependencies = [
+ "libc",
+ "windows-sys",
+]
+
+[[package]]
+name = "fastrand"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
+
 [[package]]
 name = "flatbuffers"
 version = "24.12.23"
@@ -826,16 +842,6 @@ dependencies = [
  "wasm-bindgen",
 ]
 
-[[package]]
-name = "kanal"
-version = "0.1.0-pre8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b05d55519627edaf7fd0f29981f6dc03fb52df3f5b257130eb8d0bf2801ea1d7"
-dependencies = [
- "futures-core",
- "lock_api",
-]
-
 [[package]]
 name = "lazy_static"
 version = "1.5.0"
@@ -944,6 +950,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "linux-raw-sys"
+version = "0.4.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab"
+
 [[package]]
 name = "litemap"
 version = "0.7.4"
@@ -975,18 +987,6 @@ dependencies = [
  "twox-hash",
 ]
 
-[[package]]
-name = "magnus"
-version = "0.6.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b1597ef40aa8c36be098249e82c9a20cf7199278ac1c1a1a995eeead6a184479"
-dependencies = [
- "magnus-macros",
- "rb-sys",
- "rb-sys-env",
- "seq-macro",
-]
-
 [[package]]
 name = "magnus"
 version = "0.7.1"
@@ -1203,13 +1203,11 @@ dependencies = [
  "itertools 0.14.0",
  "jemallocator",
  "jiff",
- "kanal",
- "magnus 0.7.1",
+ "magnus",
  "mimalloc",
  "parquet 54.0.0",
  "rb-sys",
- "serde",
- "serde_magnus",
+ "tempfile",
  "thiserror",
 ]
 
@@ -1402,6 +1400,19 @@ dependencies = [
  "semver",
 ]
 
+[[package]]
+name = "rustix"
+version = "0.38.43"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a78891ee6bf2340288408954ac787aa063d8e8817e9f53abb37c695c6d834ef6"
+dependencies = [
+ "bitflags 2.6.0",
+ "errno",
+ "libc",
+ "linux-raw-sys",
+ "windows-sys",
+]
+
 [[package]]
 name = "ryu"
 version = "1.0.18"
@@ -1467,17 +1478,6 @@ dependencies = [
  "serde",
 ]
 
-[[package]]
-name = "serde_magnus"
-version = "0.8.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "76c20da583b5e1016e9199ef5f3260f7a8d1b253307d232600f6b12737262dbd"
-dependencies = [
- "magnus 0.6.4",
- "serde",
- "tap",
-]
-
 [[package]]
 name = "shell-words"
 version = "1.1.0"
@@ -1567,10 +1567,18 @@ dependencies = [
 ]
 
 [[package]]
-name = "tap"
-version = "1.0.1"
+name = "tempfile"
+version = "3.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
+checksum = "9a8a559c81686f576e8cd0290cd2a24a2a9ad80c98b3478856500fcbd7acd704"
+dependencies = [
+ "cfg-if",
+ "fastrand",
+ "getrandom",
+ "once_cell",
+ "rustix",
+ "windows-sys",
+]
 
 [[package]]
 name = "thiserror"

data/Gemfile

@@ -8,7 +8,7 @@ gemspec
 
 group :development do
   gem "benchmark-ips", "~> 2.12"
-  # gem "polars-df"
+  gem "polars-df"
   gem "duckdb"
 end
 

data/README.md

@@ -4,8 +4,6 @@
 
 This project is a Ruby library wrapping the [parquet-rs](https://github.com/apache/parquet-rs) rust crate.
 
-At the moment, it only supports iterating rows as either a hash or an array.
-
 ## Usage
 
 This library provides high-level bindings to parquet-rs with two primary APIs for reading Parquet files: row-wise and column-wise iteration. The column-wise API generally offers better performance, especially when working with subset of columns.
@@ -83,3 +81,95 @@ Additional arguments for `each_column`:
 - `batch_size`: Number of rows per batch (defaults to implementation-defined value)
 
 When no block is given, both methods return an Enumerator.
+
+### Writing Row-wise Data
+
+The `write_rows` method allows you to write data row by row:
+
+```ruby
+require "parquet"
+
+# Define the schema for your data
+schema = [
+  { "id" => "int64" },
+  { "name" => "string" },
+  { "score" => "double" }
+]
+
+# Create an enumerator that yields arrays of row values
+rows = [
+  [1, "Alice", 95.5],
+  [2, "Bob", 82.3],
+  [3, "Charlie", 88.7]
+].each
+
+# Write to a file
+Parquet.write_rows(rows, schema: schema, write_to: "data.parquet")
+
+# Write to an IO object
+File.open("data.parquet", "wb") do |file|
+  Parquet.write_rows(rows, schema: schema, write_to: file)
+end
+
+# Optionally specify batch size (default is 1000)
+Parquet.write_rows(rows,
+  schema: schema,
+  write_to: "data.parquet",
+  batch_size: 500
+)
+```
+
+### Writing Column-wise Data
+
+The `write_columns` method provides a more efficient way to write data in column-oriented batches:
+
+```ruby
+require "parquet"
+
+# Define the schema
+schema = [
+  { "id" => "int64" },
+  { "name" => "string" },
+  { "score" => "double" }
+]
+
+# Create batches of column data
+batches = [
+  # First batch
+  [
+    [1, 2],          # id column
+    ["Alice", "Bob"], # name column
+    [95.5, 82.3]     # score column
+  ],
+  # Second batch
+  [
+    [3],             # id column
+    ["Charlie"],     # name column
+    [88.7]           # score column
+  ]
+]
+
+# Create an enumerator from the batches
+columns = batches.each
+
+# Write to a parquet file
+Parquet.write_columns(columns, schema: schema, write_to: "data.parquet")
+
+# Write to an IO object
+File.open("data.parquet", "wb") do |file|
+  Parquet.write_columns(columns, schema: schema, write_to: file)
+end
+```
+
+The following data types are supported in the schema:
+
+- `int8`, `int16`, `int32`, `int64`
+- `uint8`, `uint16`, `uint32`, `uint64`
+- `float`, `double`
+- `string`
+- `binary`
+- `boolean`
+- `date32`
+- `timestamp_millis`, `timestamp_micros`
+
+Note: List and Map types are currently not supported.

data/ext/parquet/Cargo.toml

@@ -8,19 +8,16 @@ crate-type = ["cdylib"]
 
 [dependencies]
 ahash = "0.8"
-parquet = { version = "^54.0", features = ["json", "object_store"] }
-arrow-schema = "54.0.0"
 arrow-array = "54.0.0"
+arrow-schema = "54.0.0"
 bytes = "^1.9"
-kanal = "0.1.0-pre8"
+itertools = "^0.14"
+jiff = "0.1.19"
 magnus = { version = "0.7", features = ["rb-sys"] }
+parquet = { version = "^54.0", features = ["json", "object_store"] }
 rb-sys = "^0.9"
-serde = { version = "1.0", features = ["derive"] }
-serde_magnus = "0.8.1"
 thiserror = "2.0"
-itertools = "^0.14"
-jiff = "0.1.19"
-
+tempfile = "^3.15"
 
 [target.'cfg(target_os = "linux")'.dependencies]
 jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }

data/ext/parquet/src/enumerator.rs

@@ -3,12 +3,12 @@ use magnus::{
     block::Yield, value::ReprValue, Error as MagnusError, KwArgs, RArray, RHash, Symbol, Value,
 };
 
-use crate::{ColumnRecord, RowRecord};
+use crate::{ColumnRecord, ParserResultType, RowRecord};
 
 pub struct RowEnumeratorArgs {
     pub rb_self: Value,
     pub to_read: Value,
-    pub result_type: String,
+    pub result_type: ParserResultType,
     pub columns: Option<Vec<String>>,
 }
 
@@ -17,7 +17,10 @@ pub fn create_row_enumerator(
     args: RowEnumeratorArgs,
 ) -> Result<Yield<Box<dyn Iterator<Item = RowRecord<RandomState>>>>, MagnusError> {
     let kwargs = RHash::new();
-    kwargs.aset(Symbol::new("result_type"), Symbol::new(args.result_type))?;
+    kwargs.aset(
+        Symbol::new("result_type"),
+        Symbol::new(args.result_type.to_string()),
+    )?;
     if let Some(columns) = args.columns {
         kwargs.aset(Symbol::new("columns"), RArray::from_vec(columns))?;
     }
@@ -30,7 +33,7 @@ pub fn create_row_enumerator(
 pub struct ColumnEnumeratorArgs {
     pub rb_self: Value,
     pub to_read: Value,
-    pub result_type: String,
+    pub result_type: ParserResultType,
     pub columns: Option<Vec<String>>,
     pub batch_size: Option<usize>,
 }
@@ -40,7 +43,10 @@ pub fn create_column_enumerator(
     args: ColumnEnumeratorArgs,
 ) -> Result<Yield<Box<dyn Iterator<Item = ColumnRecord<RandomState>>>>, MagnusError> {
     let kwargs = RHash::new();
-    kwargs.aset(Symbol::new("result_type"), Symbol::new(args.result_type))?;
+    kwargs.aset(
+        Symbol::new("result_type"),
+        Symbol::new(args.result_type.to_string()),
+    )?;
     if let Some(columns) = args.columns {
         kwargs.aset(Symbol::new("columns"), RArray::from_vec(columns))?;
     }

data/ext/parquet/src/lib.rs

@@ -6,6 +6,7 @@ mod ruby_integration;
 mod ruby_reader;
 mod types;
 mod utils;
+mod writer;
 
 use crate::enumerator::*;
 use crate::reader::*;
@@ -13,6 +14,8 @@ use crate::ruby_integration::*;
 use crate::types::*;
 
 use magnus::{Error, Ruby};
+use writer::write_columns;
+use writer::write_rows;
 
 /// Initializes the Ruby extension and defines methods.
 #[magnus::init]
@@ -20,5 +23,7 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     let module = ruby.define_module("Parquet")?;
     module.define_module_function("each_row", magnus::method!(parse_parquet_rows, -1))?;
     module.define_module_function("each_column", magnus::method!(parse_parquet_columns, -1))?;
+    module.define_module_function("write_rows", magnus::function!(write_rows, -1))?;
+    module.define_module_function("write_columns", magnus::function!(write_columns, -1))?;
     Ok(())
 }

data/ext/parquet/src/reader/mod.rs

@@ -0,0 +1,42 @@
+mod parquet_column_reader;
+mod parquet_row_reader;
+
+use std::io;
+
+use magnus::{Error as MagnusError, Ruby};
+use thiserror::Error;
+
+use crate::header_cache::CacheError;
+pub use parquet_column_reader::parse_parquet_columns;
+pub use parquet_row_reader::parse_parquet_rows;
+
+#[derive(Error, Debug)]
+pub enum ReaderError {
+    #[error("Failed to get file descriptor: {0}")]
+    FileDescriptor(String),
+    #[error("Invalid file descriptor")]
+    InvalidFileDescriptor,
+    #[error("Failed to open file: {0}")]
+    FileOpen(#[from] io::Error),
+    #[error("Failed to intern headers: {0}")]
+    HeaderIntern(#[from] CacheError),
+    #[error("Ruby error: {0}")]
+    Ruby(String),
+    #[error("Parquet error: {0}")]
+    Parquet(#[from] parquet::errors::ParquetError),
+}
+
+impl From<MagnusError> for ReaderError {
+    fn from(err: MagnusError) -> Self {
+        Self::Ruby(err.to_string())
+    }
+}
+
+impl From<ReaderError> for MagnusError {
+    fn from(err: ReaderError) -> Self {
+        MagnusError::new(
+            Ruby::get().unwrap().exception_runtime_error(),
+            err.to_string(),
+        )
+    }
+}

data/ext/parquet/src/{reader.rs → reader/parquet_column_reader.rs}

@@ -1,11 +1,7 @@
-// =============================================================================
-// Imports and Dependencies
-// =============================================================================
 use crate::header_cache::{CacheError, HeaderCacheCleanupIter, StringCache};
 use crate::{
-    create_column_enumerator, create_row_enumerator, utils::*, ColumnEnumeratorArgs, ColumnRecord,
-    ForgottenFileHandle, ParquetField, ParquetValueVec, RowEnumeratorArgs, RowRecord,
-    SeekableRubyValue,
+    create_column_enumerator, utils::*, ColumnEnumeratorArgs, ColumnRecord, ForgottenFileHandle,
+    ParquetValueVec, ParserResultType, SeekableRubyValue,
 };
 use ahash::RandomState;
 use magnus::rb_sys::AsRawValue;
@@ -14,149 +10,20 @@ use magnus::{block::Yield, Error as MagnusError, Ruby, Value};
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use parquet::arrow::ProjectionMask;
 use parquet::errors::ParquetError;
-use parquet::file::reader::FileReader;
-use parquet::file::reader::SerializedFileReader;
-use parquet::record::reader::RowIter as ParquetRowIter;
-use parquet::schema::types::{Type as SchemaType, TypePtr};
 use std::collections::HashMap;
 use std::fs::File;
-use std::io::{self};
+use std::io;
 use std::mem::ManuallyDrop;
 use std::os::fd::FromRawFd;
 use std::sync::OnceLock;
 use thiserror::Error;
 
-#[inline]
-pub fn parse_parquet_rows<'a>(
-    rb_self: Value,
-    args: &[Value],
-) -> Result<Yield<Box<dyn Iterator<Item = RowRecord<RandomState>>>>, MagnusError> {
-    let original = unsafe { Ruby::get_unchecked() };
-    let ruby: &'static Ruby = Box::leak(Box::new(original));
-
-    let ParquetRowsArgs {
-        to_read,
-        result_type,
-        columns,
-    } = parse_parquet_rows_args(&ruby, args)?;
-
-    if !ruby.block_given() {
-        return create_row_enumerator(RowEnumeratorArgs {
-            rb_self,
-            to_read,
-            result_type,
-            columns,
-        });
-    }
-
-    let (schema, mut iter) = if to_read.is_kind_of(ruby.class_string()) {
-        let path_string = to_read.to_r_string()?;
-        let file_path = unsafe { path_string.as_str()? };
-        let file = File::open(file_path).unwrap();
-        let reader = SerializedFileReader::new(file).unwrap();
-        let schema = reader.metadata().file_metadata().schema().clone();
-
-        (schema, ParquetRowIter::from_file_into(Box::new(reader)))
-    } else if to_read.is_kind_of(ruby.class_io()) {
-        let raw_value = to_read.as_raw();
-        let fd = std::panic::catch_unwind(|| unsafe { rb_sys::rb_io_descriptor(raw_value) })
-            .map_err(|_| {
-                ReaderError::FileDescriptor("Failed to get file descriptor".to_string())
-            })?;
-
-        if fd < 0 {
-            return Err(ReaderError::InvalidFileDescriptor.into());
-        }
-
-        let file = unsafe { File::from_raw_fd(fd) };
-        let file = ForgottenFileHandle(ManuallyDrop::new(file));
-        let reader = SerializedFileReader::new(file).unwrap();
-        let schema = reader.metadata().file_metadata().schema().clone();
-
-        (schema, ParquetRowIter::from_file_into(Box::new(reader)))
-    } else {
-        let readable = SeekableRubyValue(Opaque::from(to_read));
-        let reader = SerializedFileReader::new(readable).unwrap();
-        let schema = reader.metadata().file_metadata().schema().clone();
-
-        (schema, ParquetRowIter::from_file_into(Box::new(reader)))
-    };
-
-    if let Some(cols) = columns {
-        let projection = create_projection_schema(&schema, &cols);
-        iter = iter.project(Some(projection.to_owned())).map_err(|e| {
-            MagnusError::new(
-                ruby.exception_runtime_error(),
-                format!("Failed to create projection: {}", e),
-            )
-        })?;
-    }
-
-    let iter: Box<dyn Iterator<Item = RowRecord<RandomState>>> = match result_type.as_str() {
-        "hash" => {
-            let headers = OnceLock::new();
-            let headers_clone = headers.clone();
-            let iter = iter
-                .filter_map(move |row| {
-                    row.ok().map(|row| {
-                        let headers = headers_clone.get_or_init(|| {
-                            let column_count = row.get_column_iter().count();
-
-                            let mut header_string = Vec::with_capacity(column_count);
-                            for (k, _) in row.get_column_iter() {
-                                header_string.push(k.to_owned());
-                            }
-
-                            let headers = StringCache::intern_many(&header_string).unwrap();
-
-                            headers
-                        });
-
-                        let mut map =
-                            HashMap::with_capacity_and_hasher(headers.len(), Default::default());
-                        row.get_column_iter().enumerate().for_each(|(i, (_, v))| {
-                            map.insert(headers[i], ParquetField(v.clone()));
-                        });
-                        map
-                    })
-                })
-                .map(RowRecord::Map);
-
-            Box::new(HeaderCacheCleanupIter {
-                inner: iter,
-                headers,
-            })
-        }
-        "array" => Box::new(
-            iter.filter_map(|row| {
-                row.ok().map(|row| {
-                    let column_count = row.get_column_iter().count();
-                    let mut vec = Vec::with_capacity(column_count);
-                    row.get_column_iter()
-                        .for_each(|(_, v)| vec.push(ParquetField(v.clone())));
-                    vec
-                })
-            })
-            .map(RowRecord::Vec),
-        ),
-        _ => {
-            return Err(MagnusError::new(
-                ruby.exception_runtime_error(),
-                "Invalid result type",
-            ))
-        }
-    };
-
-    Ok(Yield::Iter(iter))
-}
-
 #[inline]
 pub fn parse_parquet_columns<'a>(
     rb_self: Value,
     args: &[Value],
 ) -> Result<Yield<Box<dyn Iterator<Item = ColumnRecord<RandomState>>>>, MagnusError> {
-    let original = unsafe { Ruby::get_unchecked() };
-    let ruby: &'static Ruby = Box::leak(Box::new(original));
+    let ruby = unsafe { Ruby::get_unchecked() };
 
     let ParquetColumnsArgs {
         to_read,
@@ -282,8 +149,8 @@ pub fn parse_parquet_columns<'a>(
         return Ok(Yield::Iter(Box::new(column_record.into_iter())));
     }
 
-    let iter: Box<dyn Iterator<Item = ColumnRecord<RandomState>>> = match result_type.as_str() {
-        "hash" => {
+    let iter: Box<dyn Iterator<Item = ColumnRecord<RandomState>>> = match result_type {
+        ParserResultType::Hash => {
             let headers = OnceLock::new();
             let headers_clone = headers.clone();
             let iter = batch_reader
@@ -318,7 +185,7 @@ pub fn parse_parquet_columns<'a>(
                 headers,
             })
         }
-        "array" => Box::new(
+        ParserResultType::Array => Box::new(
             batch_reader
                 .filter_map(|batch| {
                     batch.ok().map(|batch| {
@@ -334,35 +201,11 @@ pub fn parse_parquet_columns<'a>(
                 })
                 .map(ColumnRecord::Vec),
         ),
-        _ => {
-            return Err(MagnusError::new(
-                ruby.exception_runtime_error(),
-                "Invalid result type",
-            ))
-        }
     };
 
     Ok(Yield::Iter(iter))
 }
 
-fn create_projection_schema(schema: &SchemaType, columns: &[String]) -> SchemaType {
-    if let SchemaType::GroupType { fields, .. } = schema {
-        let projected_fields: Vec<TypePtr> = fields
-            .iter()
-            .filter(|field| columns.contains(&field.name().to_string()))
-            .cloned()
-            .collect();
-
-        SchemaType::GroupType {
-            basic_info: schema.get_basic_info().clone(),
-            fields: projected_fields,
-        }
-    } else {
-        // Return original schema if not a group type
-        schema.clone()
-    }
-}
-
 #[derive(Error, Debug)]
 pub enum ReaderError {
     #[error("Failed to get file descriptor: {0}")]