Skip to content

PiecewiseMergeJoin drops rows when a stream batch spans multiple output chunks (small batch_size) #24348

Description

@viirya

Describe the bug

A classic RIGHT/FULL (and any) PiecewiseMergeJoin can drop output rows when a single stream batch produces more than one batch_size output chunk.

ClassicPWMJStream::process_stream_batch returns a single completed batch from the coalescer and then advances to fetch the next stream batch. But finish_buffered_batch can leave several completed batches queued; advancing while some remain strands them for the final stream batch, so those rows are never emitted.

This only shows up at small datafusion.execution.batch_size (the default 8192 usually keeps a stream batch's output within a single chunk, hiding it). It is independent of the NULL-key fix in #24335/#24336.

To Reproduce

set datafusion.optimizer.enable_piecewise_merge_join = true;
set datafusion.execution.batch_size = 4;

create table l(id int, v int) as values
  (0,0),(1,3),(2,1),(3,-5),(4,NULL),(5,-3),(6,NULL),(7,2),(8,4),(9,NULL),
  (10,-1),(11,-2),(12,2),(13,2),(14,-3),(15,NULL),(16,-2),(17,NULL),(18,0),(19,-5),
  (20,-1),(21,NULL),(22,-5),(23,0),(24,NULL),(25,NULL),(26,-3),(27,2),(28,-3),(29,NULL);
create table r(id int, v int) as values
  (0,-5),(1,NULL),(2,1),(3,2),(4,3),(5,-5),(6,NULL),(7,NULL),(8,3),(9,-3),
  (10,NULL),(11,-2),(12,NULL),(13,0),(14,NULL),(15,2),(16,-1),(17,NULL),(18,-4),(19,-5),
  (20,-3),(21,2),(22,0),(23,-1),(24,-2),(25,NULL),(26,1),(27,-3),(28,-1),(29,-5);

select count(*) from l right join r on l.v < r.v;

With PiecewiseMergeJoin this returns 196; the correct count (as NestedLoopJoin returns with the flag off) is 198. The dropped rows are unmatched right rows.

Expected behavior

PiecewiseMergeJoin returns the same rows as NestedLoopJoin regardless of batch_size.

Additional context

Root cause: the !continue_process branch of process_stream_batch returns one completed batch and advances to FetchStreamBatch even when more completed batches are queued. Fix + regression test coming.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions