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

Code before:
        let inputs_ordered = orders.get_mut(&node).unwrap();

        macro_rules! set_unordered_output {
            () => {
                for (output, edge) in node_outputs {
                    orders.get_mut(output).unwrap()[*edge] = false;
                }
            };
        }

        // Pullup simplification rules.
        use MaintainOrderJoin as MOJ;
        match ir {
            IR::Sort { sort_options, .. } => {
                // Unordered -> _     ==>    maintain_order=false
                sort_options.maintain_order &= inputs_ordered[0];
            },
            IR::GroupBy {
                keys: _,
                maintain_order,
                ..
            } => {
                // If maintain_order is false, the outputs are considered unordered.
                // If maintain_order is true, we keep it – this was explicitly requested
                // (e.g. group_by_stable / with_maintain_order(true)), even if the input
                // itself is not guaranteed to be ordered.
                if !*maintain_order {
                    set_unordered_output!();
                }
            },
            IR::Sink { input: _, payload } => {
                if !inputs_ordered[0] {
                    // Set maintain order to false if input is unordered
                    match payload {
                        SinkTypeIR::Memory => {},
                        SinkTypeIR::File(FileSinkOptions {
                            unified_sink_args, ..
                        })
                        | SinkTypeIR::Partitioned(PartitionedSinkOptionsIR {
                            unified_sink_args,
                            ..
                        }) => unified_sink_args.maintain_order = false,
                        SinkTypeIR::Callback(s) => s.maintain_order = false,
                    }
                }
            },
            IR::Join { options, .. } => {
                let left_unordered = !inputs_ordered[0];
                let right_unordered = !inputs_ordered[1];

                let maintain_order = options.args.maintain_order;

Review comment (bug):
This makes no sense to me, why would we be maintaining the order of unordered input? This is removing a useful optimization (since order-maintaining group-by's are slower).

Why is the input unordered in the first place in the bug example? I believe the problem is that (for whatever reason) the `asof_join` output is marked as unordered when it shouldn't be.

Code after:
        let mut ir = ir_arena.get_mut(node);

        let inputs_ordered = orders.get_mut(&node).unwrap();

        macro_rules! set_unordered_output {
            () => {
                for (output, edge) in node_outputs {
                    orders.get_mut(output).unwrap()[*edge] = false;
                }
            };
        }

        // Pullup simplification rules.
        use MaintainOrderJoin as MOJ;
        match ir {
            IR::Sort { sort_options, .. } => {
                // Unordered -> _     ==>    maintain_order=false
                sort_options.maintain_order &= inputs_ordered[0];
            },
            IR::GroupBy {
                keys,
                maintain_order,
                ..
            } => {
                if !inputs_ordered[0] && *maintain_order {
                    // Unordered -> _
                    //   to
                    // maintain_order = false
                    // and
                    // Unordered -> Unordered

                    let keys_produce_order = keys
                        .iter()
                        .any(|k| is_output_ordered(expr_arena.get(k.node()), expr_arena, false));
                    if !keys_produce_order {
                        *maintain_order = false;
                    }
                }
                if !*maintain_order {
                    set_unordered_output!();
                }
            },
            IR::Sink { input: _, payload } => {
                if !inputs_ordered[0] {
                    // Set maintain order to false if input is unordered
                    match payload {
                        SinkTypeIR::Memory => {},
                        SinkTypeIR::File(FileSinkOptions {
                            unified_sink_args, ..
                        })
                        | SinkTypeIR::Partitioned(PartitionedSinkOptionsIR {