Skip to content

Enable user defined display_name for ScalarUDF #10417

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

Merged
alamb merged 3 commits into from
May 8, 2024
Merged

Enable user defined display_name for ScalarUDF #10417

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
270889a
enable user defined display_name
yyy1000 May 7, 2024
ba395a3
add display_name to get_field
yyy1000 May 8, 2024
90b61e7
add physical name
yyy1000 May 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions datafusion/core/src/physical_planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,9 +239,7 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
}
};
}
Expr::ScalarFunction(fun) => {
create_function_physical_name(fun.name(), false, &fun.args, None)
}
Expr::ScalarFunction(fun) => fun.func.display_name(&fun.args),
Expr::WindowFunction(WindowFunction {
fun,
args,
Expand Down Expand Up @@ -491,6 +489,7 @@ impl PhysicalPlanner for DefaultPhysicalPlanner {
let plan = self
.create_initial_plan(logical_plan, session_state)
.await?;

self.optimize_internal(plan, session_state, |_, _| {})
}
}
Expand Down
4 changes: 2 additions & 2 deletions datafusion/expr/src/expr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1637,7 +1637,7 @@ fn create_function_name(fun: &str, distinct: bool, args: &[Expr]) -> Result<Stri

/// Returns a readable name of an expression based on the input schema.
/// This function recursively transverses the expression for names such as "CAST(a > 2)".
fn create_name(e: &Expr) -> Result<String> {
pub(crate) fn create_name(e: &Expr) -> Result<String> {
match e {
Expr::Alias(Alias { name, .. }) => Ok(name.clone()),
Expr::Column(c) => Ok(c.flat_name()),
Expand Down Expand Up @@ -1793,7 +1793,7 @@ fn create_name(e: &Expr) -> Result<String> {
let expr_name = create_name(expr)?;
Ok(format!("unnest({expr_name})"))
}
Expr::ScalarFunction(fun) => create_function_name(fun.name(), false, &fun.args),
Expr::ScalarFunction(fun) => fun.func.display_name(&fun.args),
Expr::WindowFunction(WindowFunction {
fun,
args,
Expand Down
15 changes: 15 additions & 0 deletions datafusion/expr/src/udf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

//! [`ScalarUDF`]: Scalar User Defined Functions

use crate::expr::create_name;
use crate::simplify::{ExprSimplifyResult, SimplifyInfo};
use crate::{
ColumnarValue, Expr, FuncMonotonicity, ReturnTypeFunction,
Expand Down Expand Up @@ -133,6 +134,13 @@ impl ScalarUDF {
self.inner.name()
}

/// Returns this function's display_name.
///
/// See [`ScalarUDFImpl::display_name`] for more details
pub fn display_name(&self, args: &[Expr]) -> Result<String> {
self.inner.display_name(args)
}

/// Returns the aliases for this function.
///
/// See [`ScalarUDF::with_aliases`] for more details
Expand Down Expand Up @@ -274,6 +282,13 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
/// Returns this function's name
fn name(&self) -> &str;

/// Returns the user-defined display name of the UDF given the arguments
///
fn display_name(&self, args: &[Expr]) -> Result<String> {
let names: Vec<String> = args.iter().map(create_name).collect::<Result<_>>()?;
Ok(format!("{}({})", self.name(), names.join(",")))
}

/// Returns the function's [`Signature`] for information about what input
/// types are accepted and the function's Volatility.
fn signature(&self) -> &Signature;
Expand Down
21 changes: 21 additions & 0 deletions datafusion/functions/src/core/getfield.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,31 @@ impl ScalarUDFImpl for GetFieldFunc {
fn as_any(&self) -> &dyn Any {
self
}

fn name(&self) -> &str {
"get_field"
}

fn display_name(&self, args: &[Expr]) -> Result<String> {
if args.len() != 2 {
return exec_err!(
"get_field function requires 2 arguments, got {}",
args.len()
);
}

let name = match &args[1] {
Expr::Literal(name) => name,
_ => {
return exec_err!(
"get_field function requires the argument field_name to be a string"
);
}
};

Ok(format!("{}[{}]", args[0].display_name()?, name))
}

Copy link
Contributor

@jayzhan211 jayzhan211 May 8, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

cli in this PR

DataFusion CLI v38.0.0
> select struct('a');
+-------------------+
| struct(Utf8("a")) |
+-------------------+
| {c0: a}           |
+-------------------+
1 row(s) fetched. 
Elapsed 0.023 seconds.

> select get_field(struct('a'), 'c0');
+-----------------------+
| struct(Utf8("a"))[c0] |
+-----------------------+
| a                     |
+-----------------------+
1 row(s) fetched. 
Elapsed 0.007 seconds.

main branch:

DataFusion CLI v38.0.0
> select struct('a');
+-------------------+
| struct(Utf8("a")) |
+-------------------+
| {c0: a}           |
+-------------------+
1 row(s) fetched. 
Elapsed 0.014 seconds.

> select get_field(struct('a'), 'c0');
+-----------------------------------------+
| get_field(struct(Utf8("a")),Utf8("c0")) |
+-----------------------------------------+
| a                                       |
+-----------------------------------------+
1 row(s) fetched. 
Elapsed 0.004 seconds.

Copy link
Contributor Author

@yyy1000 yyy1000 May 8, 2024

Choose a reason for hiding this comment

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

Yeah, this example has better view. 🙌

fn signature(&self) -> &Signature {
&self.signature
}
Expand Down