Plans that have multiple output partitions at the top level fail

[adragomir]

I have a physical plan that looks like this when it enters ray

    [ output_partitions: 16]    ProjectionExec: expr=[...]
    [ output_partitions: 16]      RayStageExec[0] (output_partitioning=RoundRobinBatch(16))
    [ output_partitions: 16]        RepartitionExec: partitioning=RoundRobinBatch(16), input_partitions=1
    [ output_partitions: 1]          GlobalLimitExec: skip=0, fetch=1
    [ output_partitions: 1]            CoalescePartitionsExec
    [ output_partitions: 16]              CoalesceBatchesExec: target_batch_size=8192, fetch=1
    [ output_partitions: 16]                FilterExec: timestamp@0 >= 1736899200000000 AND timestamp@0 < 1736985600000000
    [ output_partitions: 16]                  DeltaScan
    [ output_partitions: 16]                    ParquetExec: file_groups={...]

In this case, Ray will add a RayStageExec on top of the ProjectionExec, and then fail because Last stage expected to have one partition

I managed to fix it locally by ensuring that the plan has 1 output partition

        let num_output_partitions = match plan.output_partitioning() {
            Partitioning::RoundRobinBatch(x) => x,
            Partitioning::Hash(_, x) => x,
            Partitioning::UnknownPartitioning(x) => x
        };
        if *num_output_partitions > 1 {
            plan = Arc::new(RepartitionExec::try_new(plan, Partitioning::RoundRobinBatch(1))?);
        }

[robtandy]

@adragomir can you provide the sql to reproduce the plan? Or a sql query that will produce a similar plan with the same problem?

[adragomir]

@robtandy HTH, if not, I can start working on a full test case tomorrow, harder to extract minimal query and setup.

Table schema would look like (duckdb syntax to get a parquet file):

CREATE OR REPLACE TABLE midvalues(
  timestamp TIMESTAMP_S,  -- 0
  web STRUCT(
    webPageDetails STRUCT(
      pageViews STRUCT(value INT8)
    )
  ), 
  endUserIDs STRUCT(
    _experience STRUCT(
      mcid STRUCT(
        id VARCHAR, 
        extra1 VARCHAR
      ),
      aaid STRUCT(
        id VARCHAR, 
        extra1 VARCHAR
      )
    )
  )
);
COPY midvalues TO 'midvalues.parquet' (FORMAT PARQUET);

SQL query:

        SELECT
            timestamp,
            web.webPageDetails.pageViews.value AS pageview,
            endUserIDs._experience.mcid.id AS mcid,
            endUserIDs._experience.aaid.id AS aaid
        FROM
            midvalues
        WHERE
            timestamp >= TO_TIMESTAMP('2025-01-15')
            AND timestamp < TO_TIMESTAMP('2025-01-16')