apache/datafusion-python/make-pythonic

Make Python API Functions More Pythonic

You are improving the datafusion-python API to feel more natural to Python users. The goal is to allow functions to accept native Python types (int, float, str, bool, etc.) for arguments that are contextually always or typically literal values, instead of requiring users to manually wrap them in lit().

Core principle: A Python user should be able to write split_part(col("a"), ",", 2) instead of split_part(col("a"), lit(","), lit(2)) when the arguments are contextually obvious literals.

Scope: functions vs functions.spark

Both python/datafusion/functions/__init__.py and python/datafusion/functions/spark.py are in scope. We want both to feel pythonic — accept native Python types where the argument is contextually a literal — but functions.spark carries an additional constraint: every signature must remain compatible with `pyspark.sql.functions`.

Compatibility rules for the spark namespace:

• Parameter names must match pyspark exactly. Pyspark callers pass by

keyword (spark.shiftleft(col=..., numBits=...)), so renames break them. Do NOT rename a parameter just because it would be more pythonic in the main namespace.

• Positional order must match pyspark exactly. Reordering breaks

positional pyspark calls.

• Type unions may widen the input set, never narrow it. Pyspark

accepts Column or str (column name) for most args; we accept Expr already, and widening to Expr | int / Expr | str for literal-friendly arguments is on-brand because the int/str case is exactly what a pyspark caller would also try. Just verify the widened set is a superset of what pyspark accepts for that arg.

• Extra keyword arguments are allowed as long as they default to

None and pyspark's positional/keyword form still works (e.g. the spark avg/try_sum/collect_list/collect_set retain DataFusion's distinct/filter/order_by/null_treatment kwargs).

Practical effect: in functions.spark, apply Categories A and (where pyspark exposes the same arg as a non-Expr) B normally, but cross-check each proposed signature against pyspark.sql.functions before landing it. When pyspark's own type hint is Column | str for a "column name" arg, prefer leaving the spark wrapper at Expr — Category C ("Expr | str meaning column name") is unusual in functions.py and should remain so in functions.spark.

How to Identify Candidates

The user may specify a scope via $ARGUMENTS. If no scope is given or "all" is specified, audit all functions in python/datafusion/functions/__init__.py and python/datafusion/functions/spark.py. When updating a spark-namespace function, apply the compatibility rules from "Scope" above on top of the standard analysis.

For each function, determine if any parameter can accept native Python types by evaluating two complementary signals:

Signal 1: Contextual Understanding

Some arguments are contextually always or almost always literal values based on what the function does:

Signal 2: Upstream Rust Implementation

Check the Rust binding in crates/core/src/functions.rs and the upstream DataFusion function implementation to determine type constraints. The upstream source is cached locally at:

~/.cargo/registry/src/index.crates.io-*/datafusion-functions-<VERSION>/src/

Check the DataFusion version in crates/core/Cargo.toml to find the right directory. Key subdirectories: string/, datetime/, math/, regex/.

For aggregate functions, the upstream source is in a separate crate:

~/.cargo/registry/src/index.crates.io-*/datafusion-functions-aggregate-<VERSION>/src/

There are five concrete techniques to check, in order of signal strength:

Technique 1: Check invoke_with_args() for literal-only enforcement (strongest signal)

Some functions pattern-match on ColumnarValue::Scalar in their invoke_with_args() method and return an error if the argument is a column/array. This means the argument must be a literal — passing a column expression will fail at runtime.

Example from date_trunc.rs:

let granularity_str = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) = granularity {
    v.to_lowercase()
} else {
    return exec_err!("Granularity of `date_trunc` must be non-null scalar Utf8");
};

If you find this pattern: The argument is Category B — accept only the corresponding native Python type (e.g., str), not Expr. The function will error at runtime with a column expression anyway.

Technique 1a: Check accumulator() for literal-only enforcement (aggregate functions)

Technique 1 applies to scalar UDFs. Aggregate functions do not have invoke_with_args() — instead, they enforce literal-only arguments in their accumulator() (or create_accumulator()) method, which runs at planning time before any data is processed.

Look for these patterns inside accumulator():

• get_scalar_value(expr) — evaluates the expression against an empty batch and errors if it's not a scalar
• validate_percentile_expr(expr) — specific helper used by percentile functions
• downcast_ref::<Literal>() — checks that the physical expression is a literal constant

Example from approx_percentile_cont.rs:

fn accumulator(&self, args: AccumulatorArgs) -> Result<ApproxPercentileAccumulator> {
    let percentile =
        validate_percentile_expr(&args.exprs[1], "APPROX_PERCENTILE_CONT")?;
    // ...
}

Where validate_percentile_expr calls get_scalar_value and errors with "must be a literal".

Example from string_agg.rs:

fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
    let Some(lit) = acc_args.exprs[1].as_any().downcast_ref::<Literal>() else {
        return not_impl_err!(
            "The second argument of the string_agg function must be a string literal"
        );
    };
    // ...
}

If you find this pattern: The argument is Category B — accept only the corresponding native Python type, not Expr. The function will error at planning time with a non-literal expression.

To discover which aggregate functions have literal-only arguments, search the upstream aggregate crate for get_scalar_value, validate_percentile_expr, and downcast_ref::<Literal>() inside accumulator() methods. For example, you should expect to find approx_percentile_cont (percentile) and string_agg (delimiter) among the results.

Technique 1b: Check partition_evaluator() for literal-only enforcement (window functions)

Window functions do not have invoke_with_args() or accumulator(). Instead, they enforce literal-only arguments in their partition_evaluator() method, which constructs the evaluator that processes each partition.

The upstream source is in a separate crate:

~/.cargo/registry/src/index.crates.io-*/datafusion-functions-window-<VERSION>/src/

Look for get_scalar_value_from_args() calls inside partition_evaluator(). This helper (defined in the window crate's utils.rs) calls downcast_ref::<Literal>() and errors with "There is only support Literal types for field at idx: {index} in Window Function".

Example from ntile.rs:

fn partition_evaluator(
    &self,
    partition_evaluator_args: PartitionEvaluatorArgs,
) -> Result<Box<dyn PartitionEvaluator>> {
    let scalar_n =
        get_scalar_value_from_args(partition_evaluator_args.input_exprs(), 0)?
            .ok_or_else(|| {
                exec_datafusion_err!("NTILE requires a positive integer")
            })?;
    // ...
}

If you find this pattern: The argument is Category B — accept only the corresponding native Python type, not Expr. The function will error at planning time with a non-literal expression.

To discover which window functions have literal-only arguments, search the upstream window crate for get_scalar_value_from_args inside partition_evaluator() methods. For example, you should expect to find ntile (n) and lead/lag (offset, default_value) among the results.

Technique 2: Check the Signature for data type constraints

Each function defines a Signature::coercible(...) that specifies what data types each argument accepts, using Coercion entries. This tells you the expected data type even if it doesn't enforce literal-only.

Example from repeat.rs:

signature: Signature::coercible(
    vec![
        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
        Coercion::new_implicit(
            TypeSignatureClass::Native(logical_int64()),
            vec![TypeSignatureClass::Integer],
            NativeType::Int64,
        ),
    ],
    Volatility::Immutable,
),

This tells you arg 2 (n) must be an integer type coerced to Int64. Use this to choose the correct Python type (e.g., int not str or float).

Common mappings:

Important: In Python's type system (PEP 484), float already accepts int values, so int | float is redundant and will fail the ruff linter (rule PYI041). Use float alone when the Rust side accepts a float/numeric type — Python users can still pass integer literals like log(10, col("a")) or power(col("a"), 3) without issue. Only use int when the Rust side strictly requires an integer (e.g., logical_int64()).

Technique 3: Check return_field_from_args() for scalar_arguments usage

Functions that inspect literal values at query planning time use args.scalar_arguments.get(n) in their return_field_from_args() method. This indicates the argument is expected to be a literal for optimal behavior (e.g., to determine output type precision), but may still work as a column.

Example from round.rs:

let decimal_places: Option<i32> = match args.scalar_arguments.get(1) {
    None => Some(0),
    Some(None) => None,        // argument is not a literal (column)
    Some(Some(scalar)) if scalar.is_null() => Some(0),
    Some(Some(scalar)) => Some(decimal_places_from_scalar(scalar)?),
};

If you find this pattern: The argument is Category A — accept native types AND Expr. It works as a column but is primarily used as a literal.

Decision flow