polars-df 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: a2b7504ba8b5c20977f2699f849c2f8a8c28f7912d695f17f35eb8b492699fe6
+  data.tar.gz: d38b420fa7e192e38c91b7f7b621117fbdbb981a779e2f54ae7a9d9066e355ba
+SHA512:
+  metadata.gz: 7ac246cbb2c93433cd17078e534a372ed0e2d00c1cdd27bcf454cb4685b8bd48fd6d66f5060a9132a2fece3f877503f67e3b5f8b65da467942174a753c2d889b
+  data.tar.gz: 5e3f62085f6b02708534ed3ee3e4bdfe0e608c8c6fdd688e0e534858e53b63ec0d57a255d98bfc3385e609053385e354536efcf5e5469a0e17c9aecbaa9ea879

data/CHANGELOG.md

@@ -0,0 +1,3 @@
+## 0.1.0 (2022-11-21)
+
+- First release

data/LICENSE.txt

@@ -0,0 +1,20 @@
+Copyright (c) 2020 Ritchie Vink
+Copyright (c) 2022 Andrew Kane
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/README.md

@@ -0,0 +1,93 @@
+# Polars Ruby
+
+:fire: Blazingly fast DataFrames for Ruby, powered by [Polars](https://github.com/pola-rs/polars)
+
+[![Build Status](https://github.com/ankane/polars-ruby/workflows/build/badge.svg?branch=master)](https://github.com/ankane/polars-ruby/actions)
+
+## Installation
+
+Add this line to your application’s Gemfile:
+
+```ruby
+gem "polars-df"
+```
+
+Note: Rust is currently required for installation, and it can take 15-20 minutes to compile the extension.
+
+## Getting Started
+
+This library follows the [Polars Python API](https://pola-rs.github.io/polars/py-polars/html/reference/index.html).
+
+```ruby
+Polars.read_csv("iris.csv")
+  .lazy
+  .filter(Polars.col("sepal_length") > 5)
+  .groupby("species")
+  .agg(Polars.all.sum)
+  .collect
+```
+
+You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/introduction.html) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems. Note that many methods and options are missing at the moment.
+
+## Examples
+
+### Creating DataFrames
+
+From a CSV
+
+```ruby
+Polars.read_csv("file.csv")
+```
+
+From Parquet
+
+```ruby
+Polars.read_parquet("file.parquet")
+```
+
+From Active Record
+
+```ruby
+Polars::DataFrame.new(User.all)
+```
+
+From a hash
+
+```ruby
+Polars::DataFrame.new({
+  a: [1, 2, 3],
+  b: ["one", "two", "three"]
+})
+```
+
+From an array of series
+
+```ruby
+Polars::DataFrame.new([
+  Polars::Series.new("a", [1, 2, 3]),
+  Polars::Series.new("b", ["one", "two", "three"])
+])
+```
+
+## History
+
+View the [changelog](CHANGELOG.md)
+
+## Contributing
+
+Everyone is encouraged to help improve this project. Here are a few ways you can help:
+
+- [Report bugs](https://github.com/ankane/polars-ruby/issues)
+- Fix bugs and [submit pull requests](https://github.com/ankane/polars-ruby/pulls)
+- Write, clarify, or fix documentation
+- Suggest or add new features
+
+To get started with development:
+
+```sh
+git clone https://github.com/ankane/polars-ruby.git
+cd polars-ruby
+bundle install
+bundle exec rake compile
+bundle exec rake test
+```

data/ext/polars/Cargo.toml

@@ -0,0 +1,35 @@
+[package]
+name = "polars"
+version = "0.1.0"
+authors = ["Andrew Kane <andrew@ankane.org>"]
+edition = "2021"
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+magnus = "0.4"
+serde_json = "1"
+
+[dependencies.polars]
+version = "0.25.1"
+features = [
+    "arange",
+    "csv-file",
+    "cum_agg",
+    "diagonal_concat",
+    "dtype-full",
+    "fmt",
+    "horizontal_concat",
+    "interpolate",
+    "ipc",
+    "json",
+    "lazy",
+    "lazy_regex",
+    "parquet",
+    "semi_anti_join",
+    "serde-lazy",
+    "strings",
+    "trigonometry",
+]

data/ext/polars/extconf.rb

@@ -0,0 +1,4 @@
+require "mkmf"
+require "rb_sys/mkmf"
+
+create_rust_makefile("polars/polars")

data/ext/polars/src/conversion.rs

@@ -0,0 +1,115 @@
+use magnus::{Value, QNIL};
+use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
+use polars::datatypes::AnyValue;
+use polars::frame::DataFrame;
+use polars::prelude::*;
+
+use crate::{RbDataFrame, RbResult, RbValueError};
+
+pub fn wrap(val: AnyValue) -> Value {
+    match val {
+        AnyValue::UInt8(v) => Value::from(v),
+        AnyValue::UInt16(v) => Value::from(v),
+        AnyValue::UInt32(v) => Value::from(v),
+        AnyValue::UInt64(v) => Value::from(v),
+        AnyValue::Int8(v) => Value::from(v),
+        AnyValue::Int16(v) => Value::from(v),
+        AnyValue::Int32(v) => Value::from(v),
+        AnyValue::Int64(v) => Value::from(v),
+        AnyValue::Float32(v) => Value::from(v),
+        AnyValue::Float64(v) => Value::from(v),
+        AnyValue::Null => *QNIL,
+        AnyValue::Boolean(v) => Value::from(v),
+        AnyValue::Utf8(v) => Value::from(v),
+        _ => todo!(),
+    }
+}
+
+pub fn parse_fill_null_strategy(
+    strategy: &str,
+    limit: FillNullLimit,
+) -> RbResult<FillNullStrategy> {
+    let parsed = match strategy {
+        "forward" => FillNullStrategy::Forward(limit),
+        "backward" => FillNullStrategy::Backward(limit),
+        "min" => FillNullStrategy::Min,
+        "max" => FillNullStrategy::Max,
+        "mean" => FillNullStrategy::Mean,
+        "zero" => FillNullStrategy::Zero,
+        "one" => FillNullStrategy::One,
+        e => {
+            return Err(magnus::Error::runtime_error(format!(
+                "strategy must be one of {{'forward', 'backward', 'min', 'max', 'mean', 'zero', 'one'}}, got {}",
+                e,
+            )))
+        }
+    };
+    Ok(parsed)
+}
+
+pub fn wrap_join_type(ob: &str) -> RbResult<JoinType> {
+    let parsed = match ob {
+        "inner" => JoinType::Inner,
+        "left" => JoinType::Left,
+        "outer" => JoinType::Outer,
+        "semi" => JoinType::Semi,
+        "anti" => JoinType::Anti,
+        // #[cfg(feature = "cross_join")]
+        // "cross" => JoinType::Cross,
+        v => {
+            return Err(RbValueError::new_err(format!(
+                "how must be one of {{'inner', 'left', 'outer', 'semi', 'anti', 'cross'}}, got {}",
+                v
+            )))
+        }
+    };
+    Ok(parsed)
+}
+
+pub fn get_df(obj: Value) -> RbResult<DataFrame> {
+    let rbdf = obj.funcall::<_, _, &RbDataFrame>("_df", ())?;
+    Ok(rbdf.df.borrow().clone())
+}
+
+pub fn parse_parquet_compression(
+    compression: &str,
+    compression_level: Option<i32>,
+) -> RbResult<ParquetCompression> {
+    let parsed = match compression {
+        "uncompressed" => ParquetCompression::Uncompressed,
+        "snappy" => ParquetCompression::Snappy,
+        "gzip" => ParquetCompression::Gzip(
+            compression_level
+                .map(|lvl| {
+                    GzipLevel::try_new(lvl as u8)
+                        .map_err(|e| RbValueError::new_err(format!("{:?}", e)))
+                })
+                .transpose()?,
+        ),
+        "lzo" => ParquetCompression::Lzo,
+        "brotli" => ParquetCompression::Brotli(
+            compression_level
+                .map(|lvl| {
+                    BrotliLevel::try_new(lvl as u32)
+                        .map_err(|e| RbValueError::new_err(format!("{:?}", e)))
+                })
+                .transpose()?,
+        ),
+        "lz4" => ParquetCompression::Lz4Raw,
+        "zstd" => ParquetCompression::Zstd(
+            compression_level
+                .map(|lvl| {
+                    ZstdLevel::try_new(lvl)
+                        .map_err(|e| RbValueError::new_err(format!("{:?}", e)))
+                })
+                .transpose()?,
+        ),
+        e => {
+            return Err(RbValueError::new_err(format!(
+                "compression must be one of {{'uncompressed', 'snappy', 'gzip', 'lzo', 'brotli', 'lz4', 'zstd'}}, got {}",
+                e
+            )))
+        }
+    };
+    Ok(parsed)
+}

data/ext/polars/src/dataframe.rs

@@ -0,0 +1,304 @@
+use magnus::{r_hash::ForEach, Error, RArray, RHash, RString, Value};
+use polars::io::mmap::ReaderBytes;
+use polars::prelude::*;
+use std::cell::RefCell;
+use std::fs::File;
+use std::io::{BufReader, BufWriter, Cursor};
+use std::ops::Deref;
+use std::path::PathBuf;
+
+use crate::conversion::parse_parquet_compression;
+use crate::file::{get_file_like, get_mmap_bytes_reader};
+use crate::{series, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
+
+#[magnus::wrap(class = "Polars::RbDataFrame")]
+pub struct RbDataFrame {
+    pub df: RefCell<DataFrame>,
+}
+
+impl From<DataFrame> for RbDataFrame {
+    fn from(df: DataFrame) -> Self {
+        RbDataFrame::new(df)
+    }
+}
+
+impl RbDataFrame {
+    pub fn new(df: DataFrame) -> Self {
+        RbDataFrame {
+            df: RefCell::new(df),
+        }
+    }
+
+    pub fn init(columns: RArray) -> RbResult<Self> {
+        let mut cols = Vec::new();
+        for i in columns.each() {
+            cols.push(i?.try_convert::<&RbSeries>()?.series.borrow().clone());
+        }
+        let df = DataFrame::new(cols).map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn read_csv(rb_f: Value, has_header: bool) -> RbResult<Self> {
+        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let df = CsvReader::new(mmap_bytes_r)
+            .has_header(has_header)
+            .finish()
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn read_parquet(path: PathBuf) -> RbResult<Self> {
+        let f = File::open(&path).map_err(|e| Error::runtime_error(e.to_string()))?;
+        let reader = BufReader::new(f);
+        ParquetReader::new(reader)
+            .finish()
+            .map_err(RbPolarsErr::from)
+            .map(|v| v.into())
+    }
+
+    pub fn read_json(rb_f: Value) -> RbResult<Self> {
+        // memmap the file first
+        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let mmap_read: ReaderBytes = (&mmap_bytes_r).into();
+        let bytes = mmap_read.deref();
+
+        // Happy path is our column oriented json as that is most performant
+        // on failure we try
+        match serde_json::from_slice::<DataFrame>(bytes) {
+            Ok(df) => Ok(df.into()),
+            // try arrow json reader instead
+            // this is row oriented
+            Err(_) => {
+                let out = JsonReader::new(mmap_bytes_r)
+                    .with_json_format(JsonFormat::Json)
+                    .finish()
+                    .map_err(|e| RbPolarsErr::other(format!("{:?}", e)))?;
+                Ok(out.into())
+            }
+        }
+    }
+
+    pub fn read_ndjson(rb_f: Value) -> RbResult<Self> {
+        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+
+        let out = JsonReader::new(mmap_bytes_r)
+            .with_json_format(JsonFormat::JsonLines)
+            .finish()
+            .map_err(|e| RbPolarsErr::other(format!("{:?}", e)))?;
+        Ok(out.into())
+    }
+
+    pub fn write_json(&self, rb_f: Value, pretty: bool, row_oriented: bool) -> RbResult<()> {
+        let file = BufWriter::new(get_file_like(rb_f, true)?);
+
+        let r = match (pretty, row_oriented) {
+            (_, true) => JsonWriter::new(file)
+                .with_json_format(JsonFormat::Json)
+                .finish(&mut self.df.borrow_mut()),
+            (true, _) => serde_json::to_writer_pretty(file, &*self.df.borrow())
+                .map_err(|e| PolarsError::ComputeError(format!("{:?}", e).into())),
+            (false, _) => serde_json::to_writer(file, &*self.df.borrow())
+                .map_err(|e| PolarsError::ComputeError(format!("{:?}", e).into())),
+        };
+        r.map_err(|e| RbPolarsErr::other(format!("{:?}", e)))?;
+        Ok(())
+    }
+
+    pub fn write_ndjson(&self, rb_f: Value) -> RbResult<()> {
+        let file = BufWriter::new(get_file_like(rb_f, true)?);
+
+        let r = JsonWriter::new(file)
+            .with_json_format(JsonFormat::JsonLines)
+            .finish(&mut self.df.borrow_mut());
+
+        r.map_err(|e| RbPolarsErr::other(format!("{:?}", e)))?;
+        Ok(())
+    }
+
+    pub fn read_hash(data: RHash) -> RbResult<Self> {
+        let mut cols: Vec<Series> = Vec::new();
+        data.foreach(|name: String, values: Value| {
+            let obj: Value = series().funcall("new", (name, values))?;
+            let rbseries = obj.funcall::<_, _, &RbSeries>("_s", ())?;
+            cols.push(rbseries.series.borrow().clone());
+            Ok(ForEach::Continue)
+        })?;
+        let df = DataFrame::new(cols).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn write_csv(
+        &self,
+        rb_f: Value,
+        has_header: bool,
+        sep: u8,
+        quote: u8,
+        batch_size: usize,
+        datetime_format: Option<String>,
+        date_format: Option<String>,
+        time_format: Option<String>,
+        float_precision: Option<usize>,
+        null_value: Option<String>,
+    ) -> RbResult<()> {
+        let null = null_value.unwrap_or_default();
+
+        if let Ok(s) = rb_f.try_convert::<String>() {
+            let f = std::fs::File::create(&s).unwrap();
+            // no need for a buffered writer, because the csv writer does internal buffering
+            CsvWriter::new(f)
+                .has_header(has_header)
+                .with_delimiter(sep)
+                .with_quoting_char(quote)
+                .with_batch_size(batch_size)
+                .with_datetime_format(datetime_format)
+                .with_date_format(date_format)
+                .with_time_format(time_format)
+                .with_float_precision(float_precision)
+                .with_null_value(null)
+                .finish(&mut self.df.borrow_mut())
+                .map_err(RbPolarsErr::from)?;
+        } else {
+            let mut buf = Cursor::new(Vec::new());
+            CsvWriter::new(&mut buf)
+                .has_header(has_header)
+                .with_delimiter(sep)
+                .with_quoting_char(quote)
+                .with_batch_size(batch_size)
+                .with_datetime_format(datetime_format)
+                .with_date_format(date_format)
+                .with_time_format(time_format)
+                .with_float_precision(float_precision)
+                .with_null_value(null)
+                .finish(&mut self.df.borrow_mut())
+                .map_err(RbPolarsErr::from)?;
+            // TODO less copying
+            let rb_str = RString::from_slice(&buf.into_inner());
+            rb_f.funcall::<_, _, Value>("write", (rb_str,))?;
+        }
+
+        Ok(())
+    }
+
+    pub fn write_parquet(
+        &self,
+        rb_f: Value,
+        compression: String,
+        compression_level: Option<i32>,
+        statistics: bool,
+        row_group_size: Option<usize>,
+    ) -> RbResult<()> {
+        let compression = parse_parquet_compression(&compression, compression_level)?;
+
+        if let Ok(s) = rb_f.try_convert::<String>() {
+            let f = std::fs::File::create(&s).unwrap();
+            ParquetWriter::new(f)
+                .with_compression(compression)
+                .with_statistics(statistics)
+                .with_row_group_size(row_group_size)
+                .finish(&mut self.df.borrow_mut())
+                .map_err(RbPolarsErr::from)?;
+        } else {
+            todo!();
+        }
+
+        Ok(())
+    }
+
+    pub fn rechunk(&self) -> Self {
+        self.df.borrow().agg_chunks().into()
+    }
+
+    pub fn to_s(&self) -> String {
+        format!("{}", self.df.borrow())
+    }
+
+    pub fn columns(&self) -> Vec<String> {
+        self.df
+            .borrow()
+            .get_column_names()
+            .iter()
+            .map(|v| v.to_string())
+            .collect()
+    }
+
+    pub fn dtypes(&self) -> Vec<String> {
+        self.df
+            .borrow()
+            .iter()
+            .map(|s| s.dtype().to_string())
+            .collect()
+    }
+
+    pub fn shape(&self) -> (usize, usize) {
+        self.df.borrow().shape()
+    }
+
+    pub fn height(&self) -> usize {
+        self.df.borrow().height()
+    }
+
+    pub fn width(&self) -> usize {
+        self.df.borrow().width()
+    }
+
+    pub fn select_at_idx(&self, idx: usize) -> Option<RbSeries> {
+        self.df
+            .borrow()
+            .select_at_idx(idx)
+            .map(|s| RbSeries::new(s.clone()))
+    }
+
+    // TODO remove clone
+    pub fn column(&self, name: String) -> RbResult<RbSeries> {
+        self.df
+            .borrow()
+            .column(&name)
+            .map(|v| v.clone().into())
+            .map_err(RbPolarsErr::from)
+    }
+
+    pub fn sort(&self, by_column: String, reverse: bool, nulls_last: bool) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .sort_with_options(
+                &by_column,
+                SortOptions {
+                    descending: reverse,
+                    nulls_last,
+                },
+            )
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn head(&self, length: Option<usize>) -> Self {
+        self.df.borrow().head(length).into()
+    }
+
+    pub fn tail(&self, length: Option<usize>) -> Self {
+        self.df.borrow().tail(length).into()
+    }
+
+    pub fn frame_equal(&self, other: &RbDataFrame, null_equal: bool) -> bool {
+        if null_equal {
+            self.df.borrow().frame_equal_missing(&other.df.borrow())
+        } else {
+            self.df.borrow().frame_equal(&other.df.borrow())
+        }
+    }
+
+    pub fn lazy(&self) -> RbLazyFrame {
+        self.df.borrow().clone().lazy().into()
+    }
+
+    pub fn mean(&self) -> Self {
+        self.df.borrow().mean().into()
+    }
+
+    pub fn null_count(&self) -> Self {
+        let df = self.df.borrow().null_count();
+        df.into()
+    }
+}

data/ext/polars/src/error.rs

@@ -0,0 +1,24 @@
+use magnus::exception::arg_error;
+use magnus::Error;
+use polars::prelude::PolarsError;
+
+pub struct RbPolarsErr {}
+
+impl RbPolarsErr {
+    // convert to Error instead of Self
+    pub fn from(e: PolarsError) -> Error {
+        Error::runtime_error(e.to_string())
+    }
+
+    pub fn other(message: String) -> Error {
+        Error::runtime_error(message)
+    }
+}
+
+pub struct RbValueError {}
+
+impl RbValueError {
+    pub fn new_err(message: String) -> Error {
+        Error::new(arg_error(), message)
+    }
+}

data/ext/polars/src/file.rs

@@ -0,0 +1,28 @@
+use magnus::{Error, RString, Value};
+use polars::io::mmap::MmapBytesReader;
+use std::fs::{File, OpenOptions};
+use std::io::Cursor;
+use std::path::PathBuf;
+
+use crate::RbResult;
+
+pub fn get_file_like(f: Value, truncate: bool) -> RbResult<File> {
+    OpenOptions::new()
+        .write(true)
+        .create(true)
+        .truncate(truncate)
+        .open(f.try_convert::<PathBuf>()?)
+        .map_err(|e| Error::runtime_error(e.to_string()))
+}
+
+pub fn get_mmap_bytes_reader(rb_f: Value) -> RbResult<Box<dyn MmapBytesReader>> {
+    if let Ok(bytes) = rb_f.funcall::<_, _, RString>("read", ()) {
+        let bytes = unsafe { bytes.as_slice() };
+        // TODO avoid copy
+        Ok(Box::new(Cursor::new(bytes.to_vec())))
+    } else {
+        let p = rb_f.try_convert::<PathBuf>()?;
+        let f = File::open(p).map_err(|e| Error::runtime_error(e.to_string()))?;
+        Ok(Box::new(f))
+    }
+}