Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
39 changes: 37 additions & 2 deletions datafusion/catalog/src/table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -552,29 +552,64 @@ pub trait TableProviderFactory: Debug + Sync + Send {

/// Describes arguments provided to the table function call.
pub struct TableFunctionArgs<'e, 's> {
/// Call arguments.
/// Call arguments (potentially coerced/simplified).
exprs: &'e [Expr],
/// Original raw arguments before any coercion/simplification.
raw_exprs: &'e [Expr],
/// Session within which the function is called.
session: &'s dyn Session,
}

impl<'e, 's> TableFunctionArgs<'e, 's> {
/// Make a new [`TableFunctionArgs`].
pub fn new(exprs: &'e [Expr], session: &'s dyn Session) -> Self {
Self { exprs, session }
Self {
exprs,
raw_exprs: exprs,
session,
}
}

/// Get expressions passed as the called function arguments.
pub fn exprs(&self) -> &'e [Expr] {
self.exprs
}

/// Get the original, unprocessed call arguments (before
/// type coercion and simplification).
///
/// The arguments returned by [`Self::exprs`] may have been
/// type-coerced or simplified. Use this method to access the
/// arguments exactly as the user wrote them in the SQL query.
pub fn raw_exprs(&self) -> &'e [Expr] {
self.raw_exprs
}

/// Get a session where the table function is called.
pub fn session(&self) -> &'s dyn Session {
self.session
}
}

impl<'e, 's> TableFunctionArgs<'e, 's> {
/// Make a new [`TableFunctionArgs`] with separate processed and raw expressions.
///
/// This is useful when the table function implementation needs
/// access to both the (possibly coerced/simplified) expressions
/// and the original expressions as written.
pub fn new_with_raw(
exprs: &'e [Expr],
raw_exprs: &'e [Expr],
session: &'s dyn Session,
) -> Self {
Self {
exprs,
raw_exprs,
session,
}
}
}

/// A trait for table function implementations
pub trait TableFunctionImpl: Debug + Sync + Send + Any {
/// Create a table provider
Expand Down
10 changes: 7 additions & 3 deletions datafusion/core/src/execution/session_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1984,16 +1984,20 @@ impl ContextProvider for SessionContextProvider<'_> {
.build();
let simplifier = ExprSimplifier::new(simplify_context);
let schema = DFSchema::empty();
let raw_args = args.clone();
let args = args
.into_iter()
.map(|arg| {
let backup = arg.clone();
simplifier
.coerce(arg, &schema)
.and_then(|e| simplifier.simplify(e))
.unwrap_or(backup)
})
.collect::<datafusion_common::Result<Vec<_>>>()?;
let provider = tbl_func
.create_table_provider_with_args(TableFunctionArgs::new(&args, self.state))?;
.collect::<Vec<_>>();
let provider = tbl_func.create_table_provider_with_args(
TableFunctionArgs::new_with_raw(&args, &raw_args, self.state),
)?;

Ok(provider_as_source(provider))
}
Expand Down