parquet 0.0.4 → 0.0.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b88d6751418f21c4ec032d05b6d0a6e9dbd37304983ed80e1a290508c787d118
-  data.tar.gz: 948702f38cad3c4d4e76efccbd9d7d8ad4c81366c4dcba2c71058cc4d013c237
+  metadata.gz: 398a8ab4fe6b9c8e82d63ec832aa73163e75874c39080d87291a60397756df42
+  data.tar.gz: cace20e14d0eddc6e3185b2f9294253cb57c1689ec463ff66bc903d3c780af13
 SHA512:
-  metadata.gz: 30f90ee2f597aa6e2d5a84b8ab9780af3d71fa41d3a1152f47d7a12b34bc203b8ff06b04c3f929f689c93be9e962186a3e6c305f61724b36ad4e6ad551c11f49
-  data.tar.gz: 5a83b007e0c4789c6cfde1f8037228b0b00f2f0ef7ea0f932d7eaafefb91669db422450bbfd923f4388e2bfc644cae57f514828a2e4a2868ee6a20b492af428e
+  metadata.gz: 72ae6542b367fe433016f06fa109aaa77fe360bbc1df64e5c997db8fcc0a00aa166aa19a37240a706b3f443612770b80bc387dd41b34ee4a94ab26c3b0e74832
+  data.tar.gz: f69b10c6d4c8d879cdd3fce7c3b44933a99569358d1adfa3106760bd7c66036a2fef86737cf4dc6369be46234c124b9f2ef66e82fab118e36b5b079e9d23e10b

data/Cargo.lock

@@ -826,16 +826,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"
@@ -975,18 +965,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 +1181,10 @@ dependencies = [
  "itertools 0.14.0",
  "jemallocator",
  "jiff",
- "kanal",
- "magnus 0.7.1",
+ "magnus",
  "mimalloc",
  "parquet 54.0.0",
  "rb-sys",
- "serde",
- "serde_magnus",
  "thiserror",
 ]
 
@@ -1467,17 +1442,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"
@@ -1566,12 +1530,6 @@ dependencies = [
  "syn",
 ]
 
-[[package]]
-name = "tap"
-version = "1.0.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
-
 [[package]]
 name = "thiserror"
 version = "2.0.9"

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/ext/parquet/Cargo.toml

@@ -8,19 +8,15 @@ 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"
-
 
 [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

@@ -7,6 +7,9 @@ mod ruby_reader;
 mod types;
 mod utils;
 
+mod parquet_column_reader;
+mod parquet_row_reader;
+
 use crate::enumerator::*;
 use crate::reader::*;
 use crate::ruby_integration::*;

data/ext/parquet/src/parquet_column_reader.rs

@@ -0,0 +1,238 @@
+use crate::header_cache::{CacheError, HeaderCacheCleanupIter, StringCache};
+use crate::{
+    create_column_enumerator, utils::*, ColumnEnumeratorArgs, ColumnRecord, ForgottenFileHandle,
+    ParquetValueVec, ParserResultType, SeekableRubyValue,
+};
+use ahash::RandomState;
+use magnus::rb_sys::AsRawValue;
+use magnus::value::{Opaque, ReprValue};
+use magnus::{block::Yield, Error as MagnusError, Ruby, Value};
+use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+use parquet::arrow::ProjectionMask;
+use parquet::errors::ParquetError;
+use std::collections::HashMap;
+use std::fs::File;
+use std::io;
+use std::mem::ManuallyDrop;
+use std::os::fd::FromRawFd;
+use std::sync::OnceLock;
+use thiserror::Error;
+
+#[inline]
+pub fn parse_parquet_columns<'a>(
+    rb_self: Value,
+    args: &[Value],
+) -> Result<Yield<Box<dyn Iterator<Item = ColumnRecord<RandomState>>>>, MagnusError> {
+    let ruby = unsafe { Ruby::get_unchecked() };
+
+    let ParquetColumnsArgs {
+        to_read,
+        result_type,
+        columns,
+        batch_size,
+    } = parse_parquet_columns_args(&ruby, args)?;
+
+    if !ruby.block_given() {
+        return create_column_enumerator(ColumnEnumeratorArgs {
+            rb_self,
+            to_read,
+            result_type,
+            columns,
+            batch_size,
+        });
+    }
+
+    let (batch_reader, schema, num_rows) = 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).map_err(|e| ReaderError::FileOpen(e))?;
+
+        let mut builder =
+            ParquetRecordBatchReaderBuilder::try_new(file).map_err(|e| ReaderError::Parquet(e))?;
+        let schema = builder.schema().clone();
+        let num_rows = builder.metadata().file_metadata().num_rows();
+
+        // If columns are specified, project only those columns
+        if let Some(cols) = &columns {
+            // Get the parquet schema
+            let parquet_schema = builder.parquet_schema();
+
+            // Create a projection mask from column names
+            let projection =
+                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
+
+            builder = builder.with_projection(projection);
+        }
+
+        if let Some(batch_size) = batch_size {
+            builder = builder.with_batch_size(batch_size);
+        }
+
+        let reader = builder.build().unwrap();
+
+        (reader, schema, num_rows)
+    } 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 mut builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+        let schema = builder.schema().clone();
+        let num_rows = builder.metadata().file_metadata().num_rows();
+
+        if let Some(batch_size) = batch_size {
+            builder = builder.with_batch_size(batch_size);
+        }
+
+        // If columns are specified, project only those columns
+        if let Some(cols) = &columns {
+            // Get the parquet schema
+            let parquet_schema = builder.parquet_schema();
+
+            // Create a projection mask from column names
+            let projection =
+                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
+
+            builder = builder.with_projection(projection);
+        }
+
+        let reader = builder.build().unwrap();
+
+        (reader, schema, num_rows)
+    } else {
+        let readable = SeekableRubyValue(Opaque::from(to_read));
+
+        let mut builder = ParquetRecordBatchReaderBuilder::try_new(readable).unwrap();
+        let schema = builder.schema().clone();
+        let num_rows = builder.metadata().file_metadata().num_rows();
+
+        if let Some(batch_size) = batch_size {
+            builder = builder.with_batch_size(batch_size);
+        }
+
+        // If columns are specified, project only those columns
+        if let Some(cols) = &columns {
+            // Get the parquet schema
+            let parquet_schema = builder.parquet_schema();
+
+            // Create a projection mask from column names
+            let projection =
+                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
+
+            builder = builder.with_projection(projection);
+        }
+
+        let reader = builder.build().unwrap();
+
+        (reader, schema, num_rows)
+    };
+
+    if num_rows == 0 {
+        let mut map =
+            HashMap::with_capacity_and_hasher(schema.fields().len(), RandomState::default());
+        for field in schema.fields() {
+            map.insert(
+                StringCache::intern(field.name().to_string()).unwrap(),
+                vec![],
+            );
+        }
+        let column_record = vec![ColumnRecord::Map(map)];
+        return Ok(Yield::Iter(Box::new(column_record.into_iter())));
+    }
+
+    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
+                .filter_map(move |batch| {
+                    batch.ok().map(|batch| {
+                        let headers = headers_clone.get_or_init(|| {
+                            let schema = batch.schema();
+                            let fields = schema.fields();
+                            let mut header_string = Vec::with_capacity(fields.len());
+                            for field in fields {
+                                header_string.push(field.name().to_owned());
+                            }
+                            StringCache::intern_many(&header_string).unwrap()
+                        });
+
+                        let mut map =
+                            HashMap::with_capacity_and_hasher(headers.len(), Default::default());
+
+                        batch.columns().iter().enumerate().for_each(|(i, column)| {
+                            let header = headers[i];
+                            let values = ParquetValueVec::try_from(column.clone()).unwrap();
+                            map.insert(header, values.into_inner());
+                        });
+
+                        map
+                    })
+                })
+                .map(ColumnRecord::Map);
+
+            Box::new(HeaderCacheCleanupIter {
+                inner: iter,
+                headers,
+            })
+        }
+        ParserResultType::Array => Box::new(
+            batch_reader
+                .filter_map(|batch| {
+                    batch.ok().map(|batch| {
+                        batch
+                            .columns()
+                            .into_iter()
+                            .map(|column| {
+                                let values = ParquetValueVec::try_from(column.clone()).unwrap();
+                                values.into_inner()
+                            })
+                            .collect()
+                    })
+                })
+                .map(ColumnRecord::Vec),
+        ),
+    };
+
+    Ok(Yield::Iter(iter))
+}
+
+#[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] 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/parquet_row_reader.rs

@@ -0,0 +1,152 @@
+use crate::header_cache::{HeaderCacheCleanupIter, StringCache};
+use crate::{
+    create_row_enumerator, utils::*, ForgottenFileHandle, ParquetField, ParserResultType,
+    ReaderError, RowEnumeratorArgs, RowRecord, SeekableRubyValue,
+};
+use ahash::RandomState;
+use magnus::rb_sys::AsRawValue;
+use magnus::value::{Opaque, ReprValue};
+use magnus::{block::Yield, Error as MagnusError, Ruby, Value};
+use parquet::file::reader::{FileReader, 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::mem::ManuallyDrop;
+use std::os::fd::FromRawFd;
+use std::sync::OnceLock;
+
+#[inline]
+pub fn parse_parquet_rows<'a>(
+    rb_self: Value,
+    args: &[Value],
+) -> Result<Yield<Box<dyn Iterator<Item = RowRecord<RandomState>>>>, MagnusError> {
+    let ruby = unsafe { Ruby::get_unchecked() };
+
+    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 {
+        ParserResultType::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,
+            })
+        }
+        ParserResultType::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),
+        ),
+    };
+
+    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()
+    }
+}

data/ext/parquet/src/reader.rs

@@ -1,367 +1,11 @@
-// =============================================================================
-// 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,
-};
-use ahash::RandomState;
-use magnus::rb_sys::AsRawValue;
-use magnus::value::{Opaque, ReprValue};
-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::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 ParquetColumnsArgs {
-        to_read,
-        result_type,
-        columns,
-        batch_size,
-    } = parse_parquet_columns_args(&ruby, args)?;
-
-    if !ruby.block_given() {
-        return create_column_enumerator(ColumnEnumeratorArgs {
-            rb_self,
-            to_read,
-            result_type,
-            columns,
-            batch_size,
-        });
-    }
-
-    let (batch_reader, schema, num_rows) = 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).map_err(|e| ReaderError::FileOpen(e))?;
-
-        let mut builder =
-            ParquetRecordBatchReaderBuilder::try_new(file).map_err(|e| ReaderError::Parquet(e))?;
-        let schema = builder.schema().clone();
-        let num_rows = builder.metadata().file_metadata().num_rows();
-
-        // If columns are specified, project only those columns
-        if let Some(cols) = &columns {
-            // Get the parquet schema
-            let parquet_schema = builder.parquet_schema();
-
-            // Create a projection mask from column names
-            let projection =
-                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
-
-            builder = builder.with_projection(projection);
-        }
-
-        if let Some(batch_size) = batch_size {
-            builder = builder.with_batch_size(batch_size);
-        }
+use std::io;
 
-        let reader = builder.build().unwrap();
-
-        (reader, schema, num_rows)
-    } 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 mut builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
-        let schema = builder.schema().clone();
-        let num_rows = builder.metadata().file_metadata().num_rows();
-
-        if let Some(batch_size) = batch_size {
-            builder = builder.with_batch_size(batch_size);
-        }
-
-        // If columns are specified, project only those columns
-        if let Some(cols) = &columns {
-            // Get the parquet schema
-            let parquet_schema = builder.parquet_schema();
-
-            // Create a projection mask from column names
-            let projection =
-                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
-
-            builder = builder.with_projection(projection);
-        }
-
-        let reader = builder.build().unwrap();
-
-        (reader, schema, num_rows)
-    } else {
-        let readable = SeekableRubyValue(Opaque::from(to_read));
-
-        let mut builder = ParquetRecordBatchReaderBuilder::try_new(readable).unwrap();
-        let schema = builder.schema().clone();
-        let num_rows = builder.metadata().file_metadata().num_rows();
-
-        if let Some(batch_size) = batch_size {
-            builder = builder.with_batch_size(batch_size);
-        }
-
-        // If columns are specified, project only those columns
-        if let Some(cols) = &columns {
-            // Get the parquet schema
-            let parquet_schema = builder.parquet_schema();
-
-            // Create a projection mask from column names
-            let projection =
-                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
-
-            builder = builder.with_projection(projection);
-        }
-
-        let reader = builder.build().unwrap();
-
-        (reader, schema, num_rows)
-    };
-
-    if num_rows == 0 {
-        let mut map =
-            HashMap::with_capacity_and_hasher(schema.fields().len(), RandomState::default());
-        for field in schema.fields() {
-            map.insert(
-                StringCache::intern(field.name().to_string()).unwrap(),
-                vec![],
-            );
-        }
-        let column_record = vec![ColumnRecord::Map(map)];
-        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 headers = OnceLock::new();
-            let headers_clone = headers.clone();
-            let iter = batch_reader
-                .filter_map(move |batch| {
-                    batch.ok().map(|batch| {
-                        let headers = headers_clone.get_or_init(|| {
-                            let schema = batch.schema();
-                            let fields = schema.fields();
-                            let mut header_string = Vec::with_capacity(fields.len());
-                            for field in fields {
-                                header_string.push(field.name().to_owned());
-                            }
-                            StringCache::intern_many(&header_string).unwrap()
-                        });
-
-                        let mut map =
-                            HashMap::with_capacity_and_hasher(headers.len(), Default::default());
-
-                        batch.columns().iter().enumerate().for_each(|(i, column)| {
-                            let header = headers[i];
-                            let values = ParquetValueVec::try_from(column.clone()).unwrap();
-                            map.insert(header, values.into_inner());
-                        });
-
-                        map
-                    })
-                })
-                .map(ColumnRecord::Map);
-
-            Box::new(HeaderCacheCleanupIter {
-                inner: iter,
-                headers,
-            })
-        }
-        "array" => Box::new(
-            batch_reader
-                .filter_map(|batch| {
-                    batch.ok().map(|batch| {
-                        batch
-                            .columns()
-                            .into_iter()
-                            .map(|column| {
-                                let values = ParquetValueVec::try_from(column.clone()).unwrap();
-                                values.into_inner()
-                            })
-                            .collect()
-                    })
-                })
-                .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();
+use magnus::{Error as MagnusError, Ruby};
+use thiserror::Error;
 
-        SchemaType::GroupType {
-            basic_info: schema.get_basic_info().clone(),
-            fields: projected_fields,
-        }
-    } else {
-        // Return original schema if not a group type
-        schema.clone()
-    }
-}
+use crate::header_cache::CacheError;
+pub use crate::parquet_column_reader::parse_parquet_columns;
+pub use crate::parquet_row_reader::parse_parquet_rows;
 
 #[derive(Error, Debug)]
 pub enum ReaderError {
@@ -376,7 +20,7 @@ pub enum ReaderError {
     #[error("Ruby error: {0}")]
     Ruby(String),
     #[error("Parquet error: {0}")]
-    Parquet(#[from] ParquetError),
+    Parquet(#[from] parquet::errors::ParquetError),
 }
 
 impl From<MagnusError> for ReaderError {

data/ext/parquet/src/ruby_reader.rs

@@ -14,9 +14,8 @@ pub struct RubyReader<T> {
     offset: usize,
 }
 
-pub trait SeekableRead: std::io::Read + Seek {}
-impl SeekableRead for RubyReader<Value> {}
-impl SeekableRead for RubyReader<RString> {}
+pub trait SeekableRead: Read + Seek {}
+impl<T: Read + Seek> SeekableRead for T {}
 
 pub fn build_ruby_reader(
     ruby: &Ruby,

data/ext/parquet/src/types.rs

@@ -15,6 +15,47 @@ use parquet::record::Field;
 
 use crate::header_cache::StringCacheKey;
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum ParserResultType {
+    Hash,
+    Array,
+}
+
+impl ParserResultType {
+    pub fn iter() -> impl Iterator<Item = Self> {
+        [Self::Hash, Self::Array].into_iter()
+    }
+}
+
+impl TryFrom<&str> for ParserResultType {
+    type Error = String;
+
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        match value {
+            "hash" => Ok(ParserResultType::Hash),
+            "array" => Ok(ParserResultType::Array),
+            _ => Err(format!("Invalid parser result type: {}", value)),
+        }
+    }
+}
+
+impl TryFrom<String> for ParserResultType {
+    type Error = String;
+
+    fn try_from(value: String) -> Result<Self, Self::Error> {
+        Self::try_from(value.as_str())
+    }
+}
+
+impl std::fmt::Display for ParserResultType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ParserResultType::Hash => write!(f, "hash"),
+            ParserResultType::Array => write!(f, "array"),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub enum RowRecord<S: BuildHasher + Default> {
     Vec(Vec<ParquetField>),

data/ext/parquet/src/utils.rs

@@ -4,6 +4,8 @@ use magnus::{
     Error, RString, Ruby, Symbol, Value,
 };
 
+use crate::ParserResultType;
+
 fn parse_string_or_symbol(ruby: &Ruby, value: Value) -> Result<Option<String>, Error> {
     if value.is_nil() {
         Ok(None)
@@ -28,7 +30,7 @@ fn parse_string_or_symbol(ruby: &Ruby, value: Value) -> Result<Option<String>, E
 #[derive(Debug)]
 pub struct ParquetRowsArgs {
     pub to_read: Value,
-    pub result_type: String,
+    pub result_type: ParserResultType,
     pub columns: Option<Vec<String>>,
 }
 
@@ -43,28 +45,31 @@ pub fn parse_parquet_rows_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetRow
         &["result_type", "columns"],
     )?;
 
-    let result_type = match kwargs
+    let result_type: ParserResultType = match kwargs
         .optional
         .0
         .map(|value| parse_string_or_symbol(ruby, value))
     {
-        Some(Ok(Some(parsed))) => match parsed.as_str() {
-            "hash" | "array" => parsed,
-            _ => {
-                return Err(Error::new(
-                    magnus::exception::runtime_error(),
-                    "result_type must be either 'hash' or 'array'",
-                ))
-            }
-        },
-        Some(Ok(None)) => String::from("hash"),
+        Some(Ok(Some(parsed))) => parsed.try_into().map_err(|e| {
+            Error::new(
+                magnus::exception::runtime_error(),
+                format!(
+                    "Invalid result type: {e}. Must be one of {}",
+                    ParserResultType::iter()
+                        .map(|v| v.to_string())
+                        .collect::<Vec<_>>()
+                        .join(", ")
+                ),
+            )
+        })?,
+        Some(Ok(None)) => ParserResultType::Hash,
         Some(Err(_)) => {
             return Err(Error::new(
                 magnus::exception::type_error(),
                 "result_type must be a String or Symbol",
             ))
         }
-        None => String::from("hash"),
+        None => ParserResultType::Hash,
     };
 
     Ok(ParquetRowsArgs {
@@ -77,7 +82,7 @@ pub fn parse_parquet_rows_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetRow
 #[derive(Debug)]
 pub struct ParquetColumnsArgs {
     pub to_read: Value,
-    pub result_type: String,
+    pub result_type: ParserResultType,
     pub columns: Option<Vec<String>>,
     pub batch_size: Option<usize>,
 }
@@ -96,28 +101,31 @@ pub fn parse_parquet_columns_args(
         &["result_type", "columns", "batch_size"],
     )?;
 
-    let result_type = match kwargs
+    let result_type: ParserResultType = match kwargs
         .optional
         .0
         .map(|value| parse_string_or_symbol(ruby, value))
     {
-        Some(Ok(Some(parsed))) => match parsed.as_str() {
-            "hash" | "array" => parsed,
-            _ => {
-                return Err(Error::new(
-                    magnus::exception::runtime_error(),
-                    "result_type must be either 'hash' or 'array'",
-                ))
-            }
-        },
-        Some(Ok(None)) => String::from("hash"),
+        Some(Ok(Some(parsed))) => parsed.try_into().map_err(|e| {
+            Error::new(
+                magnus::exception::runtime_error(),
+                format!(
+                    "Invalid result type: {e}. Must be one of {}",
+                    ParserResultType::iter()
+                        .map(|v| v.to_string())
+                        .collect::<Vec<_>>()
+                        .join(", ")
+                ),
+            )
+        })?,
+        Some(Ok(None)) => ParserResultType::Hash,
         Some(Err(_)) => {
             return Err(Error::new(
                 magnus::exception::type_error(),
                 "result_type must be a String or Symbol",
             ))
         }
-        None => String::from("hash"),
+        None => ParserResultType::Hash,
     };
 
     Ok(ParquetColumnsArgs {

data/lib/parquet/version.rb

@@ -1,3 +1,3 @@
 module Parquet
-  VERSION = "0.0.4"
+  VERSION = "0.0.5"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: parquet
 version: !ruby/object:Gem::Version
-  version: 0.0.4
+  version: 0.0.5
 platform: ruby
 authors:
 - Nathan Jaremko
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-01-05 00:00:00.000000000 Z
+date: 2025-01-06 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys
@@ -60,6 +60,8 @@ files:
 - ext/parquet/src/enumerator.rs
 - ext/parquet/src/header_cache.rs
 - ext/parquet/src/lib.rs
+- ext/parquet/src/parquet_column_reader.rs
+- ext/parquet/src/parquet_row_reader.rs
 - ext/parquet/src/reader.rs
 - ext/parquet/src/ruby_integration.rs
 - ext/parquet/src/ruby_reader.rs