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

Code before:
                //    expr.diff(offset, "drop")
                //      ->
                //    let expr = expr.cast(<signed datatype>) // if needed
                //    let output_len = expr.len() - offset
                //    in expr.slice(offset, output_len) - expr.slice(0, output_len)

                let mut base_expr = inner_exprs[0].clone();
                let offset_expr = &inner_exprs[1];
                let dtype =
                    base_expr.dtype(&ctx.phys_sm[input.node].output_schema, ctx.expr_arena)?;

                // IR: expr = expr.cast(<signed datatype>)
                let cast_dtype = match dtype {
                    DataType::UInt8 => Some(DataType::Int16),
                    DataType::UInt16 => Some(DataType::Int32),
                    DataType::UInt32 | DataType::UInt64 => Some(DataType::Int64),
                    _ => None,
                };
                if let Some(cast_dtype) = cast_dtype {
                    let base_node = AExprBuilder::new_from_node(base_expr.node())
                        .cast(cast_dtype, ctx.expr_arena)
                        .node();
                    base_expr = ExprIR::from_node(base_node, ctx.expr_arena);
                }

                let shifted_expr;
                let unshifted_expr;
                match null_behavior {
                    NullBehavior::Ignore => {
                        // IR: expr.shift(offset, fill_value=None)
                        let unshifted_node = AExprBuilder::function(
                            vec![base_expr.clone(), offset_expr.clone()],
                            IRFunctionExpr::Shift,
                            ctx.expr_arena,
                        )
                        .node();
                        unshifted_expr = ExprIR::from_node(unshifted_node, ctx.expr_arena);
                        shifted_expr = base_expr;
                    },
                    NullBehavior::Drop => {
                        // IR: zero_literal = 0
                        let zero_literal =
                            AExprBuilder::lit(LiteralValue::new_idxsize(0), ctx.expr_arena);

                        // IR: expr.len()
                        let len_node = AExprBuilder::dataframe_length(ctx.expr_arena).node();
                        let len_expr = ExprIR::from_node(len_node, ctx.expr_arena);

                        // IR: output_len = expr.len() - offset
                        let output_len_node = AExprBuilder::new_from_node(len_expr.node())
                            .minus(offset_expr.node(), ctx.expr_arena)

Review comment (suggestion):
I think I would prefer `lhs`/`rhs` over `shifted` and `unshifted`, considering in the `drop_nulls` arguably both are shifted.

Code after:
                //    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,
                    );
                    output = output.slice(offset_abs, null_literal, ctx.expr_arena);
                }