feat: implement into_batches operator on flotilla distrubted engine

[ohbh]

Changes Made

This PR adds support for DataFrame.into_batches() in the distributed (Flotilla/Ray) engine and wires it through the pipeline.

The node logic is as follows:

1. materialize downstream pipeline w/ IntoBatches operator of batch_size
2. collect materialized outputs from stream
3. when num rows of materialized output >= batch_size*0.8 (batch_size_treshold), submit a new task w/ IntoBatches operator
4. Send all remainder outputs in a downstream task

for (3/4), it's important to submit the next task with a IntoBatches operator, because that will force the local pipeline to combine all the RecordBatches in the give MicroPartition, which the logic done in the local execution op

Daft/src/daft-local-execution/src/intermediate_ops/into_batches.rs

Lines 35 to 43 in 048fd24

	let out = match input.concat_or_get(IOStatsContext::new("into_batches"))? {
	Some(record_batch) => Arc::new(MicroPartition::new_loaded(
	input.schema(),
	Arc::new(vec![record_batch]),
	None,
	)),
	None => Arc::new(MicroPartition::empty(Some(input.schema()))),
	};
	Ok((state, IntermediateOperatorResult::NeedMoreInput(Some(out))))

Related Issues

Native runner: #4935

Checklist

• Documented in API Docs (if applicable)
• Documented in User Guide (if applicable)
• If adding a new documentation page, doc is added to docs/mkdocs.yml navigation
• Documentation builds and is formatted properly (tag @/ccmao1130 for docs review)

[greptile-apps]

Greptile Summary

This PR implements support for DataFrame.into_batches() in the distributed (Flotilla/Ray) execution engine, achieving parity with the native execution engine. The implementation adds a new distributed pipeline node (IntoBatchesNode) that handles batching logic across distributed workers.

The changes span multiple layers of the distributed execution stack:

1. Pipeline Node Infrastructure: Adds a new IntoBatchesNode in src/daft-distributed/src/pipeline_node/into_batches.rs that implements the distributed batching strategy. The node materializes downstream pipeline outputs, accumulates them until reaching the target batch size, then submits tasks that force local pipeline combination of RecordBatches within MicroPartitions.

2. Translation Layer: Updates the pipeline translator in translate.rs to handle LogicalPlan::IntoBatches operations by creating IntoBatchesNode instances with appropriate configuration.

3. Stage Builder: Moves IntoBatches from the list of unsupported operations to supported operations in stage_builder.rs, enabling the distributed engine to process logical plans containing batching operations.

4. API Surface: Removes the runtime restriction in dataframe.py that previously prevented into_batches() from being used with the Ray runner.

5. Test Coverage: Updates tests to run on both native and distributed engines, with adjusted assertions to accommodate the "best effort" nature of distributed batching where exact batch sizes cannot be guaranteed but minimum sizes are maintained.

The distributed implementation follows a streaming approach where it collects materialized outputs and creates new tasks when accumulated data reaches the batch size threshold, ensuring efficient memory usage and proper data distribution across workers.

Confidence score: 1/5

• This PR contains a critical compilation bug that will prevent it from working
• Score reflects a serious issue in the core IntoBatches implementation that uses a moved variable
• Pay close attention to src/daft-distributed/src/pipeline_node/into_batches.rs line 124

_{6 files reviewed, 1 comment}

_{Edit Code Review Bot Settings | Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 82.05128% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.57%. Comparing base (a05ad3c) to head (a191129).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...daft-distributed/src/pipeline_node/into_batches.rs	80.95%	28 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4958      +/-   ##
==========================================
+ Coverage   78.03%   79.57%   +1.54%     
==========================================
  Files         916      917       +1     
  Lines      126917   126934      +17     
==========================================
+ Hits        99043   101014    +1971     
+ Misses      27874    25920    -1954     

Files with missing lines	Coverage Δ
daft/dataframe/dataframe.py	86.75% <ø> (+0.06%)	⬆️
src/daft-distributed/src/pipeline_node/mod.rs	89.91% <ø> (ø)
...rc/daft-distributed/src/pipeline_node/translate.rs	97.10% <100.00%> (+0.53%)	⬆️
src/daft-distributed/src/stage/stage_builder.rs	95.50% <ø> (+2.09%)	⬆️
...daft-distributed/src/pipeline_node/into_batches.rs	80.95% <80.95%> (ø)

... and 36 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[colin-ho]

Mostly looks good!

[srilman]

Great job on this!