[Regression] Query using ARRAY_AGG(DISTINCT) causes panic

[bellwether-softworks]

Describe the bug

Beginning in v37.0.0, a previously-working query is found to result in a panic:

panicked at /Users/username/.cargo/registry/src/index.crates.io-6f17d22bba15001f/datafusion-physical-expr-37.1.0/src/aggregate/array_agg_distinct.rs:158:99:
:
assertion `left == right` failed: state array should only include 1 row!
  left: 4
 right: 1

To Reproduce

The following query is known to trigger the panic, alongside the accompanying parquet payloads contained in the attached .zip file.

SELECT
        asm.floor,
        CONCAT(bi.name, '_', bi.asm_range_lower, ' - ', bi.asm_range_upper) AS name_with_range,
        ARRAY_AGG(DISTINCT s.code) AS codes,
        COUNT(DISTINCT asm.id) AS assembly_count,
        COUNT(DISTINCT s.id) AS item_qty
    FROM 'batch_items.parquet' bi
        INNER JOIN 'target_items.parquet' s
            ON s.id = bi.target_id
        INNER JOIN 'assemblies.parquet' asm
            ON s.assembly_id = asm.id
    GROUP BY
        asm.floor_ordinal,
        asm.floor,
        bi.batch_name,
        bi.name,
        bi.asm_range_lower,
        bi.asm_range_upper
    ORDER BY
        asm.floor_ordinal,
        bi.batch_name,
        bi.name,
        bi.asm_range_lower;

failing-query-assets.zip

Expected behavior

No response

Additional context

I've confirmed that the issue does not present in v36.0.0 or earlier, and is present in v37.0.0 up to v38.0.0. The issue doesn't present when omitting ARRAY_AGG(DISTINCT ...) from my queries.

[jayzhan211]

I added the assertion because I don't know if there is any case that has len > 1.

It would be nice if you had a simpler example to add to the test!

[bellwether-softworks]

@jayzhan211 I appreciate your concern regarding the complex example case; I attempted to create a simpler contrived example, but was unable to trigger the panic doing so. I don't currently know the exact conditions that are triggering the problem.

[jayzhan211]

@jayzhan211 I appreciate your concern regarding the complex example case; I attempted to create a simpler contrived example, but was unable to trigger the panic doing so. I don't currently know the exact conditions that are triggering the problem.

That's fine. I think that is the reason why I can't easily find the test case that trigger this panic, probably only the complex aggregate query that has multiple states meet the requirement.

I think we can add your example as a test case. And, find out whether we should fix the multi states issue in merge_batch or elsewhere before entering merge_batch.

[riosw]

I was playing around with this issue and found a more minimal example:

WITH A AS (
    SELECT
        1 AS id, 1 AS foo
    UNION ALL
    SELECT
        1, 5
)
SELECT
    ARRAY_AGG(DISTINCT a.foo),
    SUM(DISTINCT 1) 
FROM
    A a
GROUP BY
    a.id;

It is interesting though, using SUM(1) instead of SUM(DISTINCT 1) did not panic

[alamb]

Added to #10517