parquet 0.0.1 → 0.0.3

data/ext/parquet/src/ruby_integration.rs

@@ -0,0 +1,77 @@
+use std::{
+    fs::File,
+    io::{BufReader, SeekFrom},
+    mem::ManuallyDrop,
+};
+
+use bytes::Bytes;
+use magnus::{value::Opaque, Ruby, Value};
+use parquet::{
+    errors::ParquetError,
+    file::reader::{ChunkReader, Length},
+};
+use std::io::Read;
+
+use crate::ruby_reader::{build_ruby_reader, SeekableRead};
+
+const READ_BUFFER_SIZE: usize = 16 * 1024;
+
+pub struct SeekableRubyValue(pub Opaque<Value>);
+
+impl Length for SeekableRubyValue {
+    fn len(&self) -> u64 {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        let current_pos = reader.seek(SeekFrom::Current(0)).unwrap();
+        let file_len = reader.seek(SeekFrom::End(0)).unwrap();
+        reader.seek(SeekFrom::Start(current_pos)).unwrap();
+        file_len
+    }
+}
+
+impl ChunkReader for SeekableRubyValue {
+    type T = BufReader<Box<dyn SeekableRead>>;
+
+    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        reader.seek(SeekFrom::Start(start))?;
+        Ok(BufReader::with_capacity(READ_BUFFER_SIZE, reader))
+    }
+
+    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut buffer = Vec::with_capacity(length);
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        reader.seek(SeekFrom::Start(start))?;
+        let read = reader.take(length as _).read_to_end(&mut buffer)?;
+
+        if read != length {
+            return Err(ParquetError::EOF(format!(
+                "Expected to read {} bytes, read only {}",
+                length, read
+            )));
+        }
+        Ok(buffer.into())
+    }
+}
+
+pub struct ForgottenFileHandle(pub ManuallyDrop<File>);
+
+impl Length for ForgottenFileHandle {
+    fn len(&self) -> u64 {
+        self.0.len()
+    }
+}
+
+impl ChunkReader for ForgottenFileHandle {
+    type T = BufReader<File>;
+
+    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
+        self.0.get_read(start)
+    }
+
+    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
+        self.0.get_bytes(start, length)
+    }
+}

data/ext/parquet/src/ruby_reader.rs

@@ -2,30 +2,24 @@ use magnus::{
     value::{Opaque, ReprValue},
     RClass, RString, Ruby, Value,
 };
-use std::io::{self, Read, Seek};
+use std::io::{self, Read, Seek, SeekFrom, Write};
 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,
+pub fn build_ruby_reader(
+    ruby: &Ruby,
     input: Value,
 ) -> Result<Box<dyn SeekableRead>, magnus::Error> {
     if RubyReader::is_string_io(ruby, &input) {
@@ -39,49 +33,48 @@ pub fn build_ruby_reader<'a>(
 
 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, offset) = match pos {
+            SeekFrom::Start(i) => (0, i as i64),
+            SeekFrom::Current(i) => (1, i),
+            SeekFrom::End(i) => (2, i),
         };
 
-        let whence = match pos {
-            io::SeekFrom::Start(_) => 0,   // SEEK_SET
-            io::SeekFrom::End(_) => 2,     // SEEK_END
-            io::SeekFrom::Current(_) => 1, // SEEK_CUR
-        };
+        let new_position = self
+            .inner
+            .funcall("seek", (offset, whence))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        Ok(new_position)
+    }
+}
 
-        // Call Ruby's seek method
-        let _: u64 = self.inner.funcall("seek", (seek_to, whence)).unwrap();
+impl Write for RubyReader<Value> {
+    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
+        let ruby_bytes = RString::from_slice(buf);
 
-        // Get current position
-        let pos: u64 = self.inner.funcall("pos", ()).unwrap();
+        let bytes_written = self
+            .inner
+            .funcall::<_, _, usize>("write", (ruby_bytes,))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
 
-        Ok(pos)
+        Ok(bytes_written)
+    }
+
+    fn flush(&mut self) -> Result<(), io::Error> {
+        self.inner
+            .funcall::<_, _, Value>("flush", ())
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        Ok(())
     }
 }
 
 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;
-            }
+            io::SeekFrom::Start(offset) => self.offset = offset as usize,
+            io::SeekFrom::Current(offset) => self.offset = (self.offset as i64 + offset) as usize,
+            io::SeekFrom::End(offset) => self.offset = self.inner.len() - offset as usize,
         }
         Ok(self.offset as u64)
     }
@@ -106,59 +99,9 @@ impl RubyReader<Value> {
     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> {
@@ -176,9 +119,7 @@ impl RubyReader<RString> {
         let string_content = input.funcall::<_, _, RString>("string", ()).unwrap();
         Ok(Box::new(Self {
             inner: string_content,
-            buffer: None,
             offset: 0,
-            buffered_bytes: 0,
         }))
     }
 
@@ -197,21 +138,21 @@ impl RubyReader<RString> {
             .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)
-        }
+    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
+        let bytes = self
+            .inner
+            .funcall::<_, _, RString>("read", (buf.len(),))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        buf.write_all(unsafe { bytes.as_slice() })?;
+
+        Ok(bytes.len())
     }
 }