datafusion icon indicating copy to clipboard operation
datafusion copied to clipboard

Add a nicer way to register user defined aggregates for use in SQL

Open alamb opened this issue 3 years ago • 0 comments

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

We implement user defined aggregates in IOx called "selector functions" -- to make them visible to SQL we need to add them to the SessionState::aggregate_functions map

We use the same pattern as here: https://github.com/apache/arrow-datafusion/blob/447141c/datafusion/core/tests/user_defined_aggregates.rs#L131-L136

    // register the selector as "first"
    ctx.state
        .write()
        .aggregate_functions
        .insert(name.to_string(), Arc::new(first));
}

This is pretty ugly as it requires directly manipulating the HashMap as well as requiring duplication of name (which is already a field on AggregateUDF this allowing potential mismatches

Describe the solution you'd like I would like a nicer API to add such functions. Perhaps something like

impl SessionState {
...
  // Register the user defined aggregate function for calling via SQL
  fn register_aggregate(&mut self, func: AggregateUDF) {
...
}

Describe alternatives you've considered Leave existing API

Additional context This came up in the context of this PR in IOx: https://github.com/influxdata/influxdb_iox/pull/5628

alamb avatar Sep 13 '22 16:09 alamb