polars-df 0.1.4 → 0.1.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 6abc9619a425d8aaa0255864b063c41835349063aa4919df133ac5a4ceb972f2
-  data.tar.gz: 78372a2a9eeddb3a8080b1d615991415b9ef7752752319e250f143841bfa67f3
+  metadata.gz: 3a08e866e51227716cd3cb4454835016a7d61e30e964fe76a8b99704dcb60a12
+  data.tar.gz: 1f30c3fdd47ebf52a311909aa26ba4b6d64e426622455854b9bbc660de1229b3
 SHA512:
-  metadata.gz: e6fb27a50908c07e5f2f72c81171f07bfdf0999b5148421bdeb1ad7dc69cee1f0bae02021fa18fdad6d1740ea9273464daec513db5e3c7906d5839e77b7d6a66
-  data.tar.gz: 2eb9df841575711a057dd1ca2986403667306ead52cf540491899ffaa184d4878c1bdfc1015e3f5831c12c668de0d4126cbab7c63d1770684e10012f3d28183f
+  metadata.gz: 1531fff4fc2fab8b2dc72709a69fb2890c215ae08e4223aa32262dbb4b0debb4b6f2fbab1e8138953871f5d02d462abfaba49cc7f22a66e25aa7d60f128a89bc
+  data.tar.gz: e1041d708e2f8046c14c565a65879fa4e5c6671cf526736a3f8418a82dfa70e17692a96d383e43f393d8761e5f29f717d63185ae1ed3f0793a6876be2d946fc0

data/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.1.5 (2022-12-22)
+
+- Added `read_avro` and `write_avro` methods
+- Added more methods
+
 ## 0.1.4 (2022-12-02)
 
 - Added more methods

data/Cargo.lock

@@ -8,6 +8,12 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
+[[package]]
+name = "adler32"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aae1277d39aeec15cb388266ecc24b11c80469deae6067e17a1a7aa9e5c1f234"
+
 [[package]]
 name = "ahash"
 version = "0.7.6"
@@ -85,6 +91,7 @@ checksum = "ee6f62e41078c967a4c063fcbdfd3801a2a9632276402c045311c4d73d0845f3"
 dependencies = [
  "ahash 0.7.6",
  "arrow-format",
+ "avro-schema",
  "base64",
  "bytemuck",
  "chrono",
@@ -147,6 +154,20 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 
+[[package]]
+name = "avro-schema"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b5281855b39aba9684d2f47bf96983fbfd8f1725f12fabb0513a8ab879647bbd"
+dependencies = [
+ "crc",
+ "fallible-streaming-iterator",
+ "libflate",
+ "serde",
+ "serde_json",
+ "snap",
+]
+
 [[package]]
 name = "base64"
 version = "0.13.1"
@@ -314,6 +335,21 @@ dependencies = [
  "unicode-width",
 ]
 
+[[package]]
+name = "crc"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49fc9a695bca7f35f5f4c15cddc84415f66a74ea78eef08e90c5024f2b540e23"
+dependencies = [
+ "crc-catalog",
+]
+
+[[package]]
+name = "crc-catalog"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ccaeedb56da03b09f598226e25e80088cb4cd25f316e6e4df7d695f0feeb1403"
+
 [[package]]
 name = "crc32fast"
 version = "1.3.2"
@@ -824,6 +860,26 @@ version = "0.2.121"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "efaa7b300f3b5fe8eb6bf21ce3895e1751d9665086af2d64b42f19701015ff4f"
 
+[[package]]
+name = "libflate"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05605ab2bce11bcfc0e9c635ff29ef8b2ea83f29be257ee7d730cac3ee373093"
+dependencies = [
+ "adler32",
+ "crc32fast",
+ "libflate_lz77",
+]
+
+[[package]]
+name = "libflate_lz77"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "39a734c0493409afcd49deee13c006a04e3586b9761a03543c6272c9c51f2f5a"
+dependencies = [
+ "rle-decode-fast",
+]
+
 [[package]]
 name = "libloading"
 version = "0.7.4"
@@ -903,8 +959,7 @@ dependencies = [
 [[package]]
 name = "magnus-macros"
 version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "acc8ba6908cb0f67a4e75cb48fc81a1f0e6a6dd1501936e0c9e2c7c8f9f18e05"
+source = "git+https://github.com/matsadler/magnus#ae792419bed70107d4c930e1f8193272750b9fd2"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -1222,8 +1277,9 @@ dependencies = [
 
 [[package]]
 name = "polars"
-version = "0.1.4"
+version = "0.1.5"
 dependencies = [
+ "ahash 0.8.2",
  "jemallocator",
  "magnus",
  "mimalloc",
@@ -1500,18 +1556,18 @@ dependencies = [
 
 [[package]]
 name = "rb-sys"
-version = "0.9.44"
+version = "0.9.48"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "31f48777b8161ff5c077ad74ce486ebe963ca8a92257512bab473b405a80d69f"
+checksum = "dfc6b8f3bf2d04b0180e243ceeb033b51ca267d839aa1c12fa25f262c17d0596"
 dependencies = [
  "rb-sys-build",
 ]
 
 [[package]]
 name = "rb-sys-build"
-version = "0.9.44"
+version = "0.9.48"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a46785122aff7077527b78c2518d739c45dc0fbc410a2b8361076ff4bbf993f9"
+checksum = "2cd591ebf22c45a44e51192fbeebba473aea0fe2a708b0b24665a13010c58b8d"
 dependencies = [
  "bindgen",
  "regex",
@@ -1521,8 +1577,7 @@ dependencies = [
 [[package]]
 name = "rb-sys-env"
 version = "0.1.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "74c38752410925faeb82c400c06ba2fd9ee6aa8f719dd33994c9e53f5242d25f"
+source = "git+https://github.com/oxidize-rb/rb-sys#93c4f97a244168b9ebc2c5682275e7281421f4b8"
 
 [[package]]
 name = "redox_syscall"
@@ -1561,6 +1616,12 @@ version = "0.6.28"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "456c603be3e8d448b072f410900c09faf164fbce2d480456f50eea6e25f9c848"
 
+[[package]]
+name = "rle-decode-fast"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3582f63211428f83597b51b2ddb88e2a91a9d52d12831f9d08f5e624e8977422"
+
 [[package]]
 name = "rustc-hash"
 version = "1.1.0"

data/Cargo.toml

@@ -3,6 +3,8 @@ members = ["ext/polars"]
 
 [patch.crates-io]
 jsonpath_lib = { git = "https://github.com/ritchie46/jsonpath", rev = "24eaf0b4416edff38a4d1b6b17bc4b9f3f047b4b" }
+magnus-macros = { git = "https://github.com/matsadler/magnus" }
+rb-sys-env = { git = "https://github.com/oxidize-rb/rb-sys" }
 
 [profile.release]
 strip = true

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.1.4"
+version = "0.1.5"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
 publish = false
@@ -9,6 +9,7 @@ publish = false
 crate-type = ["cdylib"]
 
 [dependencies]
+ahash = "0.8"
 magnus = "0.4"
 polars-core = "0.25.1"
 serde_json = "1"
@@ -19,6 +20,8 @@ features = [
     "abs",
     "arange",
     "arg_where",
+    "asof_join",
+    "avro",
     "concat_str",
     "cse",
     "csv-file",
@@ -53,6 +56,7 @@ features = [
     "partition_by",
     "pct_change",
     "performant",
+    "pivot",
     "product",
     "propagate_nans",
     "random",
@@ -61,6 +65,7 @@ features = [
     "repeat_by",
     "rolling_window",
     "round_series",
+    "row_hash",
     "search_sorted",
     "semi_anti_join",
     "serde-lazy",

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

@@ -0,0 +1,292 @@
+use magnus::{class, RArray, TryConvert, Value};
+use polars::prelude::*;
+use polars_core::frame::row::{rows_to_schema_first_non_null, Row};
+
+use super::*;
+use crate::{RbDataFrame, RbPolarsErr, RbSeries, Wrap};
+
+pub fn apply_lambda_unknown<'a>(
+    df: &'a DataFrame,
+    lambda: Value,
+    inference_size: usize,
+) -> RbResult<(Value, bool)> {
+    let columns = df.get_columns();
+    let mut null_count = 0;
+
+    for idx in 0..df.height() {
+        let iter = columns.iter().map(|s: &Series| Wrap(s.get(idx)));
+        let arg = (iter.collect::<Vec<Wrap<AnyValue>>>(),);
+        let out: Value = lambda.funcall("call", arg)?;
+
+        if out.is_nil() {
+            null_count += 1;
+            continue;
+        } else if out.is_kind_of(class::true_class()) || out.is_kind_of(class::false_class()) {
+            let first_value = out.try_convert::<bool>().ok();
+            return Ok((
+                RbSeries::new(
+                    apply_lambda_with_bool_out_type(df, lambda, null_count, first_value)
+                        .into_series(),
+                )
+                .into(),
+                false,
+            ));
+        } else if out.is_kind_of(class::float()) {
+            let first_value = out.try_convert::<f64>().ok();
+
+            return Ok((
+                RbSeries::new(
+                    apply_lambda_with_primitive_out_type::<Float64Type>(
+                        df,
+                        lambda,
+                        null_count,
+                        first_value,
+                    )
+                    .into_series(),
+                )
+                .into(),
+                false,
+            ));
+        } else if out.is_kind_of(class::integer()) {
+            let first_value = out.try_convert::<i64>().ok();
+            return Ok((
+                RbSeries::new(
+                    apply_lambda_with_primitive_out_type::<Int64Type>(
+                        df,
+                        lambda,
+                        null_count,
+                        first_value,
+                    )
+                    .into_series(),
+                )
+                .into(),
+                false,
+            ));
+        // } else if out.is_kind_of(class::string()) {
+        //     let first_value = out.try_convert::<String>().ok();
+        //     return Ok((
+        //         RbSeries::new(
+        //             apply_lambda_with_utf8_out_type(df, lambda, null_count, first_value)
+        //                 .into_series(),
+        //         )
+        //         .into(),
+        //         false,
+        //     ));
+        } else if out.respond_to("_s", true)? {
+            let rb_rbseries: Value = out.funcall("_s", ()).unwrap();
+            let series = rb_rbseries
+                .try_convert::<&RbSeries>()
+                .unwrap()
+                .series
+                .borrow();
+            let dt = series.dtype();
+            return Ok((
+                RbSeries::new(
+                    apply_lambda_with_list_out_type(df, lambda, null_count, Some(&series), dt)?
+                        .into_series(),
+                )
+                .into(),
+                false,
+            ));
+        } else if out.try_convert::<Wrap<Row<'a>>>().is_ok() {
+            let first_value = out.try_convert::<Wrap<Row<'a>>>().unwrap().0;
+            return Ok((
+                RbDataFrame::from(
+                    apply_lambda_with_rows_output(
+                        df,
+                        lambda,
+                        null_count,
+                        first_value,
+                        inference_size,
+                    )
+                    .map_err(RbPolarsErr::from)?,
+                )
+                .into(),
+                true,
+            ));
+        } else if out.is_kind_of(class::array()) {
+            return Err(RbPolarsErr::other(
+                "A list output type is invalid. Do you mean to create polars List Series?\
+Then return a Series object."
+                    .into(),
+            ));
+        } else {
+            return Err(RbPolarsErr::other("Could not determine output type".into()));
+        }
+    }
+    Err(RbPolarsErr::other("Could not determine output type".into()))
+}
+
+fn apply_iter<T>(
+    df: &DataFrame,
+    lambda: Value,
+    init_null_count: usize,
+    skip: usize,
+) -> impl Iterator<Item = Option<T>> + '_
+where
+    T: TryConvert,
+{
+    let columns = df.get_columns();
+    ((init_null_count + skip)..df.height()).map(move |idx| {
+        let iter = columns.iter().map(|s: &Series| Wrap(s.get(idx)));
+        let tpl = (iter.collect::<Vec<Wrap<AnyValue>>>(),);
+        match lambda.funcall::<_, _, Value>("call", tpl) {
+            Ok(val) => val.try_convert::<T>().ok(),
+            Err(e) => panic!("ruby function failed {}", e),
+        }
+    })
+}
+
+/// Apply a lambda with a primitive output type
+pub fn apply_lambda_with_primitive_out_type<D>(
+    df: &DataFrame,
+    lambda: Value,
+    init_null_count: usize,
+    first_value: Option<D::Native>,
+) -> ChunkedArray<D>
+where
+    D: RbArrowPrimitiveType,
+    D::Native: Into<Value> + TryConvert,
+{
+    let skip = usize::from(first_value.is_some());
+    if init_null_count == df.height() {
+        ChunkedArray::full_null("apply", df.height())
+    } else {
+        let iter = apply_iter(df, lambda, init_null_count, skip);
+        iterator_to_primitive(iter, init_null_count, first_value, "apply", df.height())
+    }
+}
+
+/// Apply a lambda with a boolean output type
+pub fn apply_lambda_with_bool_out_type(
+    df: &DataFrame,
+    lambda: Value,
+    init_null_count: usize,
+    first_value: Option<bool>,
+) -> ChunkedArray<BooleanType> {
+    let skip = usize::from(first_value.is_some());
+    if init_null_count == df.height() {
+        ChunkedArray::full_null("apply", df.height())
+    } else {
+        let iter = apply_iter(df, lambda, init_null_count, skip);
+        iterator_to_bool(iter, init_null_count, first_value, "apply", df.height())
+    }
+}
+
+/// Apply a lambda with utf8 output type
+pub fn apply_lambda_with_utf8_out_type(
+    df: &DataFrame,
+    lambda: Value,
+    init_null_count: usize,
+    first_value: Option<&str>,
+) -> Utf8Chunked {
+    let skip = usize::from(first_value.is_some());
+    if init_null_count == df.height() {
+        ChunkedArray::full_null("apply", df.height())
+    } else {
+        let iter = apply_iter::<String>(df, lambda, init_null_count, skip);
+        iterator_to_utf8(iter, init_null_count, first_value, "apply", df.height())
+    }
+}
+
+/// Apply a lambda with list output type
+pub fn apply_lambda_with_list_out_type<'a>(
+    df: &'a DataFrame,
+    lambda: Value,
+    init_null_count: usize,
+    first_value: Option<&Series>,
+    dt: &DataType,
+) -> RbResult<ListChunked> {
+    let columns = df.get_columns();
+
+    let skip = usize::from(first_value.is_some());
+    if init_null_count == df.height() {
+        Ok(ChunkedArray::full_null("apply", df.height()))
+    } else {
+        let iter = ((init_null_count + skip)..df.height()).map(|idx| {
+            let iter = columns.iter().map(|s: &Series| Wrap(s.get(idx)));
+            let tpl = (iter.collect::<Vec<Wrap<AnyValue>>>(),);
+            match lambda.funcall::<_, _, Value>("call", tpl) {
+                Ok(val) => match val.funcall::<_, _, Value>("_s", ()) {
+                    Ok(val) => val
+                        .try_convert::<&RbSeries>()
+                        .ok()
+                        .map(|ps| ps.series.borrow().clone()),
+                    Err(_) => {
+                        if val.is_nil() {
+                            None
+                        } else {
+                            panic!("should return a Series, got a {:?}", val)
+                        }
+                    }
+                },
+                Err(e) => panic!("ruby function failed {}", e),
+            }
+        });
+        iterator_to_list(dt, iter, init_null_count, first_value, "apply", df.height())
+    }
+}
+
+pub fn apply_lambda_with_rows_output<'a>(
+    df: &'a DataFrame,
+    lambda: Value,
+    init_null_count: usize,
+    first_value: Row<'a>,
+    inference_size: usize,
+) -> PolarsResult<DataFrame> {
+    let columns = df.get_columns();
+    let width = first_value.0.len();
+    let null_row = Row::new(vec![AnyValue::Null; width]);
+
+    let mut row_buf = Row::default();
+
+    let skip = 1;
+    let mut row_iter = ((init_null_count + skip)..df.height()).map(|idx| {
+        let iter = columns.iter().map(|s: &Series| Wrap(s.get(idx)));
+        let tpl = (iter.collect::<Vec<Wrap<AnyValue>>>(),);
+        match lambda.funcall::<_, _, Value>("call", tpl) {
+            Ok(val) => {
+                match val.try_convert::<RArray>().ok() {
+                    Some(tuple) => {
+                        row_buf.0.clear();
+                        for v in tuple.each() {
+                            let v = v.unwrap().try_convert::<Wrap<AnyValue>>().unwrap().0;
+                            row_buf.0.push(v);
+                        }
+                        let ptr = &row_buf as *const Row;
+                        // Safety:
+                        // we know that row constructor of polars dataframe does not keep a reference
+                        // to the row. Before we mutate the row buf again, the reference is dropped.
+                        // we only cannot prove it to the compiler.
+                        // we still do this because it saves a Vec allocation in a hot loop.
+                        unsafe { &*ptr }
+                    }
+                    None => &null_row,
+                }
+            }
+            Err(e) => panic!("ruby function failed {}", e),
+        }
+    });
+
+    // first rows for schema inference
+    let mut buf = Vec::with_capacity(inference_size);
+    buf.push(first_value);
+    buf.extend((&mut row_iter).take(inference_size).cloned());
+    let schema = rows_to_schema_first_non_null(&buf, Some(50));
+
+    if init_null_count > 0 {
+        // Safety: we know the iterators size
+        let iter = unsafe {
+            (0..init_null_count)
+                .map(|_| &null_row)
+                .chain(buf.iter())
+                .chain(row_iter)
+                .trust_my_length(df.height())
+        };
+        DataFrame::from_rows_iter_and_schema(iter, &schema)
+    } else {
+        // Safety: we know the iterators size
+        let iter = unsafe { buf.iter().chain(row_iter).trust_my_length(df.height()) };
+        DataFrame::from_rows_iter_and_schema(iter, &schema)
+    }
+}