parquet 0.2.12-arm64-darwin

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

@@ -0,0 +1,575 @@
+use std::{
+    fs::File,
+    io::{self, BufReader, BufWriter},
+    mem,
+    sync::Arc,
+};
+
+use arrow_array::{Array, RecordBatch};
+use arrow_schema::{DataType, Field, Schema, TimeUnit};
+use magnus::{
+    scan_args::{get_kwargs, scan_args},
+    value::ReprValue,
+    Error as MagnusError, RArray, Ruby, TryConvert, Value,
+};
+use parquet::{
+    arrow::ArrowWriter,
+    basic::{Compression, GzipLevel, ZstdLevel},
+    file::properties::WriterProperties,
+};
+use rand::Rng;
+use tempfile::NamedTempFile;
+
+use crate::{
+    convert_ruby_array_to_arrow,
+    types::{ColumnCollector, ParquetErrorWrapper, WriterOutput},
+    IoLikeValue, ParquetSchemaType, ParquetWriteArgs, SchemaField, SendableWrite,
+};
+
+const SAMPLE_SIZE: usize = 100; // Number of rows to sample for size estimation
+const MIN_BATCH_SIZE: usize = 100; // 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 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() };
+    let parsed_args = scan_args::<(Value,), (), (), (), _, ()>(args)?;
+    let (read_from,) = parsed_args.required;
+
+    let kwargs = get_kwargs::<
+        _,
+        (Value, Value),
+        (
+            Option<Option<usize>>,
+            Option<Option<usize>>,
+            Option<Option<String>>,
+            Option<Option<usize>>,
+        ),
+        (),
+    >(
+        parsed_args.keywords,
+        &["schema", "write_to"],
+        &[
+            "batch_size",
+            "flush_threshold",
+            "compression",
+            "sample_size",
+        ],
+    )?;
+
+    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(
+                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),
+                        ))
+                    }
+                }
+            }
+
+            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,
+        });
+    }
+
+    Ok(ParquetWriteArgs {
+        read_from,
+        write_to: kwargs.required.1,
+        schema,
+        batch_size: kwargs.optional.0.flatten(),
+        flush_threshold: kwargs.optional.1.flatten(),
+        compression: kwargs.optional.2.flatten(),
+        sample_size: kwargs.optional.3.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
+                } else {
+                    16 // default estimate for binary
+                }
+            }
+            ParquetSchemaType::Boolean => 1,
+            ParquetSchemaType::List(_) | ParquetSchemaType::Map(_) => {
+                32 // rough estimate for complex types
+            }
+        };
+    }
+    Ok(row_size)
+}
+
+/// Calculate optimal batch size based on memory threshold and estimated row size
+fn calculate_batch_size(row_size: usize, memory_threshold: usize) -> usize {
+    let batch_size = memory_threshold / row_size;
+    batch_size.max(MIN_BATCH_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;
+
+                    // Recalculate batch size if we have enough samples and no user-specified size
+                    if size_samples.len() >= sample_size && user_batch_size.is_none() {
+                        let avg_row_size = size_samples.iter().sum::<usize>() / size_samples.len();
+                        current_batch_size = calculate_batch_size(avg_row_size, flush_threshold);
+                    }
+
+                    // 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);
+                }
+            }
+        }
+    } else {
+        return Err(MagnusError::new(
+            magnus::exception::type_error(),
+            "read_from must be an Enumerator",
+        ));
+    }
+
+    // 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))?;
+    }
+
+    Ok(())
+}
+
+#[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")
+                    })?;
+
+                    // 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",
+        ));
+    }
+
+    // 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))?;
+    }
+
+    Ok(())
+}
+
+fn create_writer(
+    ruby: &Ruby,
+    write_to: &Value,
+    schema: Arc<Schema>,
+    compression: Option<String>,
+) -> Result<WriterOutput, MagnusError> {
+    // Create writer properties with compression based on the option
+    let props = WriterProperties::builder()
+        .set_compression(match compression.as_deref() {
+            Some("none") | Some("uncompressed") => Compression::UNCOMPRESSED,
+            Some("snappy") => Compression::SNAPPY,
+            Some("gzip") => Compression::GZIP(GzipLevel::default()),
+            Some("lz4") => Compression::LZ4,
+            Some("zstd") => Compression::ZSTD(ZstdLevel::default()),
+            _ => Compression::UNCOMPRESSED,
+        })
+        .build();
+
+    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))?;
+        Ok(WriterOutput::File(writer))
+    } else {
+        // Create a temporary file to write to instead of directly to the IoLikeValue
+        let temp_file = NamedTempFile::new().map_err(|e| {
+            MagnusError::new(
+                magnus::exception::runtime_error(),
+                format!("Failed to create temporary file: {}", e),
+            )
+        })?;
+        let file: Box<dyn SendableWrite> = Box::new(temp_file.reopen().map_err(|e| {
+            MagnusError::new(
+                magnus::exception::runtime_error(),
+                format!("Failed to reopen temporary file: {}", e),
+            )
+        })?);
+        let writer =
+            ArrowWriter::try_new(file, schema, Some(props)).map_err(|e| ParquetErrorWrapper(e))?;
+        Ok(WriterOutput::TempFile(writer, temp_file))
+    }
+}
+
+// Helper function to copy temp file contents to IoLikeValue
+fn copy_temp_file_to_io_like(
+    temp_file: NamedTempFile,
+    io_like: IoLikeValue,
+) -> Result<(), MagnusError> {
+    let file = temp_file.reopen().map_err(|e| {
+        MagnusError::new(
+            magnus::exception::runtime_error(),
+            format!("Failed to reopen temporary file: {}", e),
+        )
+    })?;
+    let mut buf_reader = BufReader::new(file);
+    let mut buf_writer = BufWriter::new(io_like);
+
+    io::copy(&mut buf_reader, &mut buf_writer).map_err(|e| {
+        MagnusError::new(
+            magnus::exception::runtime_error(),
+            format!("Failed to copy temp file to io_like: {}", e),
+        )
+    })?;
+
+    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(())
+}

data/lib/parquet/version.rb

@@ -0,0 +1,3 @@
+module Parquet
+  VERSION = "0.2.12"
+end

data/lib/parquet.rb

@@ -0,0 +1,5 @@
+require_relative "parquet/version"
+require_relative "parquet/parquet"
+
+module Parquet
+end

data/lib/parquet.rbi

@@ -0,0 +1,113 @@
+# typed: true
+
+module Parquet
+  # Options:
+  #   - `input`: String, File, or IO object containing parquet data
+  #   - `result_type`: String specifying the output format
+  #                    ("hash" or "array" or :hash or :array)
+  #   - `columns`: When present, only the specified columns will be included in the output.
+  #                This is useful for reducing how much data is read and improving performance.
+  sig do
+    params(
+      input: T.any(String, File, StringIO, IO),
+      result_type: T.nilable(T.any(String, Symbol)),
+      columns: T.nilable(T::Array[String])
+    ).returns(T::Enumerator[T.any(T::Hash[String, T.untyped], T::Array[T.untyped])])
+  end
+  sig do
+    params(
+      input: T.any(String, File, StringIO, IO),
+      result_type: T.nilable(T.any(String, Symbol)),
+      columns: T.nilable(T::Array[String]),
+      blk: T.nilable(T.proc.params(row: T.any(T::Hash[String, T.untyped], T::Array[T.untyped])).void)
+    ).returns(NilClass)
+  end
+  def self.each_row(input, result_type: nil, columns: nil, &blk)
+  end
+
+  # Options:
+  #   - `input`: String, File, or IO object containing parquet data
+  #   - `result_type`: String specifying the output format
+  #                    ("hash" or "array" or :hash or :array)
+  #   - `columns`: When present, only the specified columns will be included in the output.
+  #   - `batch_size`: When present, specifies the number of rows per batch
+  sig do
+    params(
+      input: T.any(String, File, StringIO, IO),
+      result_type: T.nilable(T.any(String, Symbol)),
+      columns: T.nilable(T::Array[String]),
+      batch_size: T.nilable(Integer)
+    ).returns(T::Enumerator[T.any(T::Hash[String, T.untyped], T::Array[T.untyped])])
+  end
+  sig do
+    params(
+      input: T.any(String, File, StringIO, IO),
+      result_type: T.nilable(T.any(String, Symbol)),
+      columns: T.nilable(T::Array[String]),
+      batch_size: T.nilable(Integer),
+      blk:
+        T.nilable(T.proc.params(batch: T.any(T::Hash[String, T::Array[T.untyped]], T::Array[T::Array[T.untyped]])).void)
+    ).returns(NilClass)
+  end
+  def self.each_column(input, result_type: nil, columns: nil, batch_size: nil, &blk)
+  end
+
+  # Options:
+  #   - `read_from`: An Enumerator yielding arrays of values representing each row
+  #   - `schema`: Array of hashes specifying column names and types. Supported types:
+  #     - `int8`, `int16`, `int32`, `int64`
+  #     - `uint8`, `uint16`, `uint32`, `uint64`
+  #     - `float`, `double`
+  #     - `string`
+  #     - `binary`
+  #     - `boolean`
+  #     - `date32`
+  #     - `timestamp_millis`, `timestamp_micros`
+  #   - `write_to`: String path or IO object to write the parquet file to
+  #   - `batch_size`: Optional batch size for writing (defaults to 1000)
+  #   - `flush_threshold`: Optional memory threshold in bytes before flushing (defaults to 64MB)
+  #   - `compression`: Optional compression type to use (defaults to "zstd")
+  #                   Supported values: "none", "uncompressed", "snappy", "gzip", "lz4", "zstd"
+  #   - `sample_size`: Optional number of rows to sample for size estimation (defaults to 100)
+  sig do
+    params(
+      read_from: T::Enumerator[T::Array[T.untyped]],
+      schema: T::Array[T::Hash[String, String]],
+      write_to: T.any(String, IO),
+      batch_size: T.nilable(Integer),
+      flush_threshold: T.nilable(Integer),
+      compression: T.nilable(String),
+      sample_size: T.nilable(Integer)
+    ).void
+  end
+  def self.write_rows(read_from, schema:, write_to:, batch_size: nil, flush_threshold: nil, compression: nil, sample_size: nil)
+  end
+
+  # Options:
+  #   - `read_from`: An Enumerator yielding arrays of column batches
+  #   - `schema`: Array of hashes specifying column names and types. Supported types:
+  #     - `int8`, `int16`, `int32`, `int64`
+  #     - `uint8`, `uint16`, `uint32`, `uint64`
+  #     - `float`, `double`
+  #     - `string`
+  #     - `binary`
+  #     - `boolean`
+  #     - `date32`
+  #     - `timestamp_millis`, `timestamp_micros`
+  #     - Looks like [{"column_name" => {"type" => "date32", "format" => "%Y-%m-%d"}}, {"column_name" => "int8"}]
+  #   - `write_to`: String path or IO object to write the parquet file to
+  #   - `flush_threshold`: Optional memory threshold in bytes before flushing (defaults to 64MB)
+  #   - `compression`: Optional compression type to use (defaults to "zstd")
+  #                   Supported values: "none", "uncompressed", "snappy", "gzip", "lz4", "zstd"
+  sig do
+    params(
+      read_from: T::Enumerator[T::Array[T::Array[T.untyped]]],
+      schema: T::Array[T::Hash[String, String]],
+      write_to: T.any(String, IO),
+      flush_threshold: T.nilable(Integer),
+      compression: T.nilable(String)
+    ).void
+  end
+  def self.write_columns(read_from, schema:, write_to:, flush_threshold: nil, compression: nil)
+  end
+end