urosstan-db opened a new pull request, #59590:
URL: https://github.com/apache/spark/pull/59590
### What changes were proposed in this pull request?
- Allow `vector_avg` to accept `ARRAY<DOUBLE>` as well as `ARRAY<FLOAT>`,
returning an array with the same element type as its input.
- Follow-up to SPARK-60245 (#59466), which added the same support to
`vector_sum` through the `VectorAggregateKernel` abstraction: this PR adds a
`DoubleAvg` kernel next to the existing `FloatAvg` one, using the same
running-average update (`old_avg + (x - old_avg) / n`) and count-weighted
merge, computed in double precision.
- Update the SQL `ExpressionDescription`, Scala `functions.vector_avg`, and
PySpark `vector_avg` documentation (`versionchanged:: 4.4.0`).
### Why are the changes needed?
- Double-valued vectors currently have to be cast to floats before
averaging, which loses precision and limits the representable range (e.g.
averaging `16777216.0` and `16777218.0` yields `16777216.0` in float).
- As called out in #59466, `vector_sum` and `vector_avg` are used for
mathematical finance workloads that need double precision; the other vector
functions are embedding-oriented and stay float-only.
### Does this PR introduce _any_ user-facing change?
Yes. `vector_avg` now accepts double arrays, which previously raised
`DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE`, and averages them in double
precision:
```sql
SELECT vector_avg(v)
FROM VALUES (array(16777216.0D, 0.1D)), (array(16777218.0D, 0.2D)) AS t(v);
-- [1.6777217E7,0.15000000000000002] (ARRAY<DOUBLE>)
-- The same query with FLOAT inputs returns [1.6777216E7,0.15]
```
- Float inputs keep their existing result type and behavior.
- The type-mismatch error for unsupported inputs now reports `"ARRAY<FLOAT>"
or "ARRAY<DOUBLE>"` as the required type.
### How was this patch tested?
- `VectorAggSuite`: added double `VectorAvg` coverage for result type,
unsupported element types, multiple vectors, precision beyond float, merging
buffers with different counts, empty buffers, null vectors / null elements,
empty vectors, special values (NaN, -Infinity, `Double.MaxValue` without
overflow), dimension mismatches, and result independence from later buffer
updates. `vector_avg` was removed from the "rejects double vectors" grid test.
- `vector-agg.sql`: added double `vector_avg` queries (basic, precision,
GROUP BY, NULL handling, all-NULL, empty, window, dimension mismatch);
regenerated and reviewed the golden files.
- `MiscFunctionsSuite`: added a DataFrame double-precision test and a
`("double", "vector_avg")` case to the observe grid test.
- `test_functions.py`: added `test_vector_avg_double`.
Validated locally with Java 17:
- `build/sbt 'catalyst/testOnly *VectorAggSuite'`: 70 tests passed.
- `SPARK_GENERATE_GOLDEN_FILES=1 build/sbt 'sql/testOnly *SQLQueryTestSuite
-- -z vector-agg.sql'`, then reran without regeneration: passed.
- `build/sbt 'sql/testOnly *MiscFunctionsSuite' 'sql/testOnly
*ExpressionInfoSuite *ExpressionsSchemaSuite'`: passed.
- `python/run-tests --testnames "pyspark.sql.tests.test_functions
FunctionsTests.test_vector_avg_double"`: passed.
- `scalastyle` for `catalyst` (Compile/Test), `sql-api`, and `sql` (Test): 0
errors; `ruff format --check` / `ruff check` on the changed Python files:
passed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
This pull request and its description were written by Isaac.
--
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]
