fix tests

pola-rs · Nov 18, 2024 · 569a827 · 569a827
1 parent 9241c27
commit 569a827
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 21 deletions.
diff --git a/crates/polars-python/src/dataframe/general.rs b/crates/polars-python/src/dataframe/general.rs
@@ -727,7 +727,7 @@ impl PyDataFrame {
             let chunks = df
                 .get_columns()
                 .iter()
-                .map(|c| c.as_materialized_series().chunks()[0].to_boxed())
+                .map(|c| c.as_materialized_series().to_physical_repr().chunks()[0].to_boxed())
                 .collect::<Vec<_>>();
             let fields = fields
                 .into_iter()

diff --git a/crates/polars-python/src/series/general.rs b/crates/polars-python/src/series/general.rs
@@ -10,7 +10,6 @@ use pyo3::Python;
 
 use super::PySeries;
 use crate::dataframe::PyDataFrame;
-use crate::datatypes::PyDataType;
 use crate::error::PyPolarsErr;
 use crate::prelude::*;
 use crate::py_modules::POLARS;
@@ -541,14 +540,18 @@ impl PySeries {
     fn _row_decode<'py>(
         &'py self,
         py: Python<'py>,
-        dtypes: Vec<(String, PyDataType)>,
+        dtypes: Vec<(String, Wrap<DataType>)>,
         fields: Vec<(bool, bool, bool)>,
     ) -> PyResult<PyDataFrame> {
         assert_eq!(dtypes.len(), fields.len());
 
         let arrow_dtypes = dtypes
             .iter()
-            .map(|(_, dt)| DataType::from(dt.clone()).to_arrow(CompatLevel::newest()))
+            .map(|(_, dt)| {
+                DataType::from(dt.0.clone())
+                    .to_physical()
+                    .to_arrow(CompatLevel::newest())
+            })
             .collect::<Vec<_>>();
         let fields = fields
             .into_iter()
@@ -582,7 +585,7 @@ impl PySeries {
                             Series::from_chunks_and_dtype_unchecked(
                                 PlSmallStr::from(name),
                                 vec![arr],
-                                &DataType::from(dtype),
+                                &DataType::from(dtype.0),
                             )
                         }
                         .into_column()

diff --git a/py-polars/polars/series/series.py b/py-polars/polars/series/series.py
@@ -7519,8 +7519,8 @@ def plot(self) -> SeriesPlot:
 
     def _row_decode(
         self,
-        dtypes: list[tuple[str, DataType]],
-        fields: list[tuple[bool, bool, bool]],
+        dtypes: Iterable[tuple[str, DataType]],  # type: ignore[valid-type]
+        fields: Iterable[tuple[bool, bool, bool]],
     ) -> DataFrame:
         """
         Row decode the given Series.
@@ -7533,7 +7533,7 @@ def _row_decode(
         - nulls_last
         - no_order
         """
-        return pl.DataFrame._from_pydf(self._s._row_decode(dtypes, fields))
+        return pl.DataFrame._from_pydf(self._s._row_decode(list(dtypes), list(fields)))
 
 
 def _resolve_temporal_dtype(

diff --git a/py-polars/tests/unit/test_row_encoding.py b/py-polars/tests/unit/test_row_encoding.py
@@ -1,14 +1,13 @@
 from __future__ import annotations
 
-import numpy as np
 import pytest
 from hypothesis import given
 
 import polars as pl
 from polars.testing import assert_frame_equal
 from polars.testing.parametric import dataframes
 
-# @TODO: At the moment no_order=True breaks roundtripping for some reason
+# @TODO: Deal with no_order
 FIELD_COMBS = [
     (descending, nulls_last, False)
     for descending in [False, True]
@@ -37,12 +36,6 @@ def roundtrip_re(
             pl.Struct,
             pl.Categorical,
             pl.Enum,
-            pl.Time,
-            pl.Datetime,
-            pl.Date,
-            pl.Duration,
-            pl.Null,
-            pl.Decimal,
         ]
     )
 )
@@ -53,6 +46,15 @@ def test_row_encoding_parametric(
     roundtrip_re(df, [field] * df.width)
 
 
+@pytest.mark.parametrize("field", FIELD_COMBS)
+def test_nulls(field: tuple[bool, bool, bool]) -> None:
+    roundtrip_re(pl.Series("a", [], pl.Null).to_frame(), [field])
+    roundtrip_re(pl.Series("a", [None], pl.Null).to_frame(), [field])
+    roundtrip_re(pl.Series("a", [None] * 2, pl.Null).to_frame(), [field])
+    roundtrip_re(pl.Series("a", [None] * 13, pl.Null).to_frame(), [field])
+    roundtrip_re(pl.Series("a", [None] * 42, pl.Null).to_frame(), [field])
+
+
 @pytest.mark.parametrize("field", FIELD_COMBS)
 def test_bool(field: tuple[bool, bool, bool]) -> None:
     roundtrip_re(pl.Series("a", [], pl.Boolean).to_frame(), [field])
@@ -99,16 +101,17 @@ def test_int(dtype: pl.DataType, field: tuple[bool, bool, bool]) -> None:
 )
 @pytest.mark.parametrize("field", FIELD_COMBS)
 def test_float(dtype: pl.DataType, field: tuple[bool, bool, bool]) -> None:
+    inf = float("inf")
+    inf_b = float("-inf")
+
     roundtrip_re(pl.Series("a", [], dtype).to_frame(), [field])
     roundtrip_re(pl.Series("a", [0.0], dtype).to_frame(), [field])
-    roundtrip_re(pl.Series("a", [np.Infinity], dtype).to_frame(), [field])
-    roundtrip_re(pl.Series("a", [-np.Infinity], dtype).to_frame(), [field])
+    roundtrip_re(pl.Series("a", [inf], dtype).to_frame(), [field])
+    roundtrip_re(pl.Series("a", [-inf_b], dtype).to_frame(), [field])
 
     roundtrip_re(pl.Series("a", [1.0, 2.0, 3.0], dtype).to_frame(), [field])
     roundtrip_re(pl.Series("a", [0.0, 1.0, 2.0, 3.0], dtype).to_frame(), [field])
-    roundtrip_re(
-        pl.Series("a", [np.Infinity, 0, -np.Infinity], dtype).to_frame(), [field]
-    )
+    roundtrip_re(pl.Series("a", [inf, 0, -inf_b], dtype).to_frame(), [field])
 
 
 @pytest.mark.parametrize("field", FIELD_COMBS)