parquet 0.4.2 → 0.5.1

data/ext/parquet/src/types/writer_types.rs

@@ -1,32 +1,36 @@
+use super::{core_types::SchemaNode, ParquetGemError, PrimitiveType};
+use crate::{
+    types::{ListField, MapField, ParquetSchemaType},
+    utils::parse_string_or_symbol,
+};
+use arrow_array::{Array, RecordBatch};
+use magnus::{value::ReprValue, Error as MagnusError, RString, Ruby, TryConvert, Value};
+use parquet::{arrow::ArrowWriter, errors::ParquetError};
 use std::{
     io::{self, Write},
     str::FromStr,
     sync::Arc,
 };
-
-use arrow_array::{Array, RecordBatch};
-use magnus::{value::ReprValue, Error as MagnusError, RString, Ruby, Symbol, TryConvert, Value};
-use parquet::{arrow::ArrowWriter, errors::ParquetError};
 use tempfile::NamedTempFile;
 
-use crate::types::{convert_to_string, ListField, MapField, ParquetSchemaType};
-
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct SchemaField<'a> {
     pub name: String,
     pub type_: ParquetSchemaType<'a>,
     pub format: Option<String>,
+    pub nullable: bool,
 }
 
 #[derive(Debug)]
-pub struct ParquetWriteArgs<'a> {
+pub struct ParquetWriteArgs {
     pub read_from: Value,
     pub write_to: Value,
-    pub schema: Vec<SchemaField<'a>>,
+    pub schema: SchemaNode,
     pub batch_size: Option<usize>,
     pub flush_threshold: Option<usize>,
     pub compression: Option<String>,
     pub sample_size: Option<usize>,
+    pub logger: Option<Value>,
 }
 
 pub trait SendableWrite: Send + Write {}
@@ -59,31 +63,63 @@ impl<'a> FromStr for ParquetSchemaType<'a> {
     type Err = MagnusError;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
+        // Check if it's a list type
+        if let Some(inner_type_str) = s.strip_prefix("list<").and_then(|s| s.strip_suffix(">")) {
+            let inner_type = inner_type_str.parse::<ParquetSchemaType>()?;
+            return Ok(ParquetSchemaType::List(Box::new(ListField {
+                item_type: inner_type,
+                format: None,
+                nullable: true,
+            })));
+        }
+
+        // Check if it's a map type
+        if let Some(kv_types_str) = s.strip_prefix("map<").and_then(|s| s.strip_suffix(">")) {
+            let parts: Vec<&str> = kv_types_str.splitn(2, ',').collect();
+            if parts.len() != 2 {
+                return Err(MagnusError::new(
+                    magnus::exception::runtime_error(),
+                    format!(
+                        "Invalid map format. Expected 'map<keyType,valueType>', got '{}'",
+                        s
+                    ),
+                ));
+            }
+
+            let key_type = parts[0].trim().parse::<ParquetSchemaType>()?;
+            let value_type = parts[1].trim().parse::<ParquetSchemaType>()?;
+
+            return Ok(ParquetSchemaType::Map(Box::new(MapField {
+                key_type,
+                value_type,
+                key_format: None,
+                value_format: None,
+                value_nullable: true,
+            })));
+        }
+
+        // Handle primitive types
         match s {
-            "int8" => Ok(ParquetSchemaType::Int8),
-            "int16" => Ok(ParquetSchemaType::Int16),
-            "int32" => Ok(ParquetSchemaType::Int32),
-            "int64" => Ok(ParquetSchemaType::Int64),
-            "uint8" => Ok(ParquetSchemaType::UInt8),
-            "uint16" => Ok(ParquetSchemaType::UInt16),
-            "uint32" => Ok(ParquetSchemaType::UInt32),
-            "uint64" => Ok(ParquetSchemaType::UInt64),
-            "float" | "float32" => Ok(ParquetSchemaType::Float),
-            "double" | "float64" => Ok(ParquetSchemaType::Double),
-            "string" | "utf8" => Ok(ParquetSchemaType::String),
-            "binary" => Ok(ParquetSchemaType::Binary),
-            "boolean" | "bool" => Ok(ParquetSchemaType::Boolean),
-            "date32" => Ok(ParquetSchemaType::Date32),
-            "timestamp_millis" => Ok(ParquetSchemaType::TimestampMillis),
-            "timestamp_micros" => Ok(ParquetSchemaType::TimestampMicros),
+            "int8" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Int8)),
+            "int16" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Int16)),
+            "int32" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Int32)),
+            "int64" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Int64)),
+            "uint8" => Ok(ParquetSchemaType::Primitive(PrimitiveType::UInt8)),
+            "uint16" => Ok(ParquetSchemaType::Primitive(PrimitiveType::UInt16)),
+            "uint32" => Ok(ParquetSchemaType::Primitive(PrimitiveType::UInt32)),
+            "uint64" => Ok(ParquetSchemaType::Primitive(PrimitiveType::UInt64)),
+            "float" | "float32" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Float32)),
+            "double" | "float64" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Float64)),
+            "string" | "utf8" => Ok(ParquetSchemaType::Primitive(PrimitiveType::String)),
+            "binary" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Binary)),
+            "boolean" | "bool" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Boolean)),
+            "date32" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Date32)),
+            "timestamp_millis" => Ok(ParquetSchemaType::Primitive(PrimitiveType::TimestampMillis)),
+            "timestamp_micros" => Ok(ParquetSchemaType::Primitive(PrimitiveType::TimestampMicros)),
             "list" => Ok(ParquetSchemaType::List(Box::new(ListField {
-                item_type: ParquetSchemaType::Int8,
-                format: None,
-            }))),
-            "map" => Ok(ParquetSchemaType::Map(Box::new(MapField {
-                key_type: ParquetSchemaType::String,
-                value_type: ParquetSchemaType::Int8,
+                item_type: ParquetSchemaType::Primitive(PrimitiveType::String),
                 format: None,
+                nullable: true,
             }))),
             _ => Err(MagnusError::new(
                 magnus::exception::runtime_error(),
@@ -98,7 +134,11 @@ impl<'a> TryConvert for ParquetSchemaType<'a> {
         let ruby = unsafe { Ruby::get_unchecked() };
         let schema_type = parse_string_or_symbol(&ruby, value)?;
 
-        schema_type.unwrap().parse()
+        schema_type
+            .ok_or_else(|| {
+                MagnusError::new(magnus::exception::type_error(), "Invalid schema type")
+            })?
+            .parse()
     }
 }
 
@@ -106,31 +146,6 @@ impl<'a> TryConvert for ParquetSchemaType<'a> {
 // with simple primitive types and strings
 unsafe impl<'a> Send for ParquetSchemaType<'a> {}
 
-fn parse_string_or_symbol(ruby: &Ruby, value: Value) -> Result<Option<String>, MagnusError> {
-    if value.is_nil() {
-        Ok(None)
-    } else if value.is_kind_of(ruby.class_string()) {
-        RString::from_value(value)
-            .ok_or_else(|| {
-                MagnusError::new(magnus::exception::type_error(), "Invalid string value")
-            })?
-            .to_string()
-            .map(|s| Some(s))
-    } else if value.is_kind_of(ruby.class_symbol()) {
-        Symbol::from_value(value)
-            .ok_or_else(|| {
-                MagnusError::new(magnus::exception::type_error(), "Invalid symbol value")
-            })?
-            .funcall("to_s", ())
-            .map(|s| Some(s))
-    } else {
-        Err(MagnusError::new(
-            magnus::exception::type_error(),
-            "Value must be a String or Symbol",
-        ))
-    }
-}
-
 pub enum WriterOutput {
     File(ArrowWriter<Box<dyn SendableWrite>>),
     TempFile(ArrowWriter<Box<dyn SendableWrite>>, NamedTempFile),
@@ -157,126 +172,54 @@ impl WriterOutput {
     }
 }
 
-pub struct ParquetErrorWrapper(pub ParquetError);
-
-impl From<ParquetErrorWrapper> for MagnusError {
-    fn from(err: ParquetErrorWrapper) -> Self {
-        MagnusError::new(
-            magnus::exception::runtime_error(),
-            format!("Parquet error: {}", err.0),
-        )
-    }
-}
-
 pub struct ColumnCollector<'a> {
+    pub ruby: &'a Ruby,
     pub name: String,
     pub type_: ParquetSchemaType<'a>,
     pub format: Option<String>,
+    pub nullable: bool,
     pub values: Vec<crate::types::ParquetValue>,
 }
 
 impl<'a> ColumnCollector<'a> {
-    pub fn new(name: String, type_: ParquetSchemaType<'a>, format: Option<String>) -> Self {
+    pub fn new(
+        ruby: &'a Ruby,
+        name: String,
+        type_: ParquetSchemaType<'a>,
+        format: Option<String>,
+        nullable: bool,
+    ) -> Self {
         Self {
+            ruby,
             name,
             type_,
             format,
+            nullable,
             values: Vec::new(),
         }
     }
 
     pub fn push_value(&mut self, value: Value) -> Result<(), MagnusError> {
         use crate::types::ParquetValue;
-        use crate::{
-            convert_to_binary, convert_to_boolean, convert_to_date32, convert_to_list,
-            convert_to_map, convert_to_timestamp_micros, convert_to_timestamp_millis,
-            NumericConverter,
-        };
 
         if value.is_nil() {
-            self.values.push(ParquetValue::Null);
-            return Ok(());
+            if !self.nullable {
+                // For non-nullable fields, raise an error
+                return Err(MagnusError::new(
+                    magnus::exception::runtime_error(),
+                    "Cannot write nil value for non-nullable field",
+                ));
+            }
         }
 
-        let parquet_value = match &self.type_ {
-            ParquetSchemaType::Int8 => {
-                let v = NumericConverter::<i8>::convert_with_string_fallback(value)?;
-                ParquetValue::Int8(v)
-            }
-            ParquetSchemaType::Int16 => {
-                let v = NumericConverter::<i16>::convert_with_string_fallback(value)?;
-                ParquetValue::Int16(v)
-            }
-            ParquetSchemaType::Int32 => {
-                let v = NumericConverter::<i32>::convert_with_string_fallback(value)?;
-                ParquetValue::Int32(v)
-            }
-            ParquetSchemaType::Int64 => {
-                let v = NumericConverter::<i64>::convert_with_string_fallback(value)?;
-                ParquetValue::Int64(v)
-            }
-            ParquetSchemaType::UInt8 => {
-                let v = NumericConverter::<u8>::convert_with_string_fallback(value)?;
-                ParquetValue::UInt8(v)
-            }
-            ParquetSchemaType::UInt16 => {
-                let v = NumericConverter::<u16>::convert_with_string_fallback(value)?;
-                ParquetValue::UInt16(v)
-            }
-            ParquetSchemaType::UInt32 => {
-                let v = NumericConverter::<u32>::convert_with_string_fallback(value)?;
-                ParquetValue::UInt32(v)
-            }
-            ParquetSchemaType::UInt64 => {
-                let v = NumericConverter::<u64>::convert_with_string_fallback(value)?;
-                ParquetValue::UInt64(v)
-            }
-            ParquetSchemaType::Float => {
-                let v = NumericConverter::<f32>::convert_with_string_fallback(value)?;
-                ParquetValue::Float32(v)
-            }
-            ParquetSchemaType::Double => {
-                let v = NumericConverter::<f64>::convert_with_string_fallback(value)?;
-                ParquetValue::Float64(v)
-            }
-            ParquetSchemaType::String => {
-                let v = convert_to_string(value)?;
-                ParquetValue::String(v)
-            }
-            ParquetSchemaType::Binary => {
-                let v = convert_to_binary(value)?;
-                ParquetValue::Bytes(v)
-            }
-            ParquetSchemaType::Boolean => {
-                let v = convert_to_boolean(value)?;
-                ParquetValue::Boolean(v)
-            }
-            ParquetSchemaType::Date32 => {
-                let v = convert_to_date32(value, self.format.as_deref())?;
-                ParquetValue::Date32(v)
-            }
-            ParquetSchemaType::TimestampMillis => {
-                let v = convert_to_timestamp_millis(value, self.format.as_deref())?;
-                ParquetValue::TimestampMillis(v, None)
-            }
-            ParquetSchemaType::TimestampMicros => {
-                let v = convert_to_timestamp_micros(value, self.format.as_deref())?;
-                ParquetValue::TimestampMicros(v, None)
-            }
-            ParquetSchemaType::List(list_field) => {
-                let values = convert_to_list(value, list_field)?;
-                ParquetValue::List(values)
-            }
-            ParquetSchemaType::Map(map_field) => {
-                let map = convert_to_map(value, map_field)?;
-                ParquetValue::Map(map)
-            }
-        };
+        // For all other types, proceed as normal
+        let parquet_value =
+            ParquetValue::from_value(self.ruby, value, &self.type_, self.format.as_deref())?;
         self.values.push(parquet_value);
         Ok(())
     }
 
-    pub fn take_array(&mut self) -> Result<Arc<dyn Array>, MagnusError> {
+    pub fn take_array(&mut self) -> Result<Arc<dyn Array>, ParquetGemError> {
         let values = std::mem::take(&mut self.values);
         crate::convert_parquet_values_to_arrow(values, &self.type_)
     }

data/ext/parquet/src/utils.rs

@@ -1,12 +1,12 @@
+use crate::ParserResultType;
 use magnus::{
     scan_args::{get_kwargs, scan_args},
     value::ReprValue,
     Error, RString, Ruby, Symbol, Value,
 };
 
-use crate::ParserResultType;
-
-fn parse_string_or_symbol(ruby: &Ruby, value: Value) -> Result<Option<String>, Error> {
+/// Convert a Ruby Value to a String, handling both String and Symbol types
+pub fn parse_string_or_symbol(ruby: &Ruby, value: Value) -> Result<Option<String>, Error> {
     if value.is_nil() {
         Ok(None)
     } else if value.is_kind_of(ruby.class_string()) {
@@ -33,9 +33,10 @@ pub struct ParquetRowsArgs {
     pub result_type: ParserResultType,
     pub columns: Option<Vec<String>>,
     pub strict: bool,
+    pub logger: Option<Value>,
 }
 
-/// Parse common arguments for CSV parsing
+/// Parse common arguments for parquet row iteration
 pub fn parse_parquet_rows_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetRowsArgs, Error> {
     let parsed_args = scan_args::<(Value,), (), (), (), _, ()>(args)?;
     let (to_read,) = parsed_args.required;
@@ -47,12 +48,13 @@ pub fn parse_parquet_rows_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetRow
             Option<Option<Value>>,
             Option<Option<Vec<String>>>,
             Option<Option<bool>>,
+            Option<Option<Value>>,
         ),
         (),
     >(
         parsed_args.keywords,
         &[],
-        &["result_type", "columns", "strict"],
+        &["result_type", "columns", "strict", "logger"],
     )?;
 
     let result_type: ParserResultType = match kwargs
@@ -84,12 +86,14 @@ pub fn parse_parquet_rows_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetRow
     };
 
     let strict = kwargs.optional.2.flatten().unwrap_or(true);
+    let logger = kwargs.optional.3.flatten();
 
     Ok(ParquetRowsArgs {
         to_read,
         result_type,
         columns: kwargs.optional.1.flatten(),
         strict,
+        logger,
     })
 }
 
@@ -100,9 +104,10 @@ pub struct ParquetColumnsArgs {
     pub columns: Option<Vec<String>>,
     pub batch_size: Option<usize>,
     pub strict: bool,
+    pub logger: Option<Value>,
 }
 
-/// Parse common arguments for CSV parsing
+/// Parse common arguments for parquet column iteration
 pub fn parse_parquet_columns_args(
     ruby: &Ruby,
     args: &[Value],
@@ -118,12 +123,13 @@ pub fn parse_parquet_columns_args(
             Option<Option<Vec<String>>>,
             Option<Option<usize>>,
             Option<Option<bool>>,
+            Option<Option<Value>>,
         ),
         (),
     >(
         parsed_args.keywords,
         &[],
-        &["result_type", "columns", "batch_size", "strict"],
+        &["result_type", "columns", "batch_size", "strict", "logger"],
     )?;
 
     let result_type: ParserResultType = match kwargs
@@ -154,11 +160,25 @@ pub fn parse_parquet_columns_args(
         None => ParserResultType::Hash,
     };
 
+    let batch_size = kwargs.optional.2.flatten();
+    if let Some(sz) = batch_size {
+        if sz <= 0 {
+            return Err(Error::new(
+                ruby.exception_arg_error(),
+                format!("batch_size must be > 0, got {}", sz),
+            ));
+        }
+    }
+
+    let strict = kwargs.optional.3.flatten().unwrap_or(true);
+    let logger = kwargs.optional.4.flatten();
+
     Ok(ParquetColumnsArgs {
         to_read,
         result_type,
         columns: kwargs.optional.1.flatten(),
-        batch_size: kwargs.optional.2.flatten(),
-        strict: kwargs.optional.3.flatten().unwrap_or(true),
+        batch_size,
+        strict,
+        logger,
     })
 }