Fix incorrect OFFSET during LIMIT pushdown.

[wiedld]

Which issue does this PR close?

Fixes #12423

First commit demonstrates the bug.

Rationale for this change

First commit demonstrates the current, incorrect behavior where the offset is not applied correctly during limit pushdown.

Followup commits add the fix, as well as a few doc comments.

What changes are included in this PR?

Slight change in offset handling during one of the helper functions with the limit pushdown.
Also added some docs to help explain existing code.

Are these changes tested?

Yes.

Are there any user-facing changes?

No.

[alamb]

cc @itsjunetime

[wiedld]

@mertak-synnada this is a fix to what we believe is a bug from this (very excellent) change. We would appreciate your review 🙏🏼 .

[itsjunetime]

I think this may negate some of the performance improvements gained by the initial PR that introduced these bugs, but I think once this is merged, we can refactor the pushdown_limit_helper function slightly to keep the correct behavior while pulling in the improvements again. I think it's just more important to get a fix merged first since this did break existing behavior.

[alamb]

I have filed #12423 to track this issue and updated this PR description

[alamb]

Thank you @wiedld and @itsjunetime and @mertak-synnada

I think this code is looking good to me. I had a few suggestions about comments, but the code and testing seems 👍 to me

[alamb]

this might be a good comment to add to the pushdown_limit_helper function as well

[alamb]

I am going to make the comment suggestions to this PR so we can merge it in.

[alamb]

I think this may negate some of the performance improvements gained by the initial PR that introduced these bugs, but I think once this is merged, we can refactor the pushdown_limit_helper function slightly to keep the correct behavior while pulling in the improvements again

@itsjunetime I wonder if you can elaborate on this or file a ticket. In order to avoid the final GlobalLimitExec I believe we would have to add offset support into SortPreservingMerge and Sort(TopK) -- which we could do, but I think the benefit might be relatively low

The implementation of Limit is pretty straightforward:

datafusion/datafusion/physical-plan/src/limit.rs

Lines 342 to 404 in f24f2cb

	pub struct LimitStream {
	/// The remaining number of rows to skip
	skip: usize,
	/// The remaining number of rows to produce
	fetch: usize,
	/// The input to read from. This is set to None once the limit is
	/// reached to enable early termination
	input: Option<SendableRecordBatchStream>,
	/// Copy of the input schema
	schema: SchemaRef,
	/// Execution time metrics
	baseline_metrics: BaselineMetrics,
	}
	impl LimitStream {
	pub fn new(
	input: SendableRecordBatchStream,
	skip: usize,
	fetch: Option<usize>,
	baseline_metrics: BaselineMetrics,
	) -> Self {
	let schema = input.schema();
	Self {
	skip,
	fetch: fetch.unwrap_or(usize::MAX),
	input: Some(input),
	schema,
	baseline_metrics,
	}
	}
	fn poll_and_skip(
	&mut self,
	cx: &mut Context<'_>,
	) -> Poll<Option<Result<RecordBatch>>> {
	let input = self.input.as_mut().unwrap();
	loop {
	let poll = input.poll_next_unpin(cx);
	let poll = poll.map_ok(|batch| {
	if batch.num_rows() <= self.skip {
	self.skip -= batch.num_rows();
	RecordBatch::new_empty(input.schema())
	} else {
	let new_batch = batch.slice(self.skip, batch.num_rows() - self.skip);
	self.skip = 0;
	new_batch
	}
	});
	match &poll {
	Poll::Ready(Some(Ok(batch))) => {
	if batch.num_rows() > 0 {
	break poll;
	} else {
	// continue to poll input stream
	}
	}
	Poll::Ready(Some(Err(_e))) => break poll,
	Poll::Ready(None) => break poll,
	Poll::Pending => break poll,
	}
	}
	}

[mertak-synnada]

I think this may negate some of the performance improvements gained by the initial PR that introduced these bugs, but I think once this is merged, we can refactor the pushdown_limit_helper function slightly to keep the correct behavior while pulling in the improvements again

@itsjunetime I wonder if you can elaborate on this or file a ticket. In order to avoid the final GlobalLimitExec I believe we would have to add offset support into SortPreservingMerge and Sort(TopK) -- which we could do, but I think the benefit might be relatively low

The implementation of Limit is pretty straightforward:

datafusion/datafusion/physical-plan/src/limit.rs

Lines 342 to 404 in f24f2cb

	pub struct LimitStream {
	/// The remaining number of rows to skip
	skip: usize,
	/// The remaining number of rows to produce
	fetch: usize,
	/// The input to read from. This is set to None once the limit is
	/// reached to enable early termination
	input: Option<SendableRecordBatchStream>,
	/// Copy of the input schema
	schema: SchemaRef,
	/// Execution time metrics
	baseline_metrics: BaselineMetrics,
	}
	impl LimitStream {
	pub fn new(
	input: SendableRecordBatchStream,
	skip: usize,
	fetch: Option<usize>,
	baseline_metrics: BaselineMetrics,
	) -> Self {
	let schema = input.schema();
	Self {
	skip,
	fetch: fetch.unwrap_or(usize::MAX),
	input: Some(input),
	schema,
	baseline_metrics,
	}
	}
	fn poll_and_skip(
	&mut self,
	cx: &mut Context<'_>,
	) -> Poll<Option<Result<RecordBatch>>> {
	let input = self.input.as_mut().unwrap();
	loop {
	let poll = input.poll_next_unpin(cx);
	let poll = poll.map_ok(|batch| {
	if batch.num_rows() <= self.skip {
	self.skip -= batch.num_rows();
	RecordBatch::new_empty(input.schema())
	} else {
	let new_batch = batch.slice(self.skip, batch.num_rows() - self.skip);
	self.skip = 0;
	new_batch
	}
	});
	match &poll {
	Poll::Ready(Some(Ok(batch))) => {
	if batch.num_rows() > 0 {
	break poll;
	} else {
	// continue to poll input stream
	}
	}
	Poll::Ready(Some(Err(_e))) => break poll,
	Poll::Ready(None) => break poll,
	Poll::Pending => break poll,
	}
	}
	}

I believe @itsjunetime mentioned for the first commit, but after my change suggestion, the gains should be preserved, imo.

[alamb]

Thanks everyone for your help getting this done!