parquet 0.2.12-arm64-darwin → 0.3.0-arm64-darwin

data/Cargo.toml

@@ -1,3 +0,0 @@
-[workspace]
-members = ["./ext/parquet"]
-resolver = "2"

data/ext/parquet/Cargo.toml

@@ -1,28 +0,0 @@
-[package]
-name = "parquet"
-version = "0.1.0"
-edition = "2021"
-
-[lib]
-crate-type = ["cdylib"]
-
-[dependencies]
-ahash = "0.8"
-arrow-array = "54.0.0"
-arrow-schema = "54.0.0"
-bytes = "^1.9"
-itertools = "^0.14"
-jiff = "0.1.19"
-magnus = { version = "0.7", features = ["rb-sys"] }
-parquet = { version = "^54.0", features = ["json"] }
-rand = "0.9"
-rb-sys = "^0.9"
-thiserror = "2.0"
-tempfile = "^3.15"
-simdutf8 = "0.1.5"
-
-[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/extconf.rb

@@ -1,4 +0,0 @@
-require "mkmf"
-require "rb_sys/mkmf"
-
-create_rust_makefile("parquet/parquet")

data/ext/parquet/src/allocator.rs

@@ -1,13 +0,0 @@
-#[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

@@ -1,52 +0,0 @@
-use crate::ParserResultType;
-use magnus::{value::ReprValue, Error as MagnusError, KwArgs, RArray, RHash, Symbol, Value};
-
-pub struct RowEnumeratorArgs {
-    pub rb_self: Value,
-    pub to_read: Value,
-    pub result_type: ParserResultType,
-    pub columns: Option<Vec<String>>,
-}
-
-/// Creates an enumerator for lazy Parquet row parsing
-pub fn create_row_enumerator(args: RowEnumeratorArgs) -> Result<magnus::Enumerator, MagnusError> {
-    let kwargs = RHash::new();
-    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))?;
-    }
-    Ok(args
-        .rb_self
-        .enumeratorize("each_row", (args.to_read, KwArgs(kwargs))))
-}
-
-pub struct ColumnEnumeratorArgs {
-    pub rb_self: Value,
-    pub to_read: Value,
-    pub result_type: ParserResultType,
-    pub columns: Option<Vec<String>>,
-    pub batch_size: Option<usize>,
-}
-
-#[inline]
-pub fn create_column_enumerator(
-    args: ColumnEnumeratorArgs,
-) -> Result<magnus::Enumerator, MagnusError> {
-    let kwargs = RHash::new();
-    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))?;
-    }
-    if let Some(batch_size) = args.batch_size {
-        kwargs.aset(Symbol::new("batch_size"), batch_size)?;
-    }
-    Ok(args
-        .rb_self
-        .enumeratorize("each_column", (args.to_read, KwArgs(kwargs))))
-}

data/ext/parquet/src/header_cache.rs

@@ -1,100 +0,0 @@
-/// This module exists to avoid cloning header keys in returned HashMaps.
-/// Since the underlying RString creation already involves cloning,
-/// this caching layer aims to reduce redundant allocations.
-///
-/// Note: Performance testing on macOS showed minimal speed improvements,
-/// so this optimization could be removed if any issues arise.
-use std::{
-    collections::HashMap,
-    sync::{
-        atomic::{AtomicU32, Ordering},
-        LazyLock, Mutex,
-    },
-};
-
-use magnus::{IntoValue, RString, Ruby, Value};
-
-use thiserror::Error;
-
-#[derive(Debug, Error)]
-pub enum CacheError {
-    #[error("Failed to acquire lock: {0}")]
-    LockError(String),
-}
-
-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(&'static str);
-
-impl StringCacheKey {
-    pub fn new(string: &str) -> Self {
-        let rstr = RString::new(string);
-        let fstr = rstr.to_interned_str();
-        Self(fstr.as_str().unwrap())
-    }
-}
-
-impl AsRef<str> for StringCacheKey {
-    fn as_ref(&self) -> &'static str {
-        self.0
-    }
-}
-
-impl IntoValue for StringCacheKey {
-    fn into_value_with(self, handle: &Ruby) -> Value {
-        handle.into_value(self.0)
-    }
-}
-
-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.0.fmt(f)
-    }
-}
-
-impl PartialEq for StringCacheKey {
-    fn eq(&self, other: &Self) -> bool {
-        self.0 == other.0
-    }
-}
-
-impl std::cmp::Eq for StringCacheKey {}
-
-impl std::hash::Hash for StringCacheKey {
-    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        self.0.hash(state);
-    }
-}
-
-impl StringCache {
-    pub fn intern_many<AsStr: AsRef<str>>(
-        strings: &[AsStr],
-    ) -> Result<Vec<StringCacheKey>, CacheError> {
-        let mut cache = STRING_CACHE
-            .lock()
-            .map_err(|e| CacheError::LockError(e.to_string()))?;
-
-        let mut result: Vec<StringCacheKey> = Vec::with_capacity(strings.len());
-        for string in strings {
-            if let Some((_, (interned_string, counter))) = cache.get_key_value(string.as_ref()) {
-                counter.fetch_add(1, Ordering::Relaxed);
-                result.push(*interned_string);
-            } else {
-                let interned = StringCacheKey::new(string.as_ref());
-                cache.insert(interned.0, (interned, AtomicU32::new(1)));
-                result.push(interned);
-            }
-        }
-        Ok(result)
-    }
-}

data/ext/parquet/src/lib.rs

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

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

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

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

@@ -1,214 +0,0 @@
-use crate::header_cache::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::IntoValue;
-use magnus::{Error as MagnusError, Ruby, Value};
-use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
-use parquet::arrow::ProjectionMask;
-use std::collections::HashMap;
-use std::fs::File;
-use std::mem::ManuallyDrop;
-use std::os::fd::FromRawFd;
-use std::sync::OnceLock;
-
-use super::ReaderError;
-
-#[inline]
-pub fn parse_parquet_columns<'a>(rb_self: Value, args: &[Value]) -> Result<Value, 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,
-        })
-        .map(|yield_enum| yield_enum.into_value_with(&ruby));
-    }
-
-    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().map_err(|e| ReaderError::Parquet(e))?;
-
-        (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).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().map_err(|e| ReaderError::Parquet(e))?;
-
-        (reader, schema, num_rows)
-    } else {
-        let readable = SeekableRubyValue(Opaque::from(to_read));
-
-        let mut builder = ParquetRecordBatchReaderBuilder::try_new(readable)
-            .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().map_err(|e| ReaderError::Parquet(e))?;
-
-        (reader, schema, num_rows)
-    };
-
-    if num_rows == 0 {
-        let mut map =
-            HashMap::with_capacity_and_hasher(schema.fields().len(), RandomState::default());
-        let headers: Vec<String> = schema
-            .fields()
-            .iter()
-            .map(|field| field.name().to_string())
-            .collect();
-        let interned_headers =
-            StringCache::intern_many(&headers).map_err(|e| ReaderError::HeaderIntern(e))?;
-        for field in interned_headers.iter() {
-            map.insert(*field, vec![]);
-        }
-        let record = ColumnRecord::Map(map);
-        let _: Value = ruby.yield_value(record)?;
-        return Ok(ruby.qnil().into_value_with(&ruby));
-    }
-
-    match result_type {
-        ParserResultType::Hash => {
-            let headers = OnceLock::new();
-            let headers_clone = headers.clone();
-            let iter = batch_reader.map(move |batch| {
-                batch.map_err(ReaderError::Arrow).and_then(|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(), RandomState::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());
-                    });
-
-                    Ok(ColumnRecord::Map::<RandomState>(map))
-                })
-            });
-
-            for result in iter {
-                let record = result?;
-                let _: Value = ruby.yield_value(record)?;
-            }
-        }
-        ParserResultType::Array => {
-            let iter = batch_reader.map(|batch| {
-                batch.map_err(ReaderError::Arrow).and_then(|batch| {
-                    let vec = batch
-                        .columns()
-                        .into_iter()
-                        .map(|column| {
-                            let values = ParquetValueVec::try_from(column.clone()).unwrap();
-                            values.into_inner()
-                        })
-                        .collect();
-                    Ok(ColumnRecord::Vec::<RandomState>(vec))
-                })
-            });
-
-            for result in iter {
-                let record = result?;
-                let _: Value = ruby.yield_value(record)?;
-            }
-        }
-    }
-
-    Ok(ruby.qnil().into_value_with(&ruby))
-}

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

@@ -1,157 +0,0 @@
-use crate::header_cache::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::IntoValue;
-use magnus::{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<Value, 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,
-        })
-        .map(|yield_enum| yield_enum.into_value_with(&ruby));
-    }
-
-    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),
-            )
-        })?;
-    }
-
-    match result_type {
-        ParserResultType::Hash => {
-            let headers = OnceLock::new();
-            let headers_clone = headers.clone();
-            let iter = iter.map(move |row| {
-                row.and_then(|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(), RandomState::default());
-                    row.get_column_iter().enumerate().for_each(|(i, (_, v))| {
-                        map.insert(headers[i], ParquetField(v.clone()));
-                    });
-                    Ok(map)
-                })
-                .and_then(|row| Ok(RowRecord::Map::<RandomState>(row)))
-                .map_err(|e| ReaderError::Parquet(e))
-            });
-
-            for result in iter {
-                let record = result?;
-                let _: Value = ruby.yield_value(record)?;
-            }
-        }
-        ParserResultType::Array => {
-            let iter = iter.map(|row| {
-                row.and_then(|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())));
-                    Ok(vec)
-                })
-                .and_then(|row| Ok(RowRecord::Vec::<RandomState>(row)))
-                .map_err(|e| ReaderError::Parquet(e))
-            });
-
-            for result in iter {
-                let record = result?;
-                let _: Value = ruby.yield_value(record)?;
-            }
-        }
-    }
-
-    Ok(ruby.qnil().into_value_with(&ruby))
-}
-
-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()
-    }
-}