fix oob panic on expand_at_index and series from pyarrow chunkedarray (…

…#3938)
pola-rs · Jul 8, 2022 · 517fd6a · 517fd6a
1 parent ff9e9c0
commit 517fd6a
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 2 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/mod.rs b/polars/polars-core/src/chunked_array/ops/mod.rs
@@ -615,6 +615,9 @@ pub trait ChunkExpandAtIndex<T> {
 
 macro_rules! impl_chunk_expand {
     ($self:ident, $length:ident, $index:ident) => {{
+        if $self.is_empty() {
+            return $self.clone();
+        }
         let opt_val = $self.get($index);
         match opt_val {
             Some(val) => ChunkedArray::full($self.name(), val, $length),

diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -117,7 +117,13 @@ def wrap_s(s: PySeries) -> Series:
 
 
 ArrayLike = Union[
-    Sequence[Any], "Series", "pa.Array", "np.ndarray", "pd.Series", "pd.DatetimeIndex"
+    Sequence[Any],
+    "Series",
+    "pa.Array",
+    "pa.ChunkedArray",
+    "np.ndarray",
+    "pd.Series",
+    "pd.DatetimeIndex",
 ]
 
 
@@ -212,7 +218,7 @@ def __init__(
             self._s = sequence_to_pyseries(name, [], dtype=dtype)
         elif isinstance(values, Series):
             self._s = series_to_pyseries(name, values)
-        elif _PYARROW_AVAILABLE and isinstance(values, pa.Array):
+        elif _PYARROW_AVAILABLE and isinstance(values, (pa.Array, pa.ChunkedArray)):
             self._s = arrow_to_pyseries(name, values)
         elif _NUMPY_AVAILABLE and isinstance(values, np.ndarray):
             self._s = numpy_to_pyseries(name, values, strict, nan_to_null)

diff --git a/py-polars/tests/test_datelike.py b/py-polars/tests/test_datelike.py
@@ -1215,3 +1215,8 @@ def test_time_microseconds_3843() -> None:
     in_val = [time(0, 9, 11, 558332)]
     s = pl.Series(in_val)
     assert s.to_list() == in_val
+
+
+def test_year_empty_df() -> None:
+    df = pl.DataFrame(pl.Series(name="date", dtype=pl.Date))
+    assert df.select(pl.col("date").dt.year()).dtypes == [pl.Int32]
diff --git a/py-polars/tests/test_interop.py b/py-polars/tests/test_interop.py
@@ -396,3 +396,9 @@ def test_cat_int_types_3500() -> None:
         for t in [int_dict_type, uint_dict_type]:
             s = pl.from_arrow(pyarrow_array.cast(t))
             assert s.series_equal(pl.Series(["a", "a", "b"]).cast(pl.Categorical))
+
+
+def test_from_pyarrow_chunked_array() -> None:
+    column = pa.chunked_array([[1], [2]])
+    series = pl.Series("column", column)
+    assert series.to_list() == [1, 2]