osv 0.3.3 → 0.3.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 33c644fac6e61f8bf3b9f11e646d6706017ece2a386df03a568b5ed06fd91e2a
-  data.tar.gz: 0d977fb3a7eaf867663feb76161eb2d3fbe42e758523ef8cafbff51a9acfef0d
+  metadata.gz: 3b8a537d24e23250ff18da51bc59f9b4329b02e2fec9b881d4ee203a766e5514
+  data.tar.gz: dbdb8b1accd5897df7079adb98cc2ae50939f8ffeebf80c5095a06412dc69699
 SHA512:
-  metadata.gz: b6fe382c005837fbfc705bd02b1859fdfa9fa9f955c15f7b13ffacb97f1d5dc0714288c23a7e2431a21aada6aacb1952b33ae8d14acd019c96ed719f5580c02d
-  data.tar.gz: 4e7c3d783f23af709505a4c182d1ee6744c9f28226aecdf68c532ada9ffa653858bc730bfd879cfa94d8e8f7a1280c49380bb5b6f74ed16b7ac776318aa53d1b
+  metadata.gz: d32dc649748d62092414047bc5f7666cab3b1a4cbfb58a00edb2c0ec1a634d375e3c522ac716b57a6bb37292436433d867042b922383366a8da44a3500a6e2a4
+  data.tar.gz: deec4d1433d6da2b8242fae9c3ff2f17e5c56143f0c39d1f0ceaadd3351ef05641bb03f219a02c47b50a3e36c41048894963da41a0b1729cd153f579277cf448

data/Cargo.lock

@@ -273,6 +273,7 @@ dependencies = [
  "rb-sys",
  "serde",
  "serde_magnus",
+ "thiserror",
 ]
 
 [[package]]
@@ -436,6 +437,26 @@ version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
 
+[[package]]
+name = "thiserror"
+version = "1.0.69"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
+dependencies = [
+ "thiserror-impl",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "1.0.69"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.14"

data/Gemfile

@@ -2,6 +2,13 @@ source "https://rubygems.org"
 
 gem "rb_sys", "~> 0.9.56"
 gem "rake"
-gem "rake-compiler", "1.2.0"
+gem "csv"
 
-gem "minitest", "~> 5.0", group: :test
+# Use local version of osv
+gemspec
+
+group :development, :test do
+  gem "minitest", "~> 5.0"
+  gem "benchmark-ips", "~> 2.12"
+  gem "fastcsv", "~> 0.0.7"
+end

data/README.md

@@ -111,3 +111,93 @@ OSV.for_each(data) { |row| puts row["name"] }
 ## Performance
 
 This library is faster than the standard Ruby CSV library, and is comparable to the fastest CSV parser gems I've used.
+
+Here's some unscientific benchmarks. You can find the code in the [benchmark/comparison_benchmark.rb](benchmark/comparison_benchmark.rb) file.
+
+### 10,000 lines
+
+```
+Benchmarking with 10001 lines of data
+
+ruby 3.3.3 (2024-06-12 revision f1c7b6f435) [arm64-darwin23]
+Warming up --------------------------------------
+   OSV - Hash output     6.000 i/100ms
+   CSV - Hash output     1.000 i/100ms
+  OSV - Array output    18.000 i/100ms
+  CSV - Array output     2.000 i/100ms
+FastCSV - Array output
+                         9.000 i/100ms
+      OSV - StringIO     7.000 i/100ms
+      CSV - StringIO     1.000 i/100ms
+  FastCSV - StringIO    20.000 i/100ms
+       OSV - Gzipped     6.000 i/100ms
+       CSV - Gzipped     1.000 i/100ms
+Calculating -------------------------------------
+   OSV - Hash output     73.360 (± 4.1%) i/s   (13.63 ms/i) -    366.000 in   5.000390s
+   CSV - Hash output     11.937 (±25.1%) i/s   (83.78 ms/i) -     52.000 in   5.036297s
+  OSV - Array output    189.738 (± 8.4%) i/s    (5.27 ms/i) -    954.000 in   5.071018s
+  CSV - Array output     25.471 (±11.8%) i/s   (39.26 ms/i) -    120.000 in   5.015289s
+FastCSV - Array output
+                         97.867 (± 2.0%) i/s   (10.22 ms/i) -    495.000 in   5.060957s
+      OSV - StringIO     80.784 (± 6.2%) i/s   (12.38 ms/i) -    406.000 in   5.046696s
+      CSV - StringIO     15.872 (± 0.0%) i/s   (63.01 ms/i) -     80.000 in   5.043361s
+  FastCSV - StringIO    200.511 (± 2.0%) i/s    (4.99 ms/i) -      1.020k in   5.088592s
+       OSV - Gzipped     55.220 (±12.7%) i/s   (18.11 ms/i) -    258.000 in   5.030928s
+       CSV - Gzipped     12.591 (±15.9%) i/s   (79.42 ms/i) -     59.000 in   5.039709s
+
+Comparison:
+  FastCSV - StringIO:      200.5 i/s
+  OSV - Array output:      189.7 i/s - same-ish: difference falls within error
+FastCSV - Array output:       97.9 i/s - 2.05x  slower
+      OSV - StringIO:       80.8 i/s - 2.48x  slower
+   OSV - Hash output:       73.4 i/s - 2.73x  slower
+       OSV - Gzipped:       55.2 i/s - 3.63x  slower
+  CSV - Array output:       25.5 i/s - 7.87x  slower
+      CSV - StringIO:       15.9 i/s - 12.63x  slower
+       CSV - Gzipped:       12.6 i/s - 15.92x  slower
+   CSV - Hash output:       11.9 i/s - 16.80x  slower
+```
+
+### 1,000,000 lines
+
+```
+Benchmarking with 1000001 lines of data
+
+ruby 3.3.3 (2024-06-12 revision f1c7b6f435) [arm64-darwin23]
+Warming up --------------------------------------
+   OSV - Hash output     1.000 i/100ms
+   CSV - Hash output     1.000 i/100ms
+  OSV - Array output     1.000 i/100ms
+  CSV - Array output     1.000 i/100ms
+FastCSV - Array output
+                         1.000 i/100ms
+      OSV - StringIO     1.000 i/100ms
+      CSV - StringIO     1.000 i/100ms
+  FastCSV - StringIO     1.000 i/100ms
+       OSV - Gzipped     1.000 i/100ms
+       CSV - Gzipped     1.000 i/100ms
+Calculating -------------------------------------
+   OSV - Hash output      0.578 (± 0.0%) i/s     (1.73 s/i) -      3.000 in   5.287845s
+   CSV - Hash output      0.117 (± 0.0%) i/s     (8.57 s/i) -      1.000 in   8.571770s
+  OSV - Array output      1.142 (± 0.0%) i/s  (875.97 ms/i) -      5.000 in   5.234694s
+  CSV - Array output      0.235 (± 0.0%) i/s     (4.25 s/i) -      2.000 in   8.561144s
+FastCSV - Array output
+                          0.768 (± 0.0%) i/s     (1.30 s/i) -      4.000 in   6.924574s
+      OSV - StringIO      0.522 (± 0.0%) i/s     (1.91 s/i) -      3.000 in   5.803969s
+      CSV - StringIO      0.132 (± 0.0%) i/s     (7.59 s/i) -      1.000 in   7.593243s
+  FastCSV - StringIO      1.039 (± 0.0%) i/s  (962.53 ms/i) -      6.000 in   5.806644s
+       OSV - Gzipped      0.437 (± 0.0%) i/s     (2.29 s/i) -      3.000 in   6.885125s
+       CSV - Gzipped      0.115 (± 0.0%) i/s     (8.68 s/i) -      1.000 in   8.684069s
+
+Comparison:
+  OSV - Array output:        1.1 i/s
+  FastCSV - StringIO:        1.0 i/s - 1.10x  slower
+FastCSV - Array output:        0.8 i/s - 1.49x  slower
+   OSV - Hash output:        0.6 i/s - 1.98x  slower
+      OSV - StringIO:        0.5 i/s - 2.19x  slower
+       OSV - Gzipped:        0.4 i/s - 2.61x  slower
+  CSV - Array output:        0.2 i/s - 4.86x  slower
+      CSV - StringIO:        0.1 i/s - 8.67x  slower
+   CSV - Hash output:        0.1 i/s - 9.79x  slower
+       CSV - Gzipped:        0.1 i/s - 9.91x  slower
+```

data/ext/osv/Cargo.toml

@@ -14,3 +14,4 @@ magnus = { version = "0.7", features = ["rb-sys"] }
 rb-sys = "^0.9"
 serde = { version = "1.0", features = ["derive"] }
 serde_magnus = "0.8.1"
+thiserror = "1.0"

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

@@ -1,11 +1,50 @@
 use super::{
-    header_cache::StringCache,
+    header_cache::{CacheError, StringCache},
     parser::RecordParser,
-    reader::{ReadImpl, RecordReader},
+    read_impl::ReadImpl,
+    reader::RecordReader,
 };
 use flate2::read::GzDecoder;
-use magnus::{rb_sys::AsRawValue, value::ReprValue, Error, RString, Ruby, Value};
-use std::{fs::File, io::Read, marker::PhantomData, os::fd::FromRawFd, thread};
+use magnus::{rb_sys::AsRawValue, value::ReprValue, Error as MagnusError, RString, Ruby, Value};
+use std::{
+    fs::File,
+    io::{self, Read},
+    marker::PhantomData,
+    os::fd::FromRawFd,
+    thread,
+};
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum ReaderError {
+    #[error("Failed to get file descriptor: {0}")]
+    FileDescriptor(String),
+    #[error("Invalid file descriptor")]
+    InvalidFileDescriptor,
+    #[error("Failed to open file: {0}")]
+    FileOpen(#[from] io::Error),
+    #[error("Failed to intern headers: {0}")]
+    HeaderIntern(#[from] CacheError),
+    #[error("Unsupported GzipReader")]
+    UnsupportedGzipReader,
+    #[error("Ruby error: {0}")]
+    Ruby(String),
+}
+
+impl From<MagnusError> for ReaderError {
+    fn from(err: MagnusError) -> Self {
+        Self::Ruby(err.to_string())
+    }
+}
+
+impl From<ReaderError> for MagnusError {
+    fn from(err: ReaderError) -> Self {
+        MagnusError::new(
+            Ruby::get().unwrap().exception_runtime_error(),
+            err.to_string(),
+        )
+    }
+}
 
 pub struct RecordReaderBuilder<'a, T: RecordParser + Send + 'static> {
     ruby: &'a Ruby,
@@ -57,36 +96,83 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
         self
     }
 
-    fn get_reader(&self) -> Result<Box<dyn Read + Send + 'static>, Error> {
+    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) })
+            .map_err(|_| {
+                ReaderError::FileDescriptor("Failed to get file descriptor".to_string())
+            })?;
+
+        if fd < 0 {
+            return Err(ReaderError::InvalidFileDescriptor);
+        }
+
+        let file = unsafe { File::from_raw_fd(fd) };
+        Ok(Box::new(file))
+    }
+
+    fn handle_file_path(&self) -> Result<Box<dyn Read + Send + 'static>, ReaderError> {
+        let path = self.to_read.to_r_string()?.to_string()?;
+        let file = File::open(&path)?;
+
+        Ok(if path.ends_with(".gz") {
+            Box::new(GzDecoder::new(file))
+        } else {
+            Box::new(file)
+        })
+    }
+
+    fn get_reader(&self) -> Result<Box<dyn Read + Send + 'static>, 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) {
-            let string: RString = self.to_read.funcall("string", ())?;
-            let content = string.to_string()?;
-            Ok(Box::new(std::io::Cursor::new(content)))
+            self.handle_string_io()
+        } 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()) {
-            let fd = unsafe { rb_sys::rb_io_descriptor(self.to_read.as_raw()) };
-            let file = unsafe { File::from_raw_fd(fd) };
-            Ok(Box::new(file))
+            self.handle_file_descriptor()
         } else {
-            let path = self.to_read.to_r_string()?.to_string()?;
-            let file = std::fs::File::open(&path).map_err(|e| {
-                Error::new(
-                    self.ruby.exception_runtime_error(),
-                    format!("Failed to open file: {e}"),
-                )
-            })?;
-            if path.ends_with(".gz") {
-                let file = GzDecoder::new(file);
-                Ok(Box::new(file))
-            } else {
-                Ok(Box::new(file))
+            self.handle_file_path()
+        }
+    }
+
+    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) => self.build_multi_threaded(readable),
+            Err(_) => {
+                let readable = self.get_single_threaded_reader()?;
+                self.build_single_threaded(readable)
             }
         }
     }
 
-    pub fn build(self) -> Result<RecordReader<T>, Error> {
-        let readable = self.get_reader()?;
+    fn build_multi_threaded(
+        self,
+        readable: Box<dyn Read + Send + 'static>,
+    ) -> Result<RecordReader<T>, ReaderError> {
         let mut reader = csv::ReaderBuilder::new()
             .has_headers(self.has_headers)
             .delimiter(self.delimiter)
@@ -94,17 +180,12 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
             .from_reader(readable);
 
         let headers = RecordReader::<T>::get_headers(self.ruby, &mut reader, self.has_headers)?;
-        let null_string = self.null_string;
-
-        let static_headers = StringCache::intern_many(&headers).map_err(|e| {
-            Error::new(
-                self.ruby.exception_runtime_error(),
-                format!("Failed to intern headers: {e}"),
-            )
-        })?;
+        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 handle = thread::spawn(move || {
             let mut record = csv::StringRecord::new();
             while let Ok(true) = reader.read_record(&mut record) {
@@ -125,4 +206,58 @@ impl<'a, T: RecordParser + Send + 'static> RecordReaderBuilder<'a, T> {
             },
         })
     }
+
+    fn build_single_threaded(
+        self,
+        readable: Box<dyn Read>,
+    ) -> Result<RecordReader<T>, ReaderError> {
+        let mut reader = csv::ReaderBuilder::new()
+            .has_headers(self.has_headers)
+            .delimiter(self.delimiter)
+            .quote(self.quote_char)
+            .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 {
+            reader: ReadImpl::SingleThreaded {
+                reader,
+                headers: static_headers,
+                null_string: self.null_string,
+            },
+        })
+    }
+}
+
+struct RubyReader {
+    inner: Value,
+}
+
+impl RubyReader {
+    fn new(inner: Value) -> Self {
+        Self { inner }
+    }
+}
+
+impl Read for RubyReader {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let result = self.inner.funcall::<_, _, Value>("read", (buf.len(),));
+        match result {
+            Ok(data) => {
+                if data.is_nil() {
+                    return Ok(0);
+                }
+
+                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() };
+                let len = bytes.len().min(buf.len());
+                buf[..len].copy_from_slice(&bytes[..len]);
+                Ok(len)
+            }
+            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e.to_string())),
+        }
+    }
 }

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

@@ -4,22 +4,29 @@
 ///
 /// Note: Performance testing on macOS showed minimal speed improvements,
 /// so this optimization could be removed if any issues arise.
-
-
 use std::{
     collections::HashMap,
     sync::{atomic::AtomicU32, LazyLock, Mutex},
 };
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum CacheError {
+    #[error("Failed to acquire lock: {0}")]
+    LockError(String),
+}
 
 static STRING_CACHE: LazyLock<Mutex<HashMap<&'static str, AtomicU32>>> =
     LazyLock::new(|| Mutex::new(HashMap::with_capacity(100)));
 
-pub struct StringCache {}
+pub struct StringCache;
 
 impl StringCache {
     #[allow(dead_code)]
-    pub fn intern(string: String) -> Result<&'static str, String> {
-        let mut cache = STRING_CACHE.lock().map_err(|e| e.to_string())?;
+    pub fn intern(string: String) -> Result<&'static str, CacheError> {
+        let mut cache = STRING_CACHE
+            .lock()
+            .map_err(|e| CacheError::LockError(e.to_string()))?;
 
         if let Some((&existing, count)) = cache.get_key_value(string.as_str()) {
             count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
@@ -31,33 +38,36 @@ impl StringCache {
         }
     }
 
-    pub fn intern_many(strings: &[String]) -> Result<Vec<&'static str>, String> {
-        let mut cache = STRING_CACHE.lock().map_err(|e| e.to_string())?;
-        let mut result = Vec::with_capacity(strings.len());
+    pub fn intern_many(strings: &[String]) -> Result<Vec<&'static str>, CacheError> {
+        let mut cache = STRING_CACHE
+            .lock()
+            .map_err(|e| CacheError::LockError(e.to_string()))?;
 
+        let mut result = Vec::with_capacity(strings.len());
         for string in strings {
-            let static_str: &'static str =
-                if let Some((&existing, count)) = cache.get_key_value(string.as_str()) {
-                    count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-                    existing
-                } else {
-                    let leaked = Box::leak(string.clone().into_boxed_str());
-                    cache.insert(leaked, AtomicU32::new(1));
-                    leaked
-                };
-            result.push(static_str);
+            if let Some((&existing, count)) = cache.get_key_value(string.as_str()) {
+                count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                result.push(existing);
+            } else {
+                let leaked = Box::leak(string.clone().into_boxed_str());
+                cache.insert(leaked, AtomicU32::new(1));
+                result.push(leaked);
+            }
         }
-
         Ok(result)
     }
 
-    pub fn clear(headers: &[&'static str]) -> Result<(), String> {
-        let cache = STRING_CACHE.lock().map_err(|e| e.to_string())?;
+    pub fn clear(headers: &[&'static str]) -> Result<(), CacheError> {
+        let mut cache = STRING_CACHE
+            .lock()
+            .map_err(|e| CacheError::LockError(e.to_string()))?;
 
         for header in headers {
             if let Some(count) = cache.get(header) {
-                let remaining = count.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
-                if remaining == 0 {
+                // Returns the previous value of the counter
+                let was = count.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
+                if was == 1 {
+                    cache.remove(header);
                     let ptr = *header as *const str as *mut str;
                     unsafe {
                         let _ = Box::from_raw(ptr);

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

@@ -1,6 +1,7 @@
 mod builder;
 mod header_cache;
 mod parser;
+pub mod read_impl;
 mod reader;
 mod record;
 

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

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 pub trait RecordParser {
     type Output;
+
     fn parse(
         headers: &[&'static str],
         record: &csv::StringRecord,
@@ -11,41 +12,58 @@ pub trait RecordParser {
 
 impl RecordParser for HashMap<&'static str, Option<String>> {
     type Output = Self;
+
+    #[inline]
     fn parse(
         headers: &[&'static str],
         record: &csv::StringRecord,
         null_string: &str,
     ) -> Self::Output {
         let mut map = HashMap::with_capacity(headers.len());
-        for (header, field) in headers.iter().zip(record.iter()) {
-            map.insert(
-                *header,
-                if field == null_string {
-                    None
-                } else {
-                    Some(field.to_string())
-                },
-            );
-        }
+        headers
+            .iter()
+            .zip(record.iter())
+            .for_each(|(header, field)| {
+                map.insert(
+                    *header,
+                    if field == null_string {
+                        None
+                    } else {
+                        // Avoid allocating for empty strings
+                        if field.is_empty() {
+                            Some(String::new())
+                        } else {
+                            Some(field.to_string())
+                        }
+                    },
+                );
+            });
         map
     }
 }
 
 impl RecordParser for Vec<Option<String>> {
     type Output = Self;
+
+    #[inline]
     fn parse(
         _headers: &[&'static str],
         record: &csv::StringRecord,
         null_string: &str,
     ) -> Self::Output {
         let mut vec = Vec::with_capacity(record.len());
-        for field in record.iter() {
-            vec.push(if field == null_string {
+        vec.extend(record.iter().map(|field| {
+            if field == null_string {
                 None
             } else {
-                Some(field.to_string())
-            });
-        }
+                // Avoid allocating for empty strings
+                if field.is_empty() {
+                    Some(String::new())
+                } else {
+                    Some(field.to_string())
+                }
+            }
+        }));
         vec
     }
 }

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

@@ -0,0 +1,65 @@
+use super::{header_cache::StringCache, parser::RecordParser};
+use std::{io::Read, thread};
+
+pub enum ReadImpl<T: RecordParser> {
+    SingleThreaded {
+        reader: csv::Reader<Box<dyn Read>>,
+        headers: Vec<&'static str>,
+        null_string: String,
+    },
+    MultiThreaded {
+        headers: Vec<&'static str>,
+        receiver: kanal::Receiver<T::Output>,
+        handle: Option<thread::JoinHandle<()>>,
+    },
+}
+
+impl<T: RecordParser> ReadImpl<T> {
+    #[inline]
+    pub fn next(&mut self) -> Option<T::Output> {
+        match self {
+            Self::MultiThreaded {
+                receiver, handle, ..
+            } => match receiver.recv() {
+                Ok(record) => Some(record),
+                Err(_) => {
+                    if let Some(handle) = handle.take() {
+                        let _ = handle.join();
+                    }
+                    None
+                }
+            },
+            Self::SingleThreaded {
+                reader,
+                headers,
+                null_string,
+            } => {
+                let mut record = csv::StringRecord::new();
+                match reader.read_record(&mut record) {
+                    Ok(true) => Some(T::parse(headers, &record, null_string)),
+                    _ => None,
+                }
+            }
+        }
+    }
+
+    #[inline]
+    pub fn cleanup(&mut self) {
+        match self {
+            Self::MultiThreaded {
+                receiver,
+                handle,
+                headers,
+            } => {
+                receiver.close();
+                if let Some(handle) = handle.take() {
+                    let _ = handle.join();
+                }
+                let _ = StringCache::clear(headers);
+            }
+            Self::SingleThreaded { headers, .. } => {
+                let _ = StringCache::clear(headers);
+            }
+        }
+    }
+}

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

@@ -1,66 +1,35 @@
-use super::{header_cache::StringCache, parser::RecordParser};
+use super::{parser::RecordParser, read_impl::ReadImpl};
 use magnus::{Error, Ruby};
-use std::{io::Read, thread};
+use std::{borrow::Cow, io::Read};
 
 pub struct RecordReader<T: RecordParser> {
     pub(crate) reader: ReadImpl<T>,
 }
 
-impl<T: RecordParser> Drop for RecordReader<T> {
-    fn drop(&mut self) {
-        match &mut self.reader {
-            ReadImpl::MultiThreaded {
-                receiver,
-                handle,
-                headers,
-            } => {
-                receiver.close();
-                if let Some(handle) = handle.take() {
-                    let _ = handle.join();
-                }
-                StringCache::clear(headers).unwrap();
-            }
-            ReadImpl::SingleThreaded { headers, .. } => {
-                StringCache::clear(headers).unwrap();
-            }
-        }
-    }
-}
-
-#[allow(dead_code)]
-pub enum ReadImpl<T: RecordParser> {
-    SingleThreaded {
-        reader: csv::Reader<Box<dyn Read + Send + 'static>>,
-        headers: Vec<&'static str>,
-        null_string: String,
-    },
-    MultiThreaded {
-        headers: Vec<&'static str>,
-        receiver: kanal::Receiver<T::Output>,
-        handle: Option<thread::JoinHandle<()>>,
-    },
-}
-
 impl<T: RecordParser> RecordReader<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}"),
-                )
-            })?
-            .clone();
+        let first_row = reader.headers().map_err(|e| {
+            Error::new(
+                ruby.exception_runtime_error(),
+                Cow::Owned(format!("Failed to read headers: {e}")),
+            )
+        })?;
 
         Ok(if has_headers {
-            first_row.iter().map(String::from).collect()
+            // Pre-allocate the vector with exact capacity
+            let mut headers = Vec::with_capacity(first_row.len());
+            headers.extend(first_row.iter().map(String::from));
+            headers
         } else {
-            (0..first_row.len()).map(|i| format!("c{i}")).collect()
+            // Pre-allocate the vector with exact capacity
+            let mut headers = Vec::with_capacity(first_row.len());
+            headers.extend((0..first_row.len()).map(|i| format!("c{i}")));
+            headers
         })
     }
 }
@@ -68,30 +37,21 @@ impl<T: RecordParser> RecordReader<T> {
 impl<T: RecordParser> Iterator for RecordReader<T> {
     type Item = T::Output;
 
+    #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        match &mut self.reader {
-            ReadImpl::MultiThreaded {
-                receiver, handle, ..
-            } => match receiver.recv() {
-                Ok(record) => Some(record),
-                Err(_) => {
-                    if let Some(handle) = handle.take() {
-                        let _ = handle.join();
-                    }
-                    None
-                }
-            },
-            ReadImpl::SingleThreaded {
-                reader,
-                headers,
-                null_string,
-            } => {
-                let mut record = csv::StringRecord::new();
-                match reader.read_record(&mut record) {
-                    Ok(true) => Some(T::parse(headers, &record, null_string)),
-                    _ => None,
-                }
-            }
-        }
+        self.reader.next()
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // We can't know the exact size without reading the whole file
+        (0, None)
+    }
+}
+
+impl<T: RecordParser> Drop for RecordReader<T> {
+    #[inline]
+    fn drop(&mut self) {
+        self.reader.cleanup();
     }
 }

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

@@ -1,4 +1,4 @@
-use magnus::{IntoValue, RHash, Ruby, Value};
+use magnus::{IntoValue, Ruby, Value};
 use std::collections::HashMap;
 
 #[derive(Debug)]
@@ -8,14 +8,16 @@ pub enum CsvRecord {
 }
 
 impl IntoValue for CsvRecord {
+    #[inline]
     fn into_value_with(self, handle: &Ruby) -> Value {
         match self {
             CsvRecord::Vec(vec) => vec.into_value_with(handle),
             CsvRecord::Map(map) => {
-                let hash = RHash::new();
-                for (k, v) in map {
-                    hash.aset(k, v).unwrap();
-                }
+                // Pre-allocate the hash with the known size
+                let hash = handle.hash_new_capa(map.len());
+                map.into_iter()
+                    .try_for_each(|(k, v)| hash.aset(k, v))
+                    .unwrap();
                 hash.into_value_with(handle)
             }
         }

data/lib/osv/version.rb

@@ -1,3 +1,3 @@
 module OSV
-  VERSION = "0.3.3"
+  VERSION = "0.3.4"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: osv
 version: !ruby/object:Gem::Version
-  version: 0.3.3
+  version: 0.3.4
 platform: ruby
 authors:
 - Nathan Jaremko
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-12-23 00:00:00.000000000 Z
+date: 2024-12-24 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys
@@ -59,6 +59,7 @@ files:
 - ext/osv/src/csv/header_cache.rs
 - ext/osv/src/csv/mod.rs
 - ext/osv/src/csv/parser.rs
+- ext/osv/src/csv/read_impl.rs
 - ext/osv/src/csv/reader.rs
 - ext/osv/src/csv/record.rs
 - ext/osv/src/lib.rs