parquet 0.4.1 → 0.5.0

data/ext/parquet/src/logger.rs

@@ -0,0 +1,171 @@
+// Logger module for Parquet gem
+// Provides a Rust wrapper for Ruby logger objects
+
+use std::str::FromStr;
+
+use magnus::{exception::runtime_error, value::ReprValue, Error as MagnusError, Ruby, Value};
+
+use crate::{reader::ReaderError, utils::parse_string_or_symbol};
+
+/// Severity levels that match Ruby's Logger levels
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub enum LogLevel {
+    Debug,
+    Info,
+    Warn,
+    Error,
+    Fatal,
+}
+
+impl FromStr for LogLevel {
+    type Err = MagnusError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(match s {
+            "debug" => LogLevel::Debug,
+            "info" => LogLevel::Info,
+            "warn" => LogLevel::Warn,
+            "error" => LogLevel::Error,
+            "fatal" => LogLevel::Fatal,
+            _ => {
+                return Err(MagnusError::new(
+                    runtime_error(),
+                    format!("Invalid log level: {}", s),
+                ))
+            }
+        })
+    }
+}
+/// A wrapper around a Ruby logger object
+#[derive(Debug, Clone)]
+pub struct RubyLogger {
+    logger: Option<Value>,
+    level: LogLevel,
+}
+
+#[allow(dead_code)]
+impl RubyLogger {
+    pub fn new(ruby: &Ruby, logger_value: Option<Value>) -> Result<Self, ReaderError> {
+        let environment_level = std::env::var("PARQUET_GEM_LOG_LEVEL")
+            .unwrap_or_else(|_| "warn".to_string())
+            .parse::<LogLevel>()
+            .unwrap_or(LogLevel::Warn);
+
+        match logger_value {
+            Some(logger) => {
+                if logger.is_nil() {
+                    return Ok(Self {
+                        logger: None,
+                        level: environment_level,
+                    });
+                }
+
+                let level_value = logger.funcall::<_, _, Value>("level", ())?;
+                let level = parse_string_or_symbol(ruby, level_value)?;
+                let level = level
+                    .map(|s| s.parse::<LogLevel>())
+                    .transpose()?
+                    .unwrap_or(environment_level);
+
+                Ok(Self {
+                    logger: Some(logger),
+                    level,
+                })
+            }
+            None => Ok(Self {
+                logger: None,
+                level: environment_level,
+            }),
+        }
+    }
+
+    /// Log a message at the given level
+    pub fn log(&self, level: LogLevel, message: &str) -> Result<(), MagnusError> {
+        let method = match level {
+            LogLevel::Debug => "debug",
+            LogLevel::Info => "info",
+            LogLevel::Warn => "warn",
+            LogLevel::Error => "error",
+            LogLevel::Fatal => "fatal",
+        };
+
+        match self.logger {
+            Some(logger) => {
+                logger.funcall::<_, _, Value>(method, (message,))?;
+            }
+            None => eprintln!("{}", message),
+        }
+
+        Ok(())
+    }
+
+    /// Log a debug message
+    pub fn debug<F, S>(&self, message_fn: F) -> Result<(), MagnusError>
+    where
+        F: FnOnce() -> S,
+        S: AsRef<str>,
+    {
+        if self.level <= LogLevel::Debug {
+            let message = message_fn();
+            self.log(LogLevel::Debug, message.as_ref())
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Log an info message
+    pub fn info<F, S>(&self, message_fn: F) -> Result<(), MagnusError>
+    where
+        F: FnOnce() -> S,
+        S: AsRef<str>,
+    {
+        if self.level <= LogLevel::Info {
+            let message = message_fn();
+            self.log(LogLevel::Info, message.as_ref())
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Log a warning message
+    pub fn warn<F, S>(&self, message_fn: F) -> Result<(), MagnusError>
+    where
+        F: FnOnce() -> S,
+        S: AsRef<str>,
+    {
+        if self.level <= LogLevel::Warn {
+            let message = message_fn();
+            self.log(LogLevel::Warn, message.as_ref())
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Log an error message
+    pub fn error<F, S>(&self, message_fn: F) -> Result<(), MagnusError>
+    where
+        F: FnOnce() -> S,
+        S: AsRef<str>,
+    {
+        if self.level <= LogLevel::Error {
+            let message = message_fn();
+            self.log(LogLevel::Error, message.as_ref())
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Log a fatal message
+    pub fn fatal<F, S>(&self, message_fn: F) -> Result<(), MagnusError>
+    where
+        F: FnOnce() -> S,
+        S: AsRef<str>,
+    {
+        if self.level <= LogLevel::Fatal {
+            let message = message_fn();
+            self.log(LogLevel::Fatal, message.as_ref())
+        } else {
+            Ok(())
+        }
+    }
+}

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

@@ -0,0 +1,113 @@
+use ahash::RandomState;
+use arrow_schema::Schema;
+use either::Either;
+use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
+use parquet::arrow::ProjectionMask;
+use std::collections::HashMap;
+use std::fs::File;
+use std::sync::Arc;
+
+use magnus::value::ReprValue;
+use magnus::{Error as MagnusError, Value};
+
+use crate::header_cache::StringCache;
+use crate::ruby_reader::{RubyReader, ThreadSafeRubyReader};
+use crate::types::TryIntoValue;
+use crate::ColumnRecord;
+
+use super::ReaderError;
+
+/// Opens a parquet file or IO-like object for reading
+///
+/// This function handles both file paths (as strings) and IO-like objects,
+/// returning either a File or a ThreadSafeRubyReader that can be used with
+/// parquet readers.
+pub fn open_parquet_source(
+    to_read: Value,
+) -> Result<Either<File, ThreadSafeRubyReader>, ReaderError> {
+    let ruby = unsafe { magnus::Ruby::get_unchecked() };
+
+    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(ReaderError::from)?;
+        Ok(Either::Left(file))
+    } else {
+        let readable = ThreadSafeRubyReader::new(RubyReader::try_from(to_read)?);
+        Ok(Either::Right(readable))
+    }
+}
+
+/// Helper function to check if a block is given and create an appropriate enumerator
+/// if not
+pub fn handle_block_or_enum<F, T>(
+    _ruby: &magnus::Ruby,
+    block_given: bool,
+    create_enum: F,
+) -> Result<Option<T>, MagnusError>
+where
+    F: FnOnce() -> Result<T, MagnusError>,
+{
+    if !block_given {
+        let enum_value = create_enum()?;
+        return Ok(Some(enum_value));
+    }
+    Ok(None)
+}
+
+/// Creates a ParquetRecordBatchReader with the given columns and batch size configurations
+pub fn create_batch_reader<T: parquet::file::reader::ChunkReader + 'static>(
+    reader: T,
+    columns: &Option<Vec<String>>,
+    batch_size: Option<usize>,
+) -> Result<(ParquetRecordBatchReader, std::sync::Arc<Schema>, i64), ReaderError> {
+    let mut builder =
+        ParquetRecordBatchReaderBuilder::try_new(reader).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))?;
+    Ok((reader, schema, num_rows))
+}
+
+/// Handles the case of an empty parquet file (no rows) by yielding a record with empty arrays
+/// Returns true if the file was empty and was handled, false otherwise
+pub fn handle_empty_file(
+    ruby: &magnus::Ruby,
+    schema: &Arc<Schema>,
+    num_rows: i64,
+) -> Result<bool, ReaderError> {
+    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.try_into_value_with(&ruby)?)?;
+        return Ok(true);
+    }
+    Ok(false)
+}

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

@@ -1,9 +1,10 @@
+mod common;
 mod parquet_column_reader;
 mod parquet_row_reader;
 
 use std::io;
 
-use magnus::{Error as MagnusError, Ruby};
+use magnus::Error as MagnusError;
 use thiserror::Error;
 
 use crate::header_cache::CacheError;
@@ -12,16 +13,12 @@ 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),
+    Ruby(#[from] MagnusErrorWrapper),
     #[error("Parquet error: {0}")]
     Parquet(#[from] parquet::errors::ParquetError),
     #[error("Arrow error: {0}")]
@@ -32,17 +29,34 @@ pub enum ReaderError {
     Jiff(#[from] jiff::Error),
 }
 
+#[derive(Debug)]
+pub struct MagnusErrorWrapper(pub MagnusError);
+
+impl From<MagnusError> for MagnusErrorWrapper {
+    fn from(err: MagnusError) -> Self {
+        Self(err)
+    }
+}
+
+impl std::fmt::Display for MagnusErrorWrapper {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl std::error::Error for MagnusErrorWrapper {}
+
 impl From<MagnusError> for ReaderError {
     fn from(err: MagnusError) -> Self {
-        Self::Ruby(err.to_string())
+        Self::Ruby(MagnusErrorWrapper(err))
     }
 }
 
-impl From<ReaderError> for MagnusError {
-    fn from(err: ReaderError) -> Self {
-        MagnusError::new(
-            Ruby::get().unwrap().exception_runtime_error(),
-            err.to_string(),
-        )
+impl Into<MagnusError> for ReaderError {
+    fn into(self) -> MagnusError {
+        match self {
+            Self::Ruby(MagnusErrorWrapper(err)) => err.into(),
+            _ => MagnusError::new(magnus::exception::runtime_error(), self.to_string()),
+        }
     }
 }

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

@@ -1,24 +1,32 @@
 use crate::header_cache::StringCache;
-use crate::ruby_reader::{RubyReader, ThreadSafeRubyReader};
+use crate::logger::RubyLogger;
 use crate::types::{ArrayWrapper, TryIntoValue};
 use crate::{
     create_column_enumerator, utils::*, ColumnEnumeratorArgs, ColumnRecord, ParquetValueVec,
     ParserResultType,
 };
 use ahash::RandomState;
-use magnus::value::ReprValue;
+use either::Either;
 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::sync::OnceLock;
 
+use super::common::{
+    create_batch_reader, handle_block_or_enum, handle_empty_file, open_parquet_source,
+};
 use super::ReaderError;
 
 #[inline]
 pub fn parse_parquet_columns<'a>(rb_self: Value, args: &[Value]) -> Result<Value, MagnusError> {
+    Ok(parse_parquet_columns_impl(rb_self, args).map_err(|e| {
+        let z: MagnusError = e.into();
+        z
+    })?)
+}
+
+#[inline]
+fn parse_parquet_columns_impl<'a>(rb_self: Value, args: &[Value]) -> Result<Value, ReaderError> {
     let ruby = unsafe { Ruby::get_unchecked() };
 
     let ParquetColumnsArgs {
@@ -27,93 +35,45 @@ pub fn parse_parquet_columns<'a>(rb_self: Value, args: &[Value]) -> Result<Value
         columns,
         batch_size,
         strict,
+        logger,
     } = parse_parquet_columns_args(&ruby, args)?;
 
-    if !ruby.block_given() {
-        return create_column_enumerator(ColumnEnumeratorArgs {
+    // Initialize the logger if provided
+    let ruby_logger = RubyLogger::new(&ruby, logger)?;
+    if let Some(ref bs) = batch_size {
+        ruby_logger.debug(|| format!("Using batch size: {}", bs))?;
+    }
+
+    // Clone values for the closure to avoid move issues
+    let columns_clone = columns.clone();
+
+    // Handle block or create enumerator
+    if let Some(enum_value) = handle_block_or_enum(&ruby, ruby.block_given(), || {
+        create_column_enumerator(ColumnEnumeratorArgs {
             rb_self,
             to_read,
             result_type,
-            columns,
+            columns: columns_clone,
             batch_size,
             strict,
+            logger: logger.as_ref().map(|_| to_read),
         })
-        .map(|yield_enum| yield_enum.into_value_with(&ruby));
+        .map(|yield_enum| yield_enum.into_value_with(&ruby))
+    })? {
+        return Ok(enum_value);
     }
 
-    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 {
-        let readable = ThreadSafeRubyReader::new(RubyReader::try_from(to_read)?);
-
-        let mut builder =
-            ParquetRecordBatchReaderBuilder::try_new(readable).map_err(ReaderError::from)?;
-        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();
+    let source = open_parquet_source(to_read)?;
 
-            // Create a projection mask from column names
-            let projection =
-                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
+    // Use the common function to create the batch reader
 
-            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)
+    let (batch_reader, schema, num_rows) = match source {
+        Either::Left(file) => create_batch_reader(file, &columns, batch_size)?,
+        Either::Right(readable) => create_batch_reader(readable, &columns, batch_size)?,
     };
 
-    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.try_into_value_with(&ruby)?)?;
+    // Handle empty file case
+    if handle_empty_file(&ruby, &schema, num_rows)? {
         return Ok(ruby.qnil().into_value_with(&ruby));
     }
 

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

@@ -1,23 +1,32 @@
 use crate::header_cache::StringCache;
-use crate::ruby_reader::{RubyReader, ThreadSafeRubyReader};
+use crate::logger::RubyLogger;
 use crate::types::TryIntoValue;
 use crate::{
     create_row_enumerator, utils::*, ParquetField, ParserResultType, ReaderError,
     RowEnumeratorArgs, RowRecord,
 };
 use ahash::RandomState;
-use magnus::value::ReprValue;
+use either::Either;
 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::sync::OnceLock;
 
+use super::common::{handle_block_or_enum, open_parquet_source};
+
 #[inline]
 pub fn parse_parquet_rows<'a>(rb_self: Value, args: &[Value]) -> Result<Value, MagnusError> {
+    Ok(parse_parquet_rows_impl(rb_self, args).map_err(|e| {
+        let z: MagnusError = e.into();
+        z
+    })?)
+}
+
+#[inline]
+fn parse_parquet_rows_impl<'a>(rb_self: Value, args: &[Value]) -> Result<Value, ReaderError> {
     let ruby = unsafe { Ruby::get_unchecked() };
 
     let ParquetRowsArgs {
@@ -25,31 +34,44 @@ pub fn parse_parquet_rows<'a>(rb_self: Value, args: &[Value]) -> Result<Value, M
         result_type,
         columns,
         strict,
+        logger,
     } = parse_parquet_rows_args(&ruby, args)?;
 
-    if !ruby.block_given() {
-        return create_row_enumerator(RowEnumeratorArgs {
+    // Initialize the logger if provided
+    let ruby_logger = RubyLogger::new(&ruby, logger)?;
+
+    // Clone values for the closure to avoid move issues
+    let columns_clone = columns.clone();
+
+    // Handle block or create enumerator
+    if let Some(enum_value) = handle_block_or_enum(&ruby, ruby.block_given(), || {
+        create_row_enumerator(RowEnumeratorArgs {
             rb_self,
             to_read,
             result_type,
-            columns,
+            columns: columns_clone,
             strict,
+            logger,
         })
-        .map(|yield_enum| yield_enum.into_value_with(&ruby));
+        .map(|yield_enum| yield_enum.into_value_with(&ruby))
+    })? {
+        return Ok(enum_value);
     }
 
-    let reader: Box<dyn FileReader> = 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(ReaderError::from)?;
-        Box::new(SerializedFileReader::new(file).map_err(ReaderError::from)?)
-    } else {
-        let readable = ThreadSafeRubyReader::new(RubyReader::try_from(to_read)?);
-        Box::new(SerializedFileReader::new(readable).map_err(ReaderError::from)?)
+    let source = open_parquet_source(to_read)?;
+    let reader: Box<dyn FileReader> = match source {
+        Either::Left(file) => Box::new(SerializedFileReader::new(file).map_err(ReaderError::from)?),
+        Either::Right(readable) => {
+            Box::new(SerializedFileReader::new(readable).map_err(ReaderError::from)?)
+        }
     };
+
     let schema = reader.metadata().file_metadata().schema().clone();
+    ruby_logger.debug(|| format!("Schema loaded: {:?}", schema))?;
+
     let mut iter = ParquetRowIter::from_file_into(reader);
     if let Some(cols) = columns {
+        ruby_logger.debug(|| format!("Projecting columns: {:?}", cols))?;
         let projection = create_projection_schema(&schema, &cols);
         iter = iter.project(Some(projection.to_owned())).map_err(|e| {
             MagnusError::new(
@@ -81,9 +103,9 @@ pub fn parse_parquet_rows<'a>(rb_self: Value, args: &[Value]) -> Result<Value, M
 
                     let mut map =
                         HashMap::with_capacity_and_hasher(headers.len(), RandomState::default());
-                    row.get_column_iter().enumerate().for_each(|(i, (_, v))| {
+                    for (i, (_, v)) in row.get_column_iter().enumerate() {
                         map.insert(headers[i], ParquetField(v.clone(), strict));
-                    });
+                    }
                     Ok(map)
                 })
                 .and_then(|row| Ok(RowRecord::Map::<RandomState>(row)))
@@ -100,8 +122,9 @@ pub fn parse_parquet_rows<'a>(rb_self: Value, args: &[Value]) -> Result<Value, M
                 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(), strict)));
+                    for (_, v) in row.get_column_iter() {
+                        vec.push(ParquetField(v.clone(), strict));
+                    }
                     Ok(vec)
                 })
                 .and_then(|row| Ok(RowRecord::Vec::<RandomState>(row)))