feat(functions): add try_divide for PySpark parity#7210
Conversation
Greptile SummaryThis PR adds a
Confidence Score: 5/5This 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
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"]
%%{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"]
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 |
There was a problem hiding this comment.
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!
|
|
||
|
|
||
| def log(expr: Expression, base: int | float = math.e) -> Expression: | ||
| def log(expr: Expression, base: float = math.e) -> Expression: |
There was a problem hiding this comment.
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!
There was a problem hiding this comment.
nah these come from the repo's own lint config son
|
@greptile re-review |
Summary
Adds a
try_dividescalar function: true division that returns null where the divisor is 0, instead of infinity, NaN, or an error. This mirrors Spark'stry_dividesemantics, 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/0givesinf,0.0/0.0givesNaN), 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 thepmod(#6801) pattern: a RustScalarUDFinsrc/daft-functions/src/numeric/with a thin Python wrapper.Changes Made
src/daft-functions/src/numeric/try_divide.rs: newTryDivideScalarUDF. Computes the quotient with Daft's existingSeriesdivision, then nulls rows where the divisor equals zero via a broadcast comparison andif_else. Return type comes from the sameInferDataTypedivision inference used by the/operator (integer inputs produceFloat64).src/daft-functions/src/numeric/mod.rs: registerTryDivide.daft/functions/numeric.py:try_divide(dividend, divisor)wrapper.daft/functions/__init__.py: exporttry_divide.tests/recordbatch/numeric/test_numeric.py: 7 tests covering ints, floats (including a-0.0divisor), unsigned ints, output dtype, broadcasting in both directions, and the non-numeric error.Behavior
try_divide(1, 0)returns null;try_divide(6, 3)returns2.0(Float64, same inference as/).try_divide(1.0, 0.0),try_divide(-1.0, 0.0), andtry_divide(0.0, -0.0)all return null, matching Spark, where the regular operators would giveinf,-inf, andNaN. Daft's float equality is bitwise, so the implementation adds zero to the divisor before comparing (-0.0 + 0.0is+0.0in IEEE 754) to catch negative zero.TypeErrorat planning time.is_numericgatepmoduses (Daft'sis_numericexcludesDecimal128). Spark'stry_dividesupports decimals; this can be lifted in a follow-up if decimal division support is wanted here.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 passedcargo check -p daft-functions --libcleancargo fmt -p daft-functions --checkcleancargo clippy -p daft-functions --lib --no-depscleandaft.sqlRelated Issues
Part of #3793.