parquet 0.4.2 → 0.5.1

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

@@ -1,53 +1,57 @@
-use std::{
-    fs::File,
-    io::{self, BufReader, BufWriter},
-    mem,
-    sync::Arc,
-};
+mod write_columns;
+mod write_rows;
 
-use arrow_array::{Array, RecordBatch};
-use arrow_schema::{DataType, Field, Schema, TimeUnit};
+use arrow_schema::{DataType, Schema, TimeUnit};
+use itertools::Itertools;
 use magnus::{
     scan_args::{get_kwargs, scan_args},
     value::ReprValue,
-    Error as MagnusError, RArray, Ruby, TryConvert, Value,
+    Error as MagnusError, RArray, RHash, Ruby, Symbol, Value,
 };
 use parquet::{
     arrow::ArrowWriter,
     basic::{Compression, GzipLevel, ZstdLevel},
     file::properties::WriterProperties,
 };
-use rand::Rng;
+use std::{
+    fs::File,
+    io::{self, BufReader, BufWriter},
+    sync::Arc,
+};
 use tempfile::NamedTempFile;
+pub use write_columns::write_columns;
+pub use write_rows::write_rows;
 
+use crate::{types::PrimitiveType, SchemaNode};
 use crate::{
-    convert_ruby_array_to_arrow,
-    types::{ColumnCollector, ParquetErrorWrapper, WriterOutput},
-    IoLikeValue, ParquetSchemaType, ParquetWriteArgs, SchemaField, SendableWrite,
+    types::{ColumnCollector, ParquetGemError, ParquetSchemaType, WriterOutput},
+    utils::parse_string_or_symbol,
+    IoLikeValue, ParquetSchemaType as PST, ParquetWriteArgs, SchemaField, SendableWrite,
 };
 
-const MIN_SAMPLES_FOR_ESTIMATE: usize = 10; // Minimum samples needed for estimation
-const SAMPLE_SIZE: usize = 100; // Number of rows to sample for size estimation
-const MIN_BATCH_SIZE: usize = 10; // Minimum batch size to maintain efficiency
-const INITIAL_BATCH_SIZE: usize = 100; // Initial batch size while sampling
-
-// Maximum memory usage per batch (64MB by default)
+const MIN_SAMPLES_FOR_ESTIMATE: usize = 10;
+const SAMPLE_SIZE: usize = 100;
+const MIN_BATCH_SIZE: usize = 10;
+const INITIAL_BATCH_SIZE: usize = 100;
 const DEFAULT_MEMORY_THRESHOLD: usize = 64 * 1024 * 1024;
 
 /// Parse arguments for Parquet writing
-pub fn parse_parquet_write_args(args: &[Value]) -> Result<ParquetWriteArgs, MagnusError> {
-    let ruby = unsafe { Ruby::get_unchecked() };
+pub fn parse_parquet_write_args(
+    ruby: &Ruby,
+    args: &[Value],
+) -> Result<ParquetWriteArgs, MagnusError> {
     let parsed_args = scan_args::<(Value,), (), (), (), _, ()>(args)?;
     let (read_from,) = parsed_args.required;
 
     let kwargs = get_kwargs::<
         _,
-        (Option<RArray>, Value),
+        (Value, Value),
         (
             Option<Option<usize>>,
             Option<Option<usize>>,
             Option<Option<String>>,
             Option<Option<usize>>,
+            Option<Option<Value>>,
         ),
         (),
     >(
@@ -58,451 +62,358 @@ pub fn parse_parquet_write_args(args: &[Value]) -> Result<ParquetWriteArgs, Magn
             "flush_threshold",
             "compression",
             "sample_size",
+            "logger",
         ],
     )?;
 
-    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(),
-                "First value must be an array when schema is not provided",
-            )
-        })?;
-
-        // Generate field names f0, f1, f2, etc.
-        (0..array.len())
-            .map(|i| SchemaField {
-                name: format!("f{}", i),
-                type_: ParquetSchemaType::String,
-                format: None,
-            })
-            .collect()
+    // The schema value could be one of:
+    // 1. An array of hashes (legacy format)
+    // 2. A hash with type: :struct (new DSL format)
+    // 3. nil (infer from data)
+    let schema_value = kwargs.required.0;
+
+    // Check if it's the new DSL format (a hash with type: :struct)
+    // We need to handle both direct hash objects and objects created via Parquet::Schema.define
+
+    // First, try to convert it to a Hash if it's not already a Hash
+    // This handles the case where schema_value is a Schema object from Parquet::Schema.define
+    let schema_hash = if schema_value.is_kind_of(ruby.class_hash()) {
+        RHash::from_value(schema_value).ok_or_else(|| {
+            MagnusError::new(magnus::exception::type_error(), "Schema must be a hash")
+        })?
     } else {
-        let schema_array = kwargs.required.0.unwrap();
-
-        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(
-                    magnus::exception::type_error(),
-                    format!("schema[{}] must contain exactly one key-value pair", idx),
-                ));
-            }
-
-            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),
-                            ))
+        // Try to convert the object to a hash with to_h
+        match schema_value.respond_to("to_h", false) {
+            Ok(true) => {
+                match schema_value.funcall::<_, _, Value>("to_h", ()) {
+                    Ok(hash_val) => match RHash::from_value(hash_val) {
+                        Some(hash) => hash,
+                        None => {
+                            // Not a hash, continue to normal handling
+                            RHash::new()
                         }
+                    },
+                    Err(_) => {
+                        // couldn't call to_h, continue to normal handling
+                        RHash::new()
                     }
                 }
+            }
+            _ => {
+                // Doesn't respond to to_h, continue to normal handling
+                RHash::new()
+            }
+        }
+    };
+
+    // Now check if it's a schema hash with a type: :struct field
+    let type_val = schema_hash.get(Symbol::new("type"));
+
+    if let Some(type_val) = type_val {
+        // If it has a type: :struct, it's the new DSL format
+        // Use parse_string_or_symbol to handle both String and Symbol values
+        let ttype = parse_string_or_symbol(&ruby, type_val)?;
+        if let Some(ref type_str) = ttype {
+            if type_str == "struct" {
+                // Parse using the new schema approach
+                let schema_node = crate::parse_schema_node(&ruby, schema_value)?;
+
+                validate_schema_node(&ruby, &schema_node)?;
+
+                return Ok(ParquetWriteArgs {
+                    read_from,
+                    write_to: kwargs.required.1,
+                    schema: schema_node,
+                    batch_size: kwargs.optional.0.flatten(),
+                    flush_threshold: kwargs.optional.1.flatten(),
+                    compression: kwargs.optional.2.flatten(),
+                    sample_size: kwargs.optional.3.flatten(),
+                    logger: kwargs.optional.4.flatten(),
+                });
+            }
+        }
+    }
 
-                let type_str = type_str.ok_or_else(|| {
+    // If it's not a hash with type: :struct, handle as legacy format
+    let schema_fields = if schema_value.is_nil()
+        || (schema_value.is_kind_of(ruby.class_array())
+            && RArray::from_value(schema_value)
+                .ok_or_else(|| {
                     MagnusError::new(
                         magnus::exception::type_error(),
-                        "Missing 'type' in type definition",
+                        "Schema fields must be an array",
                     )
-                })?;
-
-                (ParquetSchemaType::try_convert(type_str)?, format_str)
-            } else {
-                (ParquetSchemaType::try_convert(type_value.clone())?, None)
-            };
+                })?
+                .len()
+                == 0)
+    {
+        // If schema is nil or an empty array, we need to peek at the first value to determine column count
+        let first_value = read_from.funcall::<_, _, Value>("peek", ())?;
+        // Default to nullable:true for auto-inferred fields
+        crate::infer_schema_from_first_row(&ruby, first_value, true)?
+    } else {
+        // Legacy array format - use our centralized parser
+        crate::parse_legacy_schema(&ruby, schema_value)?
+    };
 
-            schema.push(SchemaField {
-                name,
-                type_,
-                format,
-            });
-        }
+    // Convert the legacy schema fields to SchemaNode (DSL format)
+    let schema_node = crate::legacy_schema_to_dsl(&ruby, schema_fields)?;
 
-        schema
-    };
+    validate_schema_node(&ruby, &schema_node)?;
 
     Ok(ParquetWriteArgs {
         read_from,
         write_to: kwargs.required.1,
-        schema,
+        schema: schema_node,
         batch_size: kwargs.optional.0.flatten(),
         flush_threshold: kwargs.optional.1.flatten(),
         compression: kwargs.optional.2.flatten(),
         sample_size: kwargs.optional.3.flatten(),
+        logger: kwargs.optional.4.flatten(),
     })
 }
 
-/// Estimate the size of a row
-fn estimate_single_row_size(row: &RArray, schema: &[SchemaField]) -> Result<usize, MagnusError> {
-    let mut row_size = 0;
-    for (field, value) in schema.iter().zip(row.into_iter()) {
-        // Estimate size based on type and value
-        row_size += match &field.type_ {
-            // Use reference to avoid moving
-            ParquetSchemaType::Int8 | ParquetSchemaType::UInt8 => 1,
-            ParquetSchemaType::Int16 | ParquetSchemaType::UInt16 => 2,
-            ParquetSchemaType::Int32
-            | ParquetSchemaType::UInt32
-            | ParquetSchemaType::Float
-            | ParquetSchemaType::Date32 => 4,
-            ParquetSchemaType::Int64
-            | ParquetSchemaType::UInt64
-            | ParquetSchemaType::Double
-            | ParquetSchemaType::TimestampMillis
-            | ParquetSchemaType::TimestampMicros => 8,
-            ParquetSchemaType::String => {
-                if let Ok(s) = String::try_convert(value) {
-                    s.len() + mem::size_of::<usize>() // account for length prefix
-                } else {
-                    16 // default estimate for string
-                }
-            }
-            ParquetSchemaType::Binary => {
-                if let Ok(bytes) = Vec::<u8>::try_convert(value) {
-                    bytes.len() + mem::size_of::<usize>() // account for length prefix
+// -----------------------------------------------------------------------------
+// HELPER to invert arrow DataType back to our ParquetSchemaType
+// Converts Arrow DataType to our internal ParquetSchemaType representation.
+// This is essential for mapping Arrow types back to our schema representation
+// when working with column collections and schema validation.
+// -----------------------------------------------------------------------------
+fn arrow_data_type_to_parquet_schema_type(dt: &DataType) -> Result<ParquetSchemaType, MagnusError> {
+    match dt {
+        DataType::Boolean => Ok(PST::Primitive(PrimitiveType::Boolean)),
+        DataType::Int8 => Ok(PST::Primitive(PrimitiveType::Int8)),
+        DataType::Int16 => Ok(PST::Primitive(PrimitiveType::Int16)),
+        DataType::Int32 => Ok(PST::Primitive(PrimitiveType::Int32)),
+        DataType::Int64 => Ok(PST::Primitive(PrimitiveType::Int64)),
+        DataType::UInt8 => Ok(PST::Primitive(PrimitiveType::UInt8)),
+        DataType::UInt16 => Ok(PST::Primitive(PrimitiveType::UInt16)),
+        DataType::UInt32 => Ok(PST::Primitive(PrimitiveType::UInt32)),
+        DataType::UInt64 => Ok(PST::Primitive(PrimitiveType::UInt64)),
+        DataType::Float16 => {
+            // We do not have a direct ParquetSchemaType::Float16, we treat it as Float
+            Ok(PST::Primitive(PrimitiveType::Float32))
+        }
+        DataType::Float32 => Ok(PST::Primitive(PrimitiveType::Float32)),
+        DataType::Float64 => Ok(PST::Primitive(PrimitiveType::Float64)),
+        DataType::Date32 => Ok(PST::Primitive(PrimitiveType::Date32)),
+        DataType::Date64 => {
+            // Our code typically uses Date32 or Timestamp for 64. But Arrow has Date64
+            // We can store it as PST::Date64 if we want. If we don't have that, consider PST::Date32 or an error.
+            // If your existing code only handles Date32, you can error. But let's do PST::Date32 as fallback:
+            // Or define a new variant if you have one in your code. We'll show a fallback approach:
+            Err(MagnusError::new(
+                magnus::exception::runtime_error(),
+                "Arrow Date64 not directly supported in current ParquetSchemaType (use date32?).",
+            ))
+        }
+        DataType::Timestamp(TimeUnit::Second, _tz) => {
+            // We'll treat this as PST::TimestampMillis, or define PST::TimestampSecond
+            // For simplicity, let's map "second" to PST::TimestampMillis with a note:
+            Ok(PST::Primitive(PrimitiveType::TimestampMillis))
+        }
+        DataType::Timestamp(TimeUnit::Millisecond, _tz) => {
+            Ok(PST::Primitive(PrimitiveType::TimestampMillis))
+        }
+        DataType::Timestamp(TimeUnit::Microsecond, _tz) => {
+            Ok(PST::Primitive(PrimitiveType::TimestampMicros))
+        }
+        DataType::Timestamp(TimeUnit::Nanosecond, _tz) => {
+            // If you have a PST::TimestampNanos variant, use it. Otherwise, degrade to micros
+            // for demonstration:
+            Err(MagnusError::new(
+                magnus::exception::runtime_error(),
+                "TimestampNanos not supported, please adjust your schema or code.",
+            ))
+        }
+        DataType::Utf8 => Ok(PST::Primitive(PrimitiveType::String)),
+        DataType::Binary => Ok(PST::Primitive(PrimitiveType::Binary)),
+        DataType::LargeUtf8 => {
+            // If not supported, degrade or error. We'll degrade to PST::String
+            Ok(PST::Primitive(PrimitiveType::String))
+        }
+        DataType::LargeBinary => Ok(PST::Primitive(PrimitiveType::Binary)),
+        DataType::List(child_field) => {
+            // Recursively handle the item type
+            let child_type = arrow_data_type_to_parquet_schema_type(child_field.data_type())?;
+            Ok(PST::List(Box::new(crate::types::ListField {
+                item_type: child_type,
+                format: None,
+                nullable: true,
+            })))
+        }
+        DataType::Map(entry_field, _keys_sorted) => {
+            // Arrow's Map -> a struct<key, value> inside
+            let entry_type = entry_field.data_type();
+            if let DataType::Struct(fields) = entry_type {
+                if fields.len() == 2 {
+                    let key_type = arrow_data_type_to_parquet_schema_type(fields[0].data_type())?;
+                    let value_type = arrow_data_type_to_parquet_schema_type(fields[1].data_type())?;
+                    Ok(PST::Map(Box::new(crate::types::MapField {
+                        key_type,
+                        value_type,
+                        key_format: None,
+                        value_format: None,
+                        value_nullable: true,
+                    })))
                 } else {
-                    16 // default estimate for binary
+                    Err(MagnusError::new(
+                        magnus::exception::type_error(),
+                        "Map field must have exactly 2 child fields (key, value)",
+                    ))
                 }
+            } else {
+                Err(MagnusError::new(
+                    magnus::exception::type_error(),
+                    "Map field is not a struct? Unexpected Arrow schema layout",
+                ))
             }
-            ParquetSchemaType::Boolean => 1,
-            ParquetSchemaType::List(_) | ParquetSchemaType::Map(_) => {
-                32 // rough estimate for complex types
+        }
+        DataType::Struct(arrow_fields) => {
+            // We treat this as PST::Struct. We'll recursively handle subfields
+            // but for top-level collecting we only store them as one column
+            // so the user data must pass a Ruby Hash or something for that field.
+            let mut schema_fields = vec![];
+            for f in arrow_fields {
+                let sub_type = arrow_data_type_to_parquet_schema_type(f.data_type())?;
+                schema_fields.push(SchemaField {
+                    name: f.name().clone(),
+                    type_: sub_type,
+                    format: None, // We can't see the 'format' from Arrow
+                    nullable: f.is_nullable(),
+                });
             }
-        };
+            Ok(PST::Struct(Box::new(crate::types::StructField {
+                fields: schema_fields,
+            })))
+        }
+        _ => Err(MagnusError::new(
+            magnus::exception::runtime_error(),
+            format!("Unsupported or unhandled Arrow DataType: {:?}", dt),
+        )),
     }
-    Ok(row_size)
 }
 
-#[inline]
-pub fn write_rows(args: &[Value]) -> Result<(), MagnusError> {
-    let ruby = unsafe { Ruby::get_unchecked() };
-
-    let ParquetWriteArgs {
-        read_from,
-        write_to,
-        schema,
-        batch_size: user_batch_size,
-        compression,
-        flush_threshold,
-        sample_size: user_sample_size,
-    } = parse_parquet_write_args(args)?;
-
-    let flush_threshold = flush_threshold.unwrap_or(DEFAULT_MEMORY_THRESHOLD);
-
-    // Convert schema to Arrow schema
-    let arrow_fields: Vec<Field> = schema
-        .iter()
-        .map(|field| {
-            Field::new(
-                &field.name,
-                match field.type_ {
-                    ParquetSchemaType::Int8 => DataType::Int8,
-                    ParquetSchemaType::Int16 => DataType::Int16,
-                    ParquetSchemaType::Int32 => DataType::Int32,
-                    ParquetSchemaType::Int64 => DataType::Int64,
-                    ParquetSchemaType::UInt8 => DataType::UInt8,
-                    ParquetSchemaType::UInt16 => DataType::UInt16,
-                    ParquetSchemaType::UInt32 => DataType::UInt32,
-                    ParquetSchemaType::UInt64 => DataType::UInt64,
-                    ParquetSchemaType::Float => DataType::Float32,
-                    ParquetSchemaType::Double => DataType::Float64,
-                    ParquetSchemaType::String => DataType::Utf8,
-                    ParquetSchemaType::Binary => DataType::Binary,
-                    ParquetSchemaType::Boolean => DataType::Boolean,
-                    ParquetSchemaType::Date32 => DataType::Date32,
-                    ParquetSchemaType::TimestampMillis => {
-                        DataType::Timestamp(TimeUnit::Millisecond, None)
-                    }
-                    ParquetSchemaType::TimestampMicros => {
-                        DataType::Timestamp(TimeUnit::Microsecond, None)
-                    }
-                    ParquetSchemaType::List(_) => unimplemented!("List type not yet supported"),
-                    ParquetSchemaType::Map(_) => unimplemented!("Map type not yet supported"),
-                },
-                true,
-            )
-        })
-        .collect();
-    let arrow_schema = Arc::new(Schema::new(arrow_fields));
-
-    // Create the writer
-    let mut writer = create_writer(&ruby, &write_to, arrow_schema.clone(), compression)?;
-
-    if read_from.is_kind_of(ruby.class_enumerator()) {
-        // Create collectors for each column
-        let mut column_collectors: Vec<ColumnCollector> = schema
-            .iter()
-            .map(|field| {
-                // Clone the type to avoid moving from a reference
-                let type_clone = field.type_.clone();
-                ColumnCollector::new(field.name.clone(), type_clone, field.format.clone())
-            })
-            .collect();
-
-        let mut rows_in_batch = 0;
-        let mut total_rows = 0;
-        let mut rng = rand::rng();
-        let sample_size = user_sample_size.unwrap_or(SAMPLE_SIZE);
-        let mut size_samples = Vec::with_capacity(sample_size);
-        let mut current_batch_size = user_batch_size.unwrap_or(INITIAL_BATCH_SIZE);
-
-        loop {
-            match read_from.funcall::<_, _, Value>("next", ()) {
-                Ok(row) => {
-                    let row_array = RArray::from_value(row).ok_or_else(|| {
-                        MagnusError::new(ruby.exception_type_error(), "Row must be an array")
-                    })?;
-
-                    // Validate row length matches schema
-                    if row_array.len() != column_collectors.len() {
-                        return Err(MagnusError::new(
-                            magnus::exception::type_error(),
-                            format!(
-                                "Row length ({}) does not match schema length ({}). Schema expects columns: {:?}",
-                                row_array.len(),
-                                column_collectors.len(),
-                                column_collectors.iter().map(|c| c.name.as_str()).collect::<Vec<_>>()
-                            ),
-                        ));
-                    }
-
-                    // Sample row sizes using reservoir sampling
-                    if size_samples.len() < sample_size {
-                        size_samples.push(estimate_single_row_size(&row_array, &schema)?);
-                    } else if rng.random_range(0..=total_rows) < sample_size {
-                        let idx = rng.random_range(0..sample_size);
-                        size_samples[idx] = estimate_single_row_size(&row_array, &schema)?;
-                    }
-
-                    // Process each value in the row
-                    for (collector, value) in column_collectors.iter_mut().zip(row_array) {
-                        collector.push_value(value)?;
-                    }
-
-                    rows_in_batch += 1;
-                    total_rows += 1;
-
-                    // Calculate batch size progressively once we have minimum samples
-                    if size_samples.len() >= MIN_SAMPLES_FOR_ESTIMATE && user_batch_size.is_none() {
-                        let total_size = size_samples.iter().sum::<usize>();
-                        // Safe because we know we have at least MIN_SAMPLES_FOR_ESTIMATE samples
-                        let avg_row_size = total_size as f64 / size_samples.len() as f64;
-                        let avg_row_size = avg_row_size.max(1.0); // Ensure we don't divide by zero
-                        let suggested_batch_size =
-                            (flush_threshold as f64 / avg_row_size).floor() as usize;
-                        current_batch_size = suggested_batch_size.max(MIN_BATCH_SIZE);
-                    }
-
-                    // When we reach batch size, write the batch
-                    if rows_in_batch >= current_batch_size {
-                        write_batch(&mut writer, &mut column_collectors, flush_threshold)?;
-                        rows_in_batch = 0;
-                    }
-                }
-                Err(e) => {
-                    if e.is_kind_of(ruby.exception_stop_iteration()) {
-                        // Write any remaining rows
-                        if rows_in_batch > 0 {
-                            write_batch(&mut writer, &mut column_collectors, flush_threshold)?;
-                        }
-                        break;
-                    }
-                    return Err(e);
-                }
-            }
+// -----------------------------------------------------------------------------
+// HELPER to build ColumnCollectors for the DSL variant
+// This function converts a SchemaNode (from our DSL) into a collection of ColumnCollectors
+// that can accumulate values for each column in the schema.
+// - arrow_schema: The Arrow schema corresponding to our DSL schema
+// - root_node: The root SchemaNode (expected to be a Struct node) from which to build collectors
+// -----------------------------------------------------------------------------
+fn build_column_collectors_from_dsl<'a>(
+    ruby: &'a Ruby,
+    arrow_schema: &'a Arc<Schema>,
+    root_node: &'a SchemaNode,
+) -> Result<Vec<ColumnCollector<'a>>, MagnusError> {
+    // We expect the top-level schema node to be a Struct so that arrow_schema
+    // lines up with root_node.fields. If the user gave a top-level primitive, it would be 1 field, but
+    // our code calls build_arrow_schema under the assumption "top-level must be Struct."
+    let fields = match root_node {
+        SchemaNode::Struct { fields, .. } => fields,
+        _ => {
+            return Err(MagnusError::new(
+                ruby.exception_runtime_error(),
+                "Top-level schema for DSL must be a struct",
+            ))
         }
-    } else {
+    };
+
+    if fields.len() != arrow_schema.fields().len() {
         return Err(MagnusError::new(
-            magnus::exception::type_error(),
-            "read_from must be an Enumerator",
+            ruby.exception_runtime_error(),
+            format!(
+                "Mismatch between DSL field count ({}) and Arrow fields ({})",
+                fields.len(),
+                arrow_schema.fields().len()
+            ),
         ));
     }
 
-    // Ensure everything is written and get the temp file if it exists
-    if let Some(temp_file) = writer.close().map_err(|e| ParquetErrorWrapper(e))? {
-        // If we got a temp file back, we need to copy its contents to the IO-like object
-        copy_temp_file_to_io_like(temp_file, IoLikeValue(write_to))?;
+    let mut collectors = Vec::with_capacity(fields.len());
+    for (arrow_field, schema_field_node) in arrow_schema.fields().iter().zip(fields) {
+        let name = arrow_field.name().clone();
+        let parquet_type = arrow_data_type_to_parquet_schema_type(arrow_field.data_type())?;
+
+        // Extract the optional format from the schema node
+        let format = extract_format_from_schema_node(schema_field_node);
+
+        // Build the ColumnCollector
+        collectors.push(ColumnCollector::new(
+            ruby,
+            name,
+            parquet_type,
+            format,
+            arrow_field.is_nullable(),
+        ));
     }
-
-    Ok(())
+    Ok(collectors)
 }
 
-#[inline]
-pub fn write_columns(args: &[Value]) -> Result<(), MagnusError> {
-    let ruby = unsafe { Ruby::get_unchecked() };
-
-    let ParquetWriteArgs {
-        read_from,
-        write_to,
-        schema,
-        batch_size: _,
-        compression,
-        flush_threshold,
-        sample_size: _,
-    } = parse_parquet_write_args(args)?;
-
-    let flush_threshold = flush_threshold.unwrap_or(DEFAULT_MEMORY_THRESHOLD);
-
-    // Convert schema to Arrow schema
-    let arrow_fields: Vec<Field> = schema
-        .iter()
-        .map(|field| {
-            Field::new(
-                &field.name,
-                match field.type_ {
-                    ParquetSchemaType::Int8 => DataType::Int8,
-                    ParquetSchemaType::Int16 => DataType::Int16,
-                    ParquetSchemaType::Int32 => DataType::Int32,
-                    ParquetSchemaType::Int64 => DataType::Int64,
-                    ParquetSchemaType::UInt8 => DataType::UInt8,
-                    ParquetSchemaType::UInt16 => DataType::UInt16,
-                    ParquetSchemaType::UInt32 => DataType::UInt32,
-                    ParquetSchemaType::UInt64 => DataType::UInt64,
-                    ParquetSchemaType::Float => DataType::Float32,
-                    ParquetSchemaType::Double => DataType::Float64,
-                    ParquetSchemaType::String => DataType::Utf8,
-                    ParquetSchemaType::Binary => DataType::Binary,
-                    ParquetSchemaType::Boolean => DataType::Boolean,
-                    ParquetSchemaType::Date32 => DataType::Date32,
-                    ParquetSchemaType::TimestampMillis => {
-                        DataType::Timestamp(TimeUnit::Millisecond, None)
-                    }
-                    ParquetSchemaType::TimestampMicros => {
-                        DataType::Timestamp(TimeUnit::Microsecond, None)
-                    }
-                    ParquetSchemaType::List(_) => unimplemented!("List type not yet supported"),
-                    ParquetSchemaType::Map(_) => unimplemented!("Map type not yet supported"),
-                },
-                true,
-            )
-        })
-        .collect();
-    let arrow_schema = Arc::new(Schema::new(arrow_fields));
-
-    // Create the writer
-    let mut writer = create_writer(&ruby, &write_to, arrow_schema.clone(), compression)?;
-
-    if read_from.is_kind_of(ruby.class_enumerator()) {
-        loop {
-            match read_from.funcall::<_, _, Value>("next", ()) {
-                Ok(batch) => {
-                    let batch_array = RArray::from_value(batch).ok_or_else(|| {
-                        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(
-                            magnus::exception::type_error(),
-                            format!(
-                                "Batch column count ({}) does not match schema length ({}). Schema expects columns: {:?}",
-                                batch_array.len(),
-                                schema.len(),
-                                schema.iter().map(|f| f.name.as_str()).collect::<Vec<_>>()
-                            ),
-                        ));
-                    }
-
-                    // Convert each column in the batch to Arrow arrays
-                    let arrow_arrays: Vec<(String, Arc<dyn Array>)> = schema
-                        .iter()
-                        .zip(batch_array)
-                        .map(|(field, column)| {
-                            let column_array = RArray::from_value(column).ok_or_else(|| {
-                                MagnusError::new(
-                                    magnus::exception::type_error(),
-                                    format!("Column '{}' must be an array", field.name),
-                                )
-                            })?;
-
-                            Ok((
-                                field.name.clone(),
-                                convert_ruby_array_to_arrow(column_array, &field.type_)?,
-                            ))
-                        })
-                        .collect::<Result<_, MagnusError>>()?;
-
-                    // Create and write record batch
-                    let record_batch = RecordBatch::try_from_iter(arrow_arrays).map_err(|e| {
-                        MagnusError::new(
-                            magnus::exception::runtime_error(),
-                            format!("Failed to create record batch: {}", e),
-                        )
-                    })?;
-
-                    writer
-                        .write(&record_batch)
-                        .map_err(|e| ParquetErrorWrapper(e))?;
-
-                    match &mut writer {
-                        WriterOutput::File(w) | WriterOutput::TempFile(w, _) => {
-                            if w.in_progress_size() >= flush_threshold {
-                                w.flush().map_err(|e| ParquetErrorWrapper(e))?;
-                            }
-                        }
-                    }
-                }
-                Err(e) => {
-                    if e.is_kind_of(ruby.exception_stop_iteration()) {
-                        break;
-                    }
-                    return Err(e);
-                }
-            }
-        }
-    } else {
-        return Err(MagnusError::new(
-            magnus::exception::type_error(),
-            "read_from must be an Enumerator",
-        ));
+// Helper to extract the format from a SchemaNode if available
+fn extract_format_from_schema_node(node: &SchemaNode) -> Option<String> {
+    match node {
+        SchemaNode::Primitive {
+            format: f,
+            parquet_type: _,
+            ..
+        } => f.clone(),
+        // For struct, list, map, etc. there's no single "format." We ignore it.
+        _ => None,
     }
+}
 
-    // Ensure everything is written and get the temp file if it exists
-    if let Some(temp_file) = writer.close().map_err(|e| ParquetErrorWrapper(e))? {
-        // If we got a temp file back, we need to copy its contents to the IO-like object
-        copy_temp_file_to_io_like(temp_file, IoLikeValue(write_to))?;
+// Validates a SchemaNode to ensure it meets Parquet schema requirements
+// Currently checks for duplicate field names at the root level, which would
+// cause problems when writing Parquet files. Additional validation rules
+// could be added here in the future.
+//
+// This validation is important because schema errors are difficult to debug
+// once they reach the Parquet/Arrow layer, so we check proactively before
+// any data processing begins.
+fn validate_schema_node(ruby: &Ruby, schema_node: &SchemaNode) -> Result<(), MagnusError> {
+    if let SchemaNode::Struct { fields, .. } = &schema_node {
+        // if any root level schema fields have the same name, we raise an error
+        let field_names = fields
+            .iter()
+            .map(|f| match f {
+                SchemaNode::Struct { name, .. } => name.as_str(),
+                SchemaNode::List { name, .. } => name.as_str(),
+                SchemaNode::Map { name, .. } => name.as_str(),
+                SchemaNode::Primitive { name, .. } => name.as_str(),
+            })
+            .collect::<Vec<_>>();
+        let unique_field_names = field_names.iter().unique().collect::<Vec<_>>();
+        if field_names.len() != unique_field_names.len() {
+            return Err(MagnusError::new(
+                ruby.exception_arg_error(),
+                format!(
+                    "Duplicate field names in root level schema: {:?}",
+                    field_names
+                ),
+            ));
+        }
     }
-
     Ok(())
 }
 
+// Creates an appropriate Parquet writer based on the output target and compression settings
+// This function handles two main output scenarios:
+// 1. Writing directly to a file path (string)
+// 2. Writing to a Ruby IO-like object (using a temporary file as an intermediate buffer)
+//
+// For IO-like objects, the function creates a temporary file that is later copied to the
+// IO object when writing is complete. This approach is necessary because Parquet requires
+// random file access to write its footer after the data.
+//
+// The function also configures compression based on the user's preferences, with
+// several options available (none, snappy, gzip, lz4, zstd).
 fn create_writer(
     ruby: &Ruby,
     write_to: &Value,
     schema: Arc<Schema>,
     compression: Option<String>,
-) -> Result<WriterOutput, MagnusError> {
+) -> Result<WriterOutput, ParquetGemError> {
     // Create writer properties with compression based on the option
     let props = WriterProperties::builder()
         .set_compression(match compression.as_deref() {
@@ -517,9 +428,8 @@ fn create_writer(
 
     if write_to.is_kind_of(ruby.class_string()) {
         let path = write_to.to_r_string()?.to_string()?;
-        let file: Box<dyn SendableWrite> = Box::new(File::create(path).unwrap());
-        let writer =
-            ArrowWriter::try_new(file, schema, Some(props)).map_err(|e| ParquetErrorWrapper(e))?;
+        let file: Box<dyn SendableWrite> = Box::new(File::create(path)?);
+        let writer = ArrowWriter::try_new(file, schema, Some(props))?;
         Ok(WriterOutput::File(writer))
     } else {
         // Create a temporary file to write to instead of directly to the IoLikeValue
@@ -535,13 +445,22 @@ fn create_writer(
                 format!("Failed to reopen temporary file: {}", e),
             )
         })?);
-        let writer =
-            ArrowWriter::try_new(file, schema, Some(props)).map_err(|e| ParquetErrorWrapper(e))?;
+        let writer = ArrowWriter::try_new(file, schema, Some(props))?;
         Ok(WriterOutput::TempFile(writer, temp_file))
     }
 }
 
-// Helper function to copy temp file contents to IoLikeValue
+// Copies the contents of a temporary file to a Ruby IO-like object
+// This function is necessary because Parquet writing requires random file access
+// (especially for writing the footer after all data), but Ruby IO objects may not
+// support seeking. The solution is to:
+//
+// 1. Write the entire Parquet file to a temporary file first
+// 2. Once writing is complete, copy the entire contents to the Ruby IO object
+//
+// This approach enables support for a wide range of Ruby IO objects like StringIO,
+// network streams, etc., but does require enough disk space for the temporary file
+// and involves a second full-file read/write operation at the end.
 fn copy_temp_file_to_io_like(
     temp_file: NamedTempFile,
     io_like: IoLikeValue,
@@ -564,37 +483,3 @@ fn copy_temp_file_to_io_like(
 
     Ok(())
 }
-
-fn write_batch(
-    writer: &mut WriterOutput,
-    collectors: &mut [ColumnCollector],
-    flush_threshold: usize,
-) -> Result<(), MagnusError> {
-    // Convert columns to Arrow arrays
-    let arrow_arrays: Vec<(String, Arc<dyn Array>)> = collectors
-        .iter_mut()
-        .map(|collector| Ok((collector.name.clone(), collector.take_array()?)))
-        .collect::<Result<_, MagnusError>>()?;
-
-    // Create and write record batch
-    let record_batch = RecordBatch::try_from_iter(arrow_arrays).map_err(|e| {
-        MagnusError::new(
-            magnus::exception::runtime_error(),
-            format!("Failed to create record batch: {}", e),
-        )
-    })?;
-
-    writer
-        .write(&record_batch)
-        .map_err(|e| ParquetErrorWrapper(e))?;
-
-    match writer {
-        WriterOutput::File(w) | WriterOutput::TempFile(w, _) => {
-            if w.in_progress_size() >= flush_threshold || w.memory_size() >= flush_threshold {
-                w.flush().map_err(|e| ParquetErrorWrapper(e))?;
-            }
-        }
-    }
-
-    Ok(())
-}