polars-df 0.3.1 → 0.5.0

data/Cargo.toml

@@ -3,8 +3,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,14 +50,43 @@ From Parquet
 
 ```ruby
 Polars.read_parquet("file.parquet")
+
+# or lazily with
+Polars.scan_parquet("file.parquet")
 ```
 
 From Active Record
 
 ```ruby
-Polars.read_sql(User.all)
+Polars.read_database(User.all)
+# or
+Polars.read_database("SELECT * FROM users")
+```
+
+From JSON
+
+```ruby
+Polars.read_json("file.json")
 # or
-Polars.read_sql("SELECT * FROM users")
+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

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.3.1"
+version = "0.5.0"
 license = "MIT"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
@@ -12,13 +12,15 @@ crate-type = ["cdylib"]
 [dependencies]
 ahash = "0.8"
 magnus = "0.5"
-polars-core = "0.27.0"
+polars-core = "0.29.0"
 serde_json = "1"
+smartstring = "1"
 
 [dependencies.polars]
-version = "0.27.0"
+version = "0.29.0"
 features = [
     "abs",
+    "approx_unique",
     "arange",
     "arg_where",
     "asof_join",
@@ -26,7 +28,7 @@ features = [
     "binary_encoding",
     "concat_str",
     "cse",
-    "csv-file",
+    "csv",
     "cum_agg",
     "cumulative_eval",
     "dataframe_arithmetic",
@@ -44,10 +46,13 @@ features = [
     "ipc",
     "is_first",
     "is_in",
+    "is_unique",
     "json",
     "lazy",
     "lazy_regex",
+    "list_count",
     "list_eval",
+    "list_take",
     "list_to_struct",
     "log",
     "meta",
@@ -73,6 +78,7 @@ features = [
     "serde-lazy",
     "sign",
     "string_encoding",
+    "string_from_radix",
     "string_justify",
     "strings",
     "timezones",

data/ext/polars/src/apply/dataframe.rs

@@ -202,8 +202,8 @@ pub fn apply_lambda_with_utf8_out_type(
 }
 
 /// Apply a lambda with list output type
-pub fn apply_lambda_with_list_out_type<'a>(
-    df: &'a DataFrame,
+pub fn apply_lambda_with_list_out_type(
+    df: &DataFrame,
     lambda: Value,
     init_null_count: usize,
     first_value: Option<&Series>,

data/ext/polars/src/{lazy/apply.rs → apply/lazy.rs}

@@ -1,8 +1,7 @@
 use magnus::Value;
 use polars::prelude::*;
 
-use crate::lazy::dsl::RbExpr;
-use crate::Wrap;
+use crate::{RbExpr, Wrap};
 
 pub fn binary_lambda(_lambda: Value, _a: Series, _b: Series) -> PolarsResult<Option<Series>> {
     todo!();

data/ext/polars/src/apply/mod.rs

@@ -1,4 +1,5 @@
 pub mod dataframe;
+pub mod lazy;
 pub mod series;
 
 use magnus::{RHash, Value};

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()?;
@@ -62,11 +68,13 @@ impl RbBatchedCsv {
         };
 
         let overwrite_dtype = overwrite_dtype.map(|overwrite_dtype| {
-            let fields = overwrite_dtype.iter().map(|(name, dtype)| {
-                let dtype = dtype.0.clone();
-                Field::new(name, dtype)
-            });
-            Schema::from(fields)
+            overwrite_dtype
+                .iter()
+                .map(|(name, dtype)| {
+                    let dtype = dtype.0.clone();
+                    Field::new(name, dtype)
+                })
+                .collect::<Schema>()
         });
 
         let overwrite_dtype_slice = overwrite_dtype_slice.map(|overwrite_dtype| {
@@ -95,14 +103,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 +128,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,7 +1,11 @@
+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,
 };
+use magnus::encoding::{EncodingCapable, Index};
 use polars::chunked_array::object::PolarsObjectSafe;
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
 use polars::datatypes::AnyValue;
@@ -10,9 +14,9 @@ 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::rb_modules::utils;
 use crate::{RbDataFrame, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
 pub(crate) fn slice_to_wrapped<T>(slice: &[T]) -> &[Wrap<T>] {
@@ -21,6 +25,12 @@ pub(crate) fn slice_to_wrapped<T>(slice: &[T]) -> &[Wrap<T>] {
     unsafe { std::mem::transmute(slice) }
 }
 
+pub(crate) fn slice_extract_wrapped<T>(slice: &[Wrap<T>]) -> &[T] {
+    // Safety:
+    // Wrap is transparent.
+    unsafe { std::mem::transmute(slice) }
+}
+
 pub(crate) fn vec_extract_wrapped<T>(buf: Vec<Wrap<T>>) -> Vec<T> {
     // Safety:
     // Wrap is transparent.
@@ -82,6 +92,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 +124,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 +148,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,))
@@ -154,15 +188,25 @@ impl IntoValue for Wrap<AnyValue<'_>> {
                     t.funcall::<_, _, Value>("utc", ()).unwrap()
                 }
             }
-            AnyValue::Duration(_v, _tu) => todo!(),
+            AnyValue::Duration(v, tu) => {
+                let tu = tu.to_ascii();
+                utils().funcall("_to_ruby_duration", (v, tu)).unwrap()
+            }
             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!(),
         }
     }
 }
@@ -182,7 +226,7 @@ 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(),
@@ -210,7 +254,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))
@@ -228,6 +272,75 @@ impl IntoValue for Wrap<DataType> {
     }
 }
 
+impl IntoValue for Wrap<TimeUnit> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let tu = match self.0 {
+            TimeUnit::Nanoseconds => "ns",
+            TimeUnit::Microseconds => "us",
+            TimeUnit::Milliseconds => "ms",
+        };
+        tu.into_value()
+    }
+}
+
+impl IntoValue for Wrap<&Utf8Chunked> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let iter = self.0.into_iter();
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&BinaryChunked> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let iter = self
+            .0
+            .into_iter()
+            .map(|opt_bytes| opt_bytes.map(RString::from_slice));
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&StructChunked> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let s = self.0.clone().into_series();
+        // todo! iterate its chunks and flatten.
+        // make series::iter() accept a chunk index.
+        let s = s.rechunk();
+        let iter = s.iter().map(|av| {
+            if let AnyValue::Struct(_, _, flds) = av {
+                struct_dict(av._iter_struct_av(), flds)
+            } else {
+                unreachable!()
+            }
+        });
+
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&DurationChunked> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let utils = utils();
+        let tu = Wrap(self.0.time_unit()).into_value();
+        let iter = self.0.into_iter().map(|opt_v| {
+            opt_v.map(|v| {
+                utils
+                    .funcall::<_, _, Value>("_to_ruby_duration", (v, tu))
+                    .unwrap()
+            })
+        });
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl TryConvert for Wrap<Field> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let name: String = ob.funcall("name", ())?;
+        let dtype: Wrap<DataType> = ob.funcall("dtype", ())?;
+        Ok(Wrap(Field::new(&name, dtype.0)))
+    }
+}
+
 impl TryConvert for Wrap<DataType> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let dtype = if ob.is_kind_of(class::class()) {
@@ -261,6 +374,20 @@ impl TryConvert for Wrap<DataType> {
                     )))
                 }
             }
+        // TODO improve
+        } else if ob.try_convert::<String>().is_err() {
+            let name = unsafe { ob.class().name() }.into_owned();
+            match name.as_str() {
+                "Polars::Struct" => {
+                    let arr: RArray = ob.funcall("fields", ())?;
+                    let mut fields = Vec::with_capacity(arr.len());
+                    for v in arr.each() {
+                        fields.push(v?.try_convert::<Wrap<Field>>()?.0);
+                    }
+                    DataType::Struct(fields)
+                }
+                _ => todo!(),
+            }
         } else {
             match ob.try_convert::<String>()?.as_str() {
                 "u8" => DataType::UInt8,
@@ -306,7 +433,11 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
         } else if let Some(v) = RFloat::from_value(ob) {
             Ok(AnyValue::Float64(v.to_f64()).into())
         } else if let Some(v) = RString::from_value(ob) {
-            Ok(AnyValue::Utf8Owned(v.to_string()?.into()).into())
+            if v.enc_get() == Index::utf8() {
+                Ok(AnyValue::Utf8Owned(v.to_string()?.into()).into())
+            } else {
+                Ok(AnyValue::BinaryOwned(unsafe { v.as_slice() }.to_vec()).into())
+            }
         // call is_a? for ActiveSupport::TimeWithZone
         } else if ob.funcall::<_, _, bool>("is_a?", (class::time(),))? {
             let sec = ob.funcall::<_, _, i64>("to_i", ())?;
@@ -340,10 +471,14 @@ 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)))
             }
+        } else if ob.is_kind_of(crate::rb_modules::datetime()) {
+            let sec: i64 = ob.funcall("to_i", ())?;
+            let nsec: i64 = ob.funcall("nsec", ())?;
+            Ok(Wrap(AnyValue::Datetime(sec * 1_000_000_000 + nsec, TimeUnit::Nanoseconds, &None)))
         } else if ob.is_kind_of(crate::rb_modules::date()) {
             // convert to DateTime for UTC
             let v = ob
@@ -377,12 +512,12 @@ impl TryConvert for Wrap<Schema> {
 
         let mut schema = Vec::new();
         dict.foreach(|key: String, val: Wrap<DataType>| {
-            schema.push(Field::new(&key, val.0));
+            schema.push(Ok(Field::new(&key, val.0)));
             Ok(ForEach::Continue)
         })
         .unwrap();
 
-        Ok(Wrap(schema.into_iter().into()))
+        Ok(Wrap(schema.into_iter().collect::<RbResult<Schema>>()?))
     }
 }
 
@@ -870,3 +1005,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()
+}