polars-df 0.10.0 → 0.11.0

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

@@ -7,7 +7,7 @@ use std::num::NonZeroUsize;
 
 use magnus::{
     class, exception, prelude::*, r_hash::ForEach, value::Opaque, IntoValue, Module, RArray, RHash,
-    Ruby, Symbol, TryConvert, Value,
+    Ruby, TryConvert, Value,
 };
 use polars::chunked_array::object::PolarsObjectSafe;
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
@@ -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;
 
@@ -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;
@@ -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.",
@@ -546,57 +550,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,8 +686,8 @@ 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 },
+            "outer" => JoinType::Outer,
+            "outer_coalesce" => JoinType::Outer,
             "semi" => JoinType::Semi,
             "anti" => JoinType::Anti,
             "cross" => JoinType::Cross,

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

@@ -0,0 +1,186 @@
+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.each() {
+            let rb_row = RArray::try_convert(v?)?;
+            let mut row = Vec::with_capacity(rb_row.len());
+            for val in rb_row.each() {
+                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.each() {
+        let d = d?;
+        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.each().take(infer_schema_length) {
+        let d = d?;
+        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()
+    }
+}