make-pythonic

<!--- Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. --> # 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: | Context | Typical Arguments | Examples | |---------|------------------|----------| | **String position/count** | Character counts, indices, repetition counts | `left(str, n)`, `right(str, n)`, `repeat(str, n)`, `lpad(str, count, ...)` | | **Delimiters/separators** | Fixed separator characters | `split_part(str, delim, idx)`, `concat_ws(sep, ...)` | | **Search/replace patterns** | Literal search strings, replacements | `replace(str, from, to)`, `regexp_replace(str, pattern, replacement, flags)` | | **Date/time parts** | Part names from a fixed set | `date_part(part, date)`, `date_trunc(part, date)` | | **Rounding precision** | Decimal place counts | `round(val, places)`, `trunc(val, places)` | | **Fill characters** | Padding characters | `lpad(str, count, fill)`, `rpad(str, count, fill)` | ### 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`: ```rust 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`: ```rust 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`: ```rust 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`: ```rust 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`: ```rust 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: | Rust Type Constraint | Python Type | |---------------------|-------------| | `logical_int64()` / `TypeSignatureClass::Integer` | `int` | | `logical_float64()` / `TypeSignatureClass::Numeric` | `int \| float` | | `logical_string()` / `TypeSignatureClass::String` | `str` | | `LogicalType::Boolean` | `bool` | **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`: ```rust 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 ``` What kind of function is this? Scalar UDF: Is argument rejected at runtime if not a literal? (check invoke_with_args for ColumnarValue::Scalar-only match + exec_err!) → YES: Category B — accept only native type, no Expr → NO: continue below Aggregate: Is argument rejected at planning time if not a literal? (check accumulator() for get_scalar_value / validate_percentile_expr / downcast_ref::<Literal>() + error) → YES: Category B — accept only native type, no Expr → NO: continue below Window: Is argument rejected at planning time if not a literal? (check partition_evaluator() for get_scalar_value_from_args / downcast_ref::<Literal>() + error) → YES: Category B — accept only native type, no Expr → NO: continue below Does the Signature constrain it to a specific data type? → YES: Category A — accept Expr | <native type matching the constraint> → NO: Leave as Expr only ``` ## Coercion Categories When making a function more pythonic, apply the correct coercion pattern based on **what the argument represents**: ### Category A: Arguments That Should Accept Native Types AND Expr These are arguments that are *typically* literals but *could* be column references in advanced use cases. For these, accept a union type and coerce native types to `Expr.literal()`. **Type hint pattern:** `Expr | int`, `Expr | str`, `Expr | int | str`, etc. **When to use:** When the argument could plausibly come from a column in some use case (e.g., the repeat count might come from a column in a data-driven scenario). ```python def repeat(string: Expr, n: Expr | int) -> Expr: """Repeats the ``string`` to ``n`` times. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["ha"]}) >>> result = df.select(