Skip to content

Add FunctionRegistry::register_udaf and FunctionRegistry::register_udwf #9075

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
alamb merged 4 commits into from
Feb 7, 2024
Merged

Add FunctionRegistry::register_udaf and FunctionRegistry::register_udwf #9075

alamb merged 4 commits into from
Feb 7, 2024

Conversation

@alamb
Copy link
Contributor

@alamb alamb commented Jan 30, 2024 •
edited
Loading

Which issue does this PR close?

Closes #9074
Part of #8045 (making functions modular in DataFusion)

Rationale for this change

See #9074 (basically set us up for pulling aggregates and window functions out of the core)

What changes are included in this PR?

  1. Add FunctionRegistry::register_udaf and FunctionRegistry::register_udwf
  2. Implement them for SessionState
  3. Update SessionContext to use that API rather than modifying session state directly

Are these changes tested?

Covered by existing tests

Are there any user-facing changes?

@github-actions github-actions bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates core Core DataFusion crate labels Jan 30, 2024
@alamb alamb force-pushed the alamb/more_functions branch from c0d88f7 to 7eeb52e Compare January 31, 2024 22:47
@github-actions github-actions bot removed logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates labels Jan 31, 2024
alamb
alamb commented Jan 31, 2024
View reviewed changes
datafusion/core/src/execution/context/mod.rs
state
.scalar_functions
.insert(f.name().to_string(), Arc::new(f));
state.register_udf(Arc::new(f)).ok();
Copy link
Contributor Author

@alamb alamb Jan 31, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

moved into SessionState -- it is important that the alias resolution is done there so when new functions are registered via SessionState their aliases are as well. Otherwise aliases are only added when the function is defined via SessionContext

datafusion/execution/src/registry.rs
///
/// Returns an error (the default) if the function can not be registered,
/// for example if the registry is read only.
fn register_udaf(
Copy link
Contributor Author

@alamb alamb Jan 31, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the key (backwards compatible) API addition in this PR.

It sets us up for being able to pull out BuiltInAggregate and BuiltInWindowFunction

Copy link
Member

@viirya viirya Feb 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm wondering why this is a backwards compatible API? Does FunctionRegistry provide this API before?

Copy link
Contributor Author

@alamb alamb Feb 6, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is backwards compatible because it provides a default implementation (to return a Not Yet Implemented error)

Thus, any existing implementations of FunctionRegistry will continue to compile and work as it did before, without any required code changes

@alamb alamb force-pushed the alamb/more_functions branch from 7eeb52e to 4257c16 Compare January 31, 2024 22:50
@alamb alamb marked this pull request as ready for review February 1, 2024 16:39
This was referenced Feb 2, 2024
Merged
Closed
@alamb alamb requested a review from comphead February 6, 2024 02:06
@alamb
Copy link
Contributor Author

alamb commented Feb 6, 2024 •
edited
Loading

I wonder if you might have time to review this PR @comphead ?

Once it is merged, I can then start breaking out more functions (e.g. #9113)

viirya
viirya reviewed Feb 6, 2024
View reviewed changes
comphead
comphead approved these changes Feb 6, 2024
View reviewed changes
Copy link
Contributor

@comphead comphead left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm as for me @alamb
Maybe we can do a one liner in register_udf

datafusion/core/src/execution/context/mod.rs Outdated
let aliases = udf.aliases();
for alias in aliases {
self.scalar_functions.insert(alias.to_string(), udf.clone());
}
Copy link
Contributor

@comphead comphead Feb 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
}

alamb and others added 3 commits February 6, 2024 08:01
@alamb @comphead @viirya
Apply suggestions from code review
caef440 
Co-authored-by: comphead <comphead@users.noreply.github.com>
Co-authored-by: Liang-Chi Hsieh <viirya@gmail.com>
@alamb
Merge remote-tracking branch 'apache/main' into alamb/more_functions
20305be
@alamb
cleanup
d3cccc2
@alamb alamb merged commit 77f9775 into apache:main Feb 7, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@viirya viirya viirya approved these changes

@comphead comphead comphead approved these changes

Assignees

No one assigned

Labels

core Core DataFusion crate

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Add FunctionRegistry::register_udaf and FunctionRegistry::register_udwf

3 participants

@alamb @viirya @comphead