Repository: pola-rs/polars
File: crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs (Rust)

Code before:
                    aggs,
                    schema,
                    apply,
                    maintain_order,
                    options
                })
            }
            (Distinct {input, mut options}, Some(state)) => {
                // first restart optimization in inputs and get the updated LP
                let input_lp = lp_arena.take(input);
                let input_lp = self.pushdown(input_lp, None, lp_arena, expr_arena)?;
                let input= lp_arena.add(input_lp);
                options.slice = Some((state.offset, state.len as usize));
                Ok(Distinct {
                    input,
                    options,
                })
            }
            (Sort {input, by_column, slice: _,
                sort_options}, Some(state)) => {
                // first restart optimization in inputs and get the updated LP
                let input_lp = lp_arena.take(input);
                let input_lp = self.pushdown(input_lp, None, lp_arena, expr_arena)?;
                let input= lp_arena.add(input_lp);

                let slice = Some((state.offset, state.len as usize));
                Ok(Sort {
                    input,
                    by_column,
                    slice,
                    sort_options
                })
            }
            (Slice {
                input,
                offset,
                mut len
            }, Some(outer_slice)) => {
                let alp = lp_arena.take(input);

                // Both are positive, can combine into a single slice.
                if outer_slice.offset >= 0 && offset >= 0 {
                    let state = State {
                        offset: offset.checked_add(outer_slice.offset).unwrap(),
                        len: if len as i128 > outer_slice.offset as i128 {
                            (len - outer_slice.offset as IdxSize).min(outer_slice.len)
                        } else {
                            0
                        },
                    };
                    return self.pushdown(alp, Some(state), lp_arena, expr_arena);

Review comment (suggestion):
I think this should be another PR. As there are more mutable changes here, I cannot tell if we rely on this and if we change it, we should do a full pass.

Code after:
                    aggs,
                    schema,
                    apply,
                    maintain_order,
                    options
                })
            }
            (Distinct {input, mut options}, Some(state)) => {
                // first restart optimization in inputs and get the updated LP
                let input_lp = lp_arena.take(input);
                let input_lp = self.pushdown(input_lp, None, lp_arena, expr_arena)?;
                let input= lp_arena.add(input_lp);
                options.slice = Some((state.offset, state.len as usize));
                Ok(Distinct {
                    input,
                    options,
                })
            }
            (Sort {input, by_column, mut slice,
                sort_options}, Some(state)) => {
                // first restart optimization in inputs and get the updated LP
                let input_lp = lp_arena.take(input);
                let input_lp = self.pushdown(input_lp, None, lp_arena, expr_arena)?;
                let input= lp_arena.add(input_lp);

                slice = Some((state.offset, state.len as usize));
                Ok(Sort {
                    input,
                    by_column,
                    slice,
                    sort_options
                })
            }
            (Slice {
                input,
                offset,
                mut len
            }, Some(outer_slice)) => {
                let alp = lp_arena.take(input);

                // Both are positive, can combine into a single slice.
                if outer_slice.offset >= 0 && offset >= 0 {
                    let state = State {
                        offset: offset.checked_add(outer_slice.offset).unwrap(),
                        len: if len as i128 > outer_slice.offset as i128 {
                            (len - outer_slice.offset as IdxSize).min(outer_slice.len)
                        } else {
                            0
                        },
                    };
                    return self.pushdown(alp, Some(state), lp_arena, expr_arena);