polars-df 0.1.0 → 0.1.2

data/Cargo.toml

@@ -0,0 +1,5 @@
+[workspace]
+members = ["ext/polars"]
+
+[profile.release]
+strip = true

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.1.0"
+version = "0.1.2"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
 publish = false
@@ -15,21 +15,51 @@ serde_json = "1"
 [dependencies.polars]
 version = "0.25.1"
 features = [
+    "abs",
     "arange",
+    "concat_str",
     "csv-file",
     "cum_agg",
+    "cumulative_eval",
+    "date_offset",
     "diagonal_concat",
+    "diff",
+    "dot_product",
     "dtype-full",
+    "dynamic_groupby",
+    "ewma",
     "fmt",
     "horizontal_concat",
     "interpolate",
     "ipc",
+    "is_first",
+    "is_in",
     "json",
     "lazy",
     "lazy_regex",
+    "list_eval",
+    "log",
+    "meta",
+    "mode",
+    "moment",
     "parquet",
+    "partition_by",
+    "pct_change",
+    "product",
+    "random",
+    "rank",
+    "repeat_by",
+    "rolling_window",
+    "round_series",
+    "search_sorted",
     "semi_anti_join",
     "serde-lazy",
+    "sign",
+    "string_justify",
     "strings",
+    "timezones",
+    "to_dummies",
+    "top_k",
     "trigonometry",
+    "unique_counts",
 ]

data/ext/polars/src/batched_csv.rs

@@ -0,0 +1,120 @@
+use magnus::Value;
+use polars::io::mmap::MmapBytesReader;
+use polars::io::RowCount;
+use polars::prelude::read_impl::OwnedBatchedCsvReader;
+use polars::prelude::*;
+use std::cell::RefCell;
+use std::path::PathBuf;
+
+use crate::conversion::*;
+use crate::{RbDataFrame, RbPolarsErr, RbResult};
+
+#[magnus::wrap(class = "Polars::RbBatchedCsv")]
+pub struct RbBatchedCsv {
+    pub reader: RefCell<OwnedBatchedCsvReader>,
+}
+
+impl RbBatchedCsv {
+    pub fn new(arguments: &[Value]) -> RbResult<Self> {
+        // start arguments
+        // this pattern is needed for more than 16
+        let infer_schema_length: Option<usize> = arguments[0].try_convert()?;
+        let chunk_size: usize = arguments[1].try_convert()?;
+        let has_header: bool = arguments[2].try_convert()?;
+        let ignore_errors: bool = arguments[3].try_convert()?;
+        let n_rows: Option<usize> = arguments[4].try_convert()?;
+        let skip_rows: usize = arguments[5].try_convert()?;
+        let projection: Option<Vec<usize>> = arguments[6].try_convert()?;
+        let sep: String = arguments[7].try_convert()?;
+        let rechunk: bool = arguments[8].try_convert()?;
+        let columns: Option<Vec<String>> = arguments[9].try_convert()?;
+        let encoding: Wrap<CsvEncoding> = arguments[10].try_convert()?;
+        let n_threads: Option<usize> = arguments[11].try_convert()?;
+        let path: PathBuf = arguments[12].try_convert()?;
+        let overwrite_dtype: Option<Vec<(String, Wrap<DataType>)>> = arguments[13].try_convert()?;
+        // TODO fix
+        let overwrite_dtype_slice: Option<Vec<Wrap<DataType>>> = None; // arguments[14].try_convert()?;
+        let low_memory: bool = arguments[15].try_convert()?;
+        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 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()?;
+        let eol_char: String = arguments[23].try_convert()?;
+        // end arguments
+
+        let null_values = null_values.map(|w| w.0);
+        let comment_char = comment_char.map(|s| s.as_bytes()[0]);
+        let eol_char = eol_char.as_bytes()[0];
+
+        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+
+        let quote_char = if let Some(s) = quote_char {
+            if s.is_empty() {
+                None
+            } else {
+                Some(s.as_bytes()[0])
+            }
+        } else {
+            None
+        };
+
+        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)
+        });
+
+        let overwrite_dtype_slice = overwrite_dtype_slice.map(|overwrite_dtype| {
+            overwrite_dtype
+                .iter()
+                .map(|dt| dt.0.clone())
+                .collect::<Vec<_>>()
+        });
+
+        let file = std::fs::File::open(path).map_err(RbPolarsErr::io)?;
+        let reader = Box::new(file) as Box<dyn MmapBytesReader>;
+        let reader = CsvReader::new(reader)
+            .infer_schema(infer_schema_length)
+            .has_header(has_header)
+            .with_n_rows(n_rows)
+            .with_delimiter(sep.as_bytes()[0])
+            .with_skip_rows(skip_rows)
+            .with_ignore_parser_errors(ignore_errors)
+            .with_projection(projection)
+            .with_rechunk(rechunk)
+            .with_chunk_size(chunk_size)
+            .with_encoding(encoding.0)
+            .with_columns(columns)
+            .with_n_threads(n_threads)
+            .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_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)?;
+
+        Ok(RbBatchedCsv {
+            reader: RefCell::new(reader),
+        })
+    }
+
+    pub fn next_batches(&self, n: usize) -> RbResult<Option<Vec<RbDataFrame>>> {
+        let batches = self
+            .reader
+            .borrow_mut()
+            .next_batches(n)
+            .map_err(RbPolarsErr::from)?;
+        Ok(batches.map(|batches| batches.into_iter().map(|out| out.1.into()).collect()))
+    }
+}

data/ext/polars/src/conversion.rs

@@ -1,27 +1,345 @@
-use magnus::{Value, QNIL};
+use magnus::{RArray, Symbol, TryConvert, Value, QNIL};
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
 use polars::datatypes::AnyValue;
 use polars::frame::DataFrame;
 use polars::prelude::*;
+use polars::series::ops::NullBehavior;
 
-use crate::{RbDataFrame, RbResult, RbValueError};
+use crate::{RbDataFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
-pub fn wrap(val: AnyValue) -> Value {
-    match val {
-        AnyValue::UInt8(v) => Value::from(v),
-        AnyValue::UInt16(v) => Value::from(v),
-        AnyValue::UInt32(v) => Value::from(v),
-        AnyValue::UInt64(v) => Value::from(v),
-        AnyValue::Int8(v) => Value::from(v),
-        AnyValue::Int16(v) => Value::from(v),
-        AnyValue::Int32(v) => Value::from(v),
-        AnyValue::Int64(v) => Value::from(v),
-        AnyValue::Float32(v) => Value::from(v),
-        AnyValue::Float64(v) => Value::from(v),
-        AnyValue::Null => *QNIL,
-        AnyValue::Boolean(v) => Value::from(v),
-        AnyValue::Utf8(v) => Value::from(v),
-        _ => todo!(),
+pub struct Wrap<T>(pub T);
+
+impl<T> From<T> for Wrap<T> {
+    fn from(t: T) -> Self {
+        Wrap(t)
+    }
+}
+
+pub fn get_rbseq(obj: Value) -> RbResult<(RArray, usize)> {
+    let seq: RArray = obj.try_convert()?;
+    let len = seq.len();
+    Ok((seq, len))
+}
+
+pub fn get_df(obj: Value) -> RbResult<DataFrame> {
+    let rbdf = obj.funcall::<_, _, &RbDataFrame>("_df", ())?;
+    Ok(rbdf.df.borrow().clone())
+}
+
+pub fn get_series(obj: Value) -> RbResult<Series> {
+    let rbs = obj.funcall::<_, _, &RbSeries>("_s", ())?;
+    Ok(rbs.series.borrow().clone())
+}
+
+impl TryConvert for Wrap<Utf8Chunked> {
+    fn try_convert(obj: Value) -> RbResult<Self> {
+        let (seq, len) = get_rbseq(obj)?;
+        let mut builder = Utf8ChunkedBuilder::new("", len, len * 25);
+
+        for res in seq.each() {
+            let item = res?;
+            match item.try_convert::<String>() {
+                Ok(val) => builder.append_value(&val),
+                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>() {
+            Ok(Wrap(NullValues::AllColumnsSingle(s)))
+        } else if let Ok(s) = ob.try_convert::<Vec<String>>() {
+            Ok(Wrap(NullValues::AllColumns(s)))
+        } else if let Ok(s) = ob.try_convert::<Vec<(String, String)>>() {
+            Ok(Wrap(NullValues::Named(s)))
+        } else {
+            Err(RbPolarsErr::other(
+                "could not extract value from null_values argument".into(),
+            ))
+        }
+    }
+}
+
+impl From<Wrap<AnyValue<'_>>> for Value {
+    fn from(w: Wrap<AnyValue<'_>>) -> Self {
+        match w.0 {
+            AnyValue::UInt8(v) => Value::from(v),
+            AnyValue::UInt16(v) => Value::from(v),
+            AnyValue::UInt32(v) => Value::from(v),
+            AnyValue::UInt64(v) => Value::from(v),
+            AnyValue::Int8(v) => Value::from(v),
+            AnyValue::Int16(v) => Value::from(v),
+            AnyValue::Int32(v) => Value::from(v),
+            AnyValue::Int64(v) => Value::from(v),
+            AnyValue::Float32(v) => Value::from(v),
+            AnyValue::Float64(v) => Value::from(v),
+            AnyValue::Null => *QNIL,
+            AnyValue::Boolean(v) => Value::from(v),
+            AnyValue::Utf8(v) => Value::from(v),
+            _ => todo!(),
+        }
+    }
+}
+
+impl From<Wrap<DataType>> for Value {
+    fn from(w: Wrap<DataType>) -> Self {
+        Symbol::from(w.0.to_string()).into()
+    }
+}
+
+impl TryConvert for Wrap<DataType> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let dtype = match ob.try_convert::<String>()?.as_str() {
+            "u8" => DataType::UInt8,
+            "u16" => DataType::UInt16,
+            "u32" => DataType::UInt32,
+            "u64" => DataType::UInt64,
+            "i8" => DataType::Int8,
+            "i16" => DataType::Int16,
+            "i32" => DataType::Int32,
+            "i64" => DataType::Int64,
+            "str" => DataType::Utf8,
+            "bool" => DataType::Boolean,
+            "f32" => DataType::Float32,
+            "f64" => DataType::Float64,
+            "date" => DataType::Date,
+            _ => {
+                return Err(RbValueError::new_err(format!(
+                    "{} is not a supported DataType.",
+                    ob
+                )))
+            }
+        };
+        Ok(Wrap(dtype))
+    }
+}
+
+impl<'s> TryConvert for Wrap<AnyValue<'s>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        // TODO improve
+        if let Ok(v) = ob.try_convert::<i64>() {
+            Ok(AnyValue::Int64(v).into())
+        } else if let Ok(v) = ob.try_convert::<f64>() {
+            Ok(AnyValue::Float64(v).into())
+        } else {
+            Err(RbPolarsErr::other(format!(
+                "object type not supported {:?}",
+                ob
+            )))
+        }
+    }
+}
+
+impl TryConvert for Wrap<CategoricalOrdering> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "physical" => CategoricalOrdering::Physical,
+            "lexical" => CategoricalOrdering::Lexical,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "ordering must be one of {{'physical', 'lexical'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<ClosedWindow> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "left" => ClosedWindow::Left,
+            "right" => ClosedWindow::Right,
+            "both" => ClosedWindow::Both,
+            "none" => ClosedWindow::None,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "closed must be one of {{'left', 'right', 'both', 'none'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<CsvEncoding> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "utf8" => CsvEncoding::Utf8,
+            "utf8-lossy" => CsvEncoding::LossyUtf8,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "encoding must be one of {{'utf8', 'utf8-lossy'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<Option<IpcCompression>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "uncompressed" => None,
+            "lz4" => Some(IpcCompression::LZ4),
+            "zstd" => Some(IpcCompression::ZSTD),
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "compression must be one of {{'uncompressed', 'lz4', 'zstd'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<JoinType> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "inner" => JoinType::Inner,
+            "left" => JoinType::Left,
+            "outer" => JoinType::Outer,
+            "semi" => JoinType::Semi,
+            "anti" => JoinType::Anti,
+            // #[cfg(feature = "cross_join")]
+            // "cross" => JoinType::Cross,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                "how must be one of {{'inner', 'left', 'outer', 'semi', 'anti', 'cross'}}, 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() {
+            "drop" => NullBehavior::Drop,
+            "ignore" => NullBehavior::Ignore,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "null behavior must be one of {{'drop', 'ignore'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<NullStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "ignore" => NullStrategy::Ignore,
+            "propagate" => NullStrategy::Propagate,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "null strategy must be one of {{'ignore', 'propagate'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<ParallelStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "auto" => ParallelStrategy::Auto,
+            "columns" => ParallelStrategy::Columns,
+            "row_groups" => ParallelStrategy::RowGroups,
+            "none" => ParallelStrategy::None,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "parallel must be one of {{'auto', 'columns', 'row_groups', 'none'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<QuantileInterpolOptions> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "lower" => QuantileInterpolOptions::Lower,
+            "higher" => QuantileInterpolOptions::Higher,
+            "nearest" => QuantileInterpolOptions::Nearest,
+            "linear" => QuantileInterpolOptions::Linear,
+            "midpoint" => QuantileInterpolOptions::Midpoint,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "interpolation must be one of {{'lower', 'higher', 'nearest', 'linear', 'midpoint'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<RankMethod> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "min" => RankMethod::Min,
+            "max" => RankMethod::Max,
+            "average" => RankMethod::Average,
+            "dense" => RankMethod::Dense,
+            "ordinal" => RankMethod::Ordinal,
+            "random" => RankMethod::Random,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "method must be one of {{'min', 'max', 'average', 'dense', 'ordinal', 'random'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<TimeUnit> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "ns" => TimeUnit::Nanoseconds,
+            "us" => TimeUnit::Microseconds,
+            "ms" => TimeUnit::Milliseconds,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "time unit must be one of {{'ns', 'us', 'ms'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<UniqueKeepStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "first" => UniqueKeepStrategy::First,
+            "last" => UniqueKeepStrategy::Last,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "keep must be one of {{'first', 'last'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
     }
 }
 
@@ -47,30 +365,6 @@ pub fn parse_fill_null_strategy(
     Ok(parsed)
 }
 
-pub fn wrap_join_type(ob: &str) -> RbResult<JoinType> {
-    let parsed = match ob {
-        "inner" => JoinType::Inner,
-        "left" => JoinType::Left,
-        "outer" => JoinType::Outer,
-        "semi" => JoinType::Semi,
-        "anti" => JoinType::Anti,
-        // #[cfg(feature = "cross_join")]
-        // "cross" => JoinType::Cross,
-        v => {
-            return Err(RbValueError::new_err(format!(
-                "how must be one of {{'inner', 'left', 'outer', 'semi', 'anti', 'cross'}}, got {}",
-                v
-            )))
-        }
-    };
-    Ok(parsed)
-}
-
-pub fn get_df(obj: Value) -> RbResult<DataFrame> {
-    let rbdf = obj.funcall::<_, _, &RbDataFrame>("_df", ())?;
-    Ok(rbdf.df.borrow().clone())
-}
-
 pub fn parse_parquet_compression(
     compression: &str,
     compression_level: Option<i32>,