[NOT A BUG] Why comet does not convert the HashAggregate expression to native in my query?

[SemyonSinchenko]

Describe the bug

I'm running a query that do the following:

1. Read parquet files
2. Generate a lot of case-when columns
3. Run groupBy + agg on top of that columns (sum, min, max, mean)

I have the following logical plan (I manually truncated some parts):

== Parsed Logical Plan ==
'Aggregate ['customer_id], ['customer_id, sum('DC_food-and-household_7d_flag) AS DC_food-and-household_7d_count#18, ..., ..., ... 2057 more fields]
+- Project [customer_id AS customer_id#5422, CASE WHEN (((true AND (t_minus#5L <= cast(7 as bigint))) AND (card_type#1 = DC)) AND (trx_type#2 = food-and-household)) THEN 1 ELSE 0 END AS DC_food-and-household_7d_flag#14, ..., ... 1225 more fields]
   +- Relation [customer_id#0L,card_type#1,trx_type#2,channel#3,trx_amnt#4,t_minus#5L,part_col#6] parquet

It is converted to the following Comet plan:

== Physical Plan ==
HashAggregate(keys=[customer_id#10003], functions=[sum(DC_food-and-household_7d_flag#14), ..., ... 2057 more fields])
+- Exchange hashpartitioning(customer_id#10003, 11), ENSURE_REQUIREMENTS, [plan_id=35]
   +- ColumnarToRow
      +- CometHashAggregate [DC_food-and-household_7d_flag#14, DC_food-and-household_7d_or_none#15, ..., ... 2056 more fields]
         +- CometProject [DC_food-and-household_7d_flag#14, DC_food-and-household_7d_or_none#15, ..., ... 1224 more fields], [CASE WHEN (((t_minus#5L <= 7) AND (card_type#1 = DC)) AND (trx_type#2 = food-and-household)) THEN 1 ELSE 0 END AS DC_food-and-household_7d_flag#14, ..., ... 1224 more fields]
            +- CometScan parquet [...] Batched: true, DataFilters: [], Format: CometParquet, Location: InMemoryFileIndex(1 paths)[file:...], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<...>

Visualization:

Steps to reproduce

I'm running my own benchmark:

1. Generation of the dataset (link to github): generator --prefix test_data_tiny;
2. PySpark code (link to github);
3. Entry point (link to github)

Expected behavior

I expected to see full native plan, but for some reason the last HashAggregate is running on Spark. It looks to me that it is running even in "spark interpreter mode" (I guess because I want too much aggregations and it exceed the limit of the code size for the "Whole stage CodeGet" but I'm not 100% sure).

I checked the documentation of the Comet project and it looks like case-when expressions, sum/min/max/mean expressions are supported. HashAggregate is supported too. Exchange should be supported too because I turned on Comet shuffle (--conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager, --conf spark.comet.exec.shuffle.enabled=true, --conf spark.comet.exec.shuffle.mode=native).

Why if partial aggregation is in Comet the final one isn't and I have a ColumnarToRow instead?

Additional context

I'm ready to provide any additional information or to run any debug query.

Thanks in advance!

[viirya]

Could you run a simple query to verify if Comet shuffle can be triggered?

[viirya]

Oh, could you disable spark.sql.adaptive.coalescePartitions.enabled and retry? Comet shuffle will be disabled if spark.sql.adaptive.coalescePartitions.enabled is enabled. Although you disable AQE, but this config is still enabled by default (i.e, it won't be set to false even you disable AQE).

[SemyonSinchenko]

Wow! With a disabled spark.sql.adaptive.coalescePartitions.enabled it works! May I open a PR with updates to documentation? Looks like I need to update this page (source)

[viirya]

Yea, you can open a PR to update the document. Although it should be just a temporary limit and we are working on to remove it. We can update the document again once the limitation is removed.

[andygrove]

You may also want to set spark.comet.explainFallback.enabled=true so that you can see the reasons why parts of your query are not native (this would show up as logging in the driver log).

I wonder if we should default this to true.