polars-df 0.1.3 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9e6fb732e5dafe2fde285322554bd9159483cbbdf17d6e2bba9cba9a83563b47
-  data.tar.gz: 1b4249d0c0100f136973c601b8404cb6d92abc632d5ed0476bd93bc5360a11dc
+  metadata.gz: 6abc9619a425d8aaa0255864b063c41835349063aa4919df133ac5a4ceb972f2
+  data.tar.gz: 78372a2a9eeddb3a8080b1d615991415b9ef7752752319e250f143841bfa67f3
 SHA512:
-  metadata.gz: d9414d6f60c489e2b3b72885288822083ba8c04bac4053f4e34c1d53ee805d164f17fe4b8b3a8f4ff562550bcc657f374bea6e250b52985367f601ea50e3037f
-  data.tar.gz: 9e3a7cfe105f03ec20e9c26aa38c1475074ccc1ea057a170a97b7068b41943d561d50af49bb1d1f74b7705809dc1375900f542ab93683ba627dea080274f6d91
+  metadata.gz: e6fb27a50908c07e5f2f72c81171f07bfdf0999b5148421bdeb1ad7dc69cee1f0bae02021fa18fdad6d1740ea9273464daec513db5e3c7906d5839e77b7d6a66
+  data.tar.gz: 2eb9df841575711a057dd1ca2986403667306ead52cf540491899ffaa184d4878c1bdfc1015e3f5831c12c668de0d4126cbab7c63d1770684e10012f3d28183f

data/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.1.4 (2022-12-02)
+
+- Added more methods
+- Improved performance
+
 ## 0.1.3 (2022-11-27)
 
 - Added more methods

data/Cargo.lock

@@ -483,6 +483,12 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ee1b05cbd864bcaecbd3455d6d967862d446e4ebfc3c2e5e5b9841e53cba6673"
 
+[[package]]
+name = "fs_extra"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2022715d62ab30faffd124d40b76f4134a550a87792276512b18d63272333394"
+
 [[package]]
 name = "futures"
 version = "0.3.25"
@@ -646,6 +652,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "hex"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
+
 [[package]]
 name = "indexmap"
 version = "1.8.0"
@@ -663,6 +675,27 @@ version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4217ad341ebadf8d8e724e264f13e593e0648f5b3e94b3896a5df283be015ecc"
 
+[[package]]
+name = "jemalloc-sys"
+version = "0.5.2+5.3.0-patched"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "134163979b6eed9564c98637b710b40979939ba351f59952708234ea11b5f3f8"
+dependencies = [
+ "cc",
+ "fs_extra",
+ "libc",
+]
+
+[[package]]
+name = "jemallocator"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16c2514137880c52b0b4822b563fadd38257c1f380858addb74a400889696ea6"
+dependencies = [
+ "jemalloc-sys",
+ "libc",
+]
+
 [[package]]
 name = "jobserver"
 version = "0.1.25"
@@ -683,13 +716,23 @@ dependencies = [
 
 [[package]]
 name = "json-deserializer"
-version = "0.4.2"
+version = "0.4.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d784d2d481d0bace3450572391d6076dd6d10c66c0ebc1a0be037b3b420664bd"
+checksum = "daba674f7eecf80fe8bbbf196340908ad1a22510fe71fd6111bb50f441b26440"
 dependencies = [
  "indexmap",
 ]
 
+[[package]]
+name = "jsonpath_lib"
+version = "0.3.0"
+source = "git+https://github.com/ritchie46/jsonpath?rev=24eaf0b4416edff38a4d1b6b17bc4b9f3f047b4b#24eaf0b4416edff38a4d1b6b17bc4b9f3f047b4b"
+dependencies = [
+ "log",
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "lazy_static"
 version = "1.4.0"
@@ -797,6 +840,16 @@ version = "0.2.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "348108ab3fba42ec82ff6e9564fc4ca0247bdccdc68dd8af9764bbc79c3c8ffb"
 
+[[package]]
+name = "libmimalloc-sys"
+version = "0.1.28"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "04d1c67deb83e6b75fa4fe3309e09cfeade12e7721d95322af500d3814ea60c9"
+dependencies = [
+ "cc",
+ "libc",
+]
+
 [[package]]
 name = "lock_api"
 version = "0.4.9"
@@ -882,6 +935,15 @@ dependencies = [
  "autocfg",
 ]
 
+[[package]]
+name = "mimalloc"
+version = "0.1.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b2374e2999959a7b583e1811a1ddbf1d3a4b9496eceb9746f1192a59d871eca"
+dependencies = [
+ "libmimalloc-sys",
+]
+
 [[package]]
 name = "minimal-lexical"
 version = "0.2.1"
@@ -1160,10 +1222,13 @@ dependencies = [
 
 [[package]]
 name = "polars"
-version = "0.1.3"
+version = "0.1.4"
 dependencies = [
+ "jemallocator",
  "magnus",
+ "mimalloc",
  "polars 0.25.1",
+ "polars-core",
  "serde_json",
 ]
 
@@ -1202,11 +1267,13 @@ dependencies = [
  "ahash 0.8.2",
  "anyhow",
  "arrow2",
+ "base64",
  "bitflags",
  "chrono",
  "chrono-tz",
  "comfy-table",
  "hashbrown 0.12.3",
+ "hex",
  "indexmap",
  "num",
  "once_cell",
@@ -1277,9 +1344,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bfd3f6552b3e9539634c35047f372db331b6227f75c36fcbe4670ab58bbcbeb3"
 dependencies = [
  "arrow2",
+ "jsonpath_lib",
  "polars-arrow",
  "polars-core",
  "polars-utils",
+ "serde_json",
 ]
 
 [[package]]
@@ -1548,6 +1617,7 @@ version = "1.0.88"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8e8b3801309262e8184d9687fb697586833e939767aea0dda89f5a8e650e8bd7"
 dependencies = [
+ "indexmap",
  "itoa",
  "ryu",
  "serde",

data/Cargo.toml

@@ -1,5 +1,8 @@
 [workspace]
 members = ["ext/polars"]
 
+[patch.crates-io]
+jsonpath_lib = { git = "https://github.com/ritchie46/jsonpath", rev = "24eaf0b4416edff38a4d1b6b17bc4b9f3f047b4b" }
+
 [profile.release]
 strip = true

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.1.3"
+version = "0.1.4"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
 publish = false
@@ -10,6 +10,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 magnus = "0.4"
+polars-core = "0.25.1"
 serde_json = "1"
 
 [dependencies.polars]
@@ -19,6 +20,7 @@ features = [
     "arange",
     "arg_where",
     "concat_str",
+    "cse",
     "csv-file",
     "cum_agg",
     "cumulative_eval",
@@ -30,6 +32,7 @@ features = [
     "dtype-full",
     "dynamic_groupby",
     "ewma",
+    "extract_jsonpath",
     "fmt",
     "horizontal_concat",
     "interpolate",
@@ -49,6 +52,7 @@ features = [
     "parquet",
     "partition_by",
     "pct_change",
+    "performant",
     "product",
     "propagate_nans",
     "random",
@@ -61,6 +65,7 @@ features = [
     "semi_anti_join",
     "serde-lazy",
     "sign",
+    "string_encoding",
     "string_justify",
     "strings",
     "timezones",
@@ -69,3 +74,9 @@ features = [
     "trigonometry",
     "unique_counts",
 ]
+
+[target.'cfg(target_os = "linux")'.dependencies]
+jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }
+
+[target.'cfg(not(target_os = "linux"))'.dependencies]
+mimalloc = { version = "0.1", default-features = false }

data/ext/polars/src/conversion.rs

@@ -5,6 +5,8 @@ use polars::datatypes::AnyValue;
 use polars::frame::DataFrame;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
+use std::fmt::{Display, Formatter};
+use std::hash::{Hash, Hasher};
 
 use crate::{RbDataFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
@@ -87,6 +89,25 @@ impl From<Wrap<AnyValue<'_>>> for Value {
                 .unwrap()
                 .funcall::<_, _, Value>("to_date", ())
                 .unwrap(),
+            AnyValue::Datetime(v, tu, tz) => {
+                let t = match tu {
+                    TimeUnit::Nanoseconds => todo!(),
+                    TimeUnit::Microseconds => {
+                        let sec = v / 1000000;
+                        let subsec = v % 1000000;
+                        class::time()
+                            .funcall::<_, _, Value>("at", (sec, subsec, Symbol::new("usec")))
+                            .unwrap()
+                    }
+                    TimeUnit::Milliseconds => todo!(),
+                };
+
+                if tz.is_some() {
+                    todo!();
+                } else {
+                    t.funcall::<_, _, Value>("utc", ()).unwrap()
+                }
+            }
             _ => todo!(),
         }
     }
@@ -238,6 +259,22 @@ impl TryConvert for Wrap<JoinType> {
     }
 }
 
+impl TryConvert for Wrap<ListToStructWidthStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "first_non_null" => ListToStructWidthStrategy::FirstNonNull,
+            "max_width" => ListToStructWidthStrategy::MaxWidth,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "n_field_strategy must be one of {{'first_non_null', 'max_width'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<NullBehavior> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match ob.try_convert::<String>()?.as_str() {
@@ -425,10 +462,47 @@ pub fn parse_parquet_compression(
     Ok(parsed)
 }
 
+#[derive(Clone, Debug)]
 pub struct ObjectValue {
     pub inner: Value,
 }
 
+impl Hash for ObjectValue {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        let h = self
+            .inner
+            .funcall::<_, _, isize>("hash", ())
+            .expect("should be hashable");
+        state.write_isize(h)
+    }
+}
+
+impl Eq for ObjectValue {}
+
+impl PartialEq for ObjectValue {
+    fn eq(&self, other: &Self) -> bool {
+        self.inner.eql(&other.inner).unwrap_or(false)
+    }
+}
+
+impl Display for ObjectValue {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.inner)
+    }
+}
+
+impl PolarsObject for ObjectValue {
+    fn type_name() -> &'static str {
+        "object"
+    }
+}
+
+impl From<Value> for ObjectValue {
+    fn from(v: Value) -> Self {
+        Self { inner: v }
+    }
+}
+
 impl From<&dyn PolarsObjectSafe> for &ObjectValue {
     fn from(val: &dyn PolarsObjectSafe) -> Self {
         unsafe { &*(val as *const dyn PolarsObjectSafe as *const ObjectValue) }
@@ -440,3 +514,9 @@ impl ObjectValue {
         self.inner
     }
 }
+
+impl Default for ObjectValue {
+    fn default() -> Self {
+        ObjectValue { inner: *QNIL }
+    }
+}

data/ext/polars/src/error.rs

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

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

@@ -211,7 +211,7 @@ impl RbLazyFrame {
         projection_pushdown: bool,
         simplify_expr: bool,
         slice_pushdown: bool,
-        _cse: bool,
+        cse: bool,
         allow_streaming: bool,
     ) -> RbLazyFrame {
         let ldf = self.ldf.clone();
@@ -220,7 +220,7 @@ impl RbLazyFrame {
             .with_predicate_pushdown(predicate_pushdown)
             .with_simplify_expr(simplify_expr)
             .with_slice_pushdown(slice_pushdown)
-            // .with_common_subplan_elimination(cse)
+            .with_common_subplan_elimination(cse)
             .with_streaming(allow_streaming)
             .with_projection_pushdown(projection_pushdown);
         ldf.into()

data/ext/polars/src/lazy/dsl.rs

@@ -659,6 +659,65 @@ impl RbExpr {
         self.inner.clone().str().starts_with(sub).into()
     }
 
+    pub fn str_hex_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8().map(|s| s.hex_encode().into_series()),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.hex_encode")
+            .into()
+    }
+
+    pub fn str_hex_decode(&self, strict: Option<bool>) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8()?.hex_decode(strict).map(|s| s.into_series()),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.hex_decode")
+            .into()
+    }
+
+    pub fn str_base64_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8().map(|s| s.base64_encode().into_series()),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.base64_encode")
+            .into()
+    }
+
+    pub fn str_base64_decode(&self, strict: Option<bool>) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8()?.base64_decode(strict).map(|s| s.into_series()),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.base64_decode")
+            .into()
+    }
+
+    pub fn str_json_path_match(&self, pat: String) -> Self {
+        let function = move |s: Series| {
+            let ca = s.utf8()?;
+            match ca.json_path_match(&pat) {
+                Ok(ca) => Ok(ca.into_series()),
+                Err(e) => Err(PolarsError::ComputeError(format!("{:?}", e).into())),
+            }
+        };
+        self.clone()
+            .inner
+            .map(function, GetOutput::from_type(DataType::Utf8))
+            .with_fmt("str.json_path_match")
+            .into()
+    }
+
     pub fn str_extract(&self, pat: String, group_index: usize) -> Self {
         self.inner.clone().str().extract(&pat, group_index).into()
     }
@@ -887,6 +946,10 @@ impl RbExpr {
         self.inner.clone().dt().round(&every, &offset).into()
     }
 
+    pub fn dot(&self, other: &RbExpr) -> Self {
+        self.inner.clone().dot(other.inner.clone()).into()
+    }
+
     pub fn reinterpret(&self, signed: bool) -> Self {
         let function = move |s: Series| reinterpret(&s, signed);
         let dt = if signed {
@@ -1208,6 +1271,28 @@ impl RbExpr {
             .into()
     }
 
+    pub fn lst_to_struct(
+        &self,
+        width_strat: Wrap<ListToStructWidthStrategy>,
+        _name_gen: Option<Value>,
+    ) -> RbResult<Self> {
+        // TODO fix
+        let name_gen = None;
+        // let name_gen = name_gen.map(|lambda| {
+        //     Arc::new(move |idx: usize| {
+        //         let out: Value = lambda.funcall("call", (idx,)).unwrap();
+        //         out.try_convert::<String>().unwrap()
+        //     }) as NameGenerator
+        // });
+
+        Ok(self
+            .inner
+            .clone()
+            .arr()
+            .to_struct(width_strat.0, name_gen)
+            .into())
+    }
+
     pub fn rank(&self, method: Wrap<RankMethod>, reverse: bool) -> Self {
         let options = RankOptions {
             method: method.0,
@@ -1433,6 +1518,19 @@ pub fn repeat(value: Value, n_times: &RbExpr) -> RbResult<RbExpr> {
     }
 }
 
+pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    polars::lazy::dsl::pearson_corr(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
+pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool) -> RbExpr {
+    polars::lazy::dsl::spearman_rank_corr(a.inner.clone(), b.inner.clone(), ddof, propagate_nans)
+        .into()
+}
+
+pub fn cov(a: &RbExpr, b: &RbExpr) -> RbExpr {
+    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone()).into()
+}
+
 #[magnus::wrap(class = "Polars::RbWhen")]
 #[derive(Clone)]
 pub struct RbWhen {

data/ext/polars/src/lib.rs

@@ -4,6 +4,7 @@ mod dataframe;
 mod error;
 mod file;
 mod lazy;
+mod list_construction;
 mod series;
 mod set;
 mod utils;
@@ -27,6 +28,20 @@ use polars::functions::{diag_concat_df, hor_concat_df};
 use polars::prelude::{ClosedWindow, Duration, IntoSeries, TimeZone};
 use series::RbSeries;
 
+#[cfg(target_os = "linux")]
+use jemallocator::Jemalloc;
+
+#[cfg(not(target_os = "linux"))]
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+#[cfg(target_os = "linux")]
+static GLOBAL: Jemalloc = Jemalloc;
+
+#[global_allocator]
+#[cfg(not(target_os = "linux"))]
+static GLOBAL: MiMalloc = MiMalloc;
+
 type RbResult<T> = Result<T, Error>;
 
 fn module() -> RModule {
@@ -281,6 +296,14 @@ fn init() -> RbResult<()> {
     class.define_method("str_contains", method!(RbExpr::str_contains, 2))?;
     class.define_method("str_ends_with", method!(RbExpr::str_ends_with, 1))?;
     class.define_method("str_starts_with", method!(RbExpr::str_starts_with, 1))?;
+    class.define_method("str_hex_encode", method!(RbExpr::str_hex_encode, 0))?;
+    class.define_method("str_hex_decode", method!(RbExpr::str_hex_decode, 1))?;
+    class.define_method("str_base64_encode", method!(RbExpr::str_base64_encode, 0))?;
+    class.define_method("str_base64_decode", method!(RbExpr::str_base64_decode, 1))?;
+    class.define_method(
+        "str_json_path_match",
+        method!(RbExpr::str_json_path_match, 1),
+    )?;
     class.define_method("str_extract", method!(RbExpr::str_extract, 2))?;
     class.define_method("str_extract_all", method!(RbExpr::str_extract_all, 1))?;
     class.define_method("count_match", method!(RbExpr::count_match, 1))?;
@@ -338,6 +361,7 @@ fn init() -> RbResult<()> {
     class.define_method("dt_tz_localize", method!(RbExpr::dt_tz_localize, 1))?;
     class.define_method("dt_truncate", method!(RbExpr::dt_truncate, 2))?;
     class.define_method("dt_round", method!(RbExpr::dt_round, 2))?;
+    class.define_method("dot", method!(RbExpr::dot, 1))?;
     class.define_method("reinterpret", method!(RbExpr::reinterpret, 1))?;
     class.define_method("mode", method!(RbExpr::mode, 0))?;
     class.define_method("keep_name", method!(RbExpr::keep_name, 0))?;
@@ -372,6 +396,7 @@ fn init() -> RbResult<()> {
     class.define_method("lst_slice", method!(RbExpr::lst_slice, 2))?;
     class.define_method("lst_eval", method!(RbExpr::lst_eval, 2))?;
     class.define_method("cumulative_eval", method!(RbExpr::cumulative_eval, 3))?;
+    class.define_method("lst_to_struct", method!(RbExpr::lst_to_struct, 2))?;
     class.define_method("rank", method!(RbExpr::rank, 2))?;
     class.define_method("diff", method!(RbExpr::diff, 2))?;
     class.define_method("pct_change", method!(RbExpr::pct_change, 1))?;
@@ -424,6 +449,12 @@ fn init() -> RbResult<()> {
     class.define_singleton_method("lit", function!(crate::lazy::dsl::lit, 1))?;
     class.define_singleton_method("arange", function!(crate::lazy::dsl::arange, 3))?;
     class.define_singleton_method("repeat", function!(crate::lazy::dsl::repeat, 2))?;
+    class.define_singleton_method("pearson_corr", function!(crate::lazy::dsl::pearson_corr, 3))?;
+    class.define_singleton_method(
+        "spearman_rank_corr",
+        function!(crate::lazy::dsl::spearman_rank_corr, 4),
+    )?;
+    class.define_singleton_method("cov", function!(crate::lazy::dsl::cov, 2))?;
     class.define_singleton_method("when", function!(crate::lazy::dsl::when, 1))?;
     class.define_singleton_method("concat_str", function!(crate::lazy::dsl::concat_str, 2))?;
     class.define_singleton_method("concat_lst", function!(crate::lazy::dsl::concat_lst, 1))?;
@@ -507,7 +538,10 @@ fn init() -> RbResult<()> {
     class.define_singleton_method("new_opt_f32", function!(RbSeries::new_opt_f32, 3))?;
     class.define_singleton_method("new_opt_f64", function!(RbSeries::new_opt_f64, 3))?;
     class.define_singleton_method("new_str", function!(RbSeries::new_str, 3))?;
+    class.define_singleton_method("new_object", function!(RbSeries::new_object, 3))?;
+    class.define_singleton_method("new_list", function!(RbSeries::new_list, 3))?;
     class.define_singleton_method("new_opt_date", function!(RbSeries::new_opt_date, 3))?;
+    class.define_singleton_method("new_opt_datetime", function!(RbSeries::new_opt_datetime, 3))?;
     class.define_method("is_sorted_flag", method!(RbSeries::is_sorted_flag, 0))?;
     class.define_method(
         "is_sorted_reverse_flag",