parquet 0.2.10 → 0.2.13

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d7ad6471a6c32833344fa6196794733c398a8de814652cc85f1121ab85f14f86
-  data.tar.gz: edeb31c6b6683b45913e58782753678302c980165aab4764bbc5b498c1203798
+  metadata.gz: 0cf24938c23cee5bc8ed4049e2b3fee7794cb619755e26cf83d4bb8826ebccd7
+  data.tar.gz: 85f55738e3503729535de7854d7438bca69f0b82e648471c285a3eefdb51a69b
 SHA512:
-  metadata.gz: 13929dda3279394f9a8b2c25a3c605fd813393c44030b2c5fc52c815e0582b75c008bcb4146b08d6079ab3d19e8d545ffbe00b5a10d6757556d4f27122039927
-  data.tar.gz: a542706c99bf184b0833239950d7a269dad091dece68e09dc495528df14f90ae614fc76017a8af64231e081c247a36def095b31deaa454eac9874c60805c47ca
+  metadata.gz: 17eaa053e7c05605d63c84786958f2980817509a6ba165654bfe50459cc30a37553671cc8b57a70831c255e463b26fc2768afcee3621664b443f9e0e67dc4460
+  data.tar.gz: e1a90f2683fce4a10b489eba3b0d98754ebeed9c418bd274e1af38c8fd9b5ad50f1e4ce72568f9c98001db408367c9610723f5c24796c801a1eaed4c23377d42

data/Cargo.lock

@@ -841,6 +841,7 @@ dependencies = [
  "parquet 54.0.0",
  "rand",
  "rb-sys",
+ "simdutf8",
  "tempfile",
  "thiserror",
 ]
@@ -1113,6 +1114,12 @@ version = "1.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
 
+[[package]]
+name = "simdutf8"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e"
+
 [[package]]
 name = "snap"
 version = "1.1.1"

data/ext/parquet/Cargo.toml

@@ -19,6 +19,7 @@ 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"] }

data/ext/parquet/src/enumerator.rs

@@ -1,21 +1,16 @@
-use ahash::RandomState;
-use magnus::{
-    block::Yield, value::ReprValue, Error as MagnusError, KwArgs, RArray, RHash, Symbol, Value,
-};
-
-use crate::{ColumnRecord, ParserResultType, RowRecord};
+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>>,
+    pub strict: bool,
 }
 
-#[inline]
-pub fn create_row_enumerator(
-    args: RowEnumeratorArgs,
-) -> Result<Yield<Box<dyn Iterator<Item = RowRecord<RandomState>>>>, MagnusError> {
+/// 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"),
@@ -24,10 +19,12 @@ pub fn create_row_enumerator(
     if let Some(columns) = args.columns {
         kwargs.aset(Symbol::new("columns"), RArray::from_vec(columns))?;
     }
-    let enumerator = args
+    if args.strict {
+        kwargs.aset(Symbol::new("strict"), true)?;
+    }
+    Ok(args
         .rb_self
-        .enumeratorize("each_row", (args.to_read, KwArgs(kwargs)));
-    Ok(Yield::Enumerator(enumerator))
+        .enumeratorize("each_row", (args.to_read, KwArgs(kwargs))))
 }
 
 pub struct ColumnEnumeratorArgs {
@@ -36,12 +33,13 @@ pub struct ColumnEnumeratorArgs {
     pub result_type: ParserResultType,
     pub columns: Option<Vec<String>>,
     pub batch_size: Option<usize>,
+    pub strict: bool,
 }
 
 #[inline]
 pub fn create_column_enumerator(
     args: ColumnEnumeratorArgs,
-) -> Result<Yield<Box<dyn Iterator<Item = ColumnRecord<RandomState>>>>, MagnusError> {
+) -> Result<magnus::Enumerator, MagnusError> {
     let kwargs = RHash::new();
     kwargs.aset(
         Symbol::new("result_type"),
@@ -53,8 +51,10 @@ pub fn create_column_enumerator(
     if let Some(batch_size) = args.batch_size {
         kwargs.aset(Symbol::new("batch_size"), batch_size)?;
     }
-    let enumerator = args
+    if args.strict {
+        kwargs.aset(Symbol::new("strict"), true)?;
+    }
+    Ok(args
         .rb_self
-        .enumeratorize("each_column", (args.to_read, KwArgs(kwargs)));
-    Ok(Yield::Enumerator(enumerator))
+        .enumeratorize("each_column", (args.to_read, KwArgs(kwargs))))
 }

data/ext/parquet/src/header_cache.rs

@@ -8,15 +8,15 @@ use std::{
     collections::HashMap,
     sync::{
         atomic::{AtomicU32, Ordering},
-        LazyLock, Mutex, OnceLock,
+        LazyLock, Mutex,
     },
 };
 
-use magnus::{r_string::FString, value::Opaque, IntoValue, RString, Ruby, Value};
+use magnus::{IntoValue, RString, Ruby, Value};
 
 use thiserror::Error;
 
-#[derive(Debug, Error)]
+#[derive(Debug, Clone, Error)]
 pub enum CacheError {
     #[error("Failed to acquire lock: {0}")]
     LockError(String),
@@ -28,19 +28,19 @@ static STRING_CACHE: LazyLock<Mutex<HashMap<&'static str, (StringCacheKey, Atomi
 pub struct StringCache;
 
 #[derive(Copy, Clone)]
-pub struct StringCacheKey(Opaque<FString>, &'static str);
+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(Opaque::from(fstr), fstr.as_str().unwrap())
+        Self(fstr.as_str().unwrap())
     }
 }
 
 impl AsRef<str> for StringCacheKey {
     fn as_ref(&self) -> &'static str {
-        self.1
+        self.0
     }
 }
 
@@ -50,15 +50,21 @@ impl IntoValue for StringCacheKey {
     }
 }
 
+impl IntoValue for &StringCacheKey {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        handle.into_value(self.0)
+    }
+}
+
 impl std::fmt::Debug for StringCacheKey {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.1.fmt(f)
+        self.0.fmt(f)
     }
 }
 
 impl PartialEq for StringCacheKey {
     fn eq(&self, other: &Self) -> bool {
-        self.1 == other.1
+        self.0 == other.0
     }
 }
 
@@ -66,95 +72,29 @@ impl std::cmp::Eq for StringCacheKey {}
 
 impl std::hash::Hash for StringCacheKey {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        self.1.hash(state);
+        self.0.hash(state);
     }
 }
 
 impl StringCache {
-    #[allow(dead_code)]
-    pub fn intern(string: String) -> Result<StringCacheKey, CacheError> {
-        let mut cache = STRING_CACHE
-            .lock()
-            .map_err(|e| CacheError::LockError(e.to_string()))?;
-
-        if let Some((_, (interned_string, counter))) = cache.get_key_value(string.as_str()) {
-            counter.fetch_add(1, Ordering::Relaxed);
-            Ok(*interned_string)
-        } else {
-            let interned = StringCacheKey::new(string.as_str());
-            let leaked = Box::leak(string.into_boxed_str());
-            cache.insert(leaked, (interned, AtomicU32::new(1)));
-            Ok(interned)
-        }
-    }
-
-    pub fn intern_many(strings: &[String]) -> Result<Vec<StringCacheKey>, CacheError> {
+    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_str()) {
+            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);
-                let leaked = Box::leak(string.clone().into_boxed_str());
-                cache.insert(leaked, (interned, AtomicU32::new(1)));
+                let interned = StringCacheKey::new(string.as_ref());
+                cache.insert(interned.0, (interned, AtomicU32::new(1)));
                 result.push(interned);
             }
         }
         Ok(result)
     }
-
-    pub fn clear(headers: &[StringCacheKey]) -> Result<(), CacheError> {
-        let mut cache = STRING_CACHE
-            .lock()
-            .map_err(|e| CacheError::LockError(e.to_string()))?;
-
-        let to_remove: Vec<_> = headers
-            .iter()
-            .filter_map(|header| {
-                let key = header.as_ref();
-                if let Some((_, (_, counter))) = cache.get_key_value(key) {
-                    let prev_count = counter.fetch_sub(1, Ordering::Relaxed);
-                    if prev_count == 1 {
-                        Some(key)
-                    } else {
-                        None
-                    }
-                } else {
-                    None
-                }
-            })
-            .collect();
-
-        for key in to_remove {
-            cache.remove(key);
-        }
-
-        Ok(())
-    }
-}
-
-pub struct HeaderCacheCleanupIter<I> {
-    pub inner: I,
-    pub headers: OnceLock<Vec<StringCacheKey>>,
-}
-
-impl<I: Iterator> Iterator for HeaderCacheCleanupIter<I> {
-    type Item = I::Item;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.inner.next()
-    }
-}
-
-impl<I> Drop for HeaderCacheCleanupIter<I> {
-    fn drop(&mut self) {
-        if let Some(headers) = self.headers.get() {
-            StringCache::clear(&headers).unwrap();
-        }
-    }
 }

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

@@ -24,6 +24,12 @@ pub enum ReaderError {
     Ruby(String),
     #[error("Parquet error: {0}")]
     Parquet(#[from] parquet::errors::ParquetError),
+    #[error("Arrow error: {0}")]
+    Arrow(#[from] arrow_schema::ArrowError),
+    #[error("UTF-8 error: {0}")]
+    Utf8Error(#[from] simdutf8::basic::Utf8Error),
+    #[error("Jiff error: {0}")]
+    Jiff(#[from] jiff::Error),
 }
 
 impl From<MagnusError> for ReaderError {

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

@@ -1,4 +1,5 @@
-use crate::header_cache::{CacheError, HeaderCacheCleanupIter, StringCache};
+use crate::header_cache::StringCache;
+use crate::types::{ArrayWrapper, TryIntoValue};
 use crate::{
     create_column_enumerator, utils::*, ColumnEnumeratorArgs, ColumnRecord, ForgottenFileHandle,
     ParquetValueVec, ParserResultType, SeekableRubyValue,
@@ -6,23 +7,20 @@ use crate::{
 use ahash::RandomState;
 use magnus::rb_sys::AsRawValue;
 use magnus::value::{Opaque, ReprValue};
-use magnus::{block::Yield, Error as MagnusError, Ruby, Value};
+use magnus::IntoValue;
+use magnus::{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;
+
+use super::ReaderError;
 
 #[inline]
-pub fn parse_parquet_columns<'a>(
-    rb_self: Value,
-    args: &[Value],
-) -> Result<Yield<Box<dyn Iterator<Item = ColumnRecord<RandomState>>>>, MagnusError> {
+pub fn parse_parquet_columns<'a>(rb_self: Value, args: &[Value]) -> Result<Value, MagnusError> {
     let ruby = unsafe { Ruby::get_unchecked() };
 
     let ParquetColumnsArgs {
@@ -30,6 +28,7 @@ pub fn parse_parquet_columns<'a>(
         result_type,
         columns,
         batch_size,
+        strict,
     } = parse_parquet_columns_args(&ruby, args)?;
 
     if !ruby.block_given() {
@@ -39,7 +38,9 @@ pub fn parse_parquet_columns<'a>(
             result_type,
             columns,
             batch_size,
-        });
+            strict,
+        })
+        .map(|yield_enum| yield_enum.into_value_with(&ruby));
     }
 
     let (batch_reader, schema, num_rows) = if to_read.is_kind_of(ruby.class_string()) {
@@ -68,7 +69,7 @@ pub fn parse_parquet_columns<'a>(
             builder = builder.with_batch_size(batch_size);
         }
 
-        let reader = builder.build().unwrap();
+        let reader = builder.build().map_err(|e| ReaderError::Parquet(e))?;
 
         (reader, schema, num_rows)
     } else if to_read.is_kind_of(ruby.class_io()) {
@@ -85,14 +86,11 @@ pub fn parse_parquet_columns<'a>(
         let file = unsafe { File::from_raw_fd(fd) };
         let file = ForgottenFileHandle(ManuallyDrop::new(file));
 
-        let mut builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+        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 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
@@ -105,20 +103,21 @@ pub fn parse_parquet_columns<'a>(
             builder = builder.with_projection(projection);
         }
 
-        let reader = builder.build().unwrap();
+        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).unwrap();
+        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 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
@@ -131,7 +130,11 @@ pub fn parse_parquet_columns<'a>(
             builder = builder.with_projection(projection);
         }
 
-        let reader = builder.build().unwrap();
+        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)
     };
@@ -139,100 +142,92 @@ pub fn parse_parquet_columns<'a>(
     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 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 column_record = vec![ColumnRecord::Map(map)];
-        return Ok(Yield::Iter(Box::new(column_record.into_iter())));
+        let record = ColumnRecord::Map(map);
+        let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+        return Ok(ruby.qnil().into_value_with(&ruby));
     }
 
-    let iter: Box<dyn Iterator<Item = ColumnRecord<RandomState>>> = match result_type {
+    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 iter = batch_reader.map(move |batch| {
+                batch.map_err(ReaderError::Arrow).and_then(|batch| {
+                    let local_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();
+                            StringCache::intern_many(&header_string)
+                        })
+                        .as_ref()
+                        .map_err(|e| ReaderError::HeaderIntern(e.clone()))?;
+
+                    let mut map = HashMap::with_capacity_and_hasher(
+                        local_headers.len(),
+                        RandomState::default(),
+                    );
+
+                    batch
+                        .columns()
+                        .iter()
+                        .enumerate()
+                        .try_for_each(|(i, column)| {
+                            let header = local_headers[i];
+                            let values = ParquetValueVec::try_from(ArrayWrapper {
+                                array: &*column,
+                                strict: strict,
+                            })?;
                             map.insert(header, values.into_inner());
-                        });
+                            Ok::<_, ReaderError>(())
+                        })?;
 
-                        map
-                    })
+                    Ok(ColumnRecord::Map::<RandomState>(map))
                 })
-                .map(ColumnRecord::Map);
+            });
 
-            Box::new(HeaderCacheCleanupIter {
-                inner: iter,
-                headers,
-            })
+            for result in iter {
+                let record = result?;
+                let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+            }
         }
-        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()
-                    })
+        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(ArrayWrapper {
+                                array: &*column,
+                                strict: strict,
+                            })?;
+                            Ok::<_, ReaderError>(values.into_inner())
+                        })
+                        .collect::<Result<Vec<_>, _>>()?;
+                    Ok(ColumnRecord::Vec::<RandomState>(vec))
                 })
-                .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())
+            for result in iter {
+                let record = result?;
+                let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+            }
+        }
     }
-}
 
-impl From<ReaderError> for MagnusError {
-    fn from(err: ReaderError) -> Self {
-        MagnusError::new(
-            Ruby::get().unwrap().exception_runtime_error(),
-            err.to_string(),
-        )
-    }
+    Ok(ruby.qnil().into_value_with(&ruby))
 }