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

Code before:
            )
        )

    elif isinstance(value, enum.Enum):
        return lit(value.value, dtype=dtype)

    if dtype:
        return wrap_expr(plr.lit(value, allow_object, is_scalar=True)).cast(dtype)

    if _check_for_numpy(value) and isinstance(value, np.generic):
        # note: the item() is a py-native datetime/timedelta when units < 'ns'
        if isinstance(item := value.item(), (datetime, timedelta)):
            return lit(item)

        # handle 'ns' units
        if isinstance(item, int) and hasattr(value, "dtype"):
            dtype_name = value.dtype.name
            if dtype_name.startswith("datetime64["):
                time_unit = dtype_name[len("datetime64[") : -1]  # type: ignore[assignment]
                return lit(item).cast(Datetime(time_unit))
            if dtype_name.startswith("timedelta64["):
                time_unit = dtype_name[len("timedelta64[") : -1]  # type: ignore[assignment]
                return lit(item).cast(Duration(time_unit))

        # handle numeric values
        if isinstance(value, np.number):
            dtype = DataTypeMappings.NUMPY_KIND_AND_ITEMSIZE_TO_DTYPE.get(
                (value.dtype.kind, value.dtype.itemsize)
            )
            if dtype is not None:
                return lit(value, dtype=dtype)
    else:
        item = value

    return wrap_expr(plr.lit(item, allow_object, is_scalar=True))


Review comment (suggestion):
It's probably better match against `np.generic` instead of just `np.number` here


Code after:
            )
        )

    elif isinstance(value, enum.Enum):
        return lit(value.value, dtype=dtype)

    if dtype:
        return wrap_expr(plr.lit(value, allow_object, is_scalar=True)).cast(dtype)

    if _check_for_numpy(value) and isinstance(value, np.generic):
        # note: the item() is a py-native datetime/timedelta when units < 'ns'
        if isinstance(item := value.item(), (datetime, timedelta)):
            return lit(item)

        # handle 'ns' units
        if isinstance(item, int) and hasattr(value, "dtype"):
            dtype_name = value.dtype.name
            if dtype_name.startswith("datetime64["):
                time_unit = dtype_name[len("datetime64[") : -1]  # type: ignore[assignment]
                return lit(item).cast(Datetime(time_unit))
            if dtype_name.startswith("timedelta64["):
                time_unit = dtype_name[len("timedelta64[") : -1]  # type: ignore[assignment]
                return lit(item).cast(Duration(time_unit))

        # handle numeric values
        if isinstance(value, np.generic):
            dtype = DataTypeMappings.NUMPY_KIND_AND_ITEMSIZE_TO_DTYPE.get(
                (value.dtype.kind, value.dtype.itemsize)
            )
            if dtype is not None:
                return lit(value, dtype=dtype)
    else:
        item = value

    return wrap_expr(plr.lit(item, allow_object, is_scalar=True))
Repository: pola-rs/polars
File: py-polars/polars/__init__.py (Python)

Code before:
#  ____       _
# |  _ \ ___ | | __ _ _ __ ___
# | |_) / _ \| |/ _` | '__/ __|
# |  __/ (_) | | (_| | |  \__ \
# |_|   \___/|_|\__,_|_|  |___/
#

"""
Polars: Blazingly fast DataFrames in Python
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Polars is a fast, open-source library for data manipulation with an expressive, typed API.

Basic usage:

   >>> import polars as pl
   >>> df = pl.DataFrame({
   ...     "name": ["Alice", "Bob", "Charlie"],
   ...     "age": [25, 30, 35],
   ...     "city": ["New York", "London", "Tokyo"]
   ... })
   >>> df.filter(pl.col("age") > 28)
   shape: (2, 3)
   ┌─────────┬─────┬────────┐
   │ name    ┆ age ┆ city   │
   │ ---     ┆ --- ┆ ---    │
   │ str     ┆ i64 ┆ str    │
   ╞═════════╪═════╪════════╡
   │ Bob     ┆ 30  ┆ London │
   │ Charlie ┆ 35  ┆ Tokyo  │
   └─────────┴─────┴────────┘

User Guide: https://docs.pola.rs/
Python API Documentation: https://docs.pola.rs/api/python/stable/reference/

Review comment (suggestion):
I think the user is aware that they're using Python, and the bulk is actually implemented in Rust. So I'd just leave off the 'in Python'.

Code after:
"""
Polars: Blazingly fast DataFrames
=================================

Polars is a fast, open-source library for data manipulation with an expressive, typed API.

Basic usage:

   >>> import polars as pl
   >>> df = pl.DataFrame(
   ...     {
   ...         "name": ["Alice", "Bob", "Charlie"],
   ...         "age": [25, 30, 35],
   ...         "city": ["New York", "London", "Tokyo"],
   ...     }
   ... )
   >>> df.filter(pl.col("age") > 28)
   shape: (2, 3)
   ┌─────────┬─────┬────────┐
   │ name    ┆ age ┆ city   │
   │ ---     ┆ --- ┆ ---    │
   │ str     ┆ i64 ┆ str    │
   ╞═════════╪═════╪════════╡
   │ Bob     ┆ 30  ┆ London │
   │ Charlie ┆ 35  ┆ Tokyo  │
   └─────────┴─────┴────────┘

User Guide: https://docs.pola.rs/
Python API Documentation: https://docs.pola.rs/api/python/stable/
Source Code: https://github.com/pola-rs/polars
"""  # noqa: D400, W505, D205

import contextlib
Repository: pola-rs/polars
File: py-polars/tests/unit/functions/as_datatype/test_struct.py (Python)

Code before:
    ).alias("OUTER_STRUCT")

    agg_df = df.group_by("ID").agg(expr_outer_struct)

    assert agg_df.collect_schema() == agg_df.collect().schema


agg_expressions = [
    pl.lit("blah"),  # LiteralScalar
    pl.col("s"),  # NotAggregated
    pl.arange(pl.len()),  # AggregatedList
    pl.col("s").first(),  # AggregatedScalar
]


@pytest.mark.parametrize("lhs", agg_expressions)
@pytest.mark.parametrize("rhs", agg_expressions)
@pytest.mark.parametrize("n_rows", [0, 1, 2, 3])
def test_struct_schema_in_group_by_apply_expr_24168(
    lhs: pl.Expr, rhs: pl.Expr, n_rows: int
) -> None:
    df = pl.DataFrame({"g": [10, 10, 20], "s": ["foo11", "foo12", "foo21"]})
    lf = df.head(n_rows).lazy()
    expr = pl.struct(lhs.alias("lhs"), rhs.alias("rhs")).alias("expr")
    q = lf.group_by("g").agg(expr)
    assert q.collect_schema() == q.collect().schema


Review comment (suggestion):
Can you:
1. also parametrize over `maintain_order` (using the `check_row_order=maintain_order` on `assert_frame_equal`)
2. check the output with `assert_frame_equal`

Code after:
    ).alias("OUTER_STRUCT")

    agg_df = df.group_by("ID").agg(expr_outer_struct)

    assert agg_df.collect_schema() == agg_df.collect().schema


# parametric tuples: (expr, is_scalar)
agg_expressions = [
    (pl.lit(1, pl.Int64), True),  # LiteralScalar
    (pl.col("n"), False),  # NotAggregated
    (pl.int_range(1, pl.len() + 1), False),  # AggregatedList
    (pl.col("n").first(), True),  # AggregatedScalar
]


@pytest.mark.parametrize("lhs", agg_expressions)
@pytest.mark.parametrize("rhs", agg_expressions)
@pytest.mark.parametrize("n_rows", [0, 1, 2, 3])
@pytest.mark.parametrize("maintain_order", [True, False])
def test_struct_schema_in_group_by_apply_expr_24168(
    lhs: tuple[pl.Expr, bool],
    rhs: tuple[pl.Expr, bool],
    n_rows: int,
    maintain_order: bool,
) -> None:
    df = pl.DataFrame({"g": [10, 10, 20], "n": [1, 2, 3]})
    lf = df.head(n_rows).lazy()
    expr = pl.struct(lhs[0].alias("lhs"), rhs[0].alias("rhs")).alias("expr")
    q = lf.group_by("g", maintain_order=maintain_order).agg(expr)
    out = q.collect()

    # check schema
    assert q.collect_schema() == out.schema

    # check output against ground truth (single group only)
    if n_rows in [0, 1]:
        if lhs[1] and rhs[1]:
            expected = pl.DataFrame({"g": [10], "expr": [{"lhs": 1, "rhs": 1}]})
            expected = expected.head(n_rows)
            assert_frame_equal(out, expected, check_row_order=maintain_order)
        else:
            expected = pl.DataFrame({"g": [10], "expr": [[{"lhs": 1, "rhs": 1}]]})
            expected = expected.head(n_rows)
            assert_frame_equal(out, expected, check_row_order=maintain_order)

    if n_rows == 2:
        if lhs[1] and rhs[1]:
            expected = pl.DataFrame({"g": [10], "expr": [{"lhs": 1, "rhs": 1}]})
            assert_frame_equal(out, expected, check_row_order=maintain_order)
        else:
Repository: pola-rs/polars
File: crates/polars-core/src/series/implementations/decimal.rs (Rust)

Code before:

    fn median(&self) -> Option<f64> {
        self.0
            .physical()
            .median()
            .map(|v| v / self.scale_factor() as f64)
    }
    fn median_reduce(&self) -> PolarsResult<Scalar> {
        Ok(self.apply_scale(self.0.physical().median_reduce()))
    }

    fn std(&self, ddof: u8) -> Option<f64> {
        self.0
            .physical()
            .std(ddof)
            .map(|v| v / self.scale_factor() as f64)
    }
    fn std_reduce(&self, ddof: u8) -> PolarsResult<Scalar> {
        Ok(self.apply_scale(self.0.physical().std_reduce(ddof)))
    }

    fn var(&self, ddof: u8) -> Option<f64> {
        self.0
            .physical()
            .var(ddof)
            .map(|v| v / (self.scale_factor() as f64 * self.scale_factor() as f64))
    }
    fn var_reduce(&self, _ddof: u8) -> PolarsResult<Scalar> {
        Ok(self.apply_scale(self.0.physical().var_reduce(_ddof)))
    }

    fn quantile_reduce(&self, quantile: f64, method: QuantileMethod) -> PolarsResult<Scalar> {
        self.0
            .physical()
            .quantile_reduce(quantile, method)
            .map(|v| self.apply_scale(v))
    }

    fn find_validity_mismatch(&self, other: &Series, idxs: &mut Vec<IdxSize>) {
        self.0.physical().find_validity_mismatch(other, idxs)
    }

    fn as_any(&self) -> &dyn Any {
        &self.0
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        &mut self.0
    }

    fn as_phys_any(&self) -> &dyn Any {

Review comment (suggestion):
Please cast to `f64` first and then use that var implementation.

Code after:
            };
            Scalar::new(self.dtype().clone(), av)
        }))
    }

    fn _sum_as_f64(&self) -> f64 {
        self.0.physical()._sum_as_f64() / self.scale_factor() as f64
    }

    fn mean(&self) -> Option<f64> {
        self.0
            .physical()
            .mean()
            .map(|v| v / self.scale_factor() as f64)
    }
    fn mean_reduce(&self) -> PolarsResult<Scalar> {
        Ok(Scalar::new(DataType::Float64, self.mean().into()))
    }

    fn median(&self) -> Option<f64> {
        self.0
            .physical()
            .median()
            .map(|v| v / self.scale_factor() as f64)
    }

    fn median_reduce(&self) -> PolarsResult<Scalar> {
        Ok(self.apply_scale(self.0.physical().median_reduce()))
    }

    fn std(&self, ddof: u8) -> Option<f64> {
        self.0.cast(&DataType::Float64).ok()?.std(ddof)
    }

    fn std_reduce(&self, ddof: u8) -> PolarsResult<Scalar> {
        self.0.cast(&DataType::Float64)?.std_reduce(ddof)
    }

    fn var(&self, ddof: u8) -> Option<f64> {
        self.0.cast(&DataType::Float64).ok()?.var(ddof)
    }

    fn var_reduce(&self, ddof: u8) -> PolarsResult<Scalar> {
        self.0.cast(&DataType::Float64)?.var_reduce(ddof)
    }

    fn quantile_reduce(&self, quantile: f64, method: QuantileMethod) -> PolarsResult<Scalar> {
        self.0
            .physical()
            .quantile_reduce(quantile, method)
            .map(|v| self.apply_scale(v))
Repository: pola-rs/polars
File: py-polars/polars/_plr.py (Python)

Code before:
            pl64()
        elif _FORCE == "32":
            pl32()
        else:
            msg = f"Invalid value for `POLARS_FORCE_PKG` variable: '{_FORCE}'"
            raise ValueError(msg)
    else:
        if _PREFER == "64":
            try:
                pl64()
            except ImportError as _:
                try:
                    pllts()
                except ImportError as _:
                    pl32()
        elif _PREFER == "32":
            try:
                pl32()
            except ImportError as _:
                try:
                    pllts()
                except ImportError as _:
                    pl64()
        else:
            msg = f"Invalid value for `POLARS_PREFER_PKG` variable: '{_PREFER}'"
            raise ValueError(msg)


Review comment (suggestion):
I think this can be a bit neater:

```python
pkgs = {
    "lts": pllts,
    "64": pl64,
    "32": pl32
}

if _force is not None:
    try:
        pkgs[_force]()
    except KeyError:
        msg = f"Invalid value for `POLARS_FORCE_PKG` variable: '{_force}'"
        raise ValueError(msg)
else:
    preference = [pllts, pl64, pl32]
    if _prefer is not None:
        try:
            preference.insert(0, pkgs[_prefer])
        except KeyError:
            msg = f"Invalid value for `POLARS_PREFER_PKG` variable: '{_prefer}'"
            raise ValueError(msg)

    for pkg in preference:
        try:
            pkg()
            break
        except ImportError:
            pass
```

Code after:
            msg = f"Invalid value for `POLARS_FORCE_PKG` variable: '{_force}'"
            raise ValueError(msg) from None
    else:
        preference = [pllts, pl64, pl32]
        if _prefer is not None:
            try:
                preference.insert(0, pkgs[_prefer])
            except KeyError:
                msg = f"Invalid value for `POLARS_PREFER_PKG` variable: '{_prefer}'"
                raise ValueError(msg) from None

        version_warnings = []
        for pkg in preference:
            try:
                pkg()

                if sys.modules[__name__].__version__ != PKG_VERSION:
                    import warnings

                    version_warnings += [sys.modules[__name__].__version__]
                    warnings.warn(
                        f"Skipping Polars' Rust module version '{sys.modules[__name__].__version__}' did not match version of Python package '{PKG_VERSION}'.",
                        ImportWarning,
                        stacklevel=2,
                    )
                    continue

                break
            except ImportError:
                pass
        else:
            msg = "could not find Polars' Rust module"
            if len(version_warnings) > 0:
                msg += f". Skipped versions {version_warnings} which don't match Python package version"
            raise ImportError(msg)


# The version at the top here should match the version specified by the PLR.
assert sys.modules[__name__].__version__ == PKG_VERSION
Repository: pola-rs/polars
File: py-polars/polars/_plr.py (Python)

Code before:
    _force = os.environ.get("POLARS_FORCE_PKG")
    _prefer = os.environ.get("POLARS_PREFER_PKG")

    pkgs = {"lts": pllts, "64": pl64, "32": pl32}

    if _force is not None:
        try:
            pkgs[_force]()
        except KeyError:
            msg = f"Invalid value for `POLARS_FORCE_PKG` variable: '{_force}'"
            raise ValueError(msg) from None
    else:
        preference = [pllts, pl64, pl32]
        if _prefer is not None:
            try:
                preference.insert(0, pkgs[_prefer])
            except KeyError:
                msg = f"Invalid value for `POLARS_PREFER_PKG` variable: '{_prefer}'"
                raise ValueError(msg) from None

        for i, pkg in enumerate(preference):
            try:
                pkg()
                break
            except ImportError:
                if i == len(preference) - 1:
                    raise


Review comment (suggestion):
Python's `for` loops support an `else` which is entered if `break` is never called:

```python
for pkg in preference:
    try:
        pkg()
        break
    except ImportError:
        pass
else:
    raise RuntimeError("could not find Polars' Rust module")
```

Code after:
    # the PLR from.

    _force = os.environ.get("POLARS_FORCE_PKG")
    _prefer = os.environ.get("POLARS_PREFER_PKG")

    pkgs = {"lts": pllts, "64": pl64, "32": pl32}

    if _force is not None:
        try:
            pkgs[_force]()

            if sys.modules[__name__].__version__ != PKG_VERSION:
                msg = f"Polars Rust module for '{_force}' ({sys.modules[__name__].__version__}) did not match version of Python package '{PKG_VERSION}'"
                raise ImportError(msg)
        except KeyError:
            msg = f"Invalid value for `POLARS_FORCE_PKG` variable: '{_force}'"
            raise ValueError(msg) from None
    else:
        preference = [pllts, pl64, pl32]
        if _prefer is not None:
            try:
                preference.insert(0, pkgs[_prefer])
            except KeyError:
                msg = f"Invalid value for `POLARS_PREFER_PKG` variable: '{_prefer}'"
                raise ValueError(msg) from None

        version_warnings = []
        for pkg in preference:
            try:
                pkg()

                if sys.modules[__name__].__version__ != PKG_VERSION:
                    import warnings

                    version_warnings += [sys.modules[__name__].__version__]
                    warnings.warn(
                        f"Skipping Polars' Rust module version '{sys.modules[__name__].__version__}' did not match version of Python package '{PKG_VERSION}'.",
                        ImportWarning,
                        stacklevel=2,
                    )
                    continue

                break
            except ImportError:
                pass
        else:
            msg = "could not find Polars' Rust module"
            if len(version_warnings) > 0:
                msg += f". Skipped versions {version_warnings} which don't match Python package version"
            raise ImportError(msg)
Repository: pola-rs/polars
File: crates/polars-utils/src/plpath.rs (Rust)

Code before:
    }
}

pub struct AddressDisplay<'a> {
    addr: PlPathRef<'a>,
}

impl<'a> fmt::Display for AddressDisplay<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.addr {
            PlPathRef::Local(p) => p.display().fmt(f),
            PlPathRef::Cloud(p) => p.fmt(f),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
pub enum CloudScheme {
    Abfs,
    Abfss,
    Adl,
    Az,
    Azure,
    File { has_hostname: bool },
    Gcs,
    Gs,
    Hf,
    Http,
    Https,
    S3,
    S3a,
}

impl CloudScheme {
    /// Note, private function. Users should use [`CloudScheme::from_uri`], that will handle e.g.
    /// `file:/` without hostname properly.
    fn from_scheme(s: &str) -> Option<Self> {
        use CloudScheme::*;

        Some(match s {
            "abfs" => Abfs,
            "abfss" => Abfss,
            "adl" => Adl,
            "az" => Az,
            "azure" => Azure,
            "file" => File { has_hostname: true },
            "gcs" => Gcs,
            "gs" => Gs,
            "hf" => Hf,

Review comment (refactor):
I think the smaller refactor here would be to add a `FileStripped` / `FileSpark` (or named otherwise) variant.

Code after:
    }

    pub fn strip_scheme(&self) -> &str {
        self.scheme.strip_scheme_from_uri(self.uri)
    }
}

pub struct AddressDisplay<'a> {
    addr: PlPathRef<'a>,
}

impl<'a> fmt::Display for AddressDisplay<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.addr {
            PlPathRef::Local(p) => p.display().fmt(f),
            PlPathRef::Cloud(p) => p.fmt(f),
        }
    }
}

macro_rules! impl_cloud_scheme {
    ($($t:ident = $n:literal,)+) => {
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
        #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
        #[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
        pub enum CloudScheme {
            $($t,)+
        }

        impl CloudScheme {
             /// Note, private function. Users should use [`CloudScheme::from_uri`], that will handle e.g.
            /// `file:/` without hostname properly.
            #[expect(unreachable_patterns)]
            fn from_scheme(s: &str) -> Option<Self> {
                Some(match s {
                    $($n => Self::$t,)+
                    _ => return None,
                })
            }

            const fn as_str(&self) -> &'static str {
                match self {
                    $(Self::$t => $n,)+
                }
            }
        }
    };
}

impl_cloud_scheme! {
    Abfs = "abfs",
Repository: pola-rs/polars
File: crates/polars-plan/src/plans/aexpr/function_expr/sign.rs (Rust)

Code before:
use num_traits::{One, Zero};
use polars_core::with_match_physical_numeric_polars_type;

use super::*;

pub(super) fn sign(s: &Column) -> PolarsResult<Column> {
    let s = s.as_materialized_series();
    let dtype = s.dtype();
    match dtype {
        _ if dtype.is_primitive_numeric() => with_match_physical_numeric_polars_type!(dtype, |$T| {
            let ca: &ChunkedArray<$T> = s.as_ref().as_ref();
            Ok(sign_impl(ca))
        }),
        DataType::Decimal(_, scale) => {
            let ca = s.decimal()?;

            let p_one = 10i128.pow(*scale as u32);
            let n_one = -p_one;

            let out = ca
                .physical()
                .apply_values(|x| {
                    if x < 0 {
                        n_one
                    } else if x > 0 {
                        p_one
                    } else {
                        0
                    }
                })
                .into_column();
            unsafe { out.from_physical_unchecked(dtype) }
        },
        _ => polars_bail!(opq = sign, dtype),
    }
}

fn sign_impl<T>(ca: &ChunkedArray<T>) -> Column
where
    T: PolarsNumericType,
    ChunkedArray<T>: IntoColumn,

Review comment (suggestion):
Can you add this implementation to `polars-compute/src/decimal.rs` and call that instead of doing it here?

```rust
#[inline(always)]
pub fn dec128_sign(x: i128, s: usize) -> i128 {
    if x < 0 {
        -POW10_I128[s]
    } else if x > 0 {
        POW10_I128[s]
    } else {
        0
    }
}
```

I prefer it if we keep all arithmetic implementations centralized there.

Code after:
use num_traits::{One, Zero};
use polars_core::with_match_physical_numeric_polars_type;

use super::*;

pub(super) fn sign(s: &Column) -> PolarsResult<Column> {
    let s = s.as_materialized_series();
    let dtype = s.dtype();
    match dtype {
        _ if dtype.is_primitive_numeric() => with_match_physical_numeric_polars_type!(dtype, |$T| {
            let ca: &ChunkedArray<$T> = s.as_ref().as_ref();
            Ok(sign_impl(ca))
        }),
        DataType::Decimal(_, scale) => {
            let ca = s.decimal()?;
            let out = ca
                .physical()
                .apply_values(|x| polars_compute::decimal::dec128_sign(x, *scale))
                .into_column();
            unsafe { out.from_physical_unchecked(dtype) }
        },
        _ => polars_bail!(opq = sign, dtype),
    }
}

fn sign_impl<T>(ca: &ChunkedArray<T>) -> Column
where
    T: PolarsNumericType,
    ChunkedArray<T>: IntoColumn,
{
    ca.apply_values(|x| {
        if x < T::Native::zero() {
            T::Native::zero() - T::Native::one()
        } else if x > T::Native::zero() {
            T::Native::one()
        } else {
            // Returning x here ensures we return NaN for NaN input, and
            // maintain the sign for signed zeroes (although we don't really
            // care about the latter).
            x
        }