Native broadcast nested loop joins report the streamed side's order without keeping it #6898

Describe the bug

CometBroadcastNestedLoopJoinExec reports the order Spark reports for the join, but the native join does not keep it. Spark then drops sorts, and plans sort-merge joins, around rows that are not in that order. With default settings sortWithinPartitions returns unsorted partitions, and with AQE off a sort-merge join above the join returns wrong rows.

Spark's BroadcastNestedLoopJoinExec.outputOrdering is the streamed side's ordering for (InnerLike, _), (LeftOuter, BuildRight), (RightOuter, BuildLeft) and (LeftSingle | LeftSemi | LeftAnti, BuildRight) in Spark 4.1, and 3.4 and 3.5 override it as well. CometBroadcastNestedLoopJoinExec.createExec copies op.outputOrdering (operators.scala). DataFusion's NestedLoopJoinExec keeps the order of neither input: maintains_input_order returns [false, false] (datafusion-physical-plan 55.2.0, joins/nested_loop_join.rs:427), and process_probe_batch joins each probe batch with one build row, or one range of build rows, at a time (l_row x right_batch), so each batch's output is grouped by build row.

Spark plans around the reported order before Comet converts the join. RemoveRedundantSorts drops a sort above it, and EnsureRequirements adds no sort below an operator that requires the order. So reporting a different order from the Comet node would not bring those sorts back.

Steps to reproduce

Parquet tables big as (i % 100, i) for 10000 rows, two as 0 and 1, and other as (i % 100, i) for 200 rows.

With default settings:

val cached = spark.table("big").toDF("k", "v").sortWithinPartitions("k").cache()
cached.count()
val df = cached
  .join(broadcast(spark.table("two")), $"v" % 2 <= $"_1")
  .select($"k", $"v", $"_1")
  .sortWithinPartitions("k")
df.queryExecution.toRdd
  .mapPartitions { it =>
    val keys = it.map(_.getInt(0)).toArray
    Iterator(keys.sameElements(keys.sorted))
  }
  .collect()
// Array(false, false, false, false, false), where Spark's partitions are all sorted and the rows match

The plan has no sort above the join:

CometColumnarToRow
+- CometBroadcastNestedLoopJoin Inner, BuildRight, ((v#7 % 2) <= _1#5)
   :- CometFilter [k#6, v#7], isnotnull(v#7)
   :  +- CometSparkColumnarToColumnar
   :     +- TableCacheQueryStage 0
   :        +- InMemoryTableScan [k#6, v#7], [isnotnull(v#7)]
   +- ...

With spark.sql.adaptive.enabled=false, spark.sql.autoBroadcastJoinThreshold=-1 and spark.sql.shuffle.partitions=2, a sort-merge join above the same kind of join returns wrong rows:

val cached = spark.table("big").toDF("k", "v").repartition(2, $"k").sortWithinPartitions("k").cache()
cached.count()
val df = cached
  .join(broadcast(spark.table("two")), $"v" % 2 <= $"_1")
  .select($"k", $"v", $"_1".as("x"))
  .join(spark.table("other").toDF("k2", "w"), $"k" === $"k2")
df.count()
// 10600, where Spark returns 30000

The join's side of the sort-merge join has no sort:

CometSortMergeJoin [k#10], [k2#79], Inner
:- CometProject [k#10, v#11, x#78], [k#10, v#11, _1#5 AS x#78]
:  +- CometBroadcastNestedLoopJoin Inner, BuildRight, ((v#11 % 2) <= _1#5)
:     :- CometFilter [k#10, v#11], (isnotnull(v#11) AND isnotnull(k#10))
:     :  +- CometSparkColumnarToColumnar
:     :     +- InMemoryTableScan [k#10, v#11], [isnotnull(v#11), isnotnull(k#10)]
:     +- CometBroadcastExchange [_1#5]
+- CometSort [k2#79, w#80], [k2#79 ASC NULLS FIRST]
   +- CometExchange hashpartitioning(k2#79, 2), ENSURE_REQUIREMENTS, CometNativeShuffle

A window partitioned by k in place of the second join fails instead, with DataFusion's All partition by columns should have an ordering from WindowAggExec.

Reproduced on main (033cb37) with Spark 4.1.

Expected behavior

The same rows, and the same order within each partition, as Spark.

Additional context

Both cases need an input whose order Spark knows at the join, which here is a cache. The sort has to be inside the cached plan: a sortWithinPartitions below the join in the same query is removed by the optimizer, and Spark then keeps the sort above the join. From Spark 3.5 a cached plan that runs under AQE, such as one with an exchange, reports no order, so the second case needs AQE off there. Spark 3.4 runs cached plans without AQE by default, but Comet's in-memory cache is off by default on 3.4.

A fix could sort the native join's output when Spark reports an order for it, as #6810 does for the null-aware anti join, fall back to Spark for such a join, or keep the probe order in DataFusion's nested loop join.

This is the nested loop join counterpart of #6787. Found while working on #6835, whose order declaration treats a native nested loop join as one that may move its streamed rows.