datafusion.udf¶

Provides the user-defined functions for evaluation of dataframes.

Attributes¶

`_R`
`_W`

Classes¶

`Accumulator`	Defines how an `AggregateUDF` accumulates values.
`AggregateUDF`	Class for performing scalar user-defined functions (UDF).
`ScalarUDF`	Class for performing scalar user-defined functions (UDF).
`Volatility`	Defines how stable or volatile a function is.
`WindowEvaluator`	Evaluator class for user-defined window functions (UDWF).
`WindowUDF`	Class for performing window user-defined functions (UDF).

Module Contents¶

class datafusion.udf.Accumulator¶

Defines how an AggregateUDF accumulates values.

abstract evaluate() → pyarrow.Scalar¶: Return the resultant value.

abstract merge(states: List[pyarrow.Array]) → None¶: Merge a set of states.

abstract state() → List[pyarrow.Scalar]¶: Return the current state.

abstract update(*values: pyarrow.Array) → None¶: Evaluate an array of values and update state.

class datafusion.udf.AggregateUDF(name: str | None, accumulator: Callable[[], Accumulator], input_types: list[pyarrow.DataType], return_type: pyarrow.DataType, state_type: list[pyarrow.DataType], volatility: Volatility | str)¶

Class for performing scalar user-defined functions (UDF).

Aggregate UDFs operate on a group of rows and return a single value. See also ScalarUDF for operating on a row by row basis.

Instantiate a user-defined aggregate function (UDAF).

See udaf() for a convenience function and argument descriptions.

__call__(*args: datafusion.expr.Expr) → datafusion.expr.Expr¶

Execute the UDAF.

This function is not typically called by an end user. These calls will occur during the evaluation of the dataframe.

static udaf(accum: Callable[[], Accumulator], input_types: pyarrow.DataType | list[pyarrow.DataType], return_type: pyarrow.DataType, state_type: list[pyarrow.DataType], volatility: Volatility | str, name: str | None = None) → AggregateUDF¶

Create a new User-Defined Aggregate Function.

If your Accumulator can be instantiated with no arguments, you can simply pass it’s type as accum. If you need to pass additional arguments to it’s constructor, you can define a lambda or a factory method. During runtime the Accumulator will be constructed for every instance in which this UDAF is used. The following examples are all valid.

Parameters:

accum – The accumulator python function.
input_types – The data types of the arguments to accum.
return_type – The data type of the return value.
state_type – The data types of the intermediate accumulation.
volatility – See Volatility for allowed values.
name – A descriptive name for the function.

Returns:

A user-defined aggregate function, which can be used in either data aggregation or window function calls.

_udaf¶

class datafusion.udf.ScalarUDF(name: str | None, func: Callable[Ellipsis, _R], input_types: pyarrow.DataType | list[pyarrow.DataType], return_type: _R, volatility: Volatility | str)¶

Class for performing scalar user-defined functions (UDF).

Scalar UDFs operate on a row by row basis. See also AggregateUDF for operating on a group of rows.

Instantiate a scalar user-defined function (UDF).

See helper method udf() for argument details.

__call__(*args: datafusion.expr.Expr) → datafusion.expr.Expr¶

Execute the UDF.

This function is not typically called by an end user. These calls will occur during the evaluation of the dataframe.

static udf(func: Callable[Ellipsis, _R], input_types: list[pyarrow.DataType], return_type: _R, volatility: Volatility | str, name: str | None = None) → ScalarUDF¶

Create a new User-Defined Function.

Parameters:

func – A callable python function.
input_types – The data types of the arguments to func. This list must be of the same length as the number of arguments.
return_type – The data type of the return value from the python function.
volatility – See Volatility for allowed values.
name – A descriptive name for the function.

Returns:

A user-defined aggregate function, which can be used in either data: aggregation or window function calls.

_udf¶

class datafusion.udf.Volatility(*args, **kwds)¶

Bases: enum.Enum

Defines how stable or volatile a function is.

When setting the volatility of a function, you can either pass this enumeration or a str. The str equivalent is the lower case value of the name (“immutable”, “stable”, or “volatile”).

__str__()¶: Returns the string equivalent.

Immutable = 1¶

An immutable function will always return the same output when given the same input.

DataFusion will attempt to inline immutable functions during planning.

Stable = 2¶

Returns the same value for a given input within a single queries.

A stable function may return different values given the same input across different queries but must return the same value for a given input within a query. An example of this is the Now function. DataFusion will attempt to inline Stable functions during planning, when possible. For query select col1, now() from t1, it might take a while to execute but now() column will be the same for each output row, which is evaluated during planning.

Volatile = 3¶

A volatile function may change the return value from evaluation to evaluation.

Multiple invocations of a volatile function may return different results when used in the same query. An example of this is the random() function. DataFusion can not evaluate such functions during planning. In the query select col1, random() from t1, random() function will be evaluated for each output row, resulting in a unique random value for each row.

class datafusion.udf.WindowEvaluator¶

Evaluator class for user-defined window functions (UDWF).

It is up to the user to decide which evaluate function is appropriate.

`uses_window_frame`	`supports_bounded_execution`	`include_rank`	function_to_implement
False (default)	False (default)	False (default)	`evaluate_all`
False	True	False	`evaluate`
False	True/False	True	`evaluate_all_with_rank`
True	True/False	True/False	`evaluate`

evaluate(values: list[pyarrow.Array], eval_range: tuple[int, int]) → pyarrow.Scalar¶

Evaluate window function on a range of rows in an input partition.

This is the simplest and most general function to implement but also the least performant as it creates output one row at a time. It is typically much faster to implement stateful evaluation using one of the other specialized methods on this trait.

Returns a [ScalarValue] that is the value of the window function within range for the entire partition. Argument values contains the evaluation result of function arguments and evaluation results of ORDER BY expressions. If function has a single argument, values[1..] will contain ORDER BY expression results.

evaluate_all(values: list[pyarrow.Array], num_rows: int) → pyarrow.Array¶

Evaluate a window function on an entire input partition.

This function is called once per input partition for window functions that do not use values from the window frame, such as row_number(), rank(), dense_rank(), percent_rank(), cume_dist(), lead(), and lag().

It produces the result of all rows in a single pass. It expects to receive the entire partition as the value and must produce an output column with one output row for every input row.

num_rows is required to correctly compute the output in case len(values) == 0

Implementing this function is an optimization. Certain window functions are not affected by the window frame definition or the query doesn’t have a frame, and evaluate skips the (costly) window frame boundary calculation and the overhead of calling evaluate for each output row.

For example, the LAG built in window function does not use the values of its window frame (it can be computed in one shot on the entire partition with Self::evaluate_all regardless of the window defined in the OVER clause)

lag(x, 1) OVER (ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING)

However, avg() computes the average in the window and thus does use its window frame.

avg(x) OVER (PARTITION BY y ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING)

evaluate_all_with_rank(num_rows: int, ranks_in_partition: list[tuple[int, int]]) → pyarrow.Array¶

Called for window functions that only need the rank of a row.

Evaluate the partition evaluator against the partition using the row ranks. For example, rank(col("a")) produces

a | rank
- + ----
A | 1
A | 1
C | 3
D | 4
D | 4

For this case, num_rows would be 5 and the ranks_in_partition would be called with

[
    (0,1),
    (2,2),
    (3,4),
]

The user must implement this method if include_rank returns True.

get_range(idx: int, num_rows: int) → tuple[int, int]¶

Return the range for the window fuction.

If uses_window_frame flag is false. This method is used to calculate required range for the window function during stateful execution.

Generally there is no required range, hence by default this returns smallest range(current row). e.g seeing current row is enough to calculate window result (such as row_number, rank, etc)

Parameters:

idx: – Current index:
num_rows – Number of rows.

include_rank() → bool¶: Can this function be evaluated with (only) rank?

is_causal() → bool¶: Get whether evaluator needs future data for its result.

memoize() → None¶

Perform a memoize operation to improve performance.

When the window frame has a fixed beginning (e.g UNBOUNDED PRECEDING), some functions such as FIRST_VALUE and NTH_VALUE do not need the (unbounded) input once they have seen a certain amount of input.

memoize is called after each input batch is processed, and such functions can save whatever they need

supports_bounded_execution() → bool¶: Can the window function be incrementally computed using bounded memory?

uses_window_frame() → bool¶: Does the window function use the values from the window frame?

class datafusion.udf.WindowUDF(name: str | None, func: Callable[[], WindowEvaluator], input_types: list[pyarrow.DataType], return_type: pyarrow.DataType, volatility: Volatility | str)¶

Class for performing window user-defined functions (UDF).

Window UDFs operate on a partition of rows. See also ScalarUDF for operating on a row by row basis.

Instantiate a user-defined window function (UDWF).

See udwf() for a convenience function and argument descriptions.

__call__(*args: datafusion.expr.Expr) → datafusion.expr.Expr¶

Execute the UDWF.

This function is not typically called by an end user. These calls will occur during the evaluation of the dataframe.

static udwf(func: Callable[[], WindowEvaluator], input_types: pyarrow.DataType | list[pyarrow.DataType], return_type: pyarrow.DataType, volatility: Volatility | str, name: str | None = None) → WindowUDF¶

Create a new User-Defined Window Function.

If your WindowEvaluator can be instantiated with no arguments, you can simply pass it’s type as func. If you need to pass additional arguments to it’s constructor, you can define a lambda or a factory method. During runtime the WindowEvaluator will be constructed for every instance in which this UDWF is used. The following examples are all valid.

import pyarrow as pa

class BiasedNumbers(WindowEvaluator):
    def __init__(self, start: int = 0) -> None:
        self.start = start

    def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array:
        return pa.array([self.start + i for i in range(num_rows)])

def bias_10() -> BiasedNumbers:
    return BiasedNumbers(10)

udwf1 = udwf(BiasedNumbers, pa.int64(), pa.int64(), "immutable")
udwf2 = udwf(bias_10, pa.int64(), pa.int64(), "immutable")
udwf3 = udwf(lambda: BiasedNumbers(20), pa.int64(), pa.int64(), "immutable")

Parameters:

func – A callable to create the window function.
input_types – The data types of the arguments to func.
return_type – The data type of the return value.
volatility – See Volatility for allowed values.
arguments – A list of arguments to pass in to the __init__ method for accum.
name – A descriptive name for the function.

Returns:

A user-defined window function.

_udwf¶

datafusion.udf._R¶

datafusion.udf._W¶

datafusion.substrait