parquet 0.2.7 → 0.2.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c1ed4f490a4f03443598dbe1b0e110746052f613a4c5575f9b8e47c6e160bb40
-  data.tar.gz: 4db314d1707e633799e996c6fb777135ff0ea364a76c0a7d8fc5c429e2394d9f
+  metadata.gz: 2dea9b9b171070949497da37aff1888de71c0782e76968ba218f38e5dc2f1606
+  data.tar.gz: 74f4599b00a818cfca62d7fc162d02a87658da014ace361a76c998b718def9f2
 SHA512:
-  metadata.gz: b3f0a15cf467d030d3002c21bc6b64b6cd16e91e972b8de1e928abfd9bd373cfb5c4f77cdd1a6db7c620055e9657ec623866e0d8a0cb3a8e21a0c252bde3df87
-  data.tar.gz: 77f41921f5818051b597d2941688f6eca2a24d86333c58dec45d6e47e7161bfdd70e78f50a0f7ddd6cc99356c2b477451ab43adf9caa201501815c6b1a731d5c
+  metadata.gz: 209ca0339ccb11224501efc1d1adfed241097763475aa44e3997fce811123e9744f1bbfb1447e91decd1b020181b722ded94a6655630288db1f22e88aa8c09ae
+  data.tar.gz: a889e46dc8fca484043b3f1513ee6487b0f8caa8096c826cdbe4fa9ff2d6aa457c2d84e1bd95f7b05819e0ce2e33017a77a720aa331be7115cfa2ac470557a59

data/README.md

@@ -152,9 +152,16 @@ batches = [
 # Create an enumerator from the batches
 columns = batches.each
 
-# Write to a parquet file
+# Write to a parquet file with default ZSTD compression
 Parquet.write_columns(columns, schema: schema, write_to: "data.parquet")
 
+# Write to a parquet file with specific compression
+Parquet.write_columns(columns,
+  schema: schema,
+  write_to: "data.parquet",
+  compression: "snappy"  # Supported: "none", "uncompressed", "snappy", "gzip", "lz4", "zstd"
+)
+
 # Write to an IO object
 File.open("data.parquet", "wb") do |file|
   Parquet.write_columns(columns, schema: schema, write_to: file)

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

@@ -24,6 +24,7 @@ pub struct ParquetWriteArgs<'a> {
     pub write_to: Value,
     pub schema: Vec<SchemaField<'a>>,
     pub batch_size: Option<usize>,
+    pub compression: Option<String>,
 }
 
 pub trait SendableWrite: Send + Write {}

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

@@ -11,7 +11,11 @@ use magnus::{
     value::ReprValue,
     Error as MagnusError, RArray, Ruby, TryConvert, Value,
 };
-use parquet::arrow::ArrowWriter;
+use parquet::{
+    arrow::ArrowWriter,
+    basic::{Compression, GzipLevel, ZstdLevel},
+    file::properties::WriterProperties,
+};
 use tempfile::NamedTempFile;
 
 use crate::{
@@ -28,11 +32,12 @@ pub fn parse_parquet_write_args(args: &[Value]) -> Result<ParquetWriteArgs, Magn
     let parsed_args = scan_args::<(Value,), (), (), (), _, ()>(args)?;
     let (read_from,) = parsed_args.required;
 
-    let kwargs = get_kwargs::<_, (Value, Value), (Option<Option<usize>>,), ()>(
-        parsed_args.keywords,
-        &["schema", "write_to"],
-        &["batch_size"],
-    )?;
+    let kwargs =
+        get_kwargs::<_, (Value, Value), (Option<Option<usize>>, Option<Option<String>>), ()>(
+            parsed_args.keywords,
+            &["schema", "write_to"],
+            &["batch_size", "compression"],
+        )?;
 
     let schema_array = RArray::from_value(kwargs.required.0).ok_or_else(|| {
         MagnusError::new(
@@ -105,6 +110,7 @@ pub fn parse_parquet_write_args(args: &[Value]) -> Result<ParquetWriteArgs, Magn
         write_to: kwargs.required.1,
         schema,
         batch_size: kwargs.optional.0.flatten(),
+        compression: kwargs.optional.1.flatten(),
     })
 }
 
@@ -117,6 +123,7 @@ pub fn write_rows(args: &[Value]) -> Result<(), MagnusError> {
         write_to,
         schema,
         batch_size,
+        compression,
     } = parse_parquet_write_args(args)?;
 
     let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
@@ -158,7 +165,7 @@ pub fn write_rows(args: &[Value]) -> Result<(), MagnusError> {
     let arrow_schema = Arc::new(Schema::new(arrow_fields));
 
     // Create the writer
-    let mut writer = create_writer(&ruby, &write_to, arrow_schema.clone())?;
+    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
@@ -238,7 +245,8 @@ pub fn write_columns(args: &[Value]) -> Result<(), MagnusError> {
         read_from,
         write_to,
         schema,
-        batch_size: _, // Batch size is determined by the input
+        batch_size: _,
+        compression,
     } = parse_parquet_write_args(args)?;
 
     // Convert schema to Arrow schema
@@ -278,7 +286,7 @@ pub fn write_columns(args: &[Value]) -> Result<(), MagnusError> {
     let arrow_schema = Arc::new(Schema::new(arrow_fields));
 
     // Create the writer
-    let mut writer = create_writer(&ruby, &write_to, arrow_schema.clone())?;
+    let mut writer = create_writer(&ruby, &write_to, arrow_schema.clone(), compression)?;
 
     if read_from.is_kind_of(ruby.class_enumerator()) {
         loop {
@@ -360,12 +368,25 @@ 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, None).map_err(|e| ParquetErrorWrapper(e))?;
+            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
@@ -382,7 +403,7 @@ fn create_writer(
             )
         })?);
         let writer =
-            ArrowWriter::try_new(file, schema, None).map_err(|e| ParquetErrorWrapper(e))?;
+            ArrowWriter::try_new(file, schema, Some(props)).map_err(|e| ParquetErrorWrapper(e))?;
         Ok(WriterOutput::TempFile(writer, temp_file))
     }
 }

data/lib/parquet/version.rb

@@ -1,3 +1,3 @@
 module Parquet
-  VERSION = "0.2.7"
+  VERSION = "0.2.8"
 end

data/lib/parquet.rbi

@@ -65,15 +65,18 @@ module Parquet
   #     - `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)
+  #   - `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.untyped]],
       schema: T::Array[T::Hash[String, String]],
       write_to: T.any(String, IO),
-      batch_size: T.nilable(Integer)
+      batch_size: T.nilable(Integer),
+      compression: T.nilable(String)
     ).void
   end
-  def self.write_rows(read_from, schema:, write_to:, batch_size: nil)
+  def self.write_rows(read_from, schema:, write_to:, batch_size: nil, compression: nil)
   end
 
   # Options:
@@ -89,13 +92,16 @@ module Parquet
   #     - `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
+  #   - `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)
+      write_to: T.any(String, IO),
+      compression: T.nilable(String)
     ).void
   end
-  def self.write_columns(read_from, schema:, write_to:)
+  def self.write_columns(read_from, schema:, write_to:, compression: nil)
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: parquet
 version: !ruby/object:Gem::Version
-  version: 0.2.7
+  version: 0.2.8
 platform: ruby
 authors:
 - Nathan Jaremko