Skip to content

feat(functions): add try_divide for PySpark parity#7210

Open
BABTUNA wants to merge 3 commits into
Eventual-Inc:mainfrom
BABTUNA:feat/functions-try-divide
Open

feat(functions): add try_divide for PySpark parity#7210
BABTUNA wants to merge 3 commits into
Eventual-Inc:mainfrom
BABTUNA:feat/functions-try-divide

Conversation

@BABTUNA

@BABTUNA BABTUNA commented Jul 5, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds a try_divide scalar function: true division that returns null where the divisor is 0, instead of infinity, NaN, or an error. This mirrors Spark's try_divide semantics, where a zero divisor yields null for every numeric type.

Why

Part of the PySpark function parity effort tracked in #3793. Daft's / follows IEEE 754 (1/0 gives inf, 0.0/0.0 gives NaN), and integer floor division by zero currently panics a worker thread, so a null-on-zero-divisor variant gives users a safe option that matches Spark. Follows the pmod (#6801) pattern: a Rust ScalarUDF in src/daft-functions/src/numeric/ with a thin Python wrapper.

Changes Made

  • src/daft-functions/src/numeric/try_divide.rs: new TryDivide ScalarUDF. Computes the quotient with Daft's existing Series division, then nulls rows where the divisor equals zero via a broadcast comparison and if_else. Return type comes from the same InferDataType division inference used by the / operator (integer inputs produce Float64).
  • src/daft-functions/src/numeric/mod.rs: register TryDivide.
  • daft/functions/numeric.py: try_divide(dividend, divisor) wrapper.
  • daft/functions/__init__.py: export try_divide.
  • tests/recordbatch/numeric/test_numeric.py: 7 tests covering ints, floats (including a -0.0 divisor), unsigned ints, output dtype, broadcasting in both directions, and the non-numeric error.

Behavior

  • try_divide(1, 0) returns null; try_divide(6, 3) returns 2.0 (Float64, same inference as /).
  • try_divide(1.0, 0.0), try_divide(-1.0, 0.0), and try_divide(0.0, -0.0) all return null, matching Spark, where the regular operators would give inf, -inf, and NaN. Daft's float equality is bitwise, so the implementation adds zero to the divisor before comparing (-0.0 + 0.0 is +0.0 in IEEE 754) to catch negative zero.
  • Nulls propagate; non-numeric inputs raise a TypeError at planning time.
  • Decimal inputs are rejected by the same is_numeric gate pmod uses (Daft's is_numeric excludes Decimal128). Spark's try_divide supports decimals; this can be lifted in a follow-up if decimal division support is wanted here.
  • Available in SQL through the function registry: SELECT try_divide(1, 0) returns null.

Test Plan

  • pytest tests/recordbatch/numeric/test_numeric.py: 392 passed (7 new)
  • pytest --doctest-modules daft/functions/numeric.py: 5 passed
  • cargo check -p daft-functions --lib clean
  • cargo fmt -p daft-functions --check clean
  • cargo clippy -p daft-functions --lib --no-deps clean
  • End-to-end check through the DataFrame API and daft.sql

Related Issues

Part of #3793.

@github-actions github-actions Bot added the feat label Jul 5, 2026
@greptile-apps

greptile-apps Bot commented Jul 5, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds a try_divide scalar function for PySpark parity that returns NULL wherever the divisor is zero, rather than inf, NaN, or a panic. It follows the established pmod pattern with a Rust ScalarUDF in src/daft-functions/src/numeric/, a thin Python wrapper in daft/functions/numeric.py, and export wiring in daft/functions/__init__.py.

  • Rust implementation (try_divide.rs): Computes dividend / divisor (which casts integer operands to Float64 via Daft's existing type-inference, avoiding the integer floor-division panic), then constructs a boolean mask by normalising the divisor with +0 (to collapse -0.0+0.0 per IEEE 754) before comparing to zero, and uses if_else to replace zero-divisor positions with NULL.
  • Broadcast handling: The post-division divisor broadcast is correct for all combinator cases (scalar/vector, vector/scalar, scalar/scalar), with the zero series broadcast to match the final divisor length.
  • Test coverage: Seven new tests covering integer pairs, float special values (-0.0), unsigned types, output dtype, scalar-broadcast in both directions, and the non-numeric TypeError.

Confidence Score: 5/5

This PR is safe to merge — it adds a new function with no changes to existing code paths, all tests pass, and the Rust implementation correctly handles integer, float, unsigned, null, and negative-zero divisor cases.

The implementation is well-reasoned: integer operands are handled via Daft's existing float-cast division (avoiding the integer floor-division panic), the -0.0 normalisation via IEEE 754 addition is correct, broadcast alignment is consistent across all scalar/vector combinator cases, and null propagation works naturally through Daft's Series arithmetic and if_else semantics. No regressions to existing functions are introduced.

No files require special attention.

Important Files Changed

Filename Overview
src/daft-functions/src/numeric/try_divide.rs New ScalarUDF: correctly delegates to Series division (float-cast for ints avoids panic), normalises -0.0 via IEEE 754 addition, handles both scalar-divisor and scalar-dividend broadcasting, and uses if_else for null-masking. Logic is sound.
daft/functions/numeric.py Adds try_divide Python wrapper; also narrows type annotations on pre-existing log/between parameters (flagged in a prior thread and accepted by the team). No new issues.
daft/functions/init.py Exports try_divide; alphabetically placed in all but inserted between conv and is_nan in the from-import block (the block was already non-alphabetical, so no regression).
src/daft-functions/src/numeric/mod.rs Registers TryDivide module and UDF; straightforward wiring consistent with other numeric functions.
tests/recordbatch/numeric/test_numeric.py Seven new tests covering int, float, unsigned int, output dtype, broadcasting in both directions, -0.0, and TypeError. Imports moved to top-level per prior review thread.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["try_divide(dividend, divisor)"] --> B["Validate numeric types via get_return_field\nInfer output dtype with InferDataType"]
    B --> C{Both inputs numeric?}
    C -- No --> D["Return DaftError::TypeError"]
    C -- Yes --> E["try_divide_impl: compute quotient = dividend / divisor\nInts cast to Float64 so 1/0 gives inf not panic"]
    E --> F{"divisor.len == 1\nAND quotient.len != 1?"}
    F -- Yes --> G["Broadcast divisor to quotient.len"]
    F -- No --> H["Keep divisor as-is"]
    G --> I["zero = Int64 zero cast to divisor dtype\nthen broadcast to divisor.len"]
    H --> I
    I --> J["normalized = divisor + zero\nIEEE 754: -0.0 + 0.0 becomes +0.0"]
    J --> K["is_zero = normalized.equal(zero)"]
    K --> L["nulls = Series of nulls matching quotient shape"]
    L --> M["result = nulls.if_else(quotient, is_zero)\nnull where is_zero is true else quotient"]
    M --> N["Return result"]
Loading
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["try_divide(dividend, divisor)"] --> B["Validate numeric types via get_return_field\nInfer output dtype with InferDataType"]
    B --> C{Both inputs numeric?}
    C -- No --> D["Return DaftError::TypeError"]
    C -- Yes --> E["try_divide_impl: compute quotient = dividend / divisor\nInts cast to Float64 so 1/0 gives inf not panic"]
    E --> F{"divisor.len == 1\nAND quotient.len != 1?"}
    F -- Yes --> G["Broadcast divisor to quotient.len"]
    F -- No --> H["Keep divisor as-is"]
    G --> I["zero = Int64 zero cast to divisor dtype\nthen broadcast to divisor.len"]
    H --> I
    I --> J["normalized = divisor + zero\nIEEE 754: -0.0 + 0.0 becomes +0.0"]
    J --> K["is_zero = normalized.equal(zero)"]
    K --> L["nulls = Series of nulls matching quotient shape"]
    L --> M["result = nulls.if_else(quotient, is_zero)\nnull where is_zero is true else quotient"]
    M --> N["Return result"]
Loading

Reviews (2): Last reviewed commit: "test(functions): move try_divide test im..." | Re-trigger Greptile



def test_try_divide_int() -> None:
from daft.functions import try_divide

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P2 Inline imports inside test functions

All 7 new test functions import try_divide (and in one case DataType) inline inside the function body rather than at the top of the file — for example from daft.functions import try_divide at lines 1066, 1076, 1084, 1093, 1102, 1111, and 1120. The custom rule says imports must be placed at the top of the file. Moving them to the existing from daft import col, lit section at the top would bring these tests in line with that requirement.

Rule Used: Import statements should be placed at the top of t... (source)

Learned From
Eventual-Inc/Daft#5078

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

done

Comment thread daft/functions/numeric.py


def log(expr: Expression, base: int | float = math.e) -> Expression:
def log(expr: Expression, base: float = math.e) -> Expression:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P2 Unrelated type-annotation changes to existing functions

This PR also narrows the type hints on two pre-existing functions: log's base parameter changes from int | float to float, and both lower/upper parameters of between drop int from their unions. These changes aren't mentioned in the PR description. While PEP 484 specifies that float implicitly accepts int, removing int from an explicit union can still surprise users relying on strict type-checking tools. It would be helpful to either document the intent or keep these changes in a separate cleanup PR.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

nah these come from the repo's own lint config son

@BABTUNA

BABTUNA commented Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

@greptile re-review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant