parquet 0.0.1 → 0.0.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 66a87903cde02fff1e0a91627311bc1eba1fccc1d61915b7504c9fc61e10ed39
-  data.tar.gz: 3d93f8aa12b664fd0dcff05b98dc3f12ba6ce745507ff18847936dafe8958be4
+  metadata.gz: 70d9932bf622cd2647423e2519013d3a9f9256217effe9610e9aeaaebbcf1778
+  data.tar.gz: fae3767ce0d950c91b17f77b740159d863293e1288063ed15d9b9c1f82e87fe1
 SHA512:
-  metadata.gz: f042ba14e11c51849be30d22df7a1c63911bcf2ca6778b463b53878b562676f151f9a91a376baf57c3819a89861af2b76e58a80b66ea8c38870b8cf18664105a
-  data.tar.gz: f708de1716e9640bba8fc3537c76b2cfefbcf9db709ea560db619eb6a50ef2a3965d06b43cb127541ae33067f2facab8bf75df6d8fe6c3aea6cadb765d099a2f
+  metadata.gz: a03e75bcd377ce5a61cd5f17685995c420601ac5917bd3d4a99dc082686423729ee5f0913bb032fe826dd1a8bac9b52c152cfb2037a376751258c17f3b0e63b1
+  data.tar.gz: ddfbb0ee14a6b7dcce47caf41962afe9610ab175d2b829c2744d62bed67cc746e64d214f64318220f2301a9ce8dcdecf9f9f9e90786df3d18f244716724abef8

data/Cargo.lock

@@ -749,6 +749,26 @@ 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 = "jobserver"
 version = "0.1.32"
@@ -876,6 +896,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 +978,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 +1158,17 @@ dependencies = [
 name = "parquet"
 version = "0.1.0"
 dependencies = [
+ "ahash",
  "bytes",
+ "jemallocator",
  "kanal",
  "magnus 0.7.1",
+ "mimalloc",
  "parquet 54.0.0",
  "rb-sys",
  "serde",
  "serde_magnus",
  "thiserror",
- "xxhash-rust",
 ]
 
 [[package]]
@@ -1796,12 +1837,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/ext/parquet/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 crate-type = ["cdylib"]
 
 [dependencies]
+ahash = "0.8"
 parquet = { version = "^54.0", features = ["json", "object_store"] }
 bytes = "^1.9"
 kanal = "0.1.0-pre8"
@@ -15,4 +16,9 @@ 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"] }
+
+[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,28 @@
+use ahash::RandomState;
+use magnus::{
+    block::Yield, value::ReprValue, Error as MagnusError, KwArgs, RArray, RHash, Symbol, Value,
+};
+
+use crate::Record;
+
+pub struct EnumeratorArgs {
+    pub rb_self: Value,
+    pub to_read: Value,
+    pub result_type: String,
+    pub columns: Option<Vec<String>>,
+}
+
+#[inline]
+pub fn create_enumerator(
+    args: EnumeratorArgs,
+) -> Result<Yield<Box<dyn Iterator<Item = Record<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("for_each", (args.to_read, KwArgs(kwargs)));
+    Ok(Yield::Enumerator(enumerator))
+}

data/ext/parquet/src/header_cache.rs

@@ -6,7 +6,7 @@
 /// so this optimization could be removed if any issues arise.
 use std::{
     collections::HashMap,
-    sync::{atomic::AtomicU32, LazyLock, Mutex},
+    sync::{atomic::AtomicU32, LazyLock, Mutex, OnceLock},
 };
 use thiserror::Error;
 
@@ -79,3 +79,25 @@ impl StringCache {
         Ok(())
     }
 }
+
+pub struct HeaderCacheCleanupIter<I> {
+    pub inner: I,
+    pub headers: OnceLock<Vec<&'static str>>,
+}
+
+impl<I: Iterator> Iterator for HeaderCacheCleanupIter<I> {
+    type Item = I::Item;
+
+    #[inline(always)]
+    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};
 

data/ext/parquet/src/reader.rs

@@ -1,106 +1,39 @@
-use crate::header_cache::{CacheError, StringCache};
-use crate::ruby_reader::{build_ruby_reader, SeekableRead};
-use crate::utils::*;
-use bytes::Bytes;
+// =============================================================================
+// Imports and Dependencies
+// =============================================================================
+use crate::header_cache::{CacheError, HeaderCacheCleanupIter, StringCache};
+use crate::{
+    create_enumerator, utils::*, EnumeratorArgs, ForgottenFileHandle, ParquetField, Record,
+    SeekableRubyValue,
+};
+use ahash::RandomState;
 use magnus::rb_sys::AsRawValue;
 use magnus::value::{Opaque, ReprValue};
-use magnus::IntoValue;
-use magnus::{block::Yield, Error as MagnusError, KwArgs, RHash, Ruby, Symbol, Value};
-use parquet::errors::ParquetError;
-use parquet::file::reader::{ChunkReader, Length, SerializedFileReader};
-use parquet::record::Field;
+use magnus::{block::Yield, Error as MagnusError, Ruby, Value};
+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, BufReader, Read, Seek, SeekFrom};
+use std::io::{self};
 use std::mem::ManuallyDrop;
 use std::os::fd::FromRawFd;
 use std::sync::OnceLock;
-use std::{borrow::Cow, hash::BuildHasher};
 use thiserror::Error;
-use xxhash_rust::xxh3::Xxh3Builder;
-
-use parquet::record::reader::RowIter as ParquetRowIter;
-
-#[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),
-}
-
-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(),
-        )
-    }
-}
-
-struct ForgottenFileHandle(ManuallyDrop<File>);
-
-impl Length for ForgottenFileHandle {
-    fn len(&self) -> u64 {
-        self.0.len()
-    }
-}
-
-impl ChunkReader for ForgottenFileHandle {
-    type T = BufReader<File>;
-
-    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
-        self.0.get_read(start)
-    }
-
-    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
-        self.0.get_bytes(start, length)
-    }
-}
-
-struct HeaderCacheCleanupIter<I> {
-    inner: I,
-    headers: OnceLock<Vec<&'static str>>,
-}
-
-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();
-        }
-    }
-}
 
+#[inline]
 pub fn parse_parquet<'a>(
     rb_self: Value,
     args: &[Value],
-) -> Result<Yield<Box<dyn Iterator<Item = Record<Xxh3Builder>>>>, MagnusError> {
+) -> Result<Yield<Box<dyn Iterator<Item = Record<RandomState>>>>, MagnusError> {
     let original = unsafe { Ruby::get_unchecked() };
     let ruby: &'static Ruby = Box::leak(Box::new(original));
 
     let ParquetArgs {
         to_read,
         result_type,
+        columns,
     } = parse_parquet_args(&ruby, args)?;
 
     if !ruby.block_given() {
@@ -108,15 +41,18 @@ pub fn parse_parquet<'a>(
             rb_self,
             to_read,
             result_type,
+            columns,
         });
     }
 
-    let iter = if to_read.is_kind_of(ruby.class_string()) {
+    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();
-        ParquetRowIter::from_file_into(Box::new(reader))
+        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) })
@@ -131,14 +67,28 @@ pub fn parse_parquet<'a>(
         let file = unsafe { File::from_raw_fd(fd) };
         let file = ForgottenFileHandle(ManuallyDrop::new(file));
         let reader = SerializedFileReader::new(file).unwrap();
-        ParquetRowIter::from_file_into(Box::new(reader))
+        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();
-        ParquetRowIter::from_file_into(Box::new(reader))
+        let schema = reader.metadata().file_metadata().schema().clone();
+
+        (schema, ParquetRowIter::from_file_into(Box::new(reader)))
     };
 
-    let iter: Box<dyn Iterator<Item = Record<Xxh3Builder>>> = match result_type.as_str() {
+    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 = Record<RandomState>>> = match result_type.as_str() {
         "hash" => {
             let headers = OnceLock::new();
             let headers_clone = headers.clone();
@@ -146,21 +96,23 @@ pub fn parse_parquet<'a>(
                 .filter_map(move |row| {
                     row.ok().map(|row| {
                         let headers = headers_clone.get_or_init(|| {
-                            row.get_column_iter()
-                                .map(|(k, _)| StringCache::intern(k.to_owned()).unwrap())
-                                .collect::<Vec<_>>()
+                            let column_count = row.get_column_iter().count();
+                            let mut headers = Vec::with_capacity(column_count);
+                            row.get_column_iter().for_each(|(k, _)| {
+                                headers.push(StringCache::intern(k.to_owned()).unwrap())
+                            });
+                            headers
                         });
 
-                        row.get_column_iter()
-                            .enumerate()
-                            .map(|(i, (_, v))| {
-                                let key = headers[i];
-                                (key, ParquetField(v.clone()))
-                            })
-                            .collect::<HashMap<&'static str, ParquetField, Xxh3Builder>>()
+                        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(|row| Record::Map(row));
+                .map(Record::Map);
 
             Box::new(HeaderCacheCleanupIter {
                 inner: iter,
@@ -170,12 +122,14 @@ pub fn parse_parquet<'a>(
         "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()
-                        .map(|(_, v)| ParquetField(v.clone()))
-                        .collect::<Vec<ParquetField>>()
+                        .for_each(|(_, v)| vec.push(ParquetField(v.clone())));
+                    vec
                 })
             })
-            .map(|row| Record::Vec(row)),
+            .map(Record::Vec),
         ),
         _ => {
             return Err(MagnusError::new(
@@ -188,150 +142,49 @@ pub fn parse_parquet<'a>(
     Ok(Yield::Iter(iter))
 }
 
-struct EnumeratorArgs {
-    rb_self: Value,
-    to_read: Value,
-    result_type: String,
-}
-
-fn create_enumerator(
-    args: EnumeratorArgs,
-) -> Result<Yield<Box<dyn Iterator<Item = Record<Xxh3Builder>>>>, MagnusError> {
-    let kwargs = RHash::new();
-
-    kwargs.aset(Symbol::new("result_type"), Symbol::new(args.result_type))?;
-
-    let enumerator = args
-        .rb_self
-        .enumeratorize("for_each", (args.to_read, KwArgs(kwargs)));
-    Ok(Yield::Enumerator(enumerator))
-}
-
-#[derive(Debug)]
-pub enum Record<S: BuildHasher + Default> {
-    Vec(Vec<ParquetField>),
-    Map(HashMap<&'static str, ParquetField, S>),
-}
+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();
 
-impl<S: BuildHasher + Default> IntoValue for Record<S> {
-    #[inline]
-    fn into_value_with(self, handle: &Ruby) -> Value {
-        match self {
-            Record::Vec(vec) => {
-                let ary = handle.ary_new_capa(vec.len());
-                vec.into_iter().try_for_each(|v| ary.push(v)).unwrap();
-                ary.into_value_with(handle)
-            }
-            Record::Map(map) => {
-                // Pre-allocate the hash with the known size
-                let hash = handle.hash_new_capa(map.len());
-                map.into_iter()
-                    .try_for_each(|(k, v)| hash.aset(k, v))
-                    .unwrap();
-                hash.into_value_with(handle)
-            }
+        SchemaType::GroupType {
+            basic_info: schema.get_basic_info().clone(),
+            fields: projected_fields,
         }
+    } else {
+        // Return original schema if not a group type
+        schema.clone()
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct CowValue<'a>(pub Cow<'a, str>);
-
-impl<'a> IntoValue for CowValue<'a> {
-    fn into_value_with(self, handle: &Ruby) -> Value {
-        self.0.into_value_with(handle)
-    }
-}
-
-#[derive(Debug)]
-pub struct ParquetField(Field);
-
-impl<'a> IntoValue for ParquetField {
-    fn into_value_with(self, handle: &Ruby) -> Value {
-        match self.0 {
-            Field::Byte(b) => b.into_value_with(handle),
-            Field::Bool(b) => b.into_value_with(handle),
-            Field::Short(s) => s.into_value_with(handle),
-            Field::Int(i) => i.into_value_with(handle),
-            Field::Long(l) => l.into_value_with(handle),
-            Field::UByte(ub) => ub.into_value_with(handle),
-            Field::UShort(us) => us.into_value_with(handle),
-            Field::UInt(ui) => ui.into_value_with(handle),
-            Field::ULong(ul) => ul.into_value_with(handle),
-            Field::Float16(f) => f32::from(f).into_value_with(handle),
-            Field::Float(f) => f.into_value_with(handle),
-            Field::Double(d) => d.into_value_with(handle),
-
-            Field::Str(s) => s.into_value_with(handle),
-            Field::Bytes(b) => handle.str_from_slice(b.data()).as_value(),
-            Field::Date(d) => d.into_value_with(handle),
-            Field::TimestampMillis(ts) => ts.into_value_with(handle),
-            Field::TimestampMicros(ts) => ts.into_value_with(handle),
-            Field::ListInternal(list) => {
-                let ary = handle.ary_new_capa(list.elements().len());
-                list.elements()
-                    .iter()
-                    .try_for_each(|e| ary.push(ParquetField(e.clone()).into_value_with(handle)))
-                    .unwrap();
-                ary.into_value_with(handle)
-            }
-            Field::MapInternal(map) => {
-                let hash = handle.hash_new_capa(map.entries().len());
-                map.entries()
-                    .iter()
-                    .try_for_each(|(k, v)| {
-                        hash.aset(
-                            ParquetField(k.clone()).into_value_with(handle),
-                            ParquetField(v.clone()).into_value_with(handle),
-                        )
-                    })
-                    .unwrap();
-                hash.into_value_with(handle)
-            }
-            // Field::Decimal(d) => d.to_string().into_value_with(handle),
-            // Field::Group(row) => row.into_value_with(handle),
-            Field::Null => handle.qnil().as_value(),
-            _ => panic!("Unsupported field type"),
-        }
-    }
+#[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),
 }
 
-struct SeekableRubyValue(Opaque<Value>);
-
-impl Length for SeekableRubyValue {
-    fn len(&self) -> u64 {
-        let ruby = unsafe { Ruby::get_unchecked() };
-        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
-        let current_pos = reader.seek(SeekFrom::Current(0)).unwrap();
-        let file_len = reader.seek(SeekFrom::End(0)).unwrap();
-        reader.seek(SeekFrom::Start(current_pos)).unwrap();
-        file_len
+impl From<MagnusError> for ReaderError {
+    fn from(err: MagnusError) -> Self {
+        Self::Ruby(err.to_string())
     }
 }
 
-impl ChunkReader for SeekableRubyValue {
-    type T = BufReader<Box<dyn SeekableRead>>;
-
-    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
-        let ruby = unsafe { Ruby::get_unchecked() };
-        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
-        reader.seek(SeekFrom::Start(start))?;
-        Ok(BufReader::new(reader))
-    }
-
-    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
-        let ruby = unsafe { Ruby::get_unchecked() };
-        let mut buffer = Vec::with_capacity(length);
-        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
-        reader.seek(SeekFrom::Start(start))?;
-        let read = reader.take(length as _).read_to_end(&mut buffer)?;
-
-        if read != length {
-            return Err(ParquetError::EOF(format!(
-                "Expected to read {} bytes, read only {}",
-                length, read
-            )));
-        }
-        Ok(buffer.into())
+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/ruby_integration.rs

@@ -0,0 +1,77 @@
+use std::{
+    fs::File,
+    io::{BufReader, SeekFrom},
+    mem::ManuallyDrop,
+};
+
+use bytes::Bytes;
+use magnus::{value::Opaque, Ruby, Value};
+use parquet::{
+    errors::ParquetError,
+    file::reader::{ChunkReader, Length},
+};
+use std::io::Read;
+
+use crate::ruby_reader::{build_ruby_reader, SeekableRead};
+
+const READ_BUFFER_SIZE: usize = 16 * 1024;
+
+pub struct SeekableRubyValue(pub Opaque<Value>);
+
+impl Length for SeekableRubyValue {
+    fn len(&self) -> u64 {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        let current_pos = reader.seek(SeekFrom::Current(0)).unwrap();
+        let file_len = reader.seek(SeekFrom::End(0)).unwrap();
+        reader.seek(SeekFrom::Start(current_pos)).unwrap();
+        file_len
+    }
+}
+
+impl ChunkReader for SeekableRubyValue {
+    type T = BufReader<Box<dyn SeekableRead>>;
+
+    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        reader.seek(SeekFrom::Start(start))?;
+        Ok(BufReader::with_capacity(READ_BUFFER_SIZE, reader))
+    }
+
+    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut buffer = Vec::with_capacity(length);
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        reader.seek(SeekFrom::Start(start))?;
+        let read = reader.take(length as _).read_to_end(&mut buffer)?;
+
+        if read != length {
+            return Err(ParquetError::EOF(format!(
+                "Expected to read {} bytes, read only {}",
+                length, read
+            )));
+        }
+        Ok(buffer.into())
+    }
+}
+
+pub struct ForgottenFileHandle(pub ManuallyDrop<File>);
+
+impl Length for ForgottenFileHandle {
+    fn len(&self) -> u64 {
+        self.0.len()
+    }
+}
+
+impl ChunkReader for ForgottenFileHandle {
+    type T = BufReader<File>;
+
+    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
+        self.0.get_read(start)
+    }
+
+    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
+        self.0.get_bytes(start, length)
+    }
+}

data/ext/parquet/src/ruby_reader.rs

@@ -2,30 +2,24 @@ use magnus::{
     value::{Opaque, ReprValue},
     RClass, RString, Ruby, Value,
 };
-use std::io::{self, Read, Seek};
+use std::io::{self, Read, Seek, SeekFrom, Write};
 use std::sync::OnceLock;
 
 static STRING_IO_CLASS: OnceLock<Opaque<RClass>> = OnceLock::new();
 
-const READ_BUFFER_SIZE: usize = 16 * 1024;
-
 /// A reader that can handle various Ruby input types (String, StringIO, IO-like objects)
 /// and provide a standard Read implementation for them.
 pub struct RubyReader<T> {
     inner: T,
-    buffer: Option<Vec<u8>>,
     offset: usize,
-    // Number of bytes that have been read into the buffer
-    // Used as an upper bound for offset
-    buffered_bytes: usize,
 }
 
 pub trait SeekableRead: std::io::Read + Seek {}
 impl SeekableRead for RubyReader<Value> {}
 impl SeekableRead for RubyReader<RString> {}
 
-pub fn build_ruby_reader<'a>(
-    ruby: &'a Ruby,
+pub fn build_ruby_reader(
+    ruby: &Ruby,
     input: Value,
 ) -> Result<Box<dyn SeekableRead>, magnus::Error> {
     if RubyReader::is_string_io(ruby, &input) {
@@ -39,49 +33,48 @@ pub fn build_ruby_reader<'a>(
 
 impl Seek for RubyReader<Value> {
     fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
-        let seek_to = match pos {
-            io::SeekFrom::Start(offset) => {
-                // SEEK_SET - absolute position
-                offset as i64
-            }
-            io::SeekFrom::End(offset) => {
-                // SEEK_END - from end of stream
-                offset
-            }
-            io::SeekFrom::Current(offset) => {
-                // SEEK_CUR - relative to current
-                offset
-            }
+        let (whence, offset) = match pos {
+            SeekFrom::Start(i) => (0, i as i64),
+            SeekFrom::Current(i) => (1, i),
+            SeekFrom::End(i) => (2, i),
         };
 
-        let whence = match pos {
-            io::SeekFrom::Start(_) => 0,   // SEEK_SET
-            io::SeekFrom::End(_) => 2,     // SEEK_END
-            io::SeekFrom::Current(_) => 1, // SEEK_CUR
-        };
+        let new_position = self
+            .inner
+            .funcall("seek", (offset, whence))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        Ok(new_position)
+    }
+}
 
-        // Call Ruby's seek method
-        let _: u64 = self.inner.funcall("seek", (seek_to, whence)).unwrap();
+impl Write for RubyReader<Value> {
+    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
+        let ruby_bytes = RString::from_slice(buf);
 
-        // Get current position
-        let pos: u64 = self.inner.funcall("pos", ()).unwrap();
+        let bytes_written = self
+            .inner
+            .funcall::<_, _, usize>("write", (ruby_bytes,))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
 
-        Ok(pos)
+        Ok(bytes_written)
+    }
+
+    fn flush(&mut self) -> Result<(), io::Error> {
+        self.inner
+            .funcall::<_, _, Value>("flush", ())
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        Ok(())
     }
 }
 
 impl Seek for RubyReader<RString> {
     fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
         match pos {
-            io::SeekFrom::Start(offset) => {
-                self.offset = offset as usize;
-            }
-            io::SeekFrom::End(offset) => {
-                self.offset = (self.inner.len() - offset as usize) as usize;
-            }
-            io::SeekFrom::Current(offset) => {
-                self.offset = (self.offset as i64 + offset) as usize;
-            }
+            io::SeekFrom::Start(offset) => self.offset = offset as usize,
+            io::SeekFrom::Current(offset) => self.offset = (self.offset as i64 + offset) as usize,
+            io::SeekFrom::End(offset) => self.offset = self.inner.len() - offset as usize,
         }
         Ok(self.offset as u64)
     }
@@ -106,59 +99,9 @@ impl RubyReader<Value> {
     fn from_io_like(input: Value) -> Self {
         Self {
             inner: input,
-            buffer: Some(vec![0; READ_BUFFER_SIZE]),
             offset: 0,
-            buffered_bytes: 0,
         }
     }
-
-    fn read_from_buffer(&mut self, to_buf: &mut [u8]) -> Option<io::Result<usize>> {
-        if let Some(from_buf) = &self.buffer {
-            // If the offset is within the buffered bytes, copy the remaining bytes to the output buffer
-            if self.offset < self.buffered_bytes {
-                let remaining = self.buffered_bytes - self.offset;
-                let copy_size = remaining.min(to_buf.len());
-                to_buf[..copy_size]
-                    .copy_from_slice(&from_buf[self.offset..self.offset + copy_size]);
-                self.offset += copy_size;
-                Some(Ok(copy_size))
-            } else {
-                None
-            }
-        } else {
-            None
-        }
-    }
-
-    fn read_from_ruby(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        let buffer = self.buffer.as_mut().unwrap();
-        let result = self
-            .inner
-            .funcall::<_, _, RString>("read", (buffer.capacity(),))
-            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
-
-        if result.is_nil() {
-            return Ok(0); // EOF
-        }
-
-        let bytes = unsafe { result.as_slice() };
-
-        // Update internal buffer
-        let bytes_len = bytes.len();
-        if bytes_len == 0 {
-            return Ok(0);
-        }
-
-        // Only copy what we actually read
-        buffer[..bytes_len].copy_from_slice(bytes);
-        self.buffered_bytes = bytes_len;
-
-        // Copy to output buffer
-        let copy_size = bytes_len.min(buf.len());
-        buf[..copy_size].copy_from_slice(&buffer[..copy_size]);
-        self.offset = copy_size;
-        Ok(copy_size)
-    }
 }
 
 impl RubyReader<RString> {
@@ -176,9 +119,7 @@ impl RubyReader<RString> {
         let string_content = input.funcall::<_, _, RString>("string", ()).unwrap();
         Ok(Box::new(Self {
             inner: string_content,
-            buffer: None,
             offset: 0,
-            buffered_bytes: 0,
         }))
     }
 
@@ -197,21 +138,21 @@ impl RubyReader<RString> {
             .or_else(|_| input.funcall::<_, _, RString>("to_s", ()))?;
         Ok(Box::new(Self {
             inner: string_content,
-            buffer: None,
             offset: 0,
-            buffered_bytes: 0,
         }))
     }
 }
 
 impl Read for RubyReader<Value> {
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        if let Some(result) = self.read_from_buffer(buf) {
-            result
-        } else {
-            // If the buffer is empty, read from Ruby
-            self.read_from_ruby(buf)
-        }
+    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
+        let bytes = self
+            .inner
+            .funcall::<_, _, RString>("read", (buf.len(),))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        buf.write_all(unsafe { bytes.as_slice() })?;
+
+        Ok(bytes.len())
     }
 }
 

data/ext/parquet/src/types.rs

@@ -0,0 +1,90 @@
+use std::{borrow::Cow, collections::HashMap, hash::BuildHasher};
+
+use magnus::{value::ReprValue, IntoValue, Ruby, Value};
+use parquet::record::Field;
+
+#[derive(Debug)]
+pub enum Record<S: BuildHasher + Default> {
+    Vec(Vec<ParquetField>),
+    Map(HashMap<&'static str, ParquetField, S>),
+}
+
+impl<S: BuildHasher + Default> IntoValue for Record<S> {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        match self {
+            Record::Vec(vec) => {
+                let ary = handle.ary_new_capa(vec.len());
+                vec.into_iter().try_for_each(|v| ary.push(v)).unwrap();
+                ary.into_value_with(handle)
+            }
+            Record::Map(map) => {
+                let hash = handle.hash_new_capa(map.len());
+                map.into_iter()
+                    .try_for_each(|(k, v)| hash.aset(k, v))
+                    .unwrap();
+                hash.into_value_with(handle)
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct CowValue<'a>(pub Cow<'a, str>);
+
+impl<'a> IntoValue for CowValue<'a> {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        self.0.into_value_with(handle)
+    }
+}
+
+#[derive(Debug)]
+pub struct ParquetField(pub Field);
+
+impl IntoValue for ParquetField {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        match self.0 {
+            Field::Byte(b) => b.into_value_with(handle),
+            Field::Bool(b) => b.into_value_with(handle),
+            Field::Short(s) => s.into_value_with(handle),
+            Field::Int(i) => i.into_value_with(handle),
+            Field::Long(l) => l.into_value_with(handle),
+            Field::UByte(ub) => ub.into_value_with(handle),
+            Field::UShort(us) => us.into_value_with(handle),
+            Field::UInt(ui) => ui.into_value_with(handle),
+            Field::ULong(ul) => ul.into_value_with(handle),
+            Field::Float16(f) => f32::from(f).into_value_with(handle),
+            Field::Float(f) => f.into_value_with(handle),
+            Field::Double(d) => d.into_value_with(handle),
+            Field::Str(s) => s.into_value_with(handle),
+            Field::Bytes(b) => handle.str_from_slice(b.data()).as_value(),
+            Field::Date(d) => d.into_value_with(handle),
+            Field::TimestampMillis(ts) => ts.into_value_with(handle),
+            Field::TimestampMicros(ts) => ts.into_value_with(handle),
+            Field::ListInternal(list) => {
+                let elements = list.elements();
+                let ary = handle.ary_new_capa(elements.len());
+                elements
+                    .iter()
+                    .try_for_each(|e| ary.push(ParquetField(e.clone()).into_value_with(handle)))
+                    .unwrap();
+                ary.into_value_with(handle)
+            }
+            Field::MapInternal(map) => {
+                let entries = map.entries();
+                let hash = handle.hash_new_capa(entries.len());
+                entries
+                    .iter()
+                    .try_for_each(|(k, v)| {
+                        hash.aset(
+                            ParquetField(k.clone()).into_value_with(handle),
+                            ParquetField(v.clone()).into_value_with(handle),
+                        )
+                    })
+                    .unwrap();
+                hash.into_value_with(handle)
+            }
+            Field::Null => handle.qnil().as_value(),
+            _ => panic!("Unsupported field type"),
+        }
+    }
+}

data/ext/parquet/src/utils.rs

@@ -29,6 +29,7 @@ fn parse_string_or_symbol(ruby: &Ruby, value: Value) -> Result<Option<String>, E
 pub struct ParquetArgs {
     pub to_read: Value,
     pub result_type: String,
+    pub columns: Option<Vec<String>>,
 }
 
 /// Parse common arguments for CSV parsing
@@ -36,8 +37,11 @@ pub fn parse_parquet_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetArgs, Er
     let parsed_args = scan_args::<(Value,), (), (), (), _, ()>(args)?;
     let (to_read,) = parsed_args.required;
 
-    let kwargs =
-        get_kwargs::<_, (), (Option<Value>,), ()>(parsed_args.keywords, &[], &["result_type"])?;
+    let kwargs = get_kwargs::<_, (), (Option<Value>, Option<Vec<String>>), ()>(
+        parsed_args.keywords,
+        &[],
+        &["result_type", "columns"],
+    )?;
 
     let result_type = match kwargs
         .optional
@@ -66,5 +70,6 @@ pub fn parse_parquet_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetArgs, Er
     Ok(ParquetArgs {
         to_read,
         result_type,
+        columns: kwargs.optional.1,
     })
 }

data/lib/parquet/version.rb

@@ -1,3 +1,3 @@
 module Parquet
-  VERSION = "0.0.1"
+  VERSION = "0.0.2"
 end

data/lib/parquet.rbi

@@ -1,17 +1,19 @@
 # typed: strict
-
 module Parquet
   # Options:
   #   - `input`: String specifying the input file
   #   - `result_type`: String specifying the output format
   #                    ("hash" or "array" or :hash or :array)
+  #   - `columns`: When present, only the specified columns will be included in the output.
+  #                This is useful for reducing how much data is read and improving performance.
   sig do
     params(
       input: T.any(String, IO),
       result_type: T.nilable(T.any(String, Symbol)),
+      columns: T.nilable(T::Array[String]),
       blk: T.nilable(T.proc.params(row: T.any(T::Hash[String, T.untyped], T::Array[T.untyped])).void)
     ).returns(T.any(Enumerator, T.untyped))
   end
-  def self.each_row(input, result_type: nil, &blk)
+  def self.each_row(input, result_type: nil, columns: nil, &blk)
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: parquet
 version: !ruby/object:Gem::Version
-  version: 0.0.1
+  version: 0.0.2
 platform: ruby
 authors:
 - Nathan Jaremko
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-01-02 00:00:00.000000000 Z
+date: 2025-01-03 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys
@@ -56,10 +56,14 @@ files:
 - Rakefile
 - ext/parquet/Cargo.toml
 - ext/parquet/extconf.rb
+- ext/parquet/src/allocator.rs
+- ext/parquet/src/enumerator.rs
 - ext/parquet/src/header_cache.rs
 - ext/parquet/src/lib.rs
 - ext/parquet/src/reader.rs
+- ext/parquet/src/ruby_integration.rs
 - ext/parquet/src/ruby_reader.rs
+- ext/parquet/src/types.rs
 - ext/parquet/src/utils.rs
 - lib/parquet.rb
 - lib/parquet.rbi