parquet 0.0.2 → 0.0.3

data/ext/parquet/src/utils.rs

@@ -26,14 +26,14 @@ fn parse_string_or_symbol(ruby: &Ruby, value: Value) -> Result<Option<String>, E
 }
 
 #[derive(Debug)]
-pub struct ParquetArgs {
+pub struct ParquetRowsArgs {
     pub to_read: Value,
     pub result_type: String,
     pub columns: Option<Vec<String>>,
 }
 
 /// Parse common arguments for CSV parsing
-pub fn parse_parquet_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetArgs, Error> {
+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;
 
@@ -67,9 +67,63 @@ pub fn parse_parquet_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetArgs, Er
         None => String::from("hash"),
     };
 
-    Ok(ParquetArgs {
+    Ok(ParquetRowsArgs {
         to_read,
         result_type,
         columns: kwargs.optional.1,
     })
 }
+
+#[derive(Debug)]
+pub struct ParquetColumnsArgs {
+    pub to_read: Value,
+    pub result_type: String,
+    pub columns: Option<Vec<String>>,
+    pub batch_size: Option<usize>,
+}
+
+/// Parse common arguments for CSV parsing
+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<Value>, Option<Vec<String>>, Option<usize>), ()>(
+        parsed_args.keywords,
+        &[],
+        &["result_type", "columns", "batch_size"],
+    )?;
+
+    let result_type = match kwargs
+        .optional
+        .0
+        .map(|value| parse_string_or_symbol(ruby, value))
+    {
+        Some(Ok(Some(parsed))) => match parsed.as_str() {
+            "hash" | "array" => parsed,
+            _ => {
+                return Err(Error::new(
+                    magnus::exception::runtime_error(),
+                    "result_type must be either 'hash' or 'array'",
+                ))
+            }
+        },
+        Some(Ok(None)) => String::from("hash"),
+        Some(Err(_)) => {
+            return Err(Error::new(
+                magnus::exception::type_error(),
+                "result_type must be a String or Symbol",
+            ))
+        }
+        None => String::from("hash"),
+    };
+
+    Ok(ParquetColumnsArgs {
+        to_read,
+        result_type,
+        columns: kwargs.optional.1,
+        batch_size: kwargs.optional.2,
+    })
+}

data/lib/parquet/version.rb

@@ -1,3 +1,3 @@
 module Parquet
-  VERSION = "0.0.2"
+  VERSION = "0.0.3"
 end

data/lib/parquet.rbi

@@ -1,19 +1,38 @@
 # typed: strict
 module Parquet
   # Options:
-  #   - `input`: String specifying the input file
+  #   - `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, IO),
+      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(T.any(Enumerator, T.untyped))
+    ).returns(T.any(Enumerator, 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),
+      blk:
+        T.nilable(T.proc.params(batch: T.any(T::Hash[String, T::Array[T.untyped]], T::Array[T::Array[T.untyped]])).void)
+    ).returns(T.any(Enumerator, NilClass))
+  end
+  def self.each_column(input, result_type: nil, columns: nil, batch_size: nil, &blk)
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: parquet
 version: !ruby/object:Gem::Version
-  version: 0.0.2
+  version: 0.0.3
 platform: ruby
 authors:
 - Nathan Jaremko
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-01-03 00:00:00.000000000 Z
+date: 2025-01-04 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys