parquet 0.5.13 → 0.6.0

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

@@ -1,329 +0,0 @@
-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 tempfile::NamedTempFile;
-
-#[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 {
-    pub read_from: Value,
-    pub write_to: Value,
-    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 {}
-impl<T: Send + Write> SendableWrite for T {}
-
-pub struct IoLikeValue(pub(crate) Value);
-
-impl Write for IoLikeValue {
-    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
-        let ruby_bytes = RString::from_slice(buf);
-
-        let bytes_written = self
-            .0
-            .funcall::<_, _, usize>("write", (ruby_bytes,))
-            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
-
-        Ok(bytes_written)
-    }
-
-    fn flush(&mut self) -> Result<(), io::Error> {
-        self.0
-            .funcall::<_, _, Value>("flush", ())
-            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
-
-        Ok(())
-    }
-}
-
-impl FromStr for ParquetSchemaType<'_> {
-    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,
-            })));
-        }
-
-        // Check if it's a decimal type with precision and scale
-        if let Some(decimal_params) = s.strip_prefix("decimal(").and_then(|s| s.strip_suffix(")")) {
-            let parts: Vec<&str> = decimal_params.split(',').collect();
-
-            // Handle both single parameter (precision only) and two parameters (precision and scale)
-            if parts.len() == 1 {
-                // Only precision provided, scale defaults to 0
-                let precision = parts[0].trim().parse::<u8>().map_err(|_| {
-                    MagnusError::new(
-                        magnus::exception::runtime_error(),
-                        format!("Invalid precision value in decimal type: {}", parts[0]),
-                    )
-                })?;
-
-                return Ok(ParquetSchemaType::Primitive(PrimitiveType::Decimal128(
-                    precision, 0,
-                )));
-            } else if parts.len() == 2 {
-                // Both precision and scale provided
-                let precision = parts[0].trim().parse::<u8>().map_err(|_| {
-                    MagnusError::new(
-                        magnus::exception::runtime_error(),
-                        format!("Invalid precision value in decimal type: {}", parts[0]),
-                    )
-                })?;
-
-                let scale = parts[1].trim().parse::<i8>().map_err(|_| {
-                    MagnusError::new(
-                        magnus::exception::runtime_error(),
-                        format!("Invalid scale value in decimal type: {}", parts[1]),
-                    )
-                })?;
-
-                return Ok(ParquetSchemaType::Primitive(PrimitiveType::Decimal128(
-                    precision, scale,
-                )));
-            } else {
-                return Err(MagnusError::new(
-                    magnus::exception::runtime_error(),
-                    format!(
-                        "Invalid decimal format. Expected 'decimal(precision)' or 'decimal(precision,scale)', got '{}'",
-                        s
-                    ),
-                ));
-            }
-        }
-
-        // Check if it's a decimal256 type with precision and scale
-        if let Some(decimal_params) = s
-            .strip_prefix("decimal256(")
-            .and_then(|s| s.strip_suffix(")"))
-        {
-            let parts: Vec<&str> = decimal_params.split(',').collect();
-
-            // Handle both single parameter (precision only) and two parameters (precision and scale)
-            if parts.len() == 1 {
-                // Only precision provided, scale defaults to 0
-                let precision = parts[0].trim().parse::<u8>().map_err(|_| {
-                    MagnusError::new(
-                        magnus::exception::runtime_error(),
-                        format!("Invalid precision value in decimal256 type: {}", parts[0]),
-                    )
-                })?;
-
-                return Ok(ParquetSchemaType::Primitive(PrimitiveType::Decimal256(
-                    precision, 0,
-                )));
-            } else if parts.len() == 2 {
-                // Both precision and scale provided
-                let precision = parts[0].trim().parse::<u8>().map_err(|_| {
-                    MagnusError::new(
-                        magnus::exception::runtime_error(),
-                        format!("Invalid precision value in decimal256 type: {}", parts[0]),
-                    )
-                })?;
-
-                let scale = parts[1].trim().parse::<i8>().map_err(|_| {
-                    MagnusError::new(
-                        magnus::exception::runtime_error(),
-                        format!("Invalid scale value in decimal256 type: {}", parts[1]),
-                    )
-                })?;
-
-                return Ok(ParquetSchemaType::Primitive(PrimitiveType::Decimal256(
-                    precision, scale,
-                )));
-            } else {
-                return Err(MagnusError::new(
-                    magnus::exception::runtime_error(),
-                    format!(
-                        "Invalid decimal256 format. Expected 'decimal256(precision)' or 'decimal256(precision,scale)', got '{}'",
-                        s
-                    ),
-                ));
-            }
-        }
-
-        // Handle primitive types
-        match s {
-            "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)),
-            "time_millis" => Ok(ParquetSchemaType::Primitive(PrimitiveType::TimeMillis)),
-            "time_micros" => Ok(ParquetSchemaType::Primitive(PrimitiveType::TimeMicros)),
-            "decimal" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Decimal128(
-                38, 0,
-            ))),
-            "decimal256" => Ok(ParquetSchemaType::Primitive(PrimitiveType::Decimal256(
-                38, 0,
-            ))),
-            "list" => Ok(ParquetSchemaType::List(Box::new(ListField {
-                item_type: ParquetSchemaType::Primitive(PrimitiveType::String),
-                format: None,
-                nullable: true,
-            }))),
-            _ => Err(MagnusError::new(
-                magnus::exception::runtime_error(),
-                format!("Invalid schema type: {}", s),
-            )),
-        }
-    }
-}
-
-impl TryConvert for ParquetSchemaType<'_> {
-    fn try_convert(value: Value) -> Result<Self, MagnusError> {
-        let ruby = unsafe { Ruby::get_unchecked() };
-        let schema_type = parse_string_or_symbol(&ruby, value)?;
-
-        schema_type
-            .ok_or_else(|| {
-                MagnusError::new(magnus::exception::type_error(), "Invalid schema type")
-            })?
-            .parse()
-    }
-}
-
-// We know this type is safe to move between threads because it's just an enum
-// with simple primitive types and strings
-unsafe impl Send for ParquetSchemaType<'_> {}
-
-pub enum WriterOutput {
-    File(ArrowWriter<Box<dyn SendableWrite>>),
-    TempFile(ArrowWriter<Box<dyn SendableWrite>>, NamedTempFile),
-}
-
-impl WriterOutput {
-    pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ParquetError> {
-        match self {
-            WriterOutput::File(writer) | WriterOutput::TempFile(writer, _) => writer.write(batch),
-        }
-    }
-
-    pub fn close(self) -> Result<Option<NamedTempFile>, ParquetError> {
-        match self {
-            WriterOutput::File(writer) => {
-                writer.close()?;
-                Ok(None)
-            }
-            WriterOutput::TempFile(writer, temp_file) => {
-                writer.close()?;
-                Ok(Some(temp_file))
-            }
-        }
-    }
-}
-
-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(
-        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;
-
-        if value.is_nil() && !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",
-            ));
-        }
-
-        // 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>, 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,184 +0,0 @@
-use crate::ParserResultType;
-use magnus::{
-    scan_args::{get_kwargs, scan_args},
-    value::ReprValue,
-    Error, RString, Ruby, Symbol, Value,
-};
-
-/// 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()) {
-        RString::from_value(value)
-            .ok_or_else(|| Error::new(magnus::exception::type_error(), "Invalid string value"))?
-            .to_string()
-            .map(Some)
-    } else if value.is_kind_of(ruby.class_symbol()) {
-        Symbol::from_value(value)
-            .ok_or_else(|| Error::new(magnus::exception::type_error(), "Invalid symbol value"))?
-            .funcall("to_s", ())
-            .map(Some)
-    } else {
-        Err(Error::new(
-            magnus::exception::type_error(),
-            "Value must be a String or Symbol",
-        ))
-    }
-}
-
-#[derive(Debug)]
-pub struct ParquetRowsArgs {
-    pub to_read: Value,
-    pub result_type: ParserResultType,
-    pub columns: Option<Vec<String>>,
-    pub strict: bool,
-    pub logger: Option<Value>,
-}
-
-/// 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;
-
-    let kwargs = get_kwargs::<
-        _,
-        (),
-        (
-            Option<Option<Value>>,
-            Option<Option<Vec<String>>>,
-            Option<Option<bool>>,
-            Option<Option<Value>>,
-        ),
-        (),
-    >(
-        parsed_args.keywords,
-        &[],
-        &["result_type", "columns", "strict", "logger"],
-    )?;
-
-    let result_type: ParserResultType = match kwargs
-        .optional
-        .0
-        .flatten()
-        .map(|value| parse_string_or_symbol(ruby, value))
-    {
-        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 => ParserResultType::Hash,
-    };
-
-    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,
-    })
-}
-
-#[derive(Debug)]
-pub struct ParquetColumnsArgs {
-    pub to_read: Value,
-    pub result_type: ParserResultType,
-    pub columns: Option<Vec<String>>,
-    pub batch_size: Option<usize>,
-    pub strict: bool,
-    pub logger: Option<Value>,
-}
-
-/// Parse common arguments for parquet column iteration
-pub fn parse_parquet_columns_args(
-    ruby: &Ruby,
-    args: &[Value],
-) -> Result<ParquetColumnsArgs, Error> {
-    let parsed_args = scan_args::<(Value,), (), (), (), _, ()>(args)?;
-    let (to_read,) = parsed_args.required;
-
-    let kwargs = get_kwargs::<
-        _,
-        (),
-        (
-            Option<Option<Value>>,
-            Option<Option<Vec<String>>>,
-            Option<Option<usize>>,
-            Option<Option<bool>>,
-            Option<Option<Value>>,
-        ),
-        (),
-    >(
-        parsed_args.keywords,
-        &[],
-        &["result_type", "columns", "batch_size", "strict", "logger"],
-    )?;
-
-    let result_type: ParserResultType = match kwargs
-        .optional
-        .0
-        .flatten()
-        .map(|value| parse_string_or_symbol(ruby, value))
-    {
-        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 => ParserResultType::Hash,
-    };
-
-    let batch_size = kwargs.optional.2.flatten();
-    if let Some(batch_size) = batch_size {
-        if batch_size == 0 {
-            return Err(Error::new(
-                magnus::exception::arg_error(),
-                "Batch size must be greater than 0",
-            ));
-        }
-    }
-
-    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,
-        strict,
-        logger,
-    })
-}