Implement `GroupsAccumulator` for `first_value` aggregate (speed up `first_value` and `DISTINCT ON` queries)

[alamb]

Is your feature request related to a problem or challenge?

As reported in #16620 by @debajyoti-truefoundry, evaluting DISTINCT ON results in a query plan that uses first_value aggregates

The current implementation of first_value appears to have only a basic Accumulator implementation, and not the faster GroupsAccumulator: https://github.com/apache/datafusion/blob/main/datafusion/functions-aggregate/src/nth_value.rs#L93

We can very likely improve the performance of such queries significantly by implementing a GroupsAccumulator (background in Aggregating Millions of Groups Fast in Apache Arrow DataFusion )

Describe the solution you'd like

1. Add a benchmark (maybe add a query to the clickbench_extended suite)
2. Implement a GroupsAccumulator for first (and maybe nth) value

Describe alternatives you've considered

I think the accumulator could be pretty straightforward and track whatever groups were new and just copy the first row seen into the output (likely by using the take filter)

Additional context

There is a similar issue for optimizing array_agg here:

• Implement GroupsAccumulator for array_agg aggregation function #10145

[buraksenn]

I think this is kinda a straightforward issue to implement. If this is the case I would like to try fixing this if no one is working on it rn

[alamb]

The trick will be to get the GroupsAccumulator to do the right thing -- I will be happy to help review it

[buraksenn]

take

[tobixdev]

This would also be interesting to us. If I can help, feel free to ping me.

[debajyoti-truefoundry]

#16620 (comment)

[baiguoname]

FirstValue should support supports_retract_batch and retract_batch, this is straitforward.

[alamb]

Supporting retract batch is somewhat more complicated as that means the accumulator needs to hold the entire window of values.

I suggest you file a separate ticket to track supporting retract_batch

[baiguoname]

Supporting retract batch is somewhat more complicated as that means the accumulator needs to hold the entire window of values.

I suggest you file a separate ticket to track supporting retract_batch

For test, I implement MyTrivalFirstValueAccumulator in the similar of TrivalFirstValueAccumulator in this way:

#[derive(Debug, Default)]
pub struct MyTrivalFirstValueAccumulator {
    window_values: VecDeque<ScalarValue>,
}

impl Accumulator for MyTrivalFirstValueAccumulator {
    fn state(&mut self) -> Result<Vec<ScalarValue>> {
       Ok(vec![ScalarValue::from(0), ScalarValue::from(true)])
    }

    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
        let now_value = get_row_at_idx(values, 0)?.remove(0);
        self.window_values.push_back(now_value);
        Ok(())
    }

    fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> {
        unreachable!("not call");
    }

    fn retract_batch(&mut self, _values: &[ArrayRef]) -> Result<()> {
        self.window_values.pop_front();
        Ok(())
    }

    fn supports_retract_batch(&self) -> bool {
        true
    }

    fn evaluate(&mut self) -> Result<ScalarValue> {
        Ok(self.window_values[0].clone())
    }

    fn size(&self) -> usize {
        size_of_val(self) - size_of_val(&self.window_values)
    }
}

And their performance are close.

[Omega359]

First_value already has a GroupsAccumulator implementation

datafusion/datafusion/functions-aggregate/src/first_last.rs

Line 207 in d9d847e

fn create_groups_accumulator(

nth_value aggregate doesn't however.