polars-df 0.1.1 → 0.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2a39332a375d211cf0e04ea95cdb4fe315c843b3df474ee3fef135ec63eb0516
-  data.tar.gz: 41881b2aa8da2bb2850fbcaba11ec1409204ec99b7d36d1470a7536684b02509
+  metadata.gz: 9e6fb732e5dafe2fde285322554bd9159483cbbdf17d6e2bba9cba9a83563b47
+  data.tar.gz: 1b4249d0c0100f136973c601b8404cb6d92abc632d5ed0476bd93bc5360a11dc
 SHA512:
-  metadata.gz: 17ba7456808346b025d52d14d5fa6d91dc0d18385d0ee796c1b821eb6272875e15927e34da8f41c1872e4b158bfd32165ce1ee4adc5cebf4474fc7e86ad98e31
-  data.tar.gz: 4c3ef1aa3ed5c6e0620d555bab73d20e8b7cd495e26c97f4ddd41f14ed8d3883331cff36158fbe4e1dbc3b86825fa83dfa4b832f901db35a73a9c021b8dd955c
+  metadata.gz: d9414d6f60c489e2b3b72885288822083ba8c04bac4053f4e34c1d53ee805d164f17fe4b8b3a8f4ff562550bcc657f374bea6e250b52985367f601ea50e3037f
+  data.tar.gz: 9e3a7cfe105f03ec20e9c26aa38c1475074ccc1ea057a170a97b7068b41943d561d50af49bb1d1f74b7705809dc1375900f542ab93683ba627dea080274f6d91

data/.yardopts

@@ -0,0 +1,3 @@
+--no-private
+--markup markdown
+--embed-mixins

data/CHANGELOG.md

@@ -1,3 +1,11 @@
+## 0.1.3 (2022-11-27)
+
+- Added more methods
+
+## 0.1.2 (2022-11-25)
+
+- Added more methods
+
 ## 0.1.1 (2022-11-23)
 
 - Added more methods

data/Cargo.lock

@@ -1160,7 +1160,7 @@ dependencies = [
 
 [[package]]
 name = "polars"
-version = "0.1.0"
+version = "0.1.3"
 dependencies = [
  "magnus",
  "polars 0.25.1",
@@ -1217,6 +1217,7 @@ dependencies = [
  "rayon",
  "regex",
  "serde",
+ "serde_json",
  "smartstring",
  "thiserror",
 ]

data/README.md

@@ -27,7 +27,7 @@ Polars.read_csv("iris.csv")
   .collect
 ```
 
-You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/introduction.html) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems. Note that many methods and options are missing at the moment.
+You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/introduction.html) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems. Some methods are missing at the moment.
 
 ## Examples
 

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.1.1"
+version = "0.1.3"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
 publish = false
@@ -17,10 +17,12 @@ version = "0.25.1"
 features = [
     "abs",
     "arange",
+    "arg_where",
     "concat_str",
     "csv-file",
     "cum_agg",
     "cumulative_eval",
+    "dataframe_arithmetic",
     "date_offset",
     "diagonal_concat",
     "diff",
@@ -38,16 +40,20 @@ features = [
     "lazy",
     "lazy_regex",
     "list_eval",
+    "list_to_struct",
     "log",
     "meta",
     "mode",
     "moment",
+    "object",
     "parquet",
     "partition_by",
     "pct_change",
     "product",
+    "propagate_nans",
     "random",
     "rank",
+    "reinterpret",
     "repeat_by",
     "rolling_window",
     "round_series",

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,11 +1,12 @@
-use magnus::{TryConvert, Value, QNIL};
+use magnus::{class, RArray, Symbol, TryConvert, Value, QNIL};
+use polars::chunked_array::object::PolarsObjectSafe;
 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, RbPolarsErr, RbResult, RbValueError};
+use crate::{RbDataFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
 pub struct Wrap<T>(pub T);
 
@@ -15,14 +16,57 @@ impl<T> From<T> for 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())
 }
 
-impl Into<Value> for Wrap<AnyValue<'_>> {
-    fn into(self) -> Value {
-        match self.0 {
+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),
@@ -36,11 +80,24 @@ impl Into<Value> for Wrap<AnyValue<'_>> {
             AnyValue::Null => *QNIL,
             AnyValue::Boolean(v) => Value::from(v),
             AnyValue::Utf8(v) => Value::from(v),
+            AnyValue::Date(v) => class::time()
+                .funcall::<_, _, Value>("at", (v * 86400,))
+                .unwrap()
+                .funcall::<_, _, Value>("utc", ())
+                .unwrap()
+                .funcall::<_, _, Value>("to_date", ())
+                .unwrap(),
             _ => 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() {
@@ -53,10 +110,19 @@ impl TryConvert for Wrap<DataType> {
             "i32" => DataType::Int32,
             "i64" => DataType::Int64,
             "str" => DataType::Utf8,
+            "bin" => DataType::Binary,
             "bool" => DataType::Boolean,
+            "cat" => DataType::Categorical(None),
+            "date" => DataType::Date,
+            "datetime" => DataType::Datetime(TimeUnit::Microseconds, None),
             "f32" => DataType::Float32,
+            "time" => DataType::Time,
+            "dur" => DataType::Duration(TimeUnit::Microseconds),
             "f64" => DataType::Float64,
-            "date" => DataType::Date,
+            // "obj" => DataType::Object(OBJECT_NAME),
+            "list" => DataType::List(Box::new(DataType::Boolean)),
+            "null" => DataType::Null,
+            "unk" => DataType::Unknown,
             _ => {
                 return Err(RbValueError::new_err(format!(
                     "{} is not a supported DataType.",
@@ -118,6 +184,39 @@ impl TryConvert for Wrap<ClosedWindow> {
     }
 }
 
+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() {
@@ -171,6 +270,24 @@ impl TryConvert for Wrap<NullStrategy> {
     }
 }
 
+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() {
@@ -307,3 +424,19 @@ pub fn parse_parquet_compression(
     };
     Ok(parsed)
 }
+
+pub struct ObjectValue {
+    pub inner: Value,
+}
+
+impl From<&dyn PolarsObjectSafe> for &ObjectValue {
+    fn from(val: &dyn PolarsObjectSafe) -> Self {
+        unsafe { &*(val as *const dyn PolarsObjectSafe as *const ObjectValue) }
+    }
+}
+
+impl ObjectValue {
+    pub fn to_object(&self) -> Value {
+        self.inner
+    }
+}