polars-df 0.10.0 → 0.12.0

data/ext/polars/src/conversion/mod.rs

@@ -1,4 +1,4 @@
-pub(crate) mod anyvalue;
+pub(crate) mod any_value;
 mod chunked_array;
 
 use std::fmt::{Debug, Display, Formatter};
@@ -18,6 +18,7 @@ use polars::io::avro::AvroCompression;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 use polars_core::utils::arrow::array::Array;
+use polars_core::utils::materialize_dyn_int;
 use polars_utils::total_ord::{TotalEq, TotalHash};
 use smartstring::alias::String as SmartString;
 
@@ -74,7 +75,7 @@ pub(crate) fn get_df(obj: Value) -> RbResult<DataFrame> {
 
 pub(crate) fn get_lf(obj: Value) -> RbResult<LazyFrame> {
     let rbdf = obj.funcall::<_, _, &RbLazyFrame>("_ldf", ())?;
-    Ok(rbdf.ldf.clone())
+    Ok(rbdf.ldf.borrow().clone())
 }
 
 pub(crate) fn get_series(obj: Value) -> RbResult<Series> {
@@ -154,7 +155,7 @@ impl IntoValue for Wrap<DataType> {
                 let class = pl.const_get::<_, Value>("Float32").unwrap();
                 class.funcall("new", ()).unwrap()
             }
-            DataType::Float64 => {
+            DataType::Float64 | DataType::Unknown(UnknownKind::Float) => {
                 let class = pl.const_get::<_, Value>("Float64").unwrap();
                 class.funcall("new", ()).unwrap()
             }
@@ -168,7 +169,7 @@ impl IntoValue for Wrap<DataType> {
                 let class = pl.const_get::<_, Value>("Boolean").unwrap();
                 class.funcall("new", ()).unwrap()
             }
-            DataType::String => {
+            DataType::String | DataType::Unknown(UnknownKind::Str) => {
                 let class = pl.const_get::<_, Value>("String").unwrap();
                 class.funcall("new", ()).unwrap()
             }
@@ -242,7 +243,10 @@ impl IntoValue for Wrap<DataType> {
                 let class = pl.const_get::<_, Value>("Null").unwrap();
                 class.funcall("new", ()).unwrap()
             }
-            DataType::Unknown => {
+            DataType::Unknown(UnknownKind::Int(v)) => {
+                Wrap(materialize_dyn_int(v).dtype()).into_value()
+            }
+            DataType::Unknown(_) => {
                 let class = pl.const_get::<_, Value>("Unknown").unwrap();
                 class.funcall("new", ()).unwrap()
             }
@@ -310,7 +314,7 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Object" => DataType::Object(OBJECT_NAME, None),
                 "Polars::List" => DataType::List(Box::new(DataType::Null)),
                 "Polars::Null" => DataType::Null,
-                "Polars::Unknown" => DataType::Unknown,
+                "Polars::Unknown" => DataType::Unknown(Default::default()),
                 dt => {
                     return Err(RbValueError::new_err(format!(
                         "{dt} is not a correct polars DataType.",
@@ -350,7 +354,7 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Float32" => DataType::Float32,
                 "Polars::Float64" => DataType::Float64,
                 "Polars::Null" => DataType::Null,
-                "Polars::Unknown" => DataType::Unknown,
+                "Polars::Unknown" => DataType::Unknown(Default::default()),
                 "Polars::Duration" => {
                     let time_unit: Value = ob.funcall("time_unit", ()).unwrap();
                     let time_unit = Wrap::<TimeUnit>::try_convert(time_unit)?.0;
@@ -375,8 +379,8 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Struct" => {
                     let arr: RArray = ob.funcall("fields", ())?;
                     let mut fields = Vec::with_capacity(arr.len());
-                    for v in arr.each() {
-                        fields.push(Wrap::<Field>::try_convert(v?)?.0);
+                    for v in arr.into_iter() {
+                        fields.push(Wrap::<Field>::try_convert(v)?.0);
                     }
                     DataType::Struct(fields)
                 }
@@ -410,7 +414,7 @@ impl TryConvert for Wrap<DataType> {
                 "obj" => DataType::Object(OBJECT_NAME, None),
                 "list" => DataType::List(Box::new(DataType::Boolean)),
                 "null" => DataType::Null,
-                "unk" => DataType::Unknown,
+                "unk" => DataType::Unknown(Default::default()),
                 _ => {
                     return Err(RbValueError::new_err(format!(
                         "{} is not a supported DataType.",
@@ -423,11 +427,36 @@ impl TryConvert for Wrap<DataType> {
     }
 }
 
+impl TryConvert for Wrap<StatisticsOptions> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let mut statistics = StatisticsOptions::empty();
+
+        let dict = RHash::try_convert(ob)?;
+        dict.foreach(|key: Symbol, val: bool| {
+            match key.name()?.as_ref() {
+                "min" => statistics.min_value = val,
+                "max" => statistics.max_value = val,
+                "distinct_count" => statistics.distinct_count = val,
+                "null_count" => statistics.null_count = val,
+                _ => {
+                    return Err(RbTypeError::new_err(format!(
+                        "'{key}' is not a valid statistic option",
+                    )))
+                }
+            }
+            Ok(ForEach::Continue)
+        })
+        .unwrap();
+
+        Ok(Wrap(statistics))
+    }
+}
+
 impl<'s> TryConvert for Wrap<Row<'s>> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let mut vals: Vec<Wrap<AnyValue<'s>>> = Vec::new();
-        for item in RArray::try_convert(ob)?.each() {
-            vals.push(Wrap::<AnyValue<'s>>::try_convert(item?)?);
+        for item in RArray::try_convert(ob)?.into_iter() {
+            vals.push(Wrap::<AnyValue<'s>>::try_convert(item)?);
         }
         let vals: Vec<AnyValue> = unsafe { std::mem::transmute(vals) };
         Ok(Wrap(Row(vals)))
@@ -546,57 +575,6 @@ impl Default for ObjectValue {
     }
 }
 
-pub(crate) fn dicts_to_rows(
-    records: &Value,
-    infer_schema_len: Option<usize>,
-    schema_columns: PlIndexSet<String>,
-) -> RbResult<(Vec<Row>, Vec<String>)> {
-    let infer_schema_len = infer_schema_len.map(|n| std::cmp::max(1, n));
-    let (dicts, len) = get_rbseq(*records)?;
-
-    let key_names = {
-        if !schema_columns.is_empty() {
-            schema_columns
-        } else {
-            let mut inferred_keys = PlIndexSet::new();
-            for d in dicts.each().take(infer_schema_len.unwrap_or(usize::MAX)) {
-                let d = d?;
-                let d = RHash::try_convert(d)?;
-
-                d.foreach(|name: Value, _value: Value| {
-                    if let Some(v) = Symbol::from_value(name) {
-                        inferred_keys.insert(v.name()?.into());
-                    } else {
-                        inferred_keys.insert(String::try_convert(name)?);
-                    };
-                    Ok(ForEach::Continue)
-                })?;
-            }
-            inferred_keys
-        }
-    };
-
-    let mut rows = Vec::with_capacity(len);
-
-    for d in dicts.each() {
-        let d = d?;
-        let d = RHash::try_convert(d)?;
-
-        let mut row = Vec::with_capacity(key_names.len());
-
-        for k in key_names.iter() {
-            // TODO improve performance
-            let val = match d.get(k.clone()).or_else(|| d.get(Symbol::new(k))) {
-                None => AnyValue::Null,
-                Some(val) => Wrap::<AnyValue>::try_convert(val)?.0,
-            };
-            row.push(val)
-        }
-        rows.push(Row(row))
-    }
-    Ok((rows, key_names.into_iter().collect()))
-}
-
 impl TryConvert for Wrap<AsofStrategy> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match String::try_convert(ob)?.as_str() {
@@ -733,14 +711,13 @@ impl TryConvert for Wrap<JoinType> {
         let parsed = match String::try_convert(ob)?.as_str() {
             "inner" => JoinType::Inner,
             "left" => JoinType::Left,
-            "outer" => JoinType::Outer { coalesce: false },
-            "outer_coalesce" => JoinType::Outer { coalesce: true },
+            "full" => JoinType::Full,
             "semi" => JoinType::Semi,
             "anti" => JoinType::Anti,
             "cross" => JoinType::Cross,
             v => {
                 return Err(RbValueError::new_err(format!(
-                "how must be one of {{'inner', 'left', 'outer', 'semi', 'anti', 'cross'}}, got {}",
+                "how must be one of {{'inner', 'left', 'full', 'semi', 'anti', 'cross'}}, got {}",
                 v
             )))
             }
@@ -950,6 +927,23 @@ impl TryConvert for Wrap<SearchSortedSide> {
     }
 }
 
+impl TryConvert for Wrap<ClosedInterval> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match String::try_convert(ob)?.as_str() {
+            "both" => ClosedInterval::Both,
+            "left" => ClosedInterval::Left,
+            "right" => ClosedInterval::Right,
+            "none" => ClosedInterval::None,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "`closed` must be one of {{'both', 'left', 'right', 'none'}}, got {v}",
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<WindowMapping> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match String::try_convert(ob)?.as_str() {

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

@@ -0,0 +1,184 @@
+use magnus::{prelude::*, r_hash::ForEach, RArray, RHash, Symbol, Value};
+use polars::frame::row::{rows_to_schema_supertypes, rows_to_supertypes, Row};
+use polars::prelude::*;
+
+use super::*;
+use crate::conversion::*;
+use crate::{RbPolarsErr, RbResult};
+
+impl RbDataFrame {
+    pub fn from_rows(
+        rb_rows: RArray,
+        infer_schema_length: Option<usize>,
+        schema: Option<Wrap<Schema>>,
+    ) -> RbResult<Self> {
+        let mut data = Vec::with_capacity(rb_rows.len());
+        for v in rb_rows.into_iter() {
+            let rb_row = RArray::try_convert(v)?;
+            let mut row = Vec::with_capacity(rb_row.len());
+            for val in rb_row.into_iter() {
+                row.push(Wrap::<AnyValue>::try_convert(val)?.0);
+            }
+            data.push(Row(row));
+        }
+        let schema = schema.map(|wrap| wrap.0);
+        finish_from_rows(data, schema, None, infer_schema_length)
+    }
+
+    pub fn from_hashes(
+        data: Value,
+        schema: Option<Wrap<Schema>>,
+        schema_overrides: Option<Wrap<Schema>>,
+        strict: bool,
+        infer_schema_length: Option<usize>,
+    ) -> RbResult<Self> {
+        let schema = schema.map(|wrap| wrap.0);
+        let schema_overrides = schema_overrides.map(|wrap| wrap.0);
+
+        let names = get_schema_names(&data, schema.as_ref(), infer_schema_length)?;
+        let rows = dicts_to_rows(&data, &names, strict)?;
+
+        let schema = schema.or_else(|| {
+            Some(columns_names_to_empty_schema(
+                names.iter().map(String::as_str),
+            ))
+        });
+
+        finish_from_rows(rows, schema, schema_overrides, infer_schema_length)
+    }
+}
+
+fn finish_from_rows(
+    rows: Vec<Row>,
+    schema: Option<Schema>,
+    schema_overrides: Option<Schema>,
+    infer_schema_length: Option<usize>,
+) -> RbResult<RbDataFrame> {
+    // Object builder must be registered
+    crate::on_startup::register_object_builder();
+
+    let mut schema = if let Some(mut schema) = schema {
+        resolve_schema_overrides(&mut schema, schema_overrides);
+        update_schema_from_rows(&mut schema, &rows, infer_schema_length)?;
+        schema
+    } else {
+        rows_to_schema_supertypes(&rows, infer_schema_length).map_err(RbPolarsErr::from)?
+    };
+
+    // TODO: Remove this step when Decimals are supported properly.
+    // Erasing the decimal precision/scale here will just require us to infer it again later.
+    // https://github.com/pola-rs/polars/issues/14427
+    erase_decimal_precision_scale(&mut schema);
+
+    let df = DataFrame::from_rows_and_schema(&rows, &schema).map_err(RbPolarsErr::from)?;
+    Ok(df.into())
+}
+
+fn update_schema_from_rows(
+    schema: &mut Schema,
+    rows: &[Row],
+    infer_schema_length: Option<usize>,
+) -> RbResult<()> {
+    let schema_is_complete = schema.iter_dtypes().all(|dtype| dtype.is_known());
+    if schema_is_complete {
+        return Ok(());
+    }
+
+    // TODO: Only infer dtypes for columns with an unknown dtype
+    let inferred_dtypes =
+        rows_to_supertypes(rows, infer_schema_length).map_err(RbPolarsErr::from)?;
+    let inferred_dtypes_slice = inferred_dtypes.as_slice();
+
+    for (i, dtype) in schema.iter_dtypes_mut().enumerate() {
+        if !dtype.is_known() {
+            *dtype = inferred_dtypes_slice.get(i).ok_or_else(|| {
+                polars_err!(SchemaMismatch: "the number of columns in the schema does not match the data")
+            })
+            .map_err(RbPolarsErr::from)?
+            .clone();
+        }
+    }
+    Ok(())
+}
+
+fn resolve_schema_overrides(schema: &mut Schema, schema_overrides: Option<Schema>) {
+    if let Some(overrides) = schema_overrides {
+        for (name, dtype) in overrides.into_iter() {
+            schema.set_dtype(name.as_str(), dtype);
+        }
+    }
+}
+
+fn erase_decimal_precision_scale(schema: &mut Schema) {
+    for dtype in schema.iter_dtypes_mut() {
+        if let DataType::Decimal(_, _) = dtype {
+            *dtype = DataType::Decimal(None, None)
+        }
+    }
+}
+
+fn columns_names_to_empty_schema<'a, I>(column_names: I) -> Schema
+where
+    I: IntoIterator<Item = &'a str>,
+{
+    let fields = column_names
+        .into_iter()
+        .map(|c| Field::new(c, DataType::Unknown(Default::default())));
+    Schema::from_iter(fields)
+}
+
+fn dicts_to_rows<'a>(data: &Value, names: &'a [String], _strict: bool) -> RbResult<Vec<Row<'a>>> {
+    let (data, len) = get_rbseq(*data)?;
+    let mut rows = Vec::with_capacity(len);
+    for d in data.into_iter() {
+        let d = RHash::try_convert(d)?;
+
+        let mut row = Vec::with_capacity(names.len());
+        for k in names.iter() {
+            // TODO improve performance
+            let val = match d.get(k.clone()).or_else(|| d.get(Symbol::new(k))) {
+                None => AnyValue::Null,
+                Some(val) => Wrap::<AnyValue>::try_convert(val)?.0,
+            };
+            row.push(val)
+        }
+        rows.push(Row(row))
+    }
+    Ok(rows)
+}
+
+fn get_schema_names(
+    data: &Value,
+    schema: Option<&Schema>,
+    infer_schema_length: Option<usize>,
+) -> RbResult<Vec<String>> {
+    if let Some(schema) = schema {
+        Ok(schema.iter_names().map(|n| n.to_string()).collect())
+    } else {
+        infer_schema_names_from_data(data, infer_schema_length)
+    }
+}
+
+fn infer_schema_names_from_data(
+    data: &Value,
+    infer_schema_length: Option<usize>,
+) -> RbResult<Vec<String>> {
+    let (data, data_len) = get_rbseq(*data)?;
+    let infer_schema_length = infer_schema_length
+        .map(|n| std::cmp::max(1, n))
+        .unwrap_or(data_len);
+
+    let mut names = PlIndexSet::new();
+    for d in data.into_iter().take(infer_schema_length) {
+        let d = RHash::try_convert(d)?;
+        d.foreach(|name: Value, _value: Value| {
+            if let Some(v) = Symbol::from_value(name) {
+                names.insert(v.name()?.into());
+            } else {
+                names.insert(String::try_convert(name)?);
+            };
+            Ok(ForEach::Continue)
+        })?;
+    }
+    Ok(names.into_iter().collect())
+}

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

@@ -0,0 +1,48 @@
+use magnus::{prelude::*, IntoValue, RArray, Value};
+
+use super::*;
+use crate::conversion::{ObjectValue, Wrap};
+
+impl RbDataFrame {
+    pub fn row_tuple(&self, idx: i64) -> Value {
+        let idx = if idx < 0 {
+            (self.df.borrow().height() as i64 + idx) as usize
+        } else {
+            idx as usize
+        };
+        RArray::from_iter(
+            self.df
+                .borrow()
+                .get_columns()
+                .iter()
+                .map(|s| match s.dtype() {
+                    DataType::Object(_, _) => {
+                        let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
+                        obj.unwrap().to_object()
+                    }
+                    _ => Wrap(s.get(idx).unwrap()).into_value(),
+                }),
+        )
+        .as_value()
+    }
+
+    pub fn row_tuples(&self) -> Value {
+        let df = &self.df;
+        RArray::from_iter((0..df.borrow().height()).map(|idx| {
+            RArray::from_iter(
+                self.df
+                    .borrow()
+                    .get_columns()
+                    .iter()
+                    .map(|s| match s.dtype() {
+                        DataType::Object(_, _) => {
+                            let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
+                            obj.unwrap().to_object()
+                        }
+                        _ => Wrap(s.get(idx).unwrap()).into_value(),
+                    }),
+            )
+        }))
+        .as_value()
+    }
+}