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

Code before:
                    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);
                }

                if dims[1..]
                    .iter()
                    .any(|d| matches!(d, ReshapeDimension::Infer))
                {
                    polars_bail!(InvalidOperation: "can only infer the first dimension");
                }

                let num_infers = dims
                    .iter()
                    .filter(|d| matches!(d, ReshapeDimension::Infer))
                    .count();

                let mut inferred_size = 0;
                if num_infers == 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();

                // We pop the outer dimension as that is the height of the series.

Review comment (suggestion):
I think this can now be `usize::from(matches!(dims[0], ReshapeDimension::Infer))`. But, in general, if we can just have a boolean here now saying `infers_dimension` as there is now only 1 dimension to infer.

Code after:
                        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)
            },
            #[cfg(feature = "rle")]
            RLE => mapper.map_dtype(|dt| {
                DataType::Struct(vec![
                    Field::new(PlSmallStr::from_static("len"), IDX_DTYPE),
                    Field::new(PlSmallStr::from_static("value"), dt.clone()),
                ])
            }),
            #[cfg(feature = "rle")]
            RLEID => mapper.with_dtype(IDX_DTYPE),
            ToPhysical => mapper.to_physical_type(),
            #[cfg(feature = "random")]
            Random { .. } => mapper.with_same_dtype(),
            SetSortedFlag(_) => mapper.with_same_dtype(),
            #[cfg(feature = "ffi_plugin")]