osv 0.3.13 → 0.3.15

data/ext/osv/src/csv/builder.rs

@@ -1,19 +1,20 @@
 use super::{
     header_cache::{CacheError, StringCache},
     parser::RecordParser,
-    read_impl::ReadImpl,
-    reader::RecordReader,
-    READ_BUFFER_SIZE,
+    record_reader::{RecordReader, READ_BUFFER_SIZE},
+    ruby_reader::{build_ruby_reader, SeekableRead},
+    ForgottenFileHandle,
 };
 use flate2::read::GzDecoder;
-use magnus::{rb_sys::AsRawValue, value::ReprValue, Error as MagnusError, RString, Ruby, Value};
+use magnus::{rb_sys::AsRawValue, value::ReprValue, Error as MagnusError, Ruby, Value};
 use std::{
     fs::File,
     io::{self, BufReader, Read},
     marker::PhantomData,
+    mem::ManuallyDrop,
     os::fd::FromRawFd,
-    thread,
 };
+
 use thiserror::Error;
 
 pub(crate) static BUFFER_CHANNEL_SIZE: usize = 1024;
@@ -28,8 +29,6 @@ pub enum ReaderError {
     FileOpen(#[from] io::Error),
     #[error("Failed to intern headers: {0}")]
     HeaderIntern(#[from] CacheError),
-    #[error("Unsupported GzipReader")]
-    UnsupportedGzipReader,
     #[error("Ruby error: {0}")]
     Ruby(String),
 }
@@ -49,7 +48,7 @@ impl From<ReaderError> for MagnusError {
     }
 }
 
-pub struct RecordReaderBuilder<'a, T: RecordParser + Send + 'static> {
+pub struct RecordReaderBuilder<'a, T: RecordParser<'a> + Send> {
     ruby: &'a Ruby,
     to_read: Value,
     has_headers: bool,
@@ -58,12 +57,53 @@ pub struct RecordReaderBuilder<'a, T: RecordParser + Send + 'static> {
     null_string: Option<String>,
     buffer: usize,
     flexible: bool,
-    flexible_default: Option<String>,
+    flexible_default: Option<&'a str>,
     trim: csv::Trim,
     _phantom: PhantomData<T>,
 }
 
-impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
+impl<T: RecordParser<'static> + Send + 'static> RecordReaderBuilder<'static, T> {
+    fn build_multi_threaded(
+        self,
+        readable: Box<dyn Read + Send + 'static>,
+    ) -> Result<RecordReader<'static, T>, ReaderError> {
+        let flexible = self.flexible || self.flexible_default.is_some();
+        let mut reader = csv::ReaderBuilder::new()
+            .has_headers(self.has_headers)
+            .delimiter(self.delimiter)
+            .quote(self.quote_char)
+            .flexible(flexible)
+            .trim(self.trim)
+            .from_reader(readable);
+
+        let headers = RecordReader::<T>::get_headers(self.ruby, &mut reader, self.has_headers)?;
+        let static_headers = StringCache::intern_many(&headers)?;
+
+        Ok(RecordReader::new_multi_threaded(
+            reader,
+            static_headers,
+            self.buffer,
+            self.null_string,
+            self.flexible_default,
+        ))
+    }
+
+    pub fn build_threaded(self) -> Result<RecordReader<'static, T>, ReaderError> {
+        if self.to_read.is_kind_of(self.ruby.class_io()) {
+            let readable = self.handle_file_descriptor()?;
+            self.build_multi_threaded(readable)
+        } else if self.to_read.is_kind_of(self.ruby.class_string()) {
+            let readable = self.handle_file_path()?;
+            self.build_multi_threaded(readable)
+        } else {
+            let readable = build_ruby_reader(self.ruby, self.to_read)?;
+            let buffered_reader = BufReader::with_capacity(READ_BUFFER_SIZE, readable);
+            self.build_single_threaded(buffered_reader)
+        }
+    }
+}
+
+impl<'a, T: RecordParser<'a> + Send> RecordReaderBuilder<'a, T> {
     pub fn new(ruby: &'a Ruby, to_read: Value) -> Self {
         Self {
             ruby,
@@ -110,7 +150,7 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
         self
     }
 
-    pub fn flexible_default(mut self, flexible_default: Option<String>) -> Self {
+    pub fn flexible_default(mut self, flexible_default: Option<&'a str>) -> Self {
         self.flexible_default = flexible_default;
         self
     }
@@ -120,12 +160,6 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
         self
     }
 
-    fn handle_string_io(&self) -> Result<Box<dyn Read + Send + 'static>, ReaderError> {
-        let string: RString = self.to_read.funcall("string", ())?;
-        let content = string.to_string()?;
-        Ok(Box::new(std::io::Cursor::new(content)))
-    }
-
     fn handle_file_descriptor(&self) -> Result<Box<dyn Read + Send + 'static>, ReaderError> {
         let raw_value = self.to_read.as_raw();
         let fd = std::panic::catch_unwind(|| unsafe { rb_sys::rb_io_descriptor(raw_value) })
@@ -138,7 +172,11 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
         }
 
         let file = unsafe { File::from_raw_fd(fd) };
-        Ok(Box::new(BufReader::with_capacity(READ_BUFFER_SIZE, file)))
+        let forgotten = ForgottenFileHandle(ManuallyDrop::new(file));
+        Ok(Box::new(BufReader::with_capacity(
+            READ_BUFFER_SIZE,
+            forgotten,
+        )))
     }
 
     fn handle_file_path(&self) -> Result<Box<dyn Read + Send + 'static>, ReaderError> {
@@ -155,102 +193,12 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
         })
     }
 
-    fn get_reader(&self) -> Result<(Box<dyn Read + Send + 'static>, bool), ReaderError> {
-        let string_io: magnus::RClass = self.ruby.eval("StringIO")?;
-        let gzip_reader_class: magnus::RClass = self.ruby.eval("Zlib::GzipReader")?;
-
-        if self.to_read.is_kind_of(string_io) {
-            self.handle_string_io().map(|r| (r, false))
-        } else if self.to_read.is_kind_of(gzip_reader_class) {
-            Err(ReaderError::UnsupportedGzipReader)
-        } else if self.to_read.is_kind_of(self.ruby.class_io()) {
-            self.handle_file_descriptor().map(|r| (r, true))
-        } else {
-            self.handle_file_path().map(|r| (r, false))
-        }
-    }
-
-    fn get_single_threaded_reader(&self) -> Result<Box<dyn Read>, ReaderError> {
-        let string_io: magnus::RClass = self.ruby.eval("StringIO")?;
-        let gzip_reader_class: magnus::RClass = self.ruby.eval("Zlib::GzipReader")?;
-
-        if self.to_read.is_kind_of(string_io) {
-            self.handle_string_io().map(|r| -> Box<dyn Read> { r })
-        } else if self.to_read.is_kind_of(gzip_reader_class) {
-            Ok(Box::new(RubyReader::new(self.to_read)))
-        } else if self.to_read.is_kind_of(self.ruby.class_io()) {
-            self.handle_file_descriptor()
-                .map(|r| -> Box<dyn Read> { r })
-        } else {
-            self.handle_file_path().map(|r| -> Box<dyn Read> { r })
-        }
-    }
-
-    pub fn build(self) -> Result<RecordReader<T>, ReaderError> {
-        match self.get_reader() {
-            Ok((readable, should_forget)) => self.build_multi_threaded(readable, should_forget),
-            Err(_) => {
-                let readable = self.get_single_threaded_reader()?;
-                self.build_single_threaded(readable)
-            }
-        }
-    }
-
-    fn build_multi_threaded(
-        self,
-        readable: Box<dyn Read + Send + 'static>,
-        should_forget: bool,
-    ) -> Result<RecordReader<T>, ReaderError> {
-        let flexible = self.flexible || self.flexible_default.is_some();
-        let mut reader = csv::ReaderBuilder::new()
-            .has_headers(self.has_headers)
-            .delimiter(self.delimiter)
-            .quote(self.quote_char)
-            .flexible(flexible)
-            .trim(self.trim)
-            .from_reader(readable);
-
-        let headers = RecordReader::<T>::get_headers(self.ruby, &mut reader, self.has_headers)?;
-        let static_headers = StringCache::intern_many(&headers)?;
-        let headers_for_cleanup = static_headers.clone();
-
-        let (sender, receiver) = kanal::bounded(self.buffer);
-        let null_string = self.null_string.clone();
-
-        let flexible_default = self.flexible_default.clone();
-        let handle = thread::spawn(move || {
-            let mut record = csv::StringRecord::with_capacity(READ_BUFFER_SIZE, headers.len());
-            while let Ok(true) = reader.read_record(&mut record) {
-                let row = T::parse(
-                    &static_headers,
-                    &record,
-                    null_string.as_deref(),
-                    flexible_default.as_deref(),
-                );
-                if sender.send(row).is_err() {
-                    break;
-                }
-            }
-            if should_forget {
-                let file_to_forget = reader.into_inner();
-                std::mem::forget(file_to_forget);
-            }
-        });
-
-        Ok(RecordReader {
-            reader: ReadImpl::MultiThreaded {
-                headers: headers_for_cleanup,
-                receiver,
-                handle: Some(handle),
-            },
-        })
-    }
-
     fn build_single_threaded(
         self,
-        readable: Box<dyn Read>,
-    ) -> Result<RecordReader<T>, ReaderError> {
+        readable: BufReader<Box<dyn SeekableRead>>,
+    ) -> Result<RecordReader<'a, T>, ReaderError> {
         let flexible = self.flexible || self.flexible_default.is_some();
+
         let mut reader = csv::ReaderBuilder::new()
             .has_headers(self.has_headers)
             .delimiter(self.delimiter)
@@ -262,70 +210,11 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
         let headers = RecordReader::<T>::get_headers(self.ruby, &mut reader, self.has_headers)?;
         let static_headers = StringCache::intern_many(&headers)?;
 
-        Ok(RecordReader {
-            reader: ReadImpl::SingleThreaded {
-                reader,
-                headers: static_headers,
-                null_string: self.null_string,
-                flexible_default: self.flexible_default,
-            },
-        })
-    }
-}
-
-struct RubyReader {
-    inner: Value,
-    buffer: Option<Vec<u8>>,
-    offset: usize,
-}
-
-impl RubyReader {
-    fn new(inner: Value) -> Self {
-        Self {
-            inner,
-            buffer: None,
-            offset: 0,
-        }
-    }
-}
-
-// Read the entire inner into a vector and then read future reads from that vector with offset
-impl Read for RubyReader {
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        // If we have an existing buffer, read from it
-        if let Some(buffer) = self.buffer.as_ref() {
-            let remaining = buffer.len() - self.offset;
-            let copy_size = remaining.min(buf.len());
-            buf[..copy_size].copy_from_slice(&buffer[self.offset..self.offset + copy_size]);
-            self.offset += copy_size;
-            return Ok(copy_size);
-        }
-
-        // No buffer yet - read the entire content from Ruby
-        let result = self.inner.funcall::<_, _, Value>("read", ());
-        match result {
-            Ok(data) => {
-                if data.is_nil() {
-                    return Ok(0); // EOF
-                }
-
-                let string = RString::from_value(data).ok_or_else(|| {
-                    io::Error::new(io::ErrorKind::Other, "Failed to convert to RString")
-                })?;
-                let bytes = unsafe { string.as_slice() };
-
-                // Store the entire content in the buffer
-                self.buffer = Some(bytes.to_vec());
-                self.offset = 0;
-
-                // Read initial chunk
-                let copy_size = bytes.len().min(buf.len());
-                buf[..copy_size].copy_from_slice(&bytes[..copy_size]);
-                self.offset = copy_size;
-
-                Ok(copy_size)
-            }
-            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e.to_string())),
-        }
+        Ok(RecordReader::new_single_threaded(
+            reader,
+            static_headers,
+            self.null_string,
+            self.flexible_default,
+        ))
     }
 }

data/ext/osv/src/csv/mod.rs

@@ -1,11 +1,13 @@
 mod builder;
 mod header_cache;
 mod parser;
-pub mod read_impl;
-mod reader;
 mod record;
+mod record_reader;
+mod ruby_integration;
+mod ruby_reader;
 
 pub use builder::RecordReaderBuilder;
 pub(crate) use builder::BUFFER_CHANNEL_SIZE;
-pub(crate) use read_impl::READ_BUFFER_SIZE;
+pub use record::CowValue;
 pub use record::CsvRecord;
+pub use ruby_integration::*;

data/ext/osv/src/csv/parser.rs

@@ -1,18 +1,23 @@
+use std::borrow::Cow;
 use std::collections::HashMap;
 use std::hash::BuildHasher;
 
-pub trait RecordParser {
-    type Output;
+use super::CowValue;
+
+pub trait RecordParser<'a> {
+    type Output: 'a;
 
     fn parse(
         headers: &[&'static str],
         record: &csv::StringRecord,
         null_string: Option<&str>,
-        flexible_default: Option<&str>,
+        flexible_default: Option<Cow<'a, str>>,
     ) -> Self::Output;
 }
 
-impl<S: BuildHasher + Default> RecordParser for HashMap<&'static str, Option<String>, S> {
+impl<'a, S: BuildHasher + Default + 'a> RecordParser<'a>
+    for HashMap<&'static str, Option<CowValue<'a>>, S>
+{
     type Output = Self;
 
     #[inline]
@@ -20,19 +25,22 @@ impl<S: BuildHasher + Default> RecordParser for HashMap<&'static str, Option<Str
         headers: &[&'static str],
         record: &csv::StringRecord,
         null_string: Option<&str>,
-        flexible_default: Option<&str>,
+        flexible_default: Option<Cow<'a, str>>,
     ) -> Self::Output {
         let mut map = HashMap::with_capacity_and_hasher(headers.len(), S::default());
+
+        let shared_empty = Cow::Borrowed("");
+        let shared_default = flexible_default.map(CowValue);
         headers.iter().enumerate().for_each(|(i, &header)| {
             let value = record.get(i).map_or_else(
-                || flexible_default.map(ToString::to_string),
+                || shared_default.clone(),
                 |field| {
                     if null_string == Some(field) {
                         None
                     } else if field.is_empty() {
-                        Some(String::new())
+                        Some(CowValue(shared_empty.clone()))
                     } else {
-                        Some(field.into())
+                        Some(CowValue(Cow::Owned(field.to_string())))
                     }
                 },
             );
@@ -42,7 +50,7 @@ impl<S: BuildHasher + Default> RecordParser for HashMap<&'static str, Option<Str
     }
 }
 
-impl RecordParser for Vec<Option<String>> {
+impl<'a> RecordParser<'a> for Vec<Option<CowValue<'a>>> {
     type Output = Self;
 
     #[inline]
@@ -50,26 +58,94 @@ impl RecordParser for Vec<Option<String>> {
         headers: &[&'static str],
         record: &csv::StringRecord,
         null_string: Option<&str>,
-        flexible_default: Option<&str>,
+        flexible_default: Option<Cow<'a, str>>,
     ) -> Self::Output {
         let target_len = headers.len();
         let mut vec = Vec::with_capacity(target_len);
+
+        let shared_empty = Cow::Borrowed("");
+        let shared_default = flexible_default.map(CowValue);
+
         for field in record.iter() {
             let value = if Some(field) == null_string {
                 None
             } else if field.is_empty() {
-                Some(String::new())
+                Some(CowValue(shared_empty.clone()))
             } else {
-                Some(field.into())
+                Some(CowValue(Cow::Owned(field.to_string())))
             };
             vec.push(value);
         }
 
         if vec.len() < target_len {
-            if let Some(default) = flexible_default {
-                vec.resize_with(target_len, || Some(default.to_string()));
+            if let Some(default) = shared_default {
+                vec.resize_with(target_len, || Some(default.clone()));
             }
         }
         vec
     }
 }
+
+// impl<'a, S: BuildHasher + Default + 'a> RecordParser<'a>
+//     for HashMap<&'static str, Option<String>, S>
+// {
+//     type Output = Self;
+
+//     #[inline]
+//     fn parse(
+//         headers: &[&'static str],
+//         record: &csv::StringRecord,
+//         null_string: Option<&str>,
+//         flexible_default: Option<Cow<'a, str>>,
+//     ) -> Self::Output {
+//         let mut map = HashMap::with_capacity_and_hasher(headers.len(), S::default());
+//         headers.iter().enumerate().for_each(|(i, &header)| {
+//             let value = record.get(i).map_or_else(
+//                 || flexible_default.clone(),
+//                 |field| {
+//                     if null_string == Some(field) {
+//                         None
+//                     } else if field.is_empty() {
+//                         Some(String::new())
+//                     } else {
+//                         Some(field.into())
+//                     }
+//                 },
+//             );
+//             map.insert(header, value);
+//         });
+//         map
+//     }
+// }
+
+// impl<'a> RecordParser<'a> for Vec<Option<String>> {
+//     type Output = Self;
+
+//     #[inline]
+//     fn parse(
+//         headers: &[&'static str],
+//         record: &csv::StringRecord,
+//         null_string: Option<&str>,
+//         flexible_default: Option<Cow<'a, str>>,
+//     ) -> Self::Output {
+//         let target_len = headers.len();
+//         let mut vec = Vec::with_capacity(target_len);
+//         for field in record.iter() {
+//             let value = if Some(field) == null_string {
+//                 None
+//             } else if field.is_empty() {
+//                 Some(String::new())
+//             } else {
+//                 Some(field.into())
+//             };
+//             vec.push(value);
+//         }
+
+//         if vec.len() < target_len {
+//             if let Some(default) = flexible_default {
+//                 vec.resize_with(target_len, || Some(default.to_string()));
+//             }
+//         }
+//         vec
+//     }
+// }

data/ext/osv/src/csv/record.rs

@@ -1,17 +1,21 @@
 use magnus::{IntoValue, Ruby, Value};
-use std::{collections::HashMap, hash::BuildHasher};
+use std::{borrow::Cow, collections::HashMap, hash::BuildHasher};
 
 #[derive(Debug)]
-pub enum CsvRecord<S: BuildHasher + Default> {
-    Vec(Vec<Option<String>>),
-    Map(HashMap<&'static str, Option<String>, S>),
+pub enum CsvRecord<'a, S: BuildHasher + Default> {
+    Vec(Vec<Option<CowValue<'a>>>),
+    Map(HashMap<&'static str, Option<CowValue<'a>>, S>),
 }
 
-impl<S: BuildHasher + Default> IntoValue for CsvRecord<S> {
+impl<S: BuildHasher + Default> IntoValue for CsvRecord<'_, S> {
     #[inline]
     fn into_value_with(self, handle: &Ruby) -> Value {
         match self {
-            CsvRecord::Vec(vec) => vec.into_value_with(handle),
+            CsvRecord::Vec(vec) => {
+                let ary = handle.ary_new_capa(vec.len());
+                vec.into_iter().try_for_each(|v| ary.push(v)).unwrap();
+                ary.into_value_with(handle)
+            }
             CsvRecord::Map(map) => {
                 // Pre-allocate the hash with the known size
                 let hash = handle.hash_new_capa(map.len());
@@ -23,3 +27,12 @@ impl<S: BuildHasher + Default> IntoValue for CsvRecord<S> {
         }
     }
 }
+
+#[derive(Debug, Clone)]
+pub struct CowValue<'a>(pub Cow<'a, str>);
+
+impl IntoValue for CowValue<'_> {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        self.0.into_value_with(handle)
+    }
+}