sunchao commented on code in PR #6761:
URL: https://github.com/apache/datafusion-comet/pull/6761#discussion_r4252785379
##########
native/spark-expr/src/math_funcs/greatest_least.rs:
##########
@@ -185,17 +240,18 @@ impl ScalarUDFImpl for SparkGreatestLeast {
// Spark gives every argument the same type up to nullability, but
Arrow also compares
// field names and nullability, which the row selection below needs to
match.
let data_type = args.return_field.data_type();
- let to_array = |arg: &ColumnarValue| -> Result<ArrayRef> {
+ let to_array = |arg: &ColumnarValue, rows: usize| -> Result<ArrayRef> {
let array = arg.to_array(rows)?;
Ok(if array.data_type() == data_type {
array
} else {
cast(&array, data_type)?
})
};
- let mut result = to_array(&args.args[0])?;
- for arg in &args.args[1..] {
- result = self.pick(&result, &to_array(arg)?)?;
+ let args = self.merge_constants(&args.args, data_type, |arg|
to_array(arg, 1))?;
+ let mut result = to_array(&args[0], rows)?;
Review Comment:
[P2] Apply the reduced non-float constant before folding column arrays.
`merge_constants` leaves a single trailing constant after preceding columns, so
this loop copies their wide intermediate result before applying a short winning
constant. For `SELECT greatest(array_repeat(CAST(id % 21 AS INT), 256),
array_repeat(CAST((id + 1) % 21 AS INT), 256), array(21)) FROM range(8192)`,
both routes correctly return `[21]`, but cumulative allocation increases from
268,252 bytes through base routing to 16,980,060 bytes at this head. DataFusion
previously applied the scalar first and avoided that intermediate. Mirroring
the signs and using `least` reproduces the increase. This affects newly routed,
already-correct non-float queries even with only one constant. Please seed
their fold with the reduced constant, preserving argument order for types with
floating-point leaves.
Evidence: Recompiled the exact-head sources against cached Arrow
59.3.0/DataFusion 55.1.0 and measured allocations with an instrumented
allocator. Inputs were two 8,192-row ListArray<Int32> columns containing 256
repetitions of `i % 21` and `(i + 1) % 21`, followed by
ScalarValue::List([21]). Base and head outputs were asserted equal. Base
allocated 268,252 bytes, head 16,980,060, and head with the scalar moved first
268,188. Both greatest and sign-mirrored least reproduced these values. Spark
3.5.9 retained the trailing literal in its optimized plan and returned the
expected singleton array. Reproduction:
`/tmp/review6761-verified-now-cycn4vtz/tests
allocation::trailing_constant_after_wide_columns --exact --nocapture
--test-threads=1`. Source and compile arguments are saved in the same
directory. This measures native cumulative allocation, not end-to-end execution
time.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
