feat: support batched table functions

[bubulalabu]

Batched Table Functions with LATERAL Join Support

Note: This is a draft PR for gathering feedback on the approach. It's working, but not polished yet

I'm interested to hear opinions if this is the right path forward before starting to finalize the PR

Which issue does this PR close?

Closes #18121.

---

The Problem

DataFusion's current table function API (TableFunctionImpl) requires constant arguments at planning time and returns a TableProvider. This prevents LATERAL joins where table functions need to correlate with outer query rows.

-- Current limitation: Cannot do this
SELECT * FROM metrics m
CROSS JOIN LATERAL generate_series(m.start_time, m.end_time) AS t(ts);

Current Architecture Limitations

// Current API
trait TableFunctionImpl {
    fn call(&self, args: &[Expr]) -> Result<Arc<dyn TableProvider>>;
}

Problems:

• Planning-time evaluation: Arguments must be constants, can't reference table columns
• No LATERAL support: Can't correlate function calls with input rows
• Separate TableProvider per call: Inefficient for batch processing

---

Proposed Solution

Introduce a new trait that processes arrays of arguments (batched invocation) and returns output with input row mapping for correlation:

#[async_trait]
trait BatchedTableFunctionImpl {
    async fn invoke_batch(
        &self,
        args: &[ArrayRef],            // Already-evaluated arrays
        projection: Option<&[usize]>, 
        filters: &[Expr],
        limit: Option<usize>,
    ) -> Result<BatchResultStream>;
}

// Results map back to input rows for LATERAL
struct BatchResultChunk {
    output: RecordBatch,
    input_row_indices: Vec<u32>,  // contain the indexes that correspond to the input indexes for every output row 
}

Key Design Decisions

1. Arrays as Arguments

• Like ScalarUDF - receives evaluated data, not planning-time expressions
• Enables batched processing: process multiple function calls in one invoke_batch
• Natural fit for LATERAL: array elements correspond to input rows

2. Async API

• Matches TableProvider::scan() async pattern

3. Direct Pushdown Parameters

• Matches TableProvider::scan() pattern
• Function gets projection/filter/limit directly in invoke_batch()

3. Streaming Results

• Returns Stream<BatchResultChunk> for memory-bounded execution
• Chunks include input_row_indices for efficient LATERAL join implementation

---

Architecture Overview

Execution Flow Example: generate_series(start, end)

Input batch:

┌────┬───────┬─────┐
│ id │ start │ end │
├────┼───────┼─────┤
│ 1  │ 10    │ 12  │
│ 2  │ 20    │ 22  │
└────┴───────┴─────┘

Evaluate args → start=[10, 20], end=[12, 22]

invoke_batch → Stream of chunks:

Chunk {
    output: [10, 11, 12],     // generated values for start=[10] and end=[12] 
    input_row_indices: [0, 0, 0]  // all from input row 0
}
Chunk {
    output: [20, 21, 22],
    input_row_indices: [1, 1, 1]  // all from input row 1
}

Combine using indices → Final output:

┌────┬───────┬─────┬───────┐
│ id │ start │ end │ value │
├────┼───────┼─────┼───────┤
│ 1  │ 10    │ 12  │ 10    │
│ 1  │ 10    │ 12  │ 11    │
│ 1  │ 10    │ 12  │ 12    │
│ 2  │ 20    │ 22  │ 20    │
│ 2  │ 20    │ 22  │ 21    │
│ 2  │ 20    │ 22  │ 22    │
└────┴───────┴─────┴───────┘

---

What's Included

Core Components

1. New Trait (datafusion/catalog/src/table.rs)

pub trait BatchedTableFunctionImpl: Send + Sync + Debug {
    fn name(&self) -> &str;
    fn signature(&self) -> &Signature;
    fn return_type(&self, arg_types: &[DataType]) -> Result<Schema>;

    async fn invoke_batch(
        &self,
        args: &[ArrayRef],
        projection: Option<&[usize]>,
        filters: &[Expr],
        limit: Option<usize>,
    ) -> Result<BatchResultStream>;
}

2. Logical Plan Nodes (datafusion/expr/src/logical_plan/plan.rs)

• StandaloneBatchedTableFunction - for SELECT * FROM func(1, 100)
• LateralBatchedTableFunction - for SELECT * FROM t LATERAL func(t.x)

3. Physical Executor (datafusion/catalog/src/batched_function/exec.rs)

• Unified BatchedTableFunctionExec with two modes (Standalone/Lateral)
• Handles argument evaluation, function invocation, and result combination

6. Optimizer Rules (datafusion/optimizer/src/)

• Extended projection/filter/limit pushdown to new logical plan nodes

7. Reference Implementation (datafusion/functions-table/src/generate_series_batched.rs)

• Complete example showing the implementation pattern
• Demonstrates limit pushdown optimization

SQL Examples Enabled

-- Standalone: constant arguments
SELECT * FROM generate_series(1, 100);

-- LATERAL: correlate with outer table
SELECT m.metric_id, t.ts
FROM metrics m
CROSS JOIN LATERAL generate_series(m.start_time, m.end_time) AS t(ts);

-- With optimizations (limit pushdown)
SELECT * FROM generate_series(1, 1000000) LIMIT 10;
-- Only generates ~10 rows instead of 1M

---

Testing

Comprehensive test coverage:

• Unit tests in batched_function/exec.rs
• Integration tests in core/tests/batched_table_function.rs
• SQL logic tests in sqllogictest/test_files/lateral.slt
• All tests pass

---

Open Questions

1. Is a new trait the right approach?

This PR introduces BatchedTableFunctionImpl as a separate trait. An alternative would be to use TableProvider wrappers.

2. Is the API appropriate?

async fn invoke_batch(
    args: &[ArrayRef],
    projection: Option<&[usize]>,
    filters: &[Expr],
    limit: Option<usize>,
) -> Result<BatchResultStream>;

3. Is "Batched" the right terminology?

Current naming: BatchedTableFunctionImpl, StandaloneBatchedTableFunction, etc.

4. Should we keep both APIs or eventually deprecate TableFunctionImpl?

---

User-Facing Changes

New SQL capabilities:

-- LATERAL joins
SELECT t.id, s.value
FROM my_table t
CROSS JOIN LATERAL generate_series(t.start, t.end) AS s(value);

New API for function developers:

// Implement the trait
struct MyFunction { /* ... */ }

#[async_trait]
impl BatchedTableFunctionImpl for MyFunction {
    async fn invoke_batch(/* ... */) -> Result<BatchResultStream> {
        // Your implementation
    }
}

// Register in session
ctx.state_ref()
    .write()
    .register_batched_table_function("my_func", Arc::new(MyFunction));

Breaking changes:

The PR adds two new values to the LogicalPlan enum.