parquet 0.0.1

data/ext/parquet/src/ruby_reader.rs

@@ -0,0 +1,231 @@
+use magnus::{
+    value::{Opaque, ReprValue},
+    RClass, RString, Ruby, Value,
+};
+use std::io::{self, Read, Seek};
+use std::sync::OnceLock;
+
+static STRING_IO_CLASS: OnceLock<Opaque<RClass>> = OnceLock::new();
+
+const READ_BUFFER_SIZE: usize = 16 * 1024;
+
+/// A reader that can handle various Ruby input types (String, StringIO, IO-like objects)
+/// and provide a standard Read implementation for them.
+pub struct RubyReader<T> {
+    inner: T,
+    buffer: Option<Vec<u8>>,
+    offset: usize,
+    // Number of bytes that have been read into the buffer
+    // Used as an upper bound for offset
+    buffered_bytes: usize,
+}
+
+pub trait SeekableRead: std::io::Read + Seek {}
+impl SeekableRead for RubyReader<Value> {}
+impl SeekableRead for RubyReader<RString> {}
+
+pub fn build_ruby_reader<'a>(
+    ruby: &'a Ruby,
+    input: Value,
+) -> Result<Box<dyn SeekableRead>, magnus::Error> {
+    if RubyReader::is_string_io(ruby, &input) {
+        RubyReader::from_string_io(ruby, input)
+    } else if RubyReader::is_io_like(&input) {
+        RubyReader::from_io(input)
+    } else {
+        RubyReader::from_string_like(input)
+    }
+}
+
+impl Seek for RubyReader<Value> {
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        let seek_to = match pos {
+            io::SeekFrom::Start(offset) => {
+                // SEEK_SET - absolute position
+                offset as i64
+            }
+            io::SeekFrom::End(offset) => {
+                // SEEK_END - from end of stream
+                offset
+            }
+            io::SeekFrom::Current(offset) => {
+                // SEEK_CUR - relative to current
+                offset
+            }
+        };
+
+        let whence = match pos {
+            io::SeekFrom::Start(_) => 0,   // SEEK_SET
+            io::SeekFrom::End(_) => 2,     // SEEK_END
+            io::SeekFrom::Current(_) => 1, // SEEK_CUR
+        };
+
+        // Call Ruby's seek method
+        let _: u64 = self.inner.funcall("seek", (seek_to, whence)).unwrap();
+
+        // Get current position
+        let pos: u64 = self.inner.funcall("pos", ()).unwrap();
+
+        Ok(pos)
+    }
+}
+
+impl Seek for RubyReader<RString> {
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        match pos {
+            io::SeekFrom::Start(offset) => {
+                self.offset = offset as usize;
+            }
+            io::SeekFrom::End(offset) => {
+                self.offset = (self.inner.len() - offset as usize) as usize;
+            }
+            io::SeekFrom::Current(offset) => {
+                self.offset = (self.offset as i64 + offset) as usize;
+            }
+        }
+        Ok(self.offset as u64)
+    }
+}
+
+impl RubyReader<Value> {
+    fn from_io(input: Value) -> Result<Box<dyn SeekableRead>, magnus::Error> {
+        if Self::is_io_like(&input) {
+            Ok(Box::new(Self::from_io_like(input)))
+        } else {
+            Err(magnus::Error::new(
+                magnus::exception::type_error(),
+                "Input is not an IO-like object",
+            ))
+        }
+    }
+
+    fn is_io_like(input: &Value) -> bool {
+        input.respond_to("read", false).unwrap_or(false)
+    }
+
+    fn from_io_like(input: Value) -> Self {
+        Self {
+            inner: input,
+            buffer: Some(vec![0; READ_BUFFER_SIZE]),
+            offset: 0,
+            buffered_bytes: 0,
+        }
+    }
+
+    fn read_from_buffer(&mut self, to_buf: &mut [u8]) -> Option<io::Result<usize>> {
+        if let Some(from_buf) = &self.buffer {
+            // If the offset is within the buffered bytes, copy the remaining bytes to the output buffer
+            if self.offset < self.buffered_bytes {
+                let remaining = self.buffered_bytes - self.offset;
+                let copy_size = remaining.min(to_buf.len());
+                to_buf[..copy_size]
+                    .copy_from_slice(&from_buf[self.offset..self.offset + copy_size]);
+                self.offset += copy_size;
+                Some(Ok(copy_size))
+            } else {
+                None
+            }
+        } else {
+            None
+        }
+    }
+
+    fn read_from_ruby(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let buffer = self.buffer.as_mut().unwrap();
+        let result = self
+            .inner
+            .funcall::<_, _, RString>("read", (buffer.capacity(),))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        if result.is_nil() {
+            return Ok(0); // EOF
+        }
+
+        let bytes = unsafe { result.as_slice() };
+
+        // Update internal buffer
+        let bytes_len = bytes.len();
+        if bytes_len == 0 {
+            return Ok(0);
+        }
+
+        // Only copy what we actually read
+        buffer[..bytes_len].copy_from_slice(bytes);
+        self.buffered_bytes = bytes_len;
+
+        // Copy to output buffer
+        let copy_size = bytes_len.min(buf.len());
+        buf[..copy_size].copy_from_slice(&buffer[..copy_size]);
+        self.offset = copy_size;
+        Ok(copy_size)
+    }
+}
+
+impl RubyReader<RString> {
+    pub fn from_string_io(
+        ruby: &Ruby,
+        input: Value,
+    ) -> Result<Box<dyn SeekableRead>, magnus::Error> {
+        if !Self::is_string_io(ruby, &input) {
+            return Err(magnus::Error::new(
+                magnus::exception::type_error(),
+                "Input is not a StringIO",
+            ));
+        }
+
+        let string_content = input.funcall::<_, _, RString>("string", ()).unwrap();
+        Ok(Box::new(Self {
+            inner: string_content,
+            buffer: None,
+            offset: 0,
+            buffered_bytes: 0,
+        }))
+    }
+
+    fn is_string_io(ruby: &Ruby, input: &Value) -> bool {
+        let string_io_class = STRING_IO_CLASS.get_or_init(|| {
+            let class = RClass::from_value(ruby.eval("StringIO").unwrap()).unwrap();
+            Opaque::from(class)
+        });
+        input.is_kind_of(ruby.get_inner(*string_io_class))
+    }
+
+    fn from_string_like(input: Value) -> Result<Box<dyn SeekableRead>, magnus::Error> {
+        // Try calling `to_str`, and if that fails, try `to_s`
+        let string_content = input
+            .funcall::<_, _, RString>("to_str", ())
+            .or_else(|_| input.funcall::<_, _, RString>("to_s", ()))?;
+        Ok(Box::new(Self {
+            inner: string_content,
+            buffer: None,
+            offset: 0,
+            buffered_bytes: 0,
+        }))
+    }
+}
+
+impl Read for RubyReader<Value> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        if let Some(result) = self.read_from_buffer(buf) {
+            result
+        } else {
+            // If the buffer is empty, read from Ruby
+            self.read_from_ruby(buf)
+        }
+    }
+}
+
+impl Read for RubyReader<RString> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let string_buffer = unsafe { self.inner.as_slice() };
+        if self.offset >= string_buffer.len() {
+            return Ok(0); // EOF
+        }
+
+        let remaining = string_buffer.len() - self.offset;
+        let copy_size = remaining.min(buf.len());
+        buf[..copy_size].copy_from_slice(&string_buffer[self.offset..self.offset + copy_size]);
+        self.offset += copy_size;
+        Ok(copy_size)
+    }
+}

data/ext/parquet/src/utils.rs

@@ -0,0 +1,70 @@
+use magnus::{
+    scan_args::{get_kwargs, scan_args},
+    value::ReprValue,
+    Error, RString, Ruby, Symbol, Value,
+};
+
+fn parse_string_or_symbol(ruby: &Ruby, value: Value) -> Result<Option<String>, Error> {
+    if value.is_nil() {
+        Ok(None)
+    } else if value.is_kind_of(ruby.class_string()) {
+        RString::from_value(value)
+            .ok_or_else(|| Error::new(magnus::exception::type_error(), "Invalid string value"))?
+            .to_string()
+            .map(|s| Some(s))
+    } else if value.is_kind_of(ruby.class_symbol()) {
+        Symbol::from_value(value)
+            .ok_or_else(|| Error::new(magnus::exception::type_error(), "Invalid symbol value"))?
+            .funcall("to_s", ())
+            .map(|s| Some(s))
+    } else {
+        Err(Error::new(
+            magnus::exception::type_error(),
+            "Value must be a String or Symbol",
+        ))
+    }
+}
+
+#[derive(Debug)]
+pub struct ParquetArgs {
+    pub to_read: Value,
+    pub result_type: String,
+}
+
+/// Parse common arguments for CSV parsing
+pub fn parse_parquet_args(ruby: &Ruby, args: &[Value]) -> Result<ParquetArgs, Error> {
+    let parsed_args = scan_args::<(Value,), (), (), (), _, ()>(args)?;
+    let (to_read,) = parsed_args.required;
+
+    let kwargs =
+        get_kwargs::<_, (), (Option<Value>,), ()>(parsed_args.keywords, &[], &["result_type"])?;
+
+    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(ParquetArgs {
+        to_read,
+        result_type,
+    })
+}

data/lib/parquet/version.rb

@@ -0,0 +1,3 @@
+module Parquet
+  VERSION = "0.0.1"
+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,17 @@
+# typed: strict
+
+module Parquet
+  # Options:
+  #   - `input`: String specifying the input file
+  #   - `result_type`: String specifying the output format
+  #                    ("hash" or "array" or :hash or :array)
+  sig do
+    params(
+      input: T.any(String, IO),
+      result_type: T.nilable(T.any(String, Symbol)),
+      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))
+  end
+  def self.each_row(input, result_type: nil, &blk)
+  end
+end

metadata

@@ -0,0 +1,96 @@
+--- !ruby/object:Gem::Specification
+name: parquet
+version: !ruby/object:Gem::Version
+  version: 0.0.1
+platform: ruby
+authors:
+- Nathan Jaremko
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2025-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: rb_sys
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.9.39
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.9.39
+- !ruby/object:Gem::Dependency
+  name: rake-compiler
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 1.2.0
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 1.2.0
+description: |2
+      Parquet is a high-performance Parquet library for Ruby, written in Rust.
+      It wraps the official Apache Rust implementation to provide fast, correct Parquet parsing.
+email:
+- nathan@jaremko.ca
+executables: []
+extensions:
+- ext/parquet/extconf.rb
+extra_rdoc_files: []
+files:
+- Cargo.lock
+- Cargo.toml
+- Gemfile
+- LICENSE
+- README.md
+- Rakefile
+- ext/parquet/Cargo.toml
+- ext/parquet/extconf.rb
+- ext/parquet/src/header_cache.rs
+- ext/parquet/src/lib.rs
+- ext/parquet/src/reader.rs
+- ext/parquet/src/ruby_reader.rs
+- ext/parquet/src/utils.rs
+- lib/parquet.rb
+- lib/parquet.rbi
+- lib/parquet/version.rb
+homepage: https://github.com/njaremko/parquet
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/njaremko/parquet
+  source_code_uri: https://github.com/njaremko/parquet-ruby
+  readme_uri: https://github.com/njaremko/parquet-ruby/blob/main/README.md
+  changelog_uri: https://github.com/njaremko/parquet-ruby/blob/main/CHANGELOG.md
+  documentation_uri: https://www.rubydoc.info/gems/parquet
+  funding_uri: https://github.com/sponsors/njaremko
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.1.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.4.19
+signing_key:
+specification_version: 4
+summary: Parquet library for Ruby, written in Rust
+test_files: []