Repository: pola-rs/polars
File: crates/polars-plan/src/plans/aexpr/function_expr/schema.rs (Rust)

Code before:
            } => mapper.with_dtype(DataType::from_categories(Categories::global())),
            #[cfg(feature = "cutqcut")]
            Cut {
                include_breaks: true,
                ..
            } => {
                let struct_dt = DataType::Struct(vec![
                    Field::new(PlSmallStr::from_static("breakpoint"), DataType::Float64),
                    Field::new(
                        PlSmallStr::from_static("category"),
                        DataType::from_categories(Categories::global()),
                    ),
                ]);
                mapper.with_dtype(struct_dt)
            },
            #[cfg(feature = "repeat_by")]
            RepeatBy => mapper.map_dtype(|dt| DataType::List(dt.clone().into())),
            #[cfg(feature = "dtype-array")]
            Reshape(dims) => mapper.try_map_dtype(|dt: &DataType| {
                let dtype = dt.inner_dtype().unwrap_or(dt).clone();

                if dims.len() == 1 {
                    return Ok(dtype);
                }

                let infers = dims.iter().enumerate().filter(|(_, d)| matches!(d, ReshapeDimension::Infer)).collect::<Vec<_>>();

                polars_ensure!(infers.len() <= 1, InvalidOperation: "can only specify one inferred dimension");
                if infers.len() == 1 {
                    polars_ensure!(infers[0].0 == 0, InvalidOperation: "can only infer the first dimension");
                }

                let mut inferred_size = 0;
                if infers.len() == 1 {
                    let mut total_size = 1u64;
                    let mut current = dt;
                    while let DataType::Array(dt, width) = current {
                        if *width == 0 {
                            total_size = 0;
                            break;
                        }

                        current = dt.as_ref();
                        total_size *= *width as u64;
                    }

                    let current_size = dims.iter().map(|d| d.get_or_infer(1)).product::<u64>();
                    inferred_size = total_size / current_size;
                }

                let mut prev_dtype = dtype.leaf_dtype().clone();

Review comment (suggestion):
By definition, if you cannot infer anything beyond the first dimension, you cannot infer more than 1 dimension. I think you can do an `.any()` on `dims[1..]`.

Code after:
            PeakMin | PeakMax => mapper.with_dtype(DataType::Boolean),
            #[cfg(feature = "cutqcut")]
            Cut {
                include_breaks: false,
                ..
            } => mapper.with_dtype(DataType::from_categories(Categories::global())),
            #[cfg(feature = "cutqcut")]
            Cut {
                include_breaks: true,
                ..
            } => {
                let struct_dt = DataType::Struct(vec![
                    Field::new(PlSmallStr::from_static("breakpoint"), DataType::Float64),
                    Field::new(
                        PlSmallStr::from_static("category"),
                        DataType::from_categories(Categories::global()),
                    ),
                ]);
                mapper.with_dtype(struct_dt)
            },
            #[cfg(feature = "repeat_by")]
            RepeatBy => mapper.map_dtype(|dt| DataType::List(dt.clone().into())),
            #[cfg(feature = "dtype-array")]
            Reshape(dims) => mapper.try_map_dtype(|dt: &DataType| {
                let mut wrapped_dtype = dt.leaf_dtype().clone();
                for dim in dims[1..].iter().rev() {
                    let Some(array_size) = dim.get() else {
                        polars_bail!(InvalidOperation: "can only infer the first dimension");
                    };
                    wrapped_dtype = DataType::Array(Box::new(wrapped_dtype), array_size as usize);
                }
                Ok(wrapped_dtype)
            }),
            #[cfg(feature = "cutqcut")]
            QCut {
                include_breaks: false,
                ..
            } => mapper.with_dtype(DataType::from_categories(Categories::global())),
            #[cfg(feature = "cutqcut")]
            QCut {
                include_breaks: true,
                ..
            } => {
                let struct_dt = DataType::Struct(vec![
                    Field::new(PlSmallStr::from_static("breakpoint"), DataType::Float64),
                    Field::new(
                        PlSmallStr::from_static("category"),
                        DataType::from_categories(Categories::global()),
                    ),
                ]);
                mapper.with_dtype(struct_dt)