osv 0.3.13 → 0.3.14

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(|f| CowValue(f));
         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(|f| CowValue(f));
+
         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<'a, S: BuildHasher + Default> IntoValue for CsvRecord<'a, 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<'a> IntoValue for CowValue<'a> {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        self.0.into_value_with(handle)
+    }
+}

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

@@ -0,0 +1,175 @@
+use super::header_cache::StringCache;
+use super::parser::RecordParser;
+use magnus::{Error, Ruby};
+use std::{borrow::Cow, io::Read, thread};
+
+pub(crate) const READ_BUFFER_SIZE: usize = 16384;
+
+pub struct RecordReader<'a, T: RecordParser<'a>> {
+    inner: ReaderImpl<'a, T>,
+}
+
+enum ReaderImpl<'a, T: RecordParser<'a>> {
+    SingleThreaded {
+        reader: csv::Reader<Box<dyn Read + 'a>>,
+        headers: Vec<&'static str>,
+        null_string: Option<String>,
+        flexible_default: Option<Cow<'a, str>>,
+        string_record: csv::StringRecord,
+    },
+    MultiThreaded {
+        headers: Vec<&'static str>,
+        receiver: kanal::Receiver<T::Output>,
+        handle: Option<thread::JoinHandle<()>>,
+    },
+}
+
+impl<'a, T: RecordParser<'a>> RecordReader<'a, T> {
+    #[inline]
+    pub(crate) fn get_headers(
+        ruby: &Ruby,
+        reader: &mut csv::Reader<impl Read>,
+        has_headers: bool,
+    ) -> Result<Vec<String>, Error> {
+        let first_row = reader.headers().map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                format!("Failed to read headers: {e}"),
+            )
+        })?;
+
+        let mut headers = Vec::with_capacity(first_row.len());
+        if has_headers {
+            headers.extend(first_row.iter().map(String::from));
+        } else {
+            headers.extend((0..first_row.len()).map(|i| format!("c{i}")));
+        }
+        Ok(headers)
+    }
+
+    pub(crate) fn new_single_threaded(
+        reader: csv::Reader<Box<dyn Read + 'a>>,
+        headers: Vec<&'static str>,
+        null_string: Option<String>,
+        flexible_default: Option<&'a str>,
+    ) -> Self {
+        let headers_len = headers.len();
+        Self {
+            inner: ReaderImpl::SingleThreaded {
+                reader,
+                headers,
+                null_string,
+                flexible_default: flexible_default.map(|s| Cow::Borrowed(s)),
+                string_record: csv::StringRecord::with_capacity(READ_BUFFER_SIZE, headers_len),
+            },
+        }
+    }
+}
+
+impl<T: RecordParser<'static> + Send> RecordReader<'static, T> {
+    pub(crate) fn new_multi_threaded(
+        mut reader: csv::Reader<Box<dyn Read + Send + 'static>>,
+        headers: Vec<&'static str>,
+        buffer_size: usize,
+        null_string: Option<String>,
+        flexible_default: Option<&'static str>,
+        should_forget: bool,
+    ) -> Self {
+        let (sender, receiver) = kanal::bounded(buffer_size);
+        let headers_for_thread = headers.clone();
+
+        let handle = thread::spawn(move || {
+            let mut record =
+                csv::StringRecord::with_capacity(READ_BUFFER_SIZE, headers_for_thread.len());
+            while let Ok(true) = reader.read_record(&mut record) {
+                let row = T::parse(
+                    &headers_for_thread,
+                    &record,
+                    null_string.as_deref(),
+                    flexible_default.map(|s| Cow::Borrowed(s)),
+                );
+                if sender.send(row).is_err() {
+                    break;
+                }
+            }
+            if should_forget {
+                let file_to_forget = reader.into_inner();
+                std::mem::forget(file_to_forget);
+            }
+        });
+
+        Self {
+            inner: ReaderImpl::MultiThreaded {
+                headers,
+                receiver,
+                handle: Some(handle),
+            },
+        }
+    }
+}
+
+impl<'a, T: RecordParser<'a>> Iterator for RecordReader<'a, T> {
+    type Item = T::Output;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        match &mut self.inner {
+            ReaderImpl::MultiThreaded {
+                receiver, handle, ..
+            } => match receiver.recv() {
+                Ok(record) => Some(record),
+                Err(_) => {
+                    if let Some(handle) = handle.take() {
+                        let _ = handle.join();
+                    }
+                    None
+                }
+            },
+            ReaderImpl::SingleThreaded {
+                reader,
+                headers,
+                null_string,
+                flexible_default,
+                ref mut string_record,
+            } => match reader.read_record(string_record) {
+                Ok(true) => Some(T::parse(
+                    headers,
+                    &string_record,
+                    null_string.as_deref(),
+                    flexible_default.clone(),
+                )),
+                Ok(false) => None,
+                Err(_e) => None,
+            },
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // We can't know the exact size without reading the whole file
+        (0, None)
+    }
+}
+
+impl<'a, T: RecordParser<'a>> Drop for RecordReader<'a, T> {
+    #[inline]
+    fn drop(&mut self) {
+        match &mut self.inner {
+            ReaderImpl::MultiThreaded {
+                receiver,
+                handle,
+                headers,
+                ..
+            } => {
+                receiver.close();
+                if let Some(handle) = handle.take() {
+                    let _ = handle.join();
+                }
+                let _ = StringCache::clear(headers);
+            }
+            ReaderImpl::SingleThreaded { headers, .. } => {
+                let _ = StringCache::clear(headers);
+            }
+        }
+    }
+}

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

@@ -0,0 +1,181 @@
+use super::READ_BUFFER_SIZE;
+use magnus::{
+    value::{Opaque, ReprValue},
+    RClass, RString, Ruby, Value,
+};
+use std::io::{self, Read};
+use std::sync::OnceLock;
+
+static STRING_IO_CLASS: OnceLock<Opaque<RClass>> = OnceLock::new();
+
+/// 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<'a, T> {
+    #[allow(unused)]
+    ruby: &'a Ruby,
+    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 fn build_ruby_reader<'a>(
+    ruby: &'a Ruby,
+    input: Value,
+) -> Result<Box<dyn Read + 'a>, 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(ruby, input)
+    } else {
+        RubyReader::from_string_like(ruby, input)
+    }
+}
+
+impl<'a> RubyReader<'a, Value> {
+    fn from_io(ruby: &'a Ruby, input: Value) -> Result<Box<dyn Read + 'a>, magnus::Error> {
+        if Self::is_io_like(&input) {
+            Ok(Box::new(Self::from_io_like(ruby, 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(ruby: &'a Ruby, input: Value) -> Self {
+        Self {
+            ruby,
+            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<'a> RubyReader<'a, RString> {
+    pub fn from_string_io(
+        ruby: &'a Ruby,
+        input: Value,
+    ) -> Result<Box<dyn Read + 'a>, 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 {
+            ruby,
+            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(ruby: &'a Ruby, input: Value) -> Result<Box<dyn Read + 'a>, 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 {
+            ruby,
+            inner: string_content,
+            buffer: None,
+            offset: 0,
+            buffered_bytes: 0,
+        }))
+    }
+}
+
+impl<'a> Read for RubyReader<'a, 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<'a> Read for RubyReader<'a, 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/osv/src/reader.rs

@@ -1,4 +1,4 @@
-use crate::csv::{CsvRecord, RecordReaderBuilder};
+use crate::csv::{CowValue, CsvRecord, RecordReaderBuilder};
 use crate::utils::*;
 use csv::Trim;
 use magnus::value::ReprValue;
@@ -6,11 +6,12 @@ use magnus::{block::Yield, Error, KwArgs, RHash, Ruby, Symbol, Value};
 use std::collections::HashMap;
 use xxhash_rust::xxh3::Xxh3Builder;
 
-pub fn parse_csv(
+pub fn parse_csv<'a>(
     rb_self: Value,
     args: &[Value],
-) -> Result<Yield<Box<dyn Iterator<Item = CsvRecord<Xxh3Builder>>>>, Error> {
-    let ruby = unsafe { Ruby::get_unchecked() };
+) -> Result<Yield<Box<dyn Iterator<Item = CsvRecord<'static, Xxh3Builder>>>>, Error> {
+    let original = unsafe { Ruby::get_unchecked() };
+    let ruby: &'static Ruby = Box::leak(Box::new(original));
 
     let CsvArgs {
         to_read,
@@ -25,6 +26,10 @@ pub fn parse_csv(
         trim,
     } = parse_csv_args(&ruby, args)?;
 
+    let flexible_default: &'static Option<String> = Box::leak(Box::new(flexible_default));
+    let leaked_flexible_default: &'static Option<&str> =
+        Box::leak(Box::new(flexible_default.as_deref()));
+
     if !ruby.block_given() {
         return create_enumerator(EnumeratorArgs {
             rb_self,
@@ -36,7 +41,7 @@ pub fn parse_csv(
             buffer_size,
             result_type,
             flexible,
-            flexible_default,
+            flexible_default: leaked_flexible_default.as_deref(),
             trim: match trim {
                 Trim::All => Some("all".to_string()),
                 Trim::Headers => Some("headers".to_string()),
@@ -47,32 +52,32 @@ pub fn parse_csv(
     }
 
     let iter: Box<dyn Iterator<Item = CsvRecord<Xxh3Builder>>> = match result_type.as_str() {
-        "hash" => Box::new(
-            RecordReaderBuilder::<HashMap<&'static str, Option<String>, Xxh3Builder>>::new(
-                &ruby, to_read,
-            )
+        "hash" => {
+            let builder = RecordReaderBuilder::<
+                HashMap<&'static str, Option<CowValue<'static>>, Xxh3Builder>,
+            >::new(&ruby, to_read)
             .has_headers(has_headers)
             .flexible(flexible)
-            .flexible_default(flexible_default)
+            .flexible_default(flexible_default.as_deref())
             .trim(trim)
             .delimiter(delimiter)
             .quote_char(quote_char)
             .null_string(null_string)
-            .buffer(buffer_size)
-            .build()?
-            .map(CsvRecord::Map),
-        ),
+            .buffer(buffer_size);
+
+            Box::new(builder.build_threaded()?.map(CsvRecord::Map))
+        }
         "array" => Box::new(
-            RecordReaderBuilder::<Vec<Option<String>>>::new(&ruby, to_read)
+            RecordReaderBuilder::<Vec<Option<CowValue<'static>>>>::new(&ruby, to_read)
                 .has_headers(has_headers)
                 .flexible(flexible)
-                .flexible_default(flexible_default)
+                .flexible_default(flexible_default.as_deref())
                 .trim(trim)
                 .delimiter(delimiter)
                 .quote_char(quote_char)
                 .null_string(null_string)
                 .buffer(buffer_size)
-                .build()?
+                .build_threaded()?
                 .map(CsvRecord::Vec),
         ),
         _ => {
@@ -96,13 +101,13 @@ struct EnumeratorArgs {
     buffer_size: usize,
     result_type: String,
     flexible: bool,
-    flexible_default: Option<String>,
+    flexible_default: Option<&'static str>,
     trim: Option<String>,
 }
 
 fn create_enumerator(
     args: EnumeratorArgs,
-) -> Result<Yield<Box<dyn Iterator<Item = CsvRecord<Xxh3Builder>>>>, Error> {
+) -> Result<Yield<Box<dyn Iterator<Item = CsvRecord<'static, Xxh3Builder>>>>, Error> {
     let kwargs = RHash::new();
     kwargs.aset(Symbol::new("has_headers"), args.has_headers)?;
     kwargs.aset(

data/lib/osv/version.rb

@@ -1,3 +1,3 @@
 module OSV
-  VERSION = "0.3.13"
+  VERSION = "0.3.14"
 end