Support FIRST/LAST with `PartialMerge`

[comphead]

Describe the bug

Found in #4003

Steps to reproduce

test("partialMerge - FIRST/LAST with distinct aggregates falls back") {
    val numValues = 10000
    Seq(100).foreach { numGroups =>
      Seq(128).foreach { batchSize =>
        withSQLConf(
          SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
          CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
          withParquetTable(
            (0 until numValues).map(i => (i, Random.nextInt() % numGroups)),
            "tbl",
            false) {
            withView("v") {
              sql("CREATE TEMP VIEW v AS SELECT _1, _2 FROM tbl ORDER BY _1")
              checkSparkAnswerAndOperator(
                "SELECT _2, FIRST(_1), LAST(_1), COUNT(DISTINCT _1)" +
                  " FROM v GROUP BY _2 ORDER BY 1, 2, 3, 4")
            }
          }
        }
      }
    }
  }

== Results ==
!== Spark Answer - 199 ==                                              == Comet Answer - 199 ==
 struct<_2:int,first(_1):int,last(_1):int,count(DISTINCT _1):bigint>   struct<_2:int,first(_1):int,last(_1):int,count(DISTINCT _1):bigint>
![-99,169,9603,61]                                                     [-99,169,9804,61]
![-98,307,9492,48]                                                     [-98,307,9870,48]
![-97,153,9197,47]                                                     [-97,153,9175,47]
![-96,571,9486,56]                                                     [-96,571,9762,56]
![-95,243,9720,50]                                                     [-95,243,9642,50]
 [-94,234,9747,45]                                                     [-94,234,9747,45]
![-93,637,9020,45]                                                     [-93,637,9515,45]
 [-92,13,9754,51]                                                      [-92,13,9754,51]
![-91,75,9643,47]                                                      [-91,75,9872,47]
![-90,106,9971,50]                                                     [-90,106,9911,50]
 [-89,323,9937,49]                                                     [-89,323,9937,49]
![-88,49,9577,43]                                                      [-88,49,9756,43]
![-87,134,9960,45]                                                     [-87,134,9998,45]
![-86,12,9180,47]                                                      [-86,12,9607,47]
![-85,220,9565,58]                                                     [-85,220,9533,58]
![-84,815,9734,48]                                                     [-84,815,9919,48]

Expected behavior

Should return correct values

Additional context

No response

[Geethapranay1]

hi @comphead ,
I looked through the aggregate conversion path, the risky case is the mixed {Partial, PartialMerge} HashAggregate generated by distinct aggregates, where FIRST/LAST can enter the PartialMerge path and depend on merge/hash-table ordering.

if its still available i'll work on a targeted fallback in CometBaseAggregate.doConvert for FIRST/LAST only when their aggregate expression is PartialMerge, plus a regression test using the distinct-aggregate expand pattern. This should preserve the new PartialMerge support for safe aggregates while fixing the correctness issue.

[comphead]

Thanks @Geethapranay1 Lets wait #4003 to be merged

[comphead]

HI @Geethapranay1 the PartialMerge PR is merged, so if you have any capacity feel free to take this one

[Geethapranay1]

Hi @comphead

I checked the native path after your pr #4003 the fallback in operators.scala prevents incorrect results.

The root cause comes from the native side that is comet builds FirstValue and LastValue without ordering at planner.rs, datafusion defaults to trivial accumulators and their merge behavior relies on position under MergeAsPartial.

Is your goal to remove the fallback and support FIRST and LAST natively in PartialMerge? I will look into the implementation if you approve

[comphead]

yeah, fallback is to prevent users to get a corrupted data.
We need to investigate the reason for correctness issues and eventually try to find a fix, I assume this happens because of different hashing between Spark and DF and if it is, we prob need to modify DF to be able to accept different hasher, for example using cargo features

[Geethapranay1]

i don’t think this is only a Spark vs DF hashing issue.
For unordered FIRST/LAST, the intermediate state is just [value, is_set] once the partial hash aggregate emits that state, the original row order is lost. After that, PartialMerge becomes dependent on hash-table emission order, so changing the hasher alone would not make it correct

I added a native-side guard to reject FIRST/LAST in PartialMerge before MergeAsPartial, and a unit test for that planner path. This keeps the existing fallback and makes the protection explicit on the native side as well
can i open pr for this?

A full native fix would need order-preserving merge semantics or richer intermediate state.

[andygrove]

@comphead can this be closed? It looks like a regression test already exists and current behavior is correct? Or maybe we can remove bug label if there is more work planned as an enhancement to remove the fallback?