Repository: pola-rs/polars
File: crates/polars-stream/src/physical_plan/lower_expr.rs (Rust)

Code before:

                // Transform:
                //    expr.diff(offset, "ignore")
                //      ->
                //    expr - expr.shift(offset)

                let base_expr_ir = &inner_exprs[0];
                let base_dtype =
                    base_expr_ir.dtype(&ctx.phys_sm[input.node].output_schema, ctx.expr_arena)?;
                let offset_expr_ir = &inner_exprs[1];
                let offset_dtype =
                    offset_expr_ir.dtype(&ctx.phys_sm[input.node].output_schema, ctx.expr_arena)?;

                let cast_dtype = |dt: &DataType| match *dt {
                    DataType::UInt8 => Some(DataType::Int16),
                    DataType::UInt16 => Some(DataType::Int32),
                    DataType::UInt32 | DataType::UInt64 => Some(DataType::Int64),
                    _ => None,
                };

                let mut base = AExprBuilder::new_from_node(base_expr_ir.node());
                if let Some(cast_dtype) = cast_dtype(base_dtype) {
                    base = base.cast(cast_dtype, ctx.expr_arena);
                }
                let mut offset = AExprBuilder::new_from_node(offset_expr_ir.node());
                if let Some(cast_dtype) = cast_dtype(offset_dtype) {
                    offset = offset.cast(cast_dtype, ctx.expr_arena);
                }

                let shifted = base.shift(offset.node(), ctx.expr_arena);
                let mut output = base.minus(shifted.node(), ctx.expr_arena);

                if null_behavior == NullBehavior::Drop {
                    // Without the column size, slice can only remove leading nulls.
                    // So if the offset was negative, the nulls appeared and the end of the column.
                    // In that case, shift the column forward to move the nulls back to the front.
                    let zero_literal =
                        AExprBuilder::lit(LiteralValue::new_idxsize(0), ctx.expr_arena);
                    let offset_neg =
                        AExprBuilder::new_from_node(offset.node()).negate(ctx.expr_arena);
                    let offset_if_negative = AExprBuilder::function(
                        vec![offset_neg.expr_ir_unnamed(), zero_literal.expr_ir_unnamed()],
                        IRFunctionExpr::MaxHorizontal,
                        ctx.expr_arena,
                    );
                    output = output.shift(offset_if_negative, ctx.expr_arena);

                    // Remove the nulls that were introduced by the shift
                    let offset_abs = offset.abs(ctx.expr_arena);
                    let null_literal = AExprBuilder::lit(
                        LiteralValue::Scalar(Scalar::null(DataType::Int64)),

Review comment (suggestion):
I don't understand why we're passing the offset through `cast_dtype`. I believe the offset should always be simply cast to Int64.

Code after:

                // Transform:
                //    expr.diff(offset, "ignore")
                //      ->
                //    expr - expr.shift(offset)

                let base_expr_ir = &inner_exprs[0];
                let base_dtype =
                    base_expr_ir.dtype(&ctx.phys_sm[input.node].output_schema, ctx.expr_arena)?;
                let offset_expr_ir = &inner_exprs[1];
                let offset_dtype =
                    offset_expr_ir.dtype(&ctx.phys_sm[input.node].output_schema, ctx.expr_arena)?;

                let mut base = AExprBuilder::new_from_node(base_expr_ir.node());
                let cast_dtype = match base_dtype {
                    DataType::UInt8 => Some(DataType::Int16),
                    DataType::UInt16 => Some(DataType::Int32),
                    DataType::UInt32 | DataType::UInt64 => Some(DataType::Int64),
                    _ => None,
                };
                if let Some(dtype) = cast_dtype {
                    base = base.cast(dtype, ctx.expr_arena);
                }

                let mut offset = AExprBuilder::new_from_node(offset_expr_ir.node());
                if offset_dtype != &DataType::Int64 {
                    offset = offset.cast(DataType::Int64, ctx.expr_arena);
                }

                let shifted = base.shift(offset.node(), ctx.expr_arena);
                let mut output = base.minus(shifted.node(), ctx.expr_arena);

                if null_behavior == NullBehavior::Drop {
                    // Without the column size, slice can only remove leading nulls.
                    // So if the offset was negative, the nulls appeared at the end of the column.
                    // In that case, shift the column forward to move the nulls back to the front.
                    let zero_literal =
                        AExprBuilder::lit(LiteralValue::new_idxsize(0), ctx.expr_arena);
                    let offset_neg = offset.negate(ctx.expr_arena);
                    let offset_if_negative = AExprBuilder::function(
                        vec![offset_neg.expr_ir_unnamed(), zero_literal.expr_ir_unnamed()],
                        IRFunctionExpr::MaxHorizontal,
                        ctx.expr_arena,
                    );
                    output = output.shift(offset_if_negative, ctx.expr_arena);

                    // Remove the nulls that were introduced by the shift
                    let offset_abs = offset.abs(ctx.expr_arena);
                    let null_literal = AExprBuilder::lit(
                        LiteralValue::Scalar(Scalar::null(DataType::Int64)),
                        ctx.expr_arena,