ronantakizawa/codereview-bench huggingface dataset:
Repository: pola-rs/polars
File: py-polars/polars/_utils/various.py (Python)

Code before:
    except (ImportError, FileNotFoundError):
        msg = (
            "the graphviz `dot` binary should be on your PATH."
            "(If not installed you can download here: https://graphviz.org/download/)"
        )
        raise ImportError(msg) from None

    if output_path:
        Path(output_path).write_bytes(graph)

    if not show:
        return None

    if _in_notebook():
        from IPython.display import SVG, display

        return display(SVG(graph))
    else:
        if (cmd := os.environ.get("POLARS_DOT_SVG_VIEWER", None)) is not None:
            import tempfile

            with tempfile.NamedTemporaryFile(suffix=".svg") as file:
                file.write(graph)
                file.flush()
                cmd = cmd.replace("%file%", file.name)
                subprocess.check_output(["sh", "-c", cmd])
            return None

        import_optional(
            "matplotlib",
            err_prefix="",
            err_suffix="should be installed to show graphs",
        )
        import matplotlib.image as mpimg
        import matplotlib.pyplot as plt

        plt.figure(figsize=figsize)
        img = mpimg.imread(BytesIO(graph))
        plt.axis("off")
        plt.imshow(img)
        plt.show()
        return None


def qualified_type_name(obj: Any, *, qualify_polars: bool = False) -> str:
    """
    Return the module-qualified name of the given object as a string.

    Parameters
    ----------
    obj

Review comment (suggestion):
If you insist on using shell commands, at least use `subprocess.check_output(cmd, shell=True)` rather than hardbaking `sh -c`.

Code after:
    except (ImportError, FileNotFoundError):
        msg = (
            "the graphviz `dot` binary should be on your PATH."
            "(If not installed you can download here: https://graphviz.org/download/)"
        )
        raise ImportError(msg) from None

    if output_path:
        Path(output_path).write_bytes(graph)

    if not show:
        return None

    if _in_notebook():
        from IPython.display import SVG, display

        return display(SVG(graph))
    else:
        if (cmd := os.environ.get("POLARS_DOT_SVG_VIEWER", None)) is not None:
            import tempfile

            with tempfile.NamedTemporaryFile(suffix=".svg") as file:
                file.write(graph)
                file.flush()
                cmd = cmd.replace("%file%", file.name)
                subprocess.run(cmd, shell=True)
            return None

        import_optional(
            "matplotlib",
            err_prefix="",
            err_suffix="should be installed to show graphs",
        )
        import matplotlib.image as mpimg
        import matplotlib.pyplot as plt

        plt.figure(figsize=figsize)
        img = mpimg.imread(BytesIO(graph))
        plt.axis("off")
        plt.imshow(img)
        plt.show()
        return None


def qualified_type_name(obj: Any, *, qualify_polars: bool = False) -> str:
    """
    Return the module-qualified name of the given object as a string.

    Parameters
    ----------
    obj
Repository: pola-rs/polars
File: crates/polars-expr/src/expressions/mod.rs (Rust)

Code before:
    }

    pub(crate) fn is_not_aggregated(&self) -> bool {
        matches!(
            &self.state,
            AggState::NotAggregated(_) | AggState::Literal(_)
        )
    }

    pub(crate) fn is_aggregated(&self) -> bool {
        !self.is_not_aggregated()
    }

    pub(crate) fn is_literal(&self) -> bool {
        matches!(self.state, AggState::Literal(_))
    }

    /// # Arguments
    /// - `aggregated` sets if the Series is a list due to aggregation (could also be a list because its
    ///   the columns dtype)
    fn new(
        column: Column,
        groups: Cow<'a, GroupPositions>,
        aggregated: bool,
    ) -> AggregationContext<'a> {
        let series = match aggregated {
            true => {
                assert_eq!(column.len(), groups.len());
                AggState::AggregatedScalar(column)
            },
            _ => AggState::NotAggregated(column),
        };

        Self {
            state: series,
            groups,
            sorted: false,
            update_groups: UpdateGroups::No,
            original_len: true,
        }
    }

    fn with_agg_state(&mut self, agg_state: AggState) {
        self.state = agg_state;
    }

    fn from_agg_state(
        agg_state: AggState,
        groups: Cow<'a, GroupPositions>,
    ) -> AggregationContext<'a> {
        Self {

Review comment (suggestion):
I'll trust you on the details that we don't have to make an `AggregatedList` here but...

```rust
match boolean {
    true => a,
    _ => b,
}
```

This is just if/else and I have no idea how clippy isn't screaming at you 😅 

Code after:
    }

    pub(crate) fn is_not_aggregated(&self) -> bool {
        matches!(
            &self.state,
            AggState::NotAggregated(_) | AggState::Literal(_)
        )
    }

    pub(crate) fn is_aggregated(&self) -> bool {
        !self.is_not_aggregated()
    }

    pub(crate) fn is_literal(&self) -> bool {
        matches!(self.state, AggState::Literal(_))
    }

    /// # Arguments
    /// - `aggregated` sets if the Series is a list due to aggregation (could also be a list because its
    ///   the columns dtype)
    fn new(
        column: Column,
        groups: Cow<'a, GroupPositions>,
        aggregated: bool,
    ) -> AggregationContext<'a> {
        let series = if aggregated {
            assert_eq!(column.len(), groups.len());
            AggState::AggregatedScalar(column)
        } else {
            AggState::NotAggregated(column)
        };

        Self {
            state: series,
            groups,
            sorted: false,
            update_groups: UpdateGroups::No,
            original_len: true,
        }
    }

    fn with_agg_state(&mut self, agg_state: AggState) {
        self.state = agg_state;
    }

    fn from_agg_state(
        agg_state: AggState,
        groups: Cow<'a, GroupPositions>,
    ) -> AggregationContext<'a> {
        Self {
            state: agg_state,
Repository: pola-rs/polars
File: crates/polars-expr/src/expressions/mod.rs (Rust)

Code before:
    }

    pub(crate) fn is_not_aggregated(&self) -> bool {
        matches!(
            &self.state,
            AggState::NotAggregated(_) | AggState::Literal(_)
        )
    }

    pub(crate) fn is_aggregated(&self) -> bool {
        !self.is_not_aggregated()
    }

    pub(crate) fn is_literal(&self) -> bool {
        matches!(self.state, AggState::Literal(_))
    }

    /// # Arguments
    /// - `aggregated` sets if the Series is a list due to aggregation (could also be a list because its
    ///   the columns dtype)
    fn new(
        column: Column,
        groups: Cow<'a, GroupPositions>,
        aggregated: bool,
    ) -> AggregationContext<'a> {
        let series = match (aggregated, column.dtype()) {
            (true, &DataType::List(_)) => {
                assert_eq!(column.len(), groups.len());
                AggState::AggregatedList(column)
            },
            (true, _) => {
                assert_eq!(column.len(), groups.len());
                AggState::AggregatedScalar(column)
            },
            _ => AggState::NotAggregated(column),
        };

        Self {
            state: series,
            groups,
            sorted: false,
            update_groups: UpdateGroups::No,
            original_len: true,
        }
    }

    fn with_agg_state(&mut self, agg_state: AggState) {
        self.state = agg_state;
    }

    fn from_agg_state(

Review comment (suggestion):
This was only called at 2 places.
- One had `aggregated=false`
- The other was `count.rs`, had `aggregated=true`, but had an constant datatype of `IDX_DTYPE`.

Code after:
    }

    pub(crate) fn is_not_aggregated(&self) -> bool {
        matches!(
            &self.state,
            AggState::NotAggregated(_) | AggState::Literal(_)
        )
    }

    pub(crate) fn is_aggregated(&self) -> bool {
        !self.is_not_aggregated()
    }

    pub(crate) fn is_literal(&self) -> bool {
        matches!(self.state, AggState::Literal(_))
    }

    /// # Arguments
    /// - `aggregated` sets if the Series is a list due to aggregation (could also be a list because its
    ///   the columns dtype)
    fn new(
        column: Column,
        groups: Cow<'a, GroupPositions>,
        aggregated: bool,
    ) -> AggregationContext<'a> {
        let series = if aggregated {
            assert_eq!(column.len(), groups.len());
            AggState::AggregatedScalar(column)
        } else {
            AggState::NotAggregated(column)
        };

        Self {
            state: series,
            groups,
            sorted: false,
            update_groups: UpdateGroups::No,
            original_len: true,
        }
    }

    fn with_agg_state(&mut self, agg_state: AggState) {
        self.state = agg_state;
    }

    fn from_agg_state(
        agg_state: AggState,
        groups: Cow<'a, GroupPositions>,
    ) -> AggregationContext<'a> {
        Self {
            state: agg_state,
Repository: pola-rs/polars
File: crates/polars-lazy/src/frame/python.rs (Rust)

Code before:
use std::sync::Arc;

use polars_plan::dsl::python_dsl;
use pyo3::PyObject;

use self::python_dsl::{PythonOptionsDsl, PythonScanSource};
use crate::prelude::*;

impl LazyFrame {
    pub fn scan_from_python_function(
        schema: python_dsl::SchemaProvider,
        scan_fn: PyObject,
        pyarrow: bool,
        // Validate that the source gives the proper schema
        validate_schema: bool,
    ) -> Self {
        DslPlan::PythonScan {
            options: PythonOptionsDsl {
                // Should be a python function that returns a generator
                scan_fn: Some(scan_fn.into()),
                schema_provider: Some(Arc::new(schema)),
                python_source: if pyarrow {
                    PythonScanSource::Pyarrow
                } else {
                    PythonScanSource::IOPlugin
                },
                validate_schema,
            },
        }
        .into()
    }
}


Review comment (suggestion):
drive-by replace `Either` with a more descriptive enum


Code after:
use std::sync::Arc;

use either::Either;
use polars_core::schema::SchemaRef;
use pyo3::PyObject;

use self::python_dsl::{PythonOptionsDsl, PythonScanSource};
use crate::prelude::*;

impl LazyFrame {
    pub fn scan_from_python_function(
        schema: Either<PyObject, SchemaRef>,
        scan_fn: PyObject,
        pyarrow: bool,
        // Validate that the source gives the proper schema
        validate_schema: bool,
    ) -> Self {
        DslPlan::PythonScan {
            options: PythonOptionsDsl {
                // Should be a python function that returns a generator
                scan_fn: Some(scan_fn.into()),
                schema_fn: Some(SpecialEq::new(Arc::new(schema.map_left(|obj| obj.into())))),
                python_source: if pyarrow {
                    PythonScanSource::Pyarrow
                } else {
                    PythonScanSource::IOPlugin
                },
                validate_schema,
            },
        }
        .into()
    }
}
Repository: pola-rs/polars
File: crates/polars-stream/src/nodes/rle.rs (Rust)

Code before:
            seq: MorselSeq::default(),
            last_length: 0,
            last: None,
        }
    }
}

impl ComputeNode for RleNode {
    fn name(&self) -> &str {
        "rle"
    }

    fn update_state(
        &mut self,
        recv: &mut [PortState],
        send: &mut [PortState],
        _state: &StreamingExecutionState,
    ) -> PolarsResult<()> {
        assert!(recv.len() == 1 && send.len() == 1);

        if send[0] == PortState::Done {
            recv[0] = PortState::Done;
            self.last_length = 0;
            self.last.take();
        } else if self.last.is_some() {
            send[0] = PortState::Ready;
        } else if recv[0] == PortState::Done {
            send[0] = PortState::Done;
        } else {
            recv.swap_with_slice(send);
        }

        Ok(())
    }

    fn spawn<'env, 's>(
        &'env mut self,
        scope: &'s TaskScope<'s, 'env>,
        recv_ports: &mut [Option<RecvPort<'_>>],
        send_ports: &mut [Option<SendPort<'_>>],
        _state: &'s StreamingExecutionState,
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        assert_eq!(recv_ports.len(), 1);
        assert_eq!(send_ports.len(), 1);

        let recv = recv_ports[0].take();
        let mut send = send_ports[0].take().unwrap().serial();

        let fields = vec![
            Field::new(PlSmallStr::from_static(RLE_LENGTH_COLUMN_NAME), IDX_DTYPE),

Review comment (suggestion):
This isn't right. If your input transitioned from `Ready` to `Blocked` you might still have `self.last.is_some()` but you won't be ready to send stuff.

You should put this inside the `if recv[0] == PortState::Done` branch.

Code after:
            seq: MorselSeq::default(),
            last_length: 0,
            last: None,
        }
    }
}

impl ComputeNode for RleNode {
    fn name(&self) -> &str {
        "rle"
    }

    fn update_state(
        &mut self,
        recv: &mut [PortState],
        send: &mut [PortState],
        _state: &StreamingExecutionState,
    ) -> PolarsResult<()> {
        assert!(recv.len() == 1 && send.len() == 1);

        if send[0] == PortState::Done {
            recv[0] = PortState::Done;
            self.last_length = 0;
            self.last.take();
        } else if recv[0] == PortState::Done {
            if self.last.is_some() {
                send[0] = PortState::Ready;
            } else {
                send[0] = PortState::Done;
            }
        } else {
            recv.swap_with_slice(send);
        }

        Ok(())
    }

    fn spawn<'env, 's>(
        &'env mut self,
        scope: &'s TaskScope<'s, 'env>,
        recv_ports: &mut [Option<RecvPort<'_>>],
        send_ports: &mut [Option<SendPort<'_>>],
        _state: &'s StreamingExecutionState,
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        assert_eq!(recv_ports.len(), 1);
        assert_eq!(send_ports.len(), 1);

        let recv = recv_ports[0].take();
        let mut send = send_ports[0].take().unwrap().serial();
Repository: pola-rs/polars
File: py-polars/tests/unit/interop/test_interop.py (Python)

Code before:
                dtype=pl.Duration("ns"),
            ),
            "categorical": pl.Series(
                ["apple", "banana", "apple"], dtype=pl.Categorical
            ),
            "list_duration": [
                [timedelta(days=1), timedelta(hours=2)],
                [timedelta(minutes=30)],
                None,
            ],
            "struct_with_duration": [
                {"x": timedelta(days=1), "y": 1},
                {"x": timedelta(hours=2), "y": 2},
                None,
            ],
        }
    ).cast(
        {
            "list_duration": pl.List(pl.Duration("us")),
            "struct_with_duration": pl.Struct({"x": pl.Duration("ns"), "y": pl.Int32}),
        }
    )

    df_roundtrip = df.map_columns(pl.selectors.all(), roundtrip_series_pycapsule)

    assert_frame_equal(df_roundtrip, df)


Review comment (suggestion):
One more thing we can do here - could we also add a test the `DataFrame` pycapsule roundtrip here - i.e. `pl.DataFrame(PyCapsuleStreamWrap(df))`


Code after:
                dtype=pl.Duration("ns"),
            ),
            "categorical": pl.Series(
                ["apple", "banana", "apple"], dtype=pl.Categorical
            ),
            "list_duration": [
                [timedelta(days=1), timedelta(hours=2)],
                [timedelta(minutes=30)],
                None,
            ],
            "struct_with_duration": [
                {"x": timedelta(days=1), "y": 1},
                {"x": timedelta(hours=2), "y": 2},
                None,
            ],
        }
    ).cast(
        {
            "list_duration": pl.List(pl.Duration("us")),
            "struct_with_duration": pl.Struct({"x": pl.Duration("ns"), "y": pl.Int32}),
        }
    )

    df_roundtrip = df.map_columns(pl.selectors.all(), roundtrip_series_pycapsule)

    assert_frame_equal(df_roundtrip, df)

    df_roundtrip_direct = pl.DataFrame(PyCapsuleStreamWrap(df))

    assert_frame_equal(df_roundtrip_direct, df)
Repository: pola-rs/polars
File: py-polars/tests/unit/series/test_series.py (Python)

Code before:
        pl.Series(["foo", "bar", "buz", "bar"], dtype=pl.Categorical).mode().item()
        == "bar"
    )
    assert pl.Series([1.0, 2.0, 3.0, 2.0]).mode().item() == 2.0
    assert pl.Series(["a", "b", "c", "b"]).mode().item() == "b"

    # sorted data
    assert set(pl.int_range(0, 3, eager=True).mode().to_list()) == {0, 1, 2}


def test_diff() -> None:
    s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])
    df = pl.DataFrame([s])

    assert_series_equal(
        df.select(pl.col("a").diff())["a"], pl.Series("a", [None, 1, 1, -1, 0, 1, -3])
    )
    assert_series_equal(
        df.select(pl.col("a").diff(null_behavior="drop"))["a"],
        pl.Series("a", [1, 1, -1, 0, 1, -3]),
    )


def test_diff_negative() -> None:
    s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])
    df = pl.DataFrame([s])

    assert_series_equal(
        df.select(pl.col("a").diff(-1))["a"],
        pl.Series("a", [-1, -1, 1, 0, -1, 3, None]),
    )
    assert_series_equal(
        df.select(pl.col("a").diff(-1, null_behavior="drop"))["a"],
        pl.Series("a", [-1, -1, 1, 0, -1, 3]),
    )


def test_pct_change() -> None:
    s = pl.Series("a", [1, 2, 4, 8, 16, 32, 64])
    expected = pl.Series("a", [None, None, 3.0, 3.0, 3.0, 3.0, 3.0])
    assert_series_equal(s.pct_change(2), expected)
    assert_series_equal(s.pct_change(pl.Series([2])), expected)
    # negative
    assert pl.Series(range(5)).pct_change(-1).to_list() == [
        -1.0,
        -0.5,
        -0.3333333333333333,
        -0.25,
        None,
    ]


Review comment (suggestion):
I would just write
```python
df = pl.DataFrame({"a": [1, 2, 3, 2, 2, 3, 0]})
```

Code after:
        pl.Series(["foo", "bar", "buz", "bar"], dtype=pl.Categorical).mode().item()
        == "bar"
    )
    assert pl.Series([1.0, 2.0, 3.0, 2.0]).mode().item() == 2.0
    assert pl.Series(["a", "b", "c", "b"]).mode().item() == "b"

    # sorted data
    assert set(pl.int_range(0, 3, eager=True).mode().to_list()) == {0, 1, 2}


def test_diff() -> None:
    s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])

    assert_series_equal(
        s.diff(),
        pl.Series("a", [None, 1, 1, -1, 0, 1, -3]),
    )
    assert_series_equal(
        s.diff(null_behavior="drop"),
        pl.Series("a", [1, 1, -1, 0, 1, -3]),
    )


def test_diff_negative() -> None:
    s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])

    assert_series_equal(
        s.diff(-1),
        pl.Series("a", [-1, -1, 1, 0, -1, 3, None]),
    )
    assert_series_equal(
        s.diff(-1, null_behavior="drop"),
        pl.Series("a", [-1, -1, 1, 0, -1, 3]),
    )


def test_pct_change() -> None:
    s = pl.Series("a", [1, 2, 4, 8, 16, 32, 64])
    expected = pl.Series("a", [None, None, 3.0, 3.0, 3.0, 3.0, 3.0])
    assert_series_equal(s.pct_change(2), expected)
    assert_series_equal(s.pct_change(pl.Series([2])), expected)
    # negative
    assert pl.Series(range(5)).pct_change(-1).to_list() == [
        -1.0,
        -0.5,
        -0.3333333333333333,
        -0.25,
        None,
    ]
Repository: pola-rs/polars
File: py-polars/tests/unit/series/test_series.py (Python)

Code before:
        == "bar"
    )
    assert pl.Series([1.0, 2.0, 3.0, 2.0]).mode().item() == 2.0
    assert pl.Series(["a", "b", "c", "b"]).mode().item() == "b"

    # sorted data
    assert set(pl.int_range(0, 3, eager=True).mode().to_list()) == {0, 1, 2}


def test_diff() -> None:
    s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])
    df = pl.DataFrame([s])

    assert_series_equal(
        df.select(pl.col("a").diff())["a"], pl.Series("a", [None, 1, 1, -1, 0, 1, -3])
    )
    assert_series_equal(
        df.select(pl.col("a").diff(null_behavior="drop"))["a"],
        pl.Series("a", [1, 1, -1, 0, 1, -3]),
    )


def test_diff_negative() -> None:
    s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])
    df = pl.DataFrame([s])

    assert_series_equal(
        df.select(pl.col("a").diff(-1))["a"],
        pl.Series("a", [-1, -1, 1, 0, -1, 3, None]),
    )
    assert_series_equal(
        df.select(pl.col("a").diff(-1, null_behavior="drop"))["a"],
        pl.Series("a", [-1, -1, 1, 0, -1, 3]),
    )


def test_pct_change() -> None:
    s = pl.Series("a", [1, 2, 4, 8, 16, 32, 64])
    expected = pl.Series("a", [None, None, 3.0, 3.0, 3.0, 3.0, 3.0])
    assert_series_equal(s.pct_change(2), expected)
    assert_series_equal(s.pct_change(pl.Series([2])), expected)
    # negative
    assert pl.Series(range(5)).pct_change(-1).to_list() == [
        -1.0,
        -0.5,
        -0.3333333333333333,
        -0.25,
        None,
    ]



Review comment (suggestion):
Here too I would just use `assert_frame_equal` with a similar DataFrame definition to the above.

Or write the entire thing using series and `Series.diff` instead of `select.

Code after:
        == "bar"
    )
    assert pl.Series([1.0, 2.0, 3.0, 2.0]).mode().item() == 2.0
    assert pl.Series(["a", "b", "c", "b"]).mode().item() == "b"

    # sorted data
    assert set(pl.int_range(0, 3, eager=True).mode().to_list()) == {0, 1, 2}


def test_diff() -> None:
    s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])

    assert_series_equal(
        s.diff(),
        pl.Series("a", [None, 1, 1, -1, 0, 1, -3]),
    )
    assert_series_equal(
        s.diff(null_behavior="drop"),
        pl.Series("a", [1, 1, -1, 0, 1, -3]),
    )


def test_diff_negative() -> None:
    s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])

    assert_series_equal(
        s.diff(-1),
        pl.Series("a", [-1, -1, 1, 0, -1, 3, None]),
    )
    assert_series_equal(
        s.diff(-1, null_behavior="drop"),
        pl.Series("a", [-1, -1, 1, 0, -1, 3]),
    )


def test_pct_change() -> None:
    s = pl.Series("a", [1, 2, 4, 8, 16, 32, 64])
    expected = pl.Series("a", [None, None, 3.0, 3.0, 3.0, 3.0, 3.0])
    assert_series_equal(s.pct_change(2), expected)
    assert_series_equal(s.pct_change(pl.Series([2])), expected)
    # negative
    assert pl.Series(range(5)).pct_change(-1).to_list() == [
        -1.0,
        -0.5,
        -0.3333333333333333,
        -0.25,
        None,
    ]


def test_skew() -> None: