polars-df 0.9.0 → 0.11.0

data/ext/polars/src/batched_csv.rs

@@ -1,23 +1,19 @@
+use std::cell::RefCell;
+use std::path::PathBuf;
+use std::sync::Mutex;
+
 use magnus::{prelude::*, RArray, Value};
+use polars::io::csv::read::OwnedBatchedCsvReader;
 use polars::io::mmap::MmapBytesReader;
 use polars::io::RowIndex;
-use polars::prelude::read_impl::OwnedBatchedCsvReader;
 use polars::prelude::*;
-use std::cell::RefCell;
-use std::path::PathBuf;
 
 use crate::conversion::*;
-use crate::prelude::read_impl::OwnedBatchedCsvReaderMmap;
 use crate::{RbDataFrame, RbPolarsErr, RbResult};
 
-pub enum BatchedReader {
-    MMap(OwnedBatchedCsvReaderMmap),
-    Read(OwnedBatchedCsvReader),
-}
-
 #[magnus::wrap(class = "Polars::RbBatchedCsv")]
 pub struct RbBatchedCsv {
-    pub reader: RefCell<BatchedReader>,
+    pub reader: RefCell<Mutex<OwnedBatchedCsvReader>>,
 }
 
 impl RbBatchedCsv {
@@ -44,18 +40,23 @@ impl RbBatchedCsv {
         let comment_prefix = Option::<String>::try_convert(arguments[16])?;
         let quote_char = Option::<String>::try_convert(arguments[17])?;
         let null_values = Option::<Wrap<NullValues>>::try_convert(arguments[18])?;
-        let try_parse_dates = bool::try_convert(arguments[19])?;
-        let skip_rows_after_header = usize::try_convert(arguments[20])?;
-        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[21])?;
-        let sample_size = usize::try_convert(arguments[22])?;
-        let eol_char = String::try_convert(arguments[23])?;
+        let missing_utf8_is_empty_string = bool::try_convert(arguments[19])?;
+        let try_parse_dates = bool::try_convert(arguments[20])?;
+        let skip_rows_after_header = usize::try_convert(arguments[21])?;
+        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[22])?;
+        let sample_size = usize::try_convert(arguments[23])?;
+        let eol_char = String::try_convert(arguments[24])?;
+        let raise_if_empty = bool::try_convert(arguments[25])?;
+        let truncate_ragged_lines = bool::try_convert(arguments[26])?;
+        let decimal_comma = bool::try_convert(arguments[27])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
         let eol_char = eol_char.as_bytes()[0];
-
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
-
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
         let quote_char = if let Some(s) = quote_char {
             if s.is_empty() {
                 None
@@ -85,53 +86,55 @@ impl RbBatchedCsv {
 
         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)
+        let reader = CsvReadOptions::default()
+            .with_infer_schema_length(infer_schema_length)
+            .with_has_header(has_header)
             .with_n_rows(n_rows)
-            .with_separator(separator.as_bytes()[0])
             .with_skip_rows(skip_rows)
             .with_ignore_errors(ignore_errors)
-            .with_projection(projection)
+            .with_projection(projection.map(Arc::new))
             .with_rechunk(rechunk)
             .with_chunk_size(chunk_size)
-            .with_encoding(encoding.0)
-            .with_columns(columns)
+            .with_columns(columns.map(Arc::new))
             .with_n_threads(n_threads)
-            .with_dtypes_slice(overwrite_dtype_slice.as_deref())
-            .low_memory(low_memory)
-            .with_comment_prefix(comment_prefix.as_deref())
-            .with_null_values(null_values)
-            .with_try_parse_dates(try_parse_dates)
-            .with_quote_char(quote_char)
-            .with_end_of_line_char(eol_char)
+            .with_dtype_overwrite(overwrite_dtype_slice.map(Arc::new))
+            .with_low_memory(low_memory)
             .with_skip_rows_after_header(skip_rows_after_header)
             .with_row_index(row_index)
-            .sample_size(sample_size);
+            .with_sample_size(sample_size)
+            .with_raise_if_empty(raise_if_empty)
+            .with_parse_options(
+                CsvParseOptions::default()
+                    .with_separator(separator.as_bytes()[0])
+                    .with_encoding(encoding.0)
+                    .with_missing_is_null(!missing_utf8_is_empty_string)
+                    .with_comment_prefix(comment_prefix.as_deref())
+                    .with_null_values(null_values)
+                    .with_try_parse_dates(try_parse_dates)
+                    .with_quote_char(quote_char)
+                    .with_eol_char(eol_char)
+                    .with_truncate_ragged_lines(truncate_ragged_lines)
+                    .with_decimal_comma(decimal_comma),
+            )
+            .into_reader_with_file_handle(reader);
 
-        let reader = if low_memory {
-            let reader = reader
-                .batched_read(overwrite_dtype.map(Arc::new))
-                .map_err(RbPolarsErr::from)?;
-            BatchedReader::Read(reader)
-        } else {
-            let reader = reader
-                .batched_mmap(overwrite_dtype.map(Arc::new))
-                .map_err(RbPolarsErr::from)?;
-            BatchedReader::MMap(reader)
-        };
+        let reader = reader
+            .batched(overwrite_dtype.map(Arc::new))
+            .map_err(RbPolarsErr::from)?;
 
         Ok(RbBatchedCsv {
-            reader: RefCell::new(reader),
+            reader: RefCell::new(Mutex::new(reader)),
         })
     }
 
     pub fn next_batches(&self, n: usize) -> RbResult<Option<RArray>> {
-        let batches = match &mut *self.reader.borrow_mut() {
-            BatchedReader::MMap(reader) => reader.next_batches(n),
-            BatchedReader::Read(reader) => reader.next_batches(n),
-        }
-        .map_err(RbPolarsErr::from)?;
+        let reader = &self.reader;
+        let batches = reader
+            .borrow()
+            .lock()
+            .map_err(|e| RbPolarsErr::other(e.to_string()))?
+            .next_batches(n)
+            .map_err(RbPolarsErr::from)?;
 
         Ok(batches.map(|batches| RArray::from_iter(batches.into_iter().map(RbDataFrame::from))))
     }

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

@@ -3,8 +3,8 @@ use magnus::{
     class, prelude::*, r_hash::ForEach, Float, Integer, IntoValue, RArray, RHash, RString, Ruby,
     TryConvert, Value,
 };
-use polars::frame::row::any_values_to_dtype;
 use polars::prelude::*;
+use polars_core::utils::any_values_to_supertype_and_n_dtypes;
 
 use super::{struct_dict, ObjectValue, Wrap};
 
@@ -120,7 +120,8 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
                     avs.push(Wrap::<AnyValue>::try_convert(item?)?.0)
                 }
 
-                let (dtype, _n_types) = any_values_to_dtype(&avs).map_err(RbPolarsErr::from)?;
+                let (dtype, _n_types) =
+                    any_values_to_supertype_and_n_dtypes(&avs).map_err(RbPolarsErr::from)?;
 
                 // push the rest
                 avs.reserve(list.len());

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()
             }
@@ -217,10 +218,7 @@ impl IntoValue for Wrap<DataType> {
                 let class = pl.const_get::<_, Value>("Enum").unwrap();
                 let s = Series::from_arrow("category", categories.to_boxed()).unwrap();
                 let series = to_series(s.into());
-                class
-                    .funcall::<_, _, Value>("new", (series,))
-                    .unwrap()
-                    .into()
+                class.funcall::<_, _, Value>("new", (series,)).unwrap()
             }
             DataType::Time => {
                 let class = pl.const_get::<_, Value>("Time").unwrap();
@@ -245,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()
             }
@@ -313,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.",
@@ -353,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;
@@ -413,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.",
@@ -549,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() {
@@ -736,12 +686,11 @@ 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,
-            // #[cfg(feature = "cross_join")]
-            // "cross" => JoinType::Cross,
+            "cross" => JoinType::Cross,
             v => {
                 return Err(RbValueError::new_err(format!(
                 "how must be one of {{'inner', 'left', 'outer', 'semi', 'anti', 'cross'}}, got {}",
@@ -785,6 +734,21 @@ impl TryConvert for Wrap<ListToStructWidthStrategy> {
     }
 }
 
+impl TryConvert for Wrap<NonExistent> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match String::try_convert(ob)?.as_str() {
+            "null" => NonExistent::Null,
+            "raise" => NonExistent::Raise,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "`non_existent` must be one of {{'null', 'raise'}}, got {v}",
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<NullBehavior> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match String::try_convert(ob)?.as_str() {
@@ -1066,7 +1030,7 @@ impl TryConvert for Wrap<NonZeroUsize> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let v = usize::try_convert(ob)?;
         NonZeroUsize::new(v)
-            .map(|v| Wrap(v))
+            .map(Wrap)
             .ok_or(RbValueError::new_err("must be non-zero".into()))
     }
 }

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()
+    }
+}