parquet 0.0.1 → 0.0.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 66a87903cde02fff1e0a91627311bc1eba1fccc1d61915b7504c9fc61e10ed39
-  data.tar.gz: 3d93f8aa12b664fd0dcff05b98dc3f12ba6ce745507ff18847936dafe8958be4
+  metadata.gz: d46e0a95ff244189cadf71b3860b03aaf7638629d8c2eeda2800eaae57c0dbd2
+  data.tar.gz: 8bd2a8a29c7b199fcd67f7ababf241ebdcc871a0ff247fbd5320b789a6e6222e
 SHA512:
-  metadata.gz: f042ba14e11c51849be30d22df7a1c63911bcf2ca6778b463b53878b562676f151f9a91a376baf57c3819a89861af2b76e58a80b66ea8c38870b8cf18664105a
-  data.tar.gz: f708de1716e9640bba8fc3537c76b2cfefbcf9db709ea560db619eb6a50ef2a3965d06b43cb127541ae33067f2facab8bf75df6d8fe6c3aea6cadb765d099a2f
+  metadata.gz: ae074d37108a5e12369638a23fcf0962ee416ac901dfcb64daf67b43e76e1566f50f56ccc890581a4cee8b14a0b801ad65dd874dd64c4d0882715e94d35e71b6
+  data.tar.gz: 9e08475b60bf1a5ee5e296aecca6df33a9f13fbe747914eadc45b171b6745feda099a3f8b4a1bca57ce3b3fdb03d4452545170d44d4dc90a4168c7b9194d8c19

data/Cargo.lock

@@ -743,12 +743,70 @@ 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"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674"
 
+[[package]]
+name = "jemalloc-sys"
+version = "0.5.4+5.3.0-patched"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac6c1946e1cea1788cbfde01c993b52a10e2da07f4bac608228d1bed20bfebf2"
+dependencies = [
+ "cc",
+ "libc",
+]
+
+[[package]]
+name = "jemallocator"
+version = "0.5.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0de374a9f8e63150e6f5e8a60cc14c668226d7a347d8aee1a45766e3c4dd3bc"
+dependencies = [
+ "jemalloc-sys",
+ "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"
@@ -876,6 +934,16 @@ version = "0.2.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa"
 
+[[package]]
+name = "libmimalloc-sys"
+version = "0.1.39"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23aa6811d3bd4deb8a84dde645f943476d13b248d818edcf8ce0b2f37f036b44"
+dependencies = [
+ "cc",
+ "libc",
+]
+
 [[package]]
 name = "litemap"
 version = "0.7.4"
@@ -948,6 +1016,15 @@ version = "2.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
 
+[[package]]
+name = "mimalloc"
+version = "0.1.43"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "68914350ae34959d83f732418d51e2427a794055d0b9529f48259ac07af65633"
+dependencies = [
+ "libmimalloc-sys",
+]
+
 [[package]]
 name = "minimal-lexical"
 version = "0.2.1"
@@ -1119,15 +1196,21 @@ dependencies = [
 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",
  "parquet 54.0.0",
  "rb-sys",
  "serde",
  "serde_magnus",
  "thiserror",
- "xxhash-rust",
 ]
 
 [[package]]
@@ -1197,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"
@@ -1796,12 +1894,6 @@ version = "0.5.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51"
 
-[[package]]
-name = "xxhash-rust"
-version = "0.8.14"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d7d48f1b18be023c95e7b75f481cac649d74be7c507ff4a407c55cfb957f7934"
-
 [[package]]
 name = "yoke"
 version = "0.7.5"

data/Gemfile

@@ -6,7 +6,12 @@ gem "rake"
 # Use local version of parquet
 gemspec
 
-group :development, :test do
-  gem "minitest", "~> 5.0"
+group :development do
   gem "benchmark-ips", "~> 2.12"
+  # gem "polars-df"
+  gem "duckdb"
+end
+
+group :test do
+  gem "minitest", "~> 5.0"
 end

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

@@ -7,7 +7,10 @@ edition = "2021"
 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"] }
@@ -15,4 +18,12 @@ rb-sys = "^0.9"
 serde = { version = "1.0", features = ["derive"] }
 serde_magnus = "0.8.1"
 thiserror = "2.0"
-xxhash-rust = { version = "0.8.12", features = ["xxh3"] }
+itertools = "^0.14"
+jiff = "0.1.19"
+
+
+[target.'cfg(target_os = "linux")'.dependencies]
+jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }
+
+[target.'cfg(not(any(target_os = "linux", target_os = "windows")))'.dependencies]
+mimalloc = { version = "0.1", default-features = false }

data/ext/parquet/src/allocator.rs

@@ -0,0 +1,13 @@
+#[cfg(target_os = "linux")]
+use jemallocator::Jemalloc;
+
+#[cfg(not(any(target_os = "linux", target_os = "windows")))]
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+#[cfg(target_os = "linux")]
+static ALLOC: Jemalloc = Jemalloc;
+
+#[global_allocator]
+#[cfg(not(any(target_os = "linux", target_os = "windows")))]
+static ALLOC: MiMalloc = MiMalloc;

data/ext/parquet/src/enumerator.rs

@@ -0,0 +1,54 @@
+use ahash::RandomState;
+use magnus::{
+    block::Yield, value::ReprValue, Error as MagnusError, KwArgs, RArray, RHash, Symbol, Value,
+};
+
+use crate::{ColumnRecord, RowRecord};
+
+pub struct RowEnumeratorArgs {
+    pub rb_self: Value,
+    pub to_read: Value,
+    pub result_type: String,
+    pub columns: Option<Vec<String>>,
+}
+
+#[inline]
+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 {
+        kwargs.aset(Symbol::new("columns"), RArray::from_vec(columns))?;
+    }
+    let enumerator = args
+        .rb_self
+        .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},
+    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,66 +22,139 @@ 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(())
     }
 }
+
+pub struct HeaderCacheCleanupIter<I> {
+    pub inner: I,
+    pub headers: OnceLock<Vec<StringCacheKey>>,
+}
+
+impl<I: Iterator> Iterator for HeaderCacheCleanupIter<I> {
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.next()
+    }
+}
+
+impl<I> Drop for HeaderCacheCleanupIter<I> {
+    fn drop(&mut self) {
+        if let Some(headers) = self.headers.get() {
+            StringCache::clear(&headers).unwrap();
+        }
+    }
+}

data/ext/parquet/src/lib.rs

@@ -1,9 +1,16 @@
+mod allocator;
+mod enumerator;
 pub mod header_cache;
 mod reader;
+mod ruby_integration;
 mod ruby_reader;
+mod types;
 mod utils;
 
+use crate::enumerator::*;
 use crate::reader::*;
+use crate::ruby_integration::*;
+use crate::types::*;
 
 use magnus::{Error, Ruby};
 
@@ -11,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(())
 }