parquet 0.0.2 → 0.0.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 70d9932bf622cd2647423e2519013d3a9f9256217effe9610e9aeaaebbcf1778
-  data.tar.gz: fae3767ce0d950c91b17f77b740159d863293e1288063ed15d9b9c1f82e87fe1
+  metadata.gz: d46e0a95ff244189cadf71b3860b03aaf7638629d8c2eeda2800eaae57c0dbd2
+  data.tar.gz: 8bd2a8a29c7b199fcd67f7ababf241ebdcc871a0ff247fbd5320b789a6e6222e
 SHA512:
-  metadata.gz: a03e75bcd377ce5a61cd5f17685995c420601ac5917bd3d4a99dc082686423729ee5f0913bb032fe826dd1a8bac9b52c152cfb2037a376751258c17f3b0e63b1
-  data.tar.gz: ddfbb0ee14a6b7dcce47caf41962afe9610ab175d2b829c2744d62bed67cc746e64d214f64318220f2301a9ce8dcdecf9f9f9e90786df3d18f244716724abef8
+  metadata.gz: ae074d37108a5e12369638a23fcf0962ee416ac901dfcb64daf67b43e76e1566f50f56ccc890581a4cee8b14a0b801ad65dd874dd64c4d0882715e94d35e71b6
+  data.tar.gz: 9e08475b60bf1a5ee5e296aecca6df33a9f13fbe747914eadc45b171b6745feda099a3f8b4a1bca57ce3b3fdb03d4452545170d44d4dc90a4168c7b9194d8c19

data/Cargo.lock

@@ -743,6 +743,15 @@ dependencies = [
  "either",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.14"
@@ -769,6 +778,35 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "jiff"
+version = "0.1.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "943611a469f78ab9afdac9022e473a80fca16a9deca6c5be3eb566d872231e76"
+dependencies = [
+ "jiff-tzdb-platform",
+ "log",
+ "portable-atomic",
+ "portable-atomic-util",
+ "serde",
+ "windows-sys",
+]
+
+[[package]]
+name = "jiff-tzdb"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "91335e575850c5c4c673b9bd467b0e025f164ca59d0564f69d0c2ee0ffad4653"
+
+[[package]]
+name = "jiff-tzdb-platform"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9835f0060a626fe59f160437bc725491a6af23133ea906500027d1bd2f8f4329"
+dependencies = [
+ "jiff-tzdb",
+]
+
 [[package]]
 name = "jobserver"
 version = "0.1.32"
@@ -1159,8 +1197,12 @@ name = "parquet"
 version = "0.1.0"
 dependencies = [
  "ahash",
+ "arrow-array",
+ "arrow-schema",
  "bytes",
+ "itertools 0.14.0",
  "jemallocator",
+ "jiff",
  "kanal",
  "magnus 0.7.1",
  "mimalloc",
@@ -1238,6 +1280,21 @@ version = "0.3.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2"
 
+[[package]]
+name = "portable-atomic"
+version = "1.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8a2f0d8d040d7848a709caf78912debcc3f33ee4b3cac47d73d1e1069e83507"
+dependencies = [
+ "portable-atomic",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.92"

data/Gemfile

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

data/README.md

@@ -8,22 +8,78 @@ 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.
+
+### Row-wise Iteration
+
+The `each_row` method provides sequential access to individual rows:
+
 ```ruby
 require "parquet"
 
-# Read each row as a hash
-Parquet.each_row("test/data.parquet") { |row| puts row.inspect }
+# Basic usage with default hash output
+Parquet.each_row("data.parquet") do |row|
+  puts row.inspect  # {"id"=>1, "name"=>"name_1"}
+end
 
-# Read each row as an array
-Parquet.each_row("test/data.parquet", result_type: :array) { |row| puts row.inspect }
+# Array output for more efficient memory usage
+Parquet.each_row("data.parquet", result_type: :array) do |row|
+  puts row.inspect  # [1, "name_1"]
+end
 
-# Read from an IO object (like File or StringIO)
-File.open("test/data.parquet", "rb") do |file|
-  Parquet.each_row(file) { |row| puts row.inspect }
+# Select specific columns to reduce I/O
+Parquet.each_row("data.parquet", columns: ["id", "name"]) do |row|
+  puts row.inspect
 end
 
-# Or with StringIO
-io = StringIO.new(File.binread("test/data.parquet"))
-Parquet.each_row(io) { |row| puts row.inspect }
+# Reading from IO objects
+File.open("data.parquet", "rb") do |file|
+  Parquet.each_row(file) do |row|
+    puts row.inspect
+  end
+end
+```
+
+### Column-wise Iteration
+
+The `each_column` method reads data in column-oriented batches, which is typically more efficient for analytical queries:
 
+```ruby
+require "parquet"
+
+# Process columns in batches of 1024 rows
+Parquet.each_column("data.parquet", batch_size: 1024) do |batch|
+  # With result_type: :hash (default)
+  puts batch.inspect
+  # {
+  #   "id" => [1, 2, ..., 1024],
+  #   "name" => ["name_1", "name_2", ..., "name_1024"]
+  # }
+end
+
+# Array output with specific columns
+Parquet.each_column("data.parquet",
+                    columns: ["id", "name"],
+                    result_type: :array,
+                    batch_size: 1024) do |batch|
+  puts batch.inspect
+  # [
+  #   [1, 2, ..., 1024],           # id column
+  #   ["name_1", "name_2", ...]    # name column
+  # ]
+end
 ```
+
+### Arguments
+
+Both methods accept these common arguments:
+
+- `input`: Path string or IO-like object containing Parquet data
+- `result_type`: Output format (`:hash` or `:array`, defaults to `:hash`)
+- `columns`: Optional array of column names to read (improves performance)
+
+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.

data/ext/parquet/Cargo.toml

@@ -9,6 +9,8 @@ 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"
 bytes = "^1.9"
 kanal = "0.1.0-pre8"
 magnus = { version = "0.7", features = ["rb-sys"] }
@@ -16,6 +18,9 @@ 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"
+
 
 [target.'cfg(target_os = "linux")'.dependencies]
 jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }

data/ext/parquet/src/enumerator.rs

@@ -3,9 +3,9 @@ use magnus::{
     block::Yield, value::ReprValue, Error as MagnusError, KwArgs, RArray, RHash, Symbol, Value,
 };
 
-use crate::Record;
+use crate::{ColumnRecord, RowRecord};
 
-pub struct EnumeratorArgs {
+pub struct RowEnumeratorArgs {
     pub rb_self: Value,
     pub to_read: Value,
     pub result_type: String,
@@ -13,9 +13,9 @@ pub struct EnumeratorArgs {
 }
 
 #[inline]
-pub fn create_enumerator(
-    args: EnumeratorArgs,
-) -> Result<Yield<Box<dyn Iterator<Item = Record<RandomState>>>>, MagnusError> {
+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))?;
     if let Some(columns) = args.columns {
@@ -23,6 +23,32 @@ pub fn create_enumerator(
     }
     let enumerator = args
         .rb_self
-        .enumeratorize("for_each", (args.to_read, KwArgs(kwargs)));
+        .enumeratorize("each_row", (args.to_read, KwArgs(kwargs)));
+    Ok(Yield::Enumerator(enumerator))
+}
+
+pub struct ColumnEnumeratorArgs {
+    pub rb_self: Value,
+    pub to_read: Value,
+    pub result_type: String,
+    pub columns: Option<Vec<String>>,
+    pub batch_size: Option<usize>,
+}
+
+#[inline]
+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))?;
+    if let Some(columns) = args.columns {
+        kwargs.aset(Symbol::new("columns"), RArray::from_vec(columns))?;
+    }
+    if let Some(batch_size) = args.batch_size {
+        kwargs.aset(Symbol::new("batch_size"), batch_size)?;
+    }
+    let enumerator = args
+        .rb_self
+        .enumeratorize("each_column", (args.to_read, KwArgs(kwargs)));
     Ok(Yield::Enumerator(enumerator))
 }

data/ext/parquet/src/header_cache.rs

@@ -6,8 +6,14 @@
 /// so this optimization could be removed if any issues arise.
 use std::{
     collections::HashMap,
-    sync::{atomic::AtomicU32, LazyLock, Mutex, OnceLock},
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        LazyLock, Mutex, OnceLock,
+    },
 };
+
+use magnus::{r_string::FString, value::Opaque, IntoValue, RString, Ruby, Value};
+
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -16,64 +22,116 @@ pub enum CacheError {
     LockError(String),
 }
 
-static STRING_CACHE: LazyLock<Mutex<HashMap<&'static str, AtomicU32>>> =
+static STRING_CACHE: LazyLock<Mutex<HashMap<&'static str, (StringCacheKey, AtomicU32)>>> =
     LazyLock::new(|| Mutex::new(HashMap::with_capacity(100)));
 
 pub struct StringCache;
 
+#[derive(Copy, Clone)]
+pub struct StringCacheKey(Opaque<FString>, &'static str);
+
+impl StringCacheKey {
+    pub fn new(string: &str) -> Self {
+        let rstr = RString::new(string);
+        let fstr = rstr.to_interned_str();
+        Self(Opaque::from(fstr), fstr.as_str().unwrap())
+    }
+}
+
+impl AsRef<str> for StringCacheKey {
+    fn as_ref(&self) -> &'static str {
+        self.1
+    }
+}
+
+impl IntoValue for StringCacheKey {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        handle.into_value(self.0)
+    }
+}
+
+impl std::fmt::Debug for StringCacheKey {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.1.fmt(f)
+    }
+}
+
+impl PartialEq for StringCacheKey {
+    fn eq(&self, other: &Self) -> bool {
+        self.1 == other.1
+    }
+}
+
+impl std::cmp::Eq for StringCacheKey {}
+
+impl std::hash::Hash for StringCacheKey {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.1.hash(state);
+    }
+}
+
 impl StringCache {
     #[allow(dead_code)]
-    pub fn intern(string: String) -> Result<&'static str, CacheError> {
+    pub fn intern(string: String) -> Result<StringCacheKey, CacheError> {
         let mut cache = STRING_CACHE
             .lock()
             .map_err(|e| CacheError::LockError(e.to_string()))?;
 
-        if let Some((&existing, count)) = cache.get_key_value(string.as_str()) {
-            count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-            Ok(existing)
+        if let Some((_, (interned_string, counter))) = cache.get_key_value(string.as_str()) {
+            counter.fetch_add(1, Ordering::Relaxed);
+            Ok(*interned_string)
         } else {
+            let interned = StringCacheKey::new(string.as_str());
             let leaked = Box::leak(string.into_boxed_str());
-            cache.insert(leaked, AtomicU32::new(1));
-            Ok(leaked)
+            cache.insert(leaked, (interned, AtomicU32::new(1)));
+            Ok(interned)
         }
     }
 
-    pub fn intern_many(strings: &[String]) -> Result<Vec<&'static str>, CacheError> {
+    pub fn intern_many(strings: &[String]) -> Result<Vec<StringCacheKey>, CacheError> {
         let mut cache = STRING_CACHE
             .lock()
             .map_err(|e| CacheError::LockError(e.to_string()))?;
 
-        let mut result = Vec::with_capacity(strings.len());
+        let mut result: Vec<StringCacheKey> = Vec::with_capacity(strings.len());
         for string in strings {
-            if let Some((&existing, count)) = cache.get_key_value(string.as_str()) {
-                count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-                result.push(existing);
+            if let Some((_, (interned_string, counter))) = cache.get_key_value(string.as_str()) {
+                counter.fetch_add(1, Ordering::Relaxed);
+                result.push(*interned_string);
             } else {
+                let interned = StringCacheKey::new(&string);
                 let leaked = Box::leak(string.clone().into_boxed_str());
-                cache.insert(leaked, AtomicU32::new(1));
-                result.push(leaked);
+                cache.insert(leaked, (interned, AtomicU32::new(1)));
+                result.push(interned);
             }
         }
         Ok(result)
     }
 
-    pub fn clear(headers: &[&'static str]) -> Result<(), CacheError> {
+    pub fn clear(headers: &[StringCacheKey]) -> Result<(), CacheError> {
         let mut cache = STRING_CACHE
             .lock()
             .map_err(|e| CacheError::LockError(e.to_string()))?;
 
-        for header in headers {
-            if let Some(count) = cache.get(header) {
-                // Returns the previous value of the counter
-                let was = count.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
-                if was == 1 {
-                    cache.remove(header);
-                    let ptr = *header as *const str as *mut str;
-                    unsafe {
-                        let _ = Box::from_raw(ptr);
+        let to_remove: Vec<_> = headers
+            .iter()
+            .filter_map(|header| {
+                let key = header.as_ref();
+                if let Some((_, (_, counter))) = cache.get_key_value(key) {
+                    let prev_count = counter.fetch_sub(1, Ordering::Relaxed);
+                    if prev_count == 1 {
+                        Some(key)
+                    } else {
+                        None
                     }
+                } else {
+                    None
                 }
-            }
+            })
+            .collect();
+
+        for key in to_remove {
+            cache.remove(key);
         }
 
         Ok(())
@@ -82,13 +140,12 @@ impl StringCache {
 
 pub struct HeaderCacheCleanupIter<I> {
     pub inner: I,
-    pub headers: OnceLock<Vec<&'static str>>,
+    pub headers: OnceLock<Vec<StringCacheKey>>,
 }
 
 impl<I: Iterator> Iterator for HeaderCacheCleanupIter<I> {
     type Item = I::Item;
 
-    #[inline(always)]
     fn next(&mut self) -> Option<Self::Item> {
         self.inner.next()
     }

data/ext/parquet/src/lib.rs

@@ -18,6 +18,7 @@ use magnus::{Error, Ruby};
 #[magnus::init]
 fn init(ruby: &Ruby) -> Result<(), Error> {
     let module = ruby.define_module("Parquet")?;
-    module.define_module_function("each_row", magnus::method!(parse_parquet, -1))?;
+    module.define_module_function("each_row", magnus::method!(parse_parquet_rows, -1))?;
+    module.define_module_function("each_column", magnus::method!(parse_parquet_columns, -1))?;
     Ok(())
 }