polars-df 0.4.0 → 0.5.0

data/Cargo.toml

@@ -3,7 +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" }
 
 [profile.release]
 strip = true

data/README.md

@@ -58,9 +58,9 @@ Polars.scan_parquet("file.parquet")
 From Active Record
 
 ```ruby
-Polars.read_sql(User.all)
+Polars.read_database(User.all)
 # or
-Polars.read_sql("SELECT * FROM users")
+Polars.read_database("SELECT * FROM users")
 ```
 
 From JSON

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.4.0"
+version = "0.5.0"
 license = "MIT"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
@@ -12,14 +12,15 @@ crate-type = ["cdylib"]
 [dependencies]
 ahash = "0.8"
 magnus = "0.5"
-polars-core = "0.28.0"
+polars-core = "0.29.0"
 serde_json = "1"
 smartstring = "1"
 
 [dependencies.polars]
-version = "0.28.0"
+version = "0.29.0"
 features = [
     "abs",
+    "approx_unique",
     "arange",
     "arg_where",
     "asof_join",
@@ -27,7 +28,7 @@ features = [
     "binary_encoding",
     "concat_str",
     "cse",
-    "csv-file",
+    "csv",
     "cum_agg",
     "cumulative_eval",
     "dataframe_arithmetic",
@@ -49,7 +50,9 @@ features = [
     "json",
     "lazy",
     "lazy_regex",
+    "list_count",
     "list_eval",
+    "list_take",
     "list_to_struct",
     "log",
     "meta",
@@ -75,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

@@ -68,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| {

data/ext/polars/src/conversion.rs

@@ -5,6 +5,7 @@ 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;
@@ -15,6 +16,7 @@ use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 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>] {
@@ -23,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.
@@ -180,7 +188,10 @@ 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) => struct_dict(av._iter_struct_av(), flds),
@@ -261,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()) {
@@ -294,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,
@@ -339,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", ())?;
@@ -377,6 +475,10 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
                     .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
@@ -410,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>>()?))
     }
 }
 

data/ext/polars/src/dataframe.rs

@@ -48,20 +48,25 @@ impl RbDataFrame {
         crate::object::register_object_builder();
 
         let schema =
-            rows_to_schema_supertypes(&rows, infer_schema_length).map_err(RbPolarsErr::from)?;
+            rows_to_schema_supertypes(&rows, infer_schema_length.map(|n| std::cmp::max(1, n)))
+                .map_err(RbPolarsErr::from)?;
         // replace inferred nulls with boolean
         let fields = schema.iter_fields().map(|mut fld| match fld.data_type() {
             DataType::Null => {
                 fld.coerce(DataType::Boolean);
                 fld
             }
+            DataType::Decimal(_, _) => {
+                fld.coerce(DataType::Decimal(None, None));
+                fld
+            }
             _ => fld,
         });
-        let mut schema = Schema::from(fields);
+        let mut schema = Schema::from_iter(fields);
 
         if let Some(schema_overwrite) = schema_overwrite {
             for (i, (name, dtype)) in schema_overwrite.into_iter().enumerate() {
-                if let Some((name_, dtype_)) = schema.get_index_mut(i) {
+                if let Some((name_, dtype_)) = schema.get_at_index_mut(i) {
                     *name_ = name;
 
                     // if user sets dtype unknown, we use the inferred datatype
@@ -139,11 +144,13 @@ impl RbDataFrame {
         };
 
         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| {
@@ -512,12 +519,10 @@ impl RbDataFrame {
                 }
             }
         }
-        let st = st?;
+        let _st = st?;
 
-        match st {
-            // TODO
-            _ => None,
-        }
+        // TODO
+        None
     }
 
     pub fn write_parquet(
@@ -914,10 +919,7 @@ impl RbDataFrame {
             true => pivot_stable,
             false => pivot,
         };
-        let agg_expr = match aggregate_expr {
-            Some(aggregate_expr) => Some(aggregate_expr.inner.clone()),
-            None => None,
-        };
+        let agg_expr = aggregate_expr.map(|aggregate_expr| aggregate_expr.inner.clone());
         let df = fun(
             &self.df.borrow(),
             values,

data/ext/polars/src/error.rs

@@ -22,10 +22,6 @@ impl RbPolarsErr {
     pub fn other(message: String) -> Error {
         Error::new(exception::runtime_error(), message)
     }
-
-    pub fn todo() -> Error {
-        Error::new(exception::runtime_error(), "not implemented yet")
-    }
 }
 
 pub struct RbValueError {}

data/ext/polars/src/expr/binary.rs

@@ -0,0 +1,69 @@
+use polars::prelude::*;
+
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn bin_contains(&self, lit: Vec<u8>) -> Self {
+        self.inner.clone().binary().contains_literal(lit).into()
+    }
+
+    pub fn bin_ends_with(&self, sub: Vec<u8>) -> Self {
+        self.inner.clone().binary().ends_with(sub).into()
+    }
+
+    pub fn bin_starts_with(&self, sub: Vec<u8>) -> Self {
+        self.inner.clone().binary().starts_with(sub).into()
+    }
+
+    pub fn bin_hex_decode(&self, strict: bool) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| {
+                    s.binary()?
+                        .hex_decode(strict)
+                        .map(|s| Some(s.into_series()))
+                },
+                GetOutput::same_type(),
+            )
+            .with_fmt("bin.hex_decode")
+            .into()
+    }
+
+    pub fn bin_base64_decode(&self, strict: bool) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| {
+                    s.binary()?
+                        .base64_decode(strict)
+                        .map(|s| Some(s.into_series()))
+                },
+                GetOutput::same_type(),
+            )
+            .with_fmt("bin.base64_decode")
+            .into()
+    }
+
+    pub fn bin_hex_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.binary().map(|s| Some(s.hex_encode().into_series())),
+                GetOutput::same_type(),
+            )
+            .with_fmt("bin.hex_encode")
+            .into()
+    }
+
+    pub fn bin_base64_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.binary().map(|s| Some(s.base64_encode().into_series())),
+                GetOutput::same_type(),
+            )
+            .with_fmt("bin.base64_encode")
+            .into()
+    }
+}

data/ext/polars/src/expr/categorical.rs

@@ -0,0 +1,10 @@
+use polars::prelude::*;
+
+use crate::conversion::Wrap;
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn cat_set_ordering(&self, ordering: Wrap<CategoricalOrdering>) -> Self {
+        self.inner.clone().cat().set_ordering(ordering.0).into()
+    }
+}

data/ext/polars/src/expr/datetime.rs

@@ -0,0 +1,223 @@
+use polars::prelude::*;
+
+use crate::conversion::Wrap;
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn dt_to_string(&self, format: String) -> Self {
+        self.inner.clone().dt().to_string(&format).into()
+    }
+
+    pub fn dt_offset_by(&self, by: String) -> Self {
+        let by = Duration::parse(&by);
+        self.inner.clone().dt().offset_by(by).into()
+    }
+
+    pub fn dt_epoch_seconds(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                |s| {
+                    s.timestamp(TimeUnit::Milliseconds)
+                        .map(|ca| Some((ca / 1000).into_series()))
+                },
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+
+    pub fn dt_with_time_unit(&self, tu: Wrap<TimeUnit>) -> Self {
+        self.inner.clone().dt().with_time_unit(tu.0).into()
+    }
+
+    pub fn dt_convert_time_zone(&self, tz: TimeZone) -> Self {
+        self.inner.clone().dt().convert_time_zone(tz).into()
+    }
+
+    pub fn dt_cast_time_unit(&self, tu: Wrap<TimeUnit>) -> Self {
+        self.inner.clone().dt().cast_time_unit(tu.0).into()
+    }
+
+    pub fn dt_replace_time_zone(&self, tz: Option<String>, use_earliest: Option<bool>) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .replace_time_zone(tz, use_earliest)
+            .into()
+    }
+
+    #[allow(deprecated)]
+    pub fn dt_tz_localize(&self, tz: String) -> Self {
+        self.inner.clone().dt().tz_localize(tz).into()
+    }
+
+    pub fn dt_truncate(&self, every: String, offset: String) -> Self {
+        self.inner.clone().dt().truncate(&every, &offset).into()
+    }
+
+    pub fn dt_month_start(&self) -> Self {
+        self.inner.clone().dt().month_start().into()
+    }
+
+    pub fn dt_month_end(&self) -> Self {
+        self.inner.clone().dt().month_end().into()
+    }
+
+    pub fn dt_round(&self, every: String, offset: String) -> Self {
+        self.inner.clone().dt().round(&every, &offset).into()
+    }
+
+    pub fn dt_combine(&self, time: &Self, time_unit: Wrap<TimeUnit>) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .combine(time.inner.clone(), time_unit.0)
+            .into()
+    }
+
+    pub fn dt_year(&self) -> Self {
+        self.clone().inner.dt().year().into()
+    }
+
+    pub fn dt_is_leap_year(&self) -> Self {
+        self.clone().inner.dt().is_leap_year().into()
+    }
+
+    pub fn dt_iso_year(&self) -> Self {
+        self.clone().inner.dt().iso_year().into()
+    }
+
+    pub fn dt_quarter(&self) -> Self {
+        self.clone().inner.dt().quarter().into()
+    }
+
+    pub fn dt_month(&self) -> Self {
+        self.clone().inner.dt().month().into()
+    }
+
+    pub fn dt_week(&self) -> Self {
+        self.clone().inner.dt().week().into()
+    }
+
+    pub fn dt_weekday(&self) -> Self {
+        self.clone().inner.dt().weekday().into()
+    }
+
+    pub fn dt_day(&self) -> Self {
+        self.clone().inner.dt().day().into()
+    }
+
+    pub fn dt_ordinal_day(&self) -> Self {
+        self.clone().inner.dt().ordinal_day().into()
+    }
+
+    pub fn dt_time(&self) -> Self {
+        self.clone().inner.dt().time().into()
+    }
+
+    pub fn dt_date(&self) -> Self {
+        self.clone().inner.dt().date().into()
+    }
+
+    pub fn dt_datetime(&self) -> Self {
+        self.clone().inner.dt().datetime().into()
+    }
+
+    pub fn dt_hour(&self) -> Self {
+        self.clone().inner.dt().hour().into()
+    }
+
+    pub fn dt_minute(&self) -> Self {
+        self.clone().inner.dt().minute().into()
+    }
+
+    pub fn dt_second(&self) -> Self {
+        self.clone().inner.dt().second().into()
+    }
+
+    pub fn dt_millisecond(&self) -> Self {
+        self.clone().inner.dt().millisecond().into()
+    }
+
+    pub fn dt_microsecond(&self) -> Self {
+        self.clone().inner.dt().microsecond().into()
+    }
+
+    pub fn dt_nanosecond(&self) -> Self {
+        self.clone().inner.dt().nanosecond().into()
+    }
+
+    pub fn dt_timestamp(&self, tu: Wrap<TimeUnit>) -> Self {
+        self.inner.clone().dt().timestamp(tu.0).into()
+    }
+
+    pub fn duration_days(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.days().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+
+    pub fn duration_hours(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.hours().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+
+    pub fn duration_minutes(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.minutes().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+
+    pub fn duration_seconds(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.seconds().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+
+    pub fn duration_milliseconds(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.milliseconds().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+
+    pub fn duration_microseconds(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.microseconds().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+
+    pub fn duration_nanoseconds(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.nanoseconds().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+}