parquet 0.4.0 → 0.4.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bb28e49c647fcf9dddba5d18806b366e73ecaa6fb3bf8a3534eb7df4749710f0
-  data.tar.gz: 44953fd9a1f3fc89f24754a6cb555f4d32a90d4af4bb4aa3c14fc373a1c389df
+  metadata.gz: 8e05c4544ed72011e1a877befa04e9c29f9d61d6f5663095f3356735c2411299
+  data.tar.gz: 6642f5575e9802a3eb05af6a4523fcc198644a6b465a2fa911c35985f27943f2
 SHA512:
-  metadata.gz: eaa9767d45deab2c3ba3fb0fc4c7aeb222b4446494262ff01083bdc488cca3164feee0f2f07bff644a3196479da0db4049dee1395b8b89f87f5aeba24a97b60b
-  data.tar.gz: 805a9cad6698f6ade3a61b40eb514a948ed6b07682db0bd3e724aad533468a114053b9589ec6261c31e8fe6b244c4808457f9e801a6991dffa86a36326736ab0
+  metadata.gz: b68013d8de3a6454d4ebddaae549fb4c02d0a6102718259b9a43133ceed8803060fc04fd06287e041ee0afdf0cd6c3556ff2c6a41d61a6acd226ef239caab090
+  data.tar.gz: 074f9faf08f2a178322cf74667ca93c24561cb49b4b91f7f98b653ea61a561628c3b06c4afe9d0a6e13410479fbd87482b88ee09d0dfbc91c554b820350fdad3

data/ext/parquet/src/ruby_reader.rs

@@ -184,37 +184,53 @@ impl Length for RubyReader {
             }
             RubyReader::RubyIoLike { inner } => {
                 let unwrapped_inner = ruby.get_inner(*inner);
-                let current_pos = unwrapped_inner.funcall::<_, _, u64>("seek", (0, 1));
 
-                if let Err(e) = current_pos {
+                // Get current position
+                let current_pos = match unwrapped_inner.funcall::<_, _, u64>("pos", ()) {
+                    Ok(pos) => pos,
+                    Err(e) => {
+                        eprintln!("Error seeking: {}", e);
+                        return 0;
+                    }
+                };
+
+                // Seek to end
+                if let Err(e) = unwrapped_inner.funcall::<_, _, u64>("seek", (0, 2)) {
                     eprintln!("Error seeking: {}", e);
                     return 0;
                 }
 
-                if let Err(e) = unwrapped_inner.funcall::<_, _, u64>("seek", (0, 2)) {
+                // Offset at the end of the file is the length of the file
+                let size = match unwrapped_inner.funcall::<_, _, u64>("pos", ()) {
+                    Ok(pos) => pos,
+                    Err(e) => {
+                        eprintln!("Error seeking: {}", e);
+                        return 0;
+                    }
+                };
+
+                // Restore original position
+                if let Err(e) = unwrapped_inner.funcall::<_, _, u64>("seek", (current_pos, 0)) {
                     eprintln!("Error seeking: {}", e);
                     return 0;
                 }
 
-                let size = unwrapped_inner.funcall::<_, _, u64>("pos", ());
-
-                match size {
-                    Ok(size) => {
-                        // Restore original position
-                        if let Err(e) = unwrapped_inner.funcall::<_, _, u64>(
-                            "seek",
-                            (current_pos.expect("Current position is not set!"), 0),
-                        ) {
-                            eprintln!("Error seeking: {}", e);
-                            return 0;
-                        }
-                        size
-                    }
+                let final_pos = match unwrapped_inner.funcall::<_, _, u64>("pos", ()) {
+                    Ok(pos) => pos,
                     Err(e) => {
                         eprintln!("Error seeking: {}", e);
                         return 0;
                     }
-                }
+                };
+
+                assert_eq!(
+                    current_pos, final_pos,
+                    "Failed to restore original position in seekable IO object. Started at position {}, but ended at {}",
+                    current_pos,
+                    final_pos
+                );
+
+                size
             }
         }
     }

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

@@ -22,7 +22,7 @@ use arrow_array::{
     TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
 };
 use arrow_schema::{DataType, TimeUnit};
-use magnus::{value::ReprValue, Error as MagnusError, IntoValue, Ruby, TryConvert, Value};
+use magnus::{value::ReprValue, Error as MagnusError, IntoValue, Ruby, Value};
 use parquet::data_type::Decimal;
 use parquet::record::Field;
 use std::{collections::HashMap, hash::BuildHasher, sync::Arc};

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

@@ -208,7 +208,7 @@ impl ParquetValue {
                 Ok(ParquetValue::Float64(v))
             }
             ParquetSchemaType::String => {
-                let v = String::try_convert(value)?;
+                let v = convert_to_string(value)?;
                 Ok(ParquetValue::String(v))
             }
             ParquetSchemaType::Binary => {

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

@@ -3,7 +3,7 @@ use std::str::FromStr;
 use super::*;
 use arrow_array::builder::*;
 use jiff::tz::{Offset, TimeZone};
-use magnus::{RArray, TryConvert};
+use magnus::{RArray, RString, TryConvert};
 
 pub struct NumericConverter<T> {
     _phantom: std::marker::PhantomData<T>,
@@ -194,6 +194,21 @@ pub fn convert_to_boolean(value: Value) -> Result<bool, MagnusError> {
     }
 }
 
+pub fn convert_to_string(value: Value) -> Result<String, MagnusError> {
+    String::try_convert(value).or_else(|_| {
+        if value.respond_to("to_s", false)? {
+            value.funcall::<_, _, RString>("to_s", ())?.to_string()
+        } else if value.respond_to("to_str", false)? {
+            value.funcall::<_, _, RString>("to_str", ())?.to_string()
+        } else {
+            Err(MagnusError::new(
+                magnus::exception::type_error(),
+                format!("Not able to convert {:?} to String", value),
+            ))
+        }
+    })
+}
+
 pub fn convert_to_list(
     value: Value,
     list_field: &ListField,

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

@@ -9,7 +9,7 @@ use magnus::{value::ReprValue, Error as MagnusError, RString, Ruby, Symbol, TryC
 use parquet::{arrow::ArrowWriter, errors::ParquetError};
 use tempfile::NamedTempFile;
 
-use crate::types::{ListField, MapField, ParquetSchemaType};
+use crate::types::{convert_to_string, ListField, MapField, ParquetSchemaType};
 
 #[derive(Debug)]
 pub struct SchemaField<'a> {
@@ -240,7 +240,7 @@ impl<'a> ColumnCollector<'a> {
                 ParquetValue::Float64(v)
             }
             ParquetSchemaType::String => {
-                let v = String::try_convert(value)?;
+                let v = convert_to_string(value)?;
                 ParquetValue::String(v)
             }
             ParquetSchemaType::Binary => {

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

@@ -42,7 +42,7 @@ pub fn parse_parquet_write_args(args: &[Value]) -> Result<ParquetWriteArgs, Magn
 
     let kwargs = get_kwargs::<
         _,
-        (Value, Value),
+        (Option<RArray>, Value),
         (
             Option<Option<usize>>,
             Option<Option<usize>>,
@@ -61,71 +61,88 @@ pub fn parse_parquet_write_args(args: &[Value]) -> Result<ParquetWriteArgs, Magn
         ],
     )?;
 
-    let schema_array = RArray::from_value(kwargs.required.0).ok_or_else(|| {
-        MagnusError::new(
-            magnus::exception::type_error(),
-            "schema must be an array of hashes",
-        )
-    })?;
-
-    let mut schema = Vec::with_capacity(schema_array.len());
-
-    for (idx, field_hash) in schema_array.into_iter().enumerate() {
-        if !field_hash.is_kind_of(ruby.class_hash()) {
-            return Err(MagnusError::new(
-                magnus::exception::type_error(),
-                format!("schema[{}] must be a hash", idx),
-            ));
-        }
-
-        let entries: Vec<(Value, Value)> = field_hash.funcall("to_a", ())?;
-        if entries.len() != 1 {
-            return Err(MagnusError::new(
+    let schema = if kwargs.required.0.is_none() || kwargs.required.0.unwrap().is_empty() {
+        // If schema is nil, we need to peek at the first value to determine column count
+        let first_value = read_from.funcall::<_, _, Value>("peek", ())?;
+        let array = RArray::from_value(first_value).ok_or_else(|| {
+            MagnusError::new(
                 magnus::exception::type_error(),
-                format!("schema[{}] must contain exactly one key-value pair", idx),
-            ));
-        }
+                "First value must be an array when schema is not provided",
+            )
+        })?;
 
-        let (name, type_value) = &entries[0];
-        let name = String::try_convert(name.clone())?;
+        // Generate field names f0, f1, f2, etc.
+        (0..array.len())
+            .map(|i| SchemaField {
+                name: format!("f{}", i),
+                type_: ParquetSchemaType::String,
+                format: None,
+            })
+            .collect()
+    } else {
+        let schema_array = kwargs.required.0.unwrap();
 
-        let (type_, format) = if type_value.is_kind_of(ruby.class_hash()) {
-            let type_hash: Vec<(Value, Value)> = type_value.funcall("to_a", ())?;
-            let mut type_str = None;
-            let mut format_str = None;
+        let mut schema = Vec::with_capacity(schema_array.len());
 
-            for (key, value) in type_hash {
-                let key = String::try_convert(key)?;
-                match key.as_str() {
-                    "type" => type_str = Some(value),
-                    "format" => format_str = Some(String::try_convert(value)?),
-                    _ => {
-                        return Err(MagnusError::new(
-                            magnus::exception::type_error(),
-                            format!("Unknown key '{}' in type definition", key),
-                        ))
-                    }
-                }
+        for (idx, field_hash) in schema_array.into_iter().enumerate() {
+            if !field_hash.is_kind_of(ruby.class_hash()) {
+                return Err(MagnusError::new(
+                    magnus::exception::type_error(),
+                    format!("schema[{}] must be a hash", idx),
+                ));
             }
 
-            let type_str = type_str.ok_or_else(|| {
-                MagnusError::new(
+            let entries: Vec<(Value, Value)> = field_hash.funcall("to_a", ())?;
+            if entries.len() != 1 {
+                return Err(MagnusError::new(
                     magnus::exception::type_error(),
-                    "Missing 'type' in type definition",
-                )
-            })?;
+                    format!("schema[{}] must contain exactly one key-value pair", idx),
+                ));
+            }
 
-            (ParquetSchemaType::try_convert(type_str)?, format_str)
-        } else {
-            (ParquetSchemaType::try_convert(type_value.clone())?, None)
-        };
+            let (name, type_value) = &entries[0];
+            let name = String::try_convert(name.clone())?;
+
+            let (type_, format) = if type_value.is_kind_of(ruby.class_hash()) {
+                let type_hash: Vec<(Value, Value)> = type_value.funcall("to_a", ())?;
+                let mut type_str = None;
+                let mut format_str = None;
+
+                for (key, value) in type_hash {
+                    let key = String::try_convert(key)?;
+                    match key.as_str() {
+                        "type" => type_str = Some(value),
+                        "format" => format_str = Some(String::try_convert(value)?),
+                        _ => {
+                            return Err(MagnusError::new(
+                                magnus::exception::type_error(),
+                                format!("Unknown key '{}' in type definition", key),
+                            ))
+                        }
+                    }
+                }
 
-        schema.push(SchemaField {
-            name,
-            type_,
-            format,
-        });
-    }
+                let type_str = type_str.ok_or_else(|| {
+                    MagnusError::new(
+                        magnus::exception::type_error(),
+                        "Missing 'type' in type definition",
+                    )
+                })?;
+
+                (ParquetSchemaType::try_convert(type_str)?, format_str)
+            } else {
+                (ParquetSchemaType::try_convert(type_value.clone())?, None)
+            };
+
+            schema.push(SchemaField {
+                name,
+                type_,
+                format,
+            });
+        }
+
+        schema
+    };
 
     Ok(ParquetWriteArgs {
         read_from,
@@ -396,6 +413,14 @@ pub fn write_columns(args: &[Value]) -> Result<(), MagnusError> {
                         MagnusError::new(ruby.exception_type_error(), "Batch must be an array")
                     })?;
 
+                    // Batch array must be an array of arrays. Check that the first value in `batch_array` is an array.
+                    batch_array.entry::<RArray>(0).map_err(|_| {
+                        MagnusError::new(
+                            ruby.exception_type_error(),
+                            "When writing columns, data must be formatted as batches of columns: [[batch1_col1, batch1_col2], [batch2_col1, batch2_col2]].",
+                        )
+                    })?;
+
                     // Validate batch length matches schema
                     if batch_array.len() != schema.len() {
                         return Err(MagnusError::new(

data/lib/parquet/version.rb

@@ -1,3 +1,3 @@
 module Parquet
-  VERSION = "0.4.0"
+  VERSION = "0.4.2"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: parquet
 version: !ruby/object:Gem::Version
-  version: 0.4.0
+  version: 0.4.2
 platform: ruby
 authors:
 - Nathan Jaremko
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-02-20 00:00:00.000000000 Z
+date: 2025-02-21 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys