polars-df 0.3.0 → 0.4.0

data/Cargo.toml

@@ -4,7 +4,6 @@ members = ["ext/polars"]
 [patch.crates-io]
 jsonpath_lib = { git = "https://github.com/ritchie46/jsonpath", rev = "24eaf0b4416edff38a4d1b6b17bc4b9f3f047b4b" }
 halfbrown = { git = "https://github.com/Licenser/halfbrown", rev = "952023c5dd6461b009bb5ba66b9aa979bd75949f" }
-arrow2 = { git = "https://github.com/ankane/arrow2", rev = "ef0270922a217070ba9942567c0ff3263ae8c531" }
 
 [profile.release]
 strip = true

data/README.md

@@ -50,6 +50,9 @@ From Parquet
 
 ```ruby
 Polars.read_parquet("file.parquet")
+
+# or lazily with
+Polars.scan_parquet("file.parquet")
 ```
 
 From Active Record
@@ -60,6 +63,32 @@ Polars.read_sql(User.all)
 Polars.read_sql("SELECT * FROM users")
 ```
 
+From JSON
+
+```ruby
+Polars.read_json("file.json")
+# or
+Polars.read_ndjson("file.ndjson")
+
+# or lazily with
+Polars.scan_ndjson("file.ndjson")
+```
+
+From Feather / Arrow IPC
+
+```ruby
+Polars.read_ipc("file.arrow")
+
+# or lazily with
+Polars.scan_ipc("file.arrow")
+```
+
+From Avro
+
+```ruby
+Polars.read_avro("file.avro")
+```
+
 From a hash
 
 ```ruby
@@ -282,10 +311,10 @@ df.to_dummies
 
 ## Conversion
 
-Array of rows
+Array of hashes
 
 ```ruby
-df.rows
+df.rows(named: true)
 ```
 
 Hash of series
@@ -308,6 +337,12 @@ Parquet
 df.write_parquet("file.parquet")
 ```
 
+Numo array
+
+```ruby
+df.to_numo
+```
+
 ## Types
 
 You can specify column types when creating a data frame
@@ -343,6 +378,38 @@ Cast a column
 df["a"].cast(Polars::Int32)
 ```
 
+## Visualization
+
+Add [Vega](https://github.com/ankane/vega-ruby) to your application’s Gemfile:
+
+```ruby
+gem "vega"
+```
+
+And use:
+
+```ruby
+df.plot("a", "b")
+```
+
+Specify the chart type (`line`, `pie`, `column`, `bar`, `area`, or `scatter`)
+
+```ruby
+df.plot("a", "b", type: "pie")
+```
+
+Group data
+
+```ruby
+df.groupby("c").plot("a", "b")
+```
+
+Stacked columns or bars
+
+```ruby
+df.groupby("c").plot("a", "b", stacked: true)
+```
+
 ## History
 
 View the [changelog](CHANGELOG.md)

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.3.0"
+version = "0.4.0"
 license = "MIT"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
@@ -12,11 +12,12 @@ crate-type = ["cdylib"]
 [dependencies]
 ahash = "0.8"
 magnus = "0.5"
-polars-core = "0.27.0"
+polars-core = "0.28.0"
 serde_json = "1"
+smartstring = "1"
 
 [dependencies.polars]
-version = "0.27.0"
+version = "0.28.0"
 features = [
     "abs",
     "arange",
@@ -44,6 +45,7 @@ features = [
     "ipc",
     "is_first",
     "is_in",
+    "is_unique",
     "json",
     "lazy",
     "lazy_regex",

data/ext/polars/src/batched_csv.rs

@@ -7,11 +7,17 @@ use std::cell::RefCell;
 use std::path::PathBuf;
 
 use crate::conversion::*;
+use crate::prelude::read_impl::OwnedBatchedCsvReaderMmap;
 use crate::{RbDataFrame, RbPolarsErr, RbResult};
 
+pub enum BatchedReader {
+    MMap(OwnedBatchedCsvReaderMmap),
+    Read(OwnedBatchedCsvReader),
+}
+
 #[magnus::wrap(class = "Polars::RbBatchedCsv")]
 pub struct RbBatchedCsv {
-    pub reader: RefCell<OwnedBatchedCsvReader>,
+    pub reader: RefCell<BatchedReader>,
 }
 
 impl RbBatchedCsv {
@@ -38,7 +44,7 @@ impl RbBatchedCsv {
         let comment_char: Option<String> = arguments[16].try_convert()?;
         let quote_char: Option<String> = arguments[17].try_convert()?;
         let null_values: Option<Wrap<NullValues>> = arguments[18].try_convert()?;
-        let parse_dates: bool = arguments[19].try_convert()?;
+        let try_parse_dates: bool = arguments[19].try_convert()?;
         let skip_rows_after_header: usize = arguments[20].try_convert()?;
         let row_count: Option<(String, IdxSize)> = arguments[21].try_convert()?;
         let sample_size: usize = arguments[22].try_convert()?;
@@ -95,14 +101,24 @@ impl RbBatchedCsv {
             .low_memory(low_memory)
             .with_comment_char(comment_char)
             .with_null_values(null_values)
-            .with_parse_dates(parse_dates)
+            .with_try_parse_dates(try_parse_dates)
             .with_quote_char(quote_char)
             .with_end_of_line_char(eol_char)
             .with_skip_rows_after_header(skip_rows_after_header)
             .with_row_count(row_count)
-            .sample_size(sample_size)
-            .batched(overwrite_dtype.map(Arc::new))
-            .map_err(RbPolarsErr::from)?;
+            .sample_size(sample_size);
+
+        let reader = if low_memory {
+            let reader = reader
+                .batched_read(overwrite_dtype.map(Arc::new))
+                .map_err(RbPolarsErr::from)?;
+            BatchedReader::Read(reader)
+        } else {
+            let reader = reader
+                .batched_mmap(overwrite_dtype.map(Arc::new))
+                .map_err(RbPolarsErr::from)?;
+            BatchedReader::MMap(reader)
+        };
 
         Ok(RbBatchedCsv {
             reader: RefCell::new(reader),
@@ -110,13 +126,12 @@ impl RbBatchedCsv {
     }
 
     pub fn next_batches(&self, n: usize) -> RbResult<Option<RArray>> {
-        let batches = self
-            .reader
-            .borrow_mut()
-            .next_batches(n)
-            .map_err(RbPolarsErr::from)?;
-        Ok(batches.map(|batches| {
-            RArray::from_iter(batches.into_iter().map(|out| RbDataFrame::from(out.1)))
-        }))
+        let batches = match &mut *self.reader.borrow_mut() {
+            BatchedReader::MMap(reader) => reader.next_batches(n),
+            BatchedReader::Read(reader) => reader.next_batches(n),
+        }
+        .map_err(RbPolarsErr::from)?;
+
+        Ok(batches.map(|batches| RArray::from_iter(batches.into_iter().map(RbDataFrame::from))))
     }
 }

data/ext/polars/src/conversion.rs

@@ -1,3 +1,6 @@
+use std::fmt::{Display, Formatter};
+use std::hash::{Hash, Hasher};
+
 use magnus::{
     class, exception, r_hash::ForEach, ruby_handle::RubyHandle, Integer, IntoValue, Module, RArray,
     RFloat, RHash, RString, Symbol, TryConvert, Value, QNIL,
@@ -10,8 +13,7 @@ use polars::frame::NullStrategy;
 use polars::io::avro::AvroCompression;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
-use std::fmt::{Display, Formatter};
-use std::hash::{Hash, Hasher};
+use smartstring::alias::String as SmartString;
 
 use crate::{RbDataFrame, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
@@ -82,6 +84,22 @@ impl TryConvert for Wrap<Utf8Chunked> {
     }
 }
 
+impl TryConvert for Wrap<BinaryChunked> {
+    fn try_convert(obj: Value) -> RbResult<Self> {
+        let (seq, len) = get_rbseq(obj)?;
+        let mut builder = BinaryChunkedBuilder::new("", len, len * 25);
+
+        for res in seq.each() {
+            let item = res?;
+            match item.try_convert::<RString>() {
+                Ok(val) => builder.append_value(unsafe { val.as_slice() }),
+                Err(_) => builder.append_null(),
+            }
+        }
+        Ok(Wrap(builder.finish()))
+    }
+}
+
 impl TryConvert for Wrap<NullValues> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         if let Ok(s) = ob.try_convert::<String>() {
@@ -98,6 +116,14 @@ impl TryConvert for Wrap<NullValues> {
     }
 }
 
+fn struct_dict<'a>(vals: impl Iterator<Item = AnyValue<'a>>, flds: &[Field]) -> Value {
+    let dict = RHash::new();
+    for (fld, val) in flds.iter().zip(vals) {
+        dict.aset(fld.name().as_str(), Wrap(val)).unwrap()
+    }
+    dict.into_value()
+}
+
 impl IntoValue for Wrap<AnyValue<'_>> {
     fn into_value_with(self, _: &RubyHandle) -> Value {
         match self.0 {
@@ -114,7 +140,7 @@ impl IntoValue for Wrap<AnyValue<'_>> {
             AnyValue::Null => *QNIL,
             AnyValue::Boolean(v) => Value::from(v),
             AnyValue::Utf8(v) => Value::from(v),
-            AnyValue::Utf8Owned(_v) => todo!(),
+            AnyValue::Utf8Owned(v) => Value::from(v.as_str()),
             AnyValue::Categorical(_idx, _rev, _arr) => todo!(),
             AnyValue::Date(v) => class::time()
                 .funcall::<_, _, Value>("at", (v * 86400,))
@@ -125,7 +151,13 @@ impl IntoValue for Wrap<AnyValue<'_>> {
                 .unwrap(),
             AnyValue::Datetime(v, tu, tz) => {
                 let t = match tu {
-                    TimeUnit::Nanoseconds => todo!(),
+                    TimeUnit::Nanoseconds => {
+                        let sec = v / 1000000000;
+                        let subsec = v % 1000000000;
+                        class::time()
+                            .funcall::<_, _, Value>("at", (sec, subsec, Symbol::new("nsec")))
+                            .unwrap()
+                    }
                     TimeUnit::Microseconds => {
                         let sec = v / 1000000;
                         let subsec = v % 1000000;
@@ -133,7 +165,13 @@ impl IntoValue for Wrap<AnyValue<'_>> {
                             .funcall::<_, _, Value>("at", (sec, subsec, Symbol::new("usec")))
                             .unwrap()
                     }
-                    TimeUnit::Milliseconds => todo!(),
+                    TimeUnit::Milliseconds => {
+                        let sec = v / 1000;
+                        let subsec = v % 1000;
+                        class::time()
+                            .funcall::<_, _, Value>("at", (sec, subsec, Symbol::new("millisecond")))
+                            .unwrap()
+                    }
                 };
 
                 if tz.is_some() {
@@ -145,12 +183,19 @@ impl IntoValue for Wrap<AnyValue<'_>> {
             AnyValue::Duration(_v, _tu) => todo!(),
             AnyValue::Time(_v) => todo!(),
             AnyValue::List(v) => RbSeries::new(v).to_a().into_value(),
-            ref _av @ AnyValue::Struct(_, _, _flds) => todo!(),
-            AnyValue::StructOwned(_payload) => todo!(),
-            AnyValue::Object(_v) => todo!(),
-            AnyValue::ObjectOwned(_v) => todo!(),
-            AnyValue::Binary(_v) => todo!(),
-            AnyValue::BinaryOwned(_v) => todo!(),
+            ref av @ AnyValue::Struct(_, _, flds) => struct_dict(av._iter_struct_av(), flds),
+            AnyValue::StructOwned(payload) => struct_dict(payload.0.into_iter(), &payload.1),
+            AnyValue::Object(v) => {
+                let object = v.as_any().downcast_ref::<ObjectValue>().unwrap();
+                object.inner
+            }
+            AnyValue::ObjectOwned(v) => {
+                let object = v.0.as_any().downcast_ref::<ObjectValue>().unwrap();
+                object.inner
+            }
+            AnyValue::Binary(v) => RString::from_slice(v).into_value(),
+            AnyValue::BinaryOwned(v) => RString::from_slice(&v).into_value(),
+            AnyValue::Decimal(_v, _scale) => todo!(),
         }
     }
 }
@@ -170,12 +215,12 @@ impl IntoValue for Wrap<DataType> {
             DataType::UInt64 => pl.const_get::<_, Value>("UInt64").unwrap(),
             DataType::Float32 => pl.const_get::<_, Value>("Float32").unwrap(),
             DataType::Float64 => pl.const_get::<_, Value>("Float64").unwrap(),
-            DataType::Decimal128(_) => todo!(),
+            DataType::Decimal(_precision, _scale) => todo!(),
             DataType::Boolean => pl.const_get::<_, Value>("Boolean").unwrap(),
             DataType::Utf8 => pl.const_get::<_, Value>("Utf8").unwrap(),
             DataType::Binary => pl.const_get::<_, Value>("Binary").unwrap(),
             DataType::List(inner) => {
-                let inner = Wrap(*inner.clone());
+                let inner = Wrap(*inner);
                 let list_class = pl.const_get::<_, Value>("List").unwrap();
                 list_class.funcall::<_, _, Value>("new", (inner,)).unwrap()
             }
@@ -183,7 +228,7 @@ impl IntoValue for Wrap<DataType> {
             DataType::Datetime(tu, tz) => {
                 let datetime_class = pl.const_get::<_, Value>("Datetime").unwrap();
                 datetime_class
-                    .funcall::<_, _, Value>("new", (tu.to_ascii(), tz.clone()))
+                    .funcall::<_, _, Value>("new", (tu.to_ascii(), tz))
                     .unwrap()
             }
             DataType::Duration(tu) => {
@@ -198,7 +243,7 @@ impl IntoValue for Wrap<DataType> {
             DataType::Struct(fields) => {
                 let field_class = pl.const_get::<_, Value>("Field").unwrap();
                 let iter = fields.iter().map(|fld| {
-                    let name = fld.name().clone();
+                    let name = fld.name().as_str();
                     let dtype = Wrap(fld.data_type().clone());
                     field_class
                         .funcall::<_, _, Value>("new", (name, dtype))
@@ -328,7 +373,7 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
                 let n = 25;
                 let dtype = any_values_to_dtype(&avs[..std::cmp::min(avs.len(), n)])
                     .map_err(RbPolarsErr::from)?;
-                let s = Series::from_any_values_and_dtype("", &avs, &dtype)
+                let s = Series::from_any_values_and_dtype("", &avs, &dtype, true)
                     .map_err(RbPolarsErr::from)?;
                 Ok(Wrap(AnyValue::List(s)))
             }
@@ -858,3 +903,11 @@ pub fn parse_parquet_compression(
     };
     Ok(parsed)
 }
+
+pub(crate) fn strings_to_smartstrings<I, S>(container: I) -> Vec<SmartString>
+where
+    I: IntoIterator<Item = S>,
+    S: AsRef<str>,
+{
+    container.into_iter().map(|s| s.as_ref().into()).collect()
+}

data/ext/polars/src/dataframe.rs

@@ -6,6 +6,7 @@ use polars::io::mmap::ReaderBytes;
 use polars::io::RowCount;
 use polars::prelude::pivot::{pivot, pivot_stable};
 use polars::prelude::*;
+use polars_core::utils::try_get_supertype;
 use std::cell::RefCell;
 use std::io::{BufWriter, Cursor};
 use std::ops::Deref;
@@ -114,7 +115,7 @@ impl RbDataFrame {
         let comment_char: Option<String> = arguments[17].try_convert()?;
         let quote_char: Option<String> = arguments[18].try_convert()?;
         let null_values: Option<Wrap<NullValues>> = arguments[19].try_convert()?;
-        let parse_dates: bool = arguments[20].try_convert()?;
+        let try_parse_dates: bool = arguments[20].try_convert()?;
         let skip_rows_after_header: usize = arguments[21].try_convert()?;
         let row_count: Option<(String, IdxSize)> = arguments[22].try_convert()?;
         let sample_size: usize = arguments[23].try_convert()?;
@@ -167,12 +168,12 @@ impl RbDataFrame {
             .with_columns(columns)
             .with_n_threads(n_threads)
             .with_path(path)
-            .with_dtypes(overwrite_dtype.as_ref())
+            .with_dtypes(overwrite_dtype.map(Arc::new))
             .with_dtypes_slice(overwrite_dtype_slice.as_deref())
             .low_memory(low_memory)
             .with_comment_char(comment_char)
             .with_null_values(null_values)
-            .with_parse_dates(parse_dates)
+            .with_try_parse_dates(try_parse_dates)
             .with_quote_char(quote_char)
             .with_end_of_line_char(eol_char)
             .with_skip_rows_after_header(skip_rows_after_header)
@@ -183,6 +184,7 @@ impl RbDataFrame {
         Ok(df.into())
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn read_parquet(
         rb_f: Value,
         columns: Option<Vec<String>>,
@@ -191,6 +193,8 @@ impl RbDataFrame {
         parallel: Wrap<ParallelStrategy>,
         row_count: Option<(String, IdxSize)>,
         low_memory: bool,
+        use_statistics: bool,
+        rechunk: bool,
     ) -> RbResult<Self> {
         let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
         let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
@@ -201,6 +205,8 @@ impl RbDataFrame {
             .with_n_rows(n_rows)
             .with_row_count(row_count)
             .set_low_memory(low_memory)
+            .use_statistics(use_statistics)
+            .set_rechunk(rechunk)
             .finish()
             .map_err(RbPolarsErr::from)?;
         Ok(RbDataFrame::new(df))
@@ -253,7 +259,7 @@ impl RbDataFrame {
         use polars::io::avro::AvroWriter;
 
         if let Ok(s) = rb_f.try_convert::<String>() {
-            let f = std::fs::File::create(&s).unwrap();
+            let f = std::fs::File::create(s).unwrap();
             AvroWriter::new(f)
                 .with_compression(compression.0)
                 .finish(&mut self.df.borrow_mut())
@@ -338,7 +344,7 @@ impl RbDataFrame {
         // ensure the new names are used
         if let Some(schema) = &schema_overwrite {
             for (new_name, name) in schema.0.iter_names().zip(names.iter_mut()) {
-                *name = new_name.clone();
+                *name = new_name.to_string();
             }
         }
         let rbdf = Self::finish_from_rows(
@@ -347,17 +353,19 @@ impl RbDataFrame {
             schema_overwrite.map(|wrap| wrap.0),
         )?;
 
-        rbdf.df
-            .borrow_mut()
-            .get_columns_mut()
-            .iter_mut()
-            .zip(&names)
-            .for_each(|(s, name)| {
-                s.rename(name);
-            });
+        unsafe {
+            rbdf.df
+                .borrow_mut()
+                .get_columns_mut()
+                .iter_mut()
+                .zip(&names)
+                .for_each(|(s, name)| {
+                    s.rename(name);
+                });
+        }
         let length = names.len();
         if names.into_iter().collect::<PlHashSet<_>>().len() != length {
-            let err = PolarsError::SchemaMisMatch("duplicate column names found".into());
+            let err = PolarsError::SchemaMismatch("duplicate column names found".into());
             Err(RbPolarsErr::from(err))?;
         }
 
@@ -393,7 +401,7 @@ impl RbDataFrame {
         let null = null_value.unwrap_or_default();
 
         if let Ok(s) = rb_f.try_convert::<String>() {
-            let f = std::fs::File::create(&s).unwrap();
+            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)
@@ -435,7 +443,7 @@ impl RbDataFrame {
         compression: Wrap<Option<IpcCompression>>,
     ) -> RbResult<()> {
         if let Ok(s) = rb_f.try_convert::<String>() {
-            let f = std::fs::File::create(&s).unwrap();
+            let f = std::fs::File::create(s).unwrap();
             IpcWriter::new(f)
                 .with_compression(compression.0)
                 .finish(&mut self.df.borrow_mut())
@@ -493,6 +501,25 @@ impl RbDataFrame {
         .into()
     }
 
+    pub fn to_numo(&self) -> Option<Value> {
+        let mut st = None;
+        for s in self.df.borrow().iter() {
+            let dt_i = s.dtype();
+            match st {
+                None => st = Some(dt_i.clone()),
+                Some(ref mut st) => {
+                    *st = try_get_supertype(st, dt_i).ok()?;
+                }
+            }
+        }
+        let st = st?;
+
+        match st {
+            // TODO
+            _ => None,
+        }
+    }
+
     pub fn write_parquet(
         &self,
         rb_f: Value,
@@ -504,7 +531,7 @@ impl RbDataFrame {
         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();
+            let f = std::fs::File::create(s).unwrap();
             ParquetWriter::new(f)
                 .with_compression(compression)
                 .with_statistics(statistics)
@@ -607,7 +634,7 @@ impl RbDataFrame {
     }
 
     pub fn get_columns(&self) -> RArray {
-        let cols = self.df.borrow().get_columns().clone();
+        let cols = self.df.borrow().get_columns().to_vec();
         to_rbseries_collection(cols)
     }
 
@@ -861,10 +888,11 @@ impl RbDataFrame {
         variable_name: Option<String>,
     ) -> RbResult<Self> {
         let args = MeltArgs {
-            id_vars,
-            value_vars,
-            value_name,
-            variable_name,
+            id_vars: strings_to_smartstrings(id_vars),
+            value_vars: strings_to_smartstrings(value_vars),
+            value_name: value_name.map(|s| s.into()),
+            variable_name: variable_name.map(|s| s.into()),
+            streamable: false,
         };
 
         let df = self.df.borrow().melt2(args).map_err(RbPolarsErr::from)?;
@@ -877,22 +905,26 @@ impl RbDataFrame {
         values: Vec<String>,
         index: Vec<String>,
         columns: Vec<String>,
-        aggregate_expr: &RbExpr,
         maintain_order: bool,
         sort_columns: bool,
+        aggregate_expr: Option<&RbExpr>,
         separator: Option<String>,
     ) -> RbResult<Self> {
         let fun = match maintain_order {
             true => pivot_stable,
             false => pivot,
         };
+        let agg_expr = match aggregate_expr {
+            Some(aggregate_expr) => Some(aggregate_expr.inner.clone()),
+            None => None,
+        };
         let df = fun(
             &self.df.borrow(),
             values,
             index,
             columns,
-            aggregate_expr.inner.clone(),
             sort_columns,
+            agg_expr,
             separator.as_deref(),
         )
         .map_err(RbPolarsErr::from)?;
@@ -913,21 +945,6 @@ impl RbDataFrame {
         self.df.borrow().shift(periods).into()
     }
 
-    pub fn unique(
-        &self,
-        maintain_order: bool,
-        subset: Option<Vec<String>>,
-        keep: Wrap<UniqueKeepStrategy>,
-    ) -> RbResult<Self> {
-        let subset = subset.as_ref().map(|v| v.as_ref());
-        let df = match maintain_order {
-            true => self.df.borrow().unique_stable(subset, keep.0),
-            false => self.df.borrow().unique(subset, keep.0),
-        }
-        .map_err(RbPolarsErr::from)?;
-        Ok(df.into())
-    }
-
     pub fn lazy(&self) -> RbLazyFrame {
         self.df.borrow().clone().lazy().into()
     }

data/ext/polars/src/error.rs

@@ -43,3 +43,11 @@ impl ComputeError {
         Error::new(exception::runtime_error(), message)
     }
 }
+
+#[macro_export]
+macro_rules! raise_err(
+    ($msg:expr, $err:ident) => {{
+        Err(PolarsError::$err($msg.into())).map_err(RbPolarsErr::from)?;
+        unreachable!()
+    }}
+);