Conversation
WalkthroughThe changes integrate evaluation mode handling into the AVG aggregation function across the entire stack. The protobuf definition for AVG is updated, replacing a boolean ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
| // cast to the result data type of AVG if the result data type is different | ||
| // from the input type, e.g. AVG(Int32). We should not expect a cast | ||
| // failure since it should have already been checked at Spark side. | ||
| // For all other numeric types (Int8/16/32/64, Float32/64): |
There was a problem hiding this comment.
Bug: Decimal AVG ignores eval_mode parameter
The eval_mode is extracted from the protobuf at line 1896 but not passed to AvgDecimal::new for Decimal128 types. This means decimal averages won't respect ANSI mode settings for overflow handling, while non-decimal averages correctly receive the eval_mode parameter. The AvgDecimal implementation needs to be updated to accept and use eval_mode for consistent behavior across all numeric types.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Bugbot AI reviewer is correct that the eval_mode is ignored for the AvgDecimal implementation. https://github.com/apache/spark/blob/211dd995b221f135340375159672dcb77ef90ef4/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala#L113 shows that it is used in the Spark implementation. Prevents wrong behavior in DataFusion Comet compared to Spark.
Code Review: Support ANSI mode avg exprSummaryThis PR adds support for ANSI evaluation mode to the AVG aggregate expression. The changes look good overall and follow the established patterns in the codebase. ✅ Strengths
🔍 Key Observations & Questions1. Eval mode is stored but not used in Float64 pathIn native/spark-expr/src/agg_funcs/avg.rs lines 143, 225: eval_mode field is marked with #[allow(dead_code)] Question: Is INFINITY handling different between ANSI and LEGACY modes? Should ANSI mode throw an error on overflow instead of returning INFINITY? If the behavior should differ, the eval_mode needs to be used in the accumulation logic. 2. Decimal AVG doesn't receive eval_modeIn native/core/src/execution/planner.rs:1899-1901, the decimal path doesn't pass eval_mode to AvgDecimal::new(), while the Float64 path does. Question: Should decimal AVG also support eval_mode for consistency? If decimals can overflow during averaging, ANSI mode should handle this differently than LEGACY mode. 🐛 Potential Issues1. Overflow behavior not implementedThe PR removes the restriction on ANSI mode (getSupportLevel check removed from aggregates.scala:153-162) but doesn't implement ANSI-specific overflow checking in the Rust code. Recommendation: Either add overflow checking in ANSI mode, document that Float64 AVG treats overflow as INFINITY in all modes, or add integration tests specifically for overflow cases. 2. Missing test for overflow/error casesCurrent tests don't cover what happens with extreme values, error propagation in ANSI mode vs NULL in LEGACY mode. 📝 Code Quality Suggestions
🔒 Security ConsiderationsNo security concerns identified. The changes don't introduce new attack vectors. ⚡ Performance ConsiderationsNo performance impact expected - the eval_mode parameter is stored but doesn't add runtime branches. 📋 Recommendations SummaryBefore merging:
Nice to have:
🎯 Overall AssessmentThis is a solid PR that makes good progress on ANSI mode support. The main concern is ensuring that the removal of the ANSI mode restriction is backed by proper behavior implementation or clear documentation of current limitations. Great work on following the project's conventions and adding comprehensive tests! 🚀 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala (1)
1474-1508: New AVG/try_avg test covers basic modes; consider adding an error-path case laterThis exercise of
avg/try_avgwith and without GROUP BY under bothANSI_ENABLEDsettings is a good sanity check that eval_mode is plumbed correctly and that Comet still executes. Once stricter ANSI/TRY semantics for AVG are implemented, it may be worth adding a follow-up test that triggers an actual eval-mode-sensitive condition (e.g., overflow) to distinguish the behaviors.native/core/src/execution/planner.rs (1)
1892-1924: Planner AVG path correctly rebinds to Float64 and carries eval_modeThe
AggExprStruct::Avgbranch now (a) readseval_modefrom the proto, (b) casts non-decimal inputs toFloat64, and (c) constructsAvg::new("avg", DataType::Float64, eval_mode). This keeps the partial/final state types coherent ((Float64 sum, Int64 count) for non-decimal AVG) while threading the evaluation mode into the UDF, and leaves the decimalAvgDecimalpath untouched.native/spark-expr/src/agg_funcs/avg.rs (1)
18-24: EvalMode is correctly threaded through AVG UDF and accumulators (currently stored only)
Avgnow accepts and stores anEvalMode, and both the scalarAvgAccumulatorand groupedAvgGroupsAccumulatorare constructed with that mode, even though it’s not yet used in their logic. The Float64-based accumulation, state layout (sum as Float64, count as Int64), and null handling viavalues()+is_nullremain unchanged and correct, so behavior matches the prior implementation while giving you a clean hook to implement ANSI/TRY-specific semantics later. Based on learnings, the array access pattern also correctly avoids reading values from null slots.Also applies to: 39-137, 139-245, 251-347
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
native/core/src/execution/planner.rs(1 hunks)native/proto/src/proto/expr.proto(1 hunks)native/spark-expr/src/agg_funcs/avg.rs(10 hunks)spark/src/main/scala/org/apache/comet/serde/aggregates.scala(2 hunks)spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala(1 hunks)
🧰 Additional context used
🧠 Learnings (2)
📚 Learning: 2025-11-11T20:44:05.014Z
Learnt from: martin-augment
Repo: martin-augment/datafusion-comet PR: 19
File: native/spark-expr/src/array_funcs/array_insert.rs:221-238
Timestamp: 2025-11-11T20:44:05.014Z
Learning: In Rust code using Apache Arrow arrays, always check `is_null(index)` before calling `value(index)` on `PrimitiveArray` types (such as `Int32Array`, `Int64Array`, etc.), because `value()` does not check for nulls and returns arbitrary values for null slots. This applies to functions in `native/spark-expr/src` that process Arrow arrays.
Applied to files:
native/spark-expr/src/agg_funcs/avg.rs
📚 Learning: 2025-11-04T14:26:48.750Z
Learnt from: martin-augment
Repo: martin-augment/datafusion-comet PR: 7
File: native/spark-expr/src/math_funcs/abs.rs:201-302
Timestamp: 2025-11-04T14:26:48.750Z
Learning: In the abs function in native/spark-expr/src/math_funcs/abs.rs (Rust), NULL values for signed integers (Int8, Int16, Int32, Int64) and decimals (Decimal128, Decimal256) should return the argument as-is (e.g., ColumnarValue::Scalar(ScalarValue::Int8(None))) rather than panicking on unwrap().
Applied to files:
native/spark-expr/src/agg_funcs/avg.rs
🧬 Code graph analysis (2)
native/core/src/execution/planner.rs (2)
native/spark-expr/src/agg_funcs/avg.rs (2)
new(50-60)new(148-154)native/spark-expr/src/conversion_funcs/cast.rs (2)
new(791-801)new(827-837)
spark/src/main/scala/org/apache/comet/serde/aggregates.scala (1)
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala (2)
QueryPlanSerde(45-717)evalModeToProto(439-446)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Cursor Bugbot
- GitHub Check: claude-review
🔇 Additional comments (2)
spark/src/main/scala/org/apache/comet/serde/aggregates.scala (1)
32-34: AVG serde now correctly threads eval_mode into the protoThe new imports and
builder.setEvalMode(evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(avg.evalMode)))cleanly propagate Spark’sevalModeforAVG/try_avginto the wire format, aligning with how other expressions handleEvalMode. No further serde changes seem necessary here.Also applies to: 181-187
native/proto/src/proto/expr.proto (1)
136-141: AVG proto now uses EvalMode at tag 4, consistent with other expressionsSwitching
Avg’s field 4 frombool fail_on_errortoEvalMode eval_modebrings AVG in line withMathExprandCast, matching the new Scala serde and Rust planner paths that pass an explicit evaluation mode. This is a wire-incompatible change across versions, but is fine as long as all Comet components are upgraded together, as in this repo.Also applies to: 206-210, 212-224
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| let builder = match datatype { | ||
| DataType::Decimal128(_, _) => { | ||
| let func = | ||
| AggregateUDF::new_from_impl(AvgDecimal::new(datatype, input_datatype)); |
There was a problem hiding this comment.
eval_mode is computed for AVG but not used in the Decimal branch (AvgDecimal::new), so ANSI/TRY mode won’t affect Decimal AVG (e.g., overflow behavior). Consider propagating eval_mode to the decimal implementation to ensure semantics match Spark’s modes (also applies given getSupportLevel removal in Scala).
🤖 Was this useful? React with 👍 or 👎
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct that the eval_mode is ignored for the AvgDecimal implementation. https://github.com/apache/spark/blob/211dd995b221f135340375159672dcb77ef90ef4/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala#L113 shows that it is used in the Spark implementation. Prevents wrong behavior in DataFusion Comet compared to Spark.
value:useful; category:bug; feedback: The Claude AI reviewer is correct that the eval_mode is ignored for the AvgDecimal implementation. https://github.com/apache/spark/blob/211dd995b221f135340375159672dcb77ef90ef4/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala#L113 shows that it is used in the Spark implementation. Prevents wrong behavior in DataFusion Comet compared to Spark. |
2 participants
2817: To review by AI