fix and test estimated_size (#3113)

pola-rs · Apr 11, 2022 · 861ac9c · 861ac9c
1 parent 9c948b7
commit 861ac9c
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/py-polars/polars/internals/frame.py b/py-polars/polars/internals/frame.py
@@ -337,7 +337,7 @@ def estimated_size(self) -> int:
 
         FFI buffers are included in this estimation.
         """
-        return self._s.estimated_size()
+        return self._df.estimated_size()
 
     @classmethod
     def _from_pydf(cls: Type[DF], py_df: "PyDataFrame") -> DF:

diff --git a/py-polars/src/dataframe.rs b/py-polars/src/dataframe.rs
@@ -41,10 +41,6 @@ impl PyDataFrame {
         PyDataFrame { df }
     }
 
-    pub fn estimated_size(&self) -> usize {
-        self.df.estimated_size()
-    }
-
     fn finish_from_rows(rows: Vec<Row>) -> PyResult<Self> {
         // replace inferred nulls with boolean
         let schema = rows_to_schema(&rows);
@@ -82,6 +78,10 @@ impl PyDataFrame {
         Ok(PyDataFrame::new(df))
     }
 
+    pub fn estimated_size(&self) -> usize {
+        self.df.estimated_size()
+    }
+
     #[staticmethod]
     #[allow(clippy::too_many_arguments)]
     pub fn read_csv(

diff --git a/py-polars/tests/test_utils.py b/py-polars/tests/test_utils.py
@@ -1,5 +1,6 @@
 from datetime import date, datetime, timedelta
 
+import polars as pl
 from polars.utils import (
     _date_to_pl_date,
     _datetime_to_pl_timestamp,
@@ -33,3 +34,8 @@ def test_timedelta_to_pl_timedelta() -> None:
     assert out == 86_400_000
     out = _timedelta_to_pl_timedelta(timedelta(days=1), tu=None)
     assert out == 86_400_000_000_000
+
+
+def test_estimated_size() -> None:
+    a = pl.Series([1, 2, 3])
+    assert a.estimated_size() == a.to_frame().estimated_size()