Skip to content

Commit

Permalink
working version with fixed test
Browse files Browse the repository at this point in the history
  • Loading branch information
barak1412 committed Oct 6, 2024
1 parent f7de80c commit a32d256
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 6 deletions.
4 changes: 2 additions & 2 deletions crates/polars-core/src/series/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -528,13 +528,13 @@ impl Series {
Ok(max)
}

/// Explode a list Series. This expands every item to a new row..
/// Explode a list/array Series. This expands every item to a new row.
pub fn explode(&self) -> PolarsResult<Series> {
match self.dtype() {
DataType::List(_) => self.list().unwrap().explode(),
#[cfg(feature = "dtype-array")]
DataType::Array(_, _) => self.array().unwrap().explode(),
_ => Ok(self.clone()),
dt => polars_bail!(opq = explode, dt),
}
}

Expand Down
6 changes: 5 additions & 1 deletion crates/polars-expr/src/expressions/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -467,7 +467,11 @@ impl<'a> AggregationContext<'a> {
AggState::AggregatedScalar(s) => (s, groups),
AggState::Literal(s) => (s, groups),
AggState::AggregatedList(s) => {
let flattened = s.explode().unwrap();
// flat only if we have a list
let flattened = match s.dtype() {
DataType::List(_) => s.explode().unwrap(),
_ => s.clone(),
};
let groups = groups.into_owned();
// unroll the possible flattened state
// say we have groups with overlapping windows:
Expand Down
8 changes: 8 additions & 0 deletions py-polars/tests/unit/dataframe/test_df.py
Original file line number Diff line number Diff line change
Expand Up @@ -621,6 +621,14 @@ def test_explode() -> None:
assert out["nrs"].to_list() == [1, 2, 1, 3]


def test_illegal_explode_19049() -> None:
df = pl.DataFrame({"a": [0]}, schema={"a": pl.Int64})
with pytest.raises(
InvalidOperationError, match="`explode` operation not supported for dtype `i64`"
):
df.select(pl.col.a.explode())


@pytest.mark.parametrize(
("stack", "exp_shape", "exp_columns"),
[
Expand Down
6 changes: 3 additions & 3 deletions py-polars/tests/unit/functions/range/test_time_range.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ def test_time_range_lit_eager() -> None:
).alias("tm")
)
if not eager:
tm = tm.select(pl.col("tm").explode())
tm = tm.select(pl.col("tm"))
assert tm["tm"].to_list() == [
time(6, 47, 13, 333000),
time(12, 32, 23, 666000),
Expand All @@ -178,7 +178,7 @@ def test_time_range_lit_eager() -> None:
).alias("tm")
)
if not eager:
tm = tm.select(pl.col("tm").explode())
tm = tm.select(pl.col("tm"))
assert tm["tm"].to_list() == [
time(0, 0),
time(5, 45, 10, 333000),
Expand All @@ -194,7 +194,7 @@ def test_time_range_lit_eager() -> None:
eager=eager,
).alias("tm")
)
tm = tm.select(pl.col("tm").explode())
tm = tm.select(pl.col("tm"))
assert tm["tm"].to_list() == [
time(23, 59, 59, 999980),
time(23, 59, 59, 999990),
Expand Down

0 comments on commit a32d256

Please sign in to comment.