fix: make vector aggregates work with GROUP BY and partial aggregation (#9338)

* fix: make vector aggregates work with GROUP BY and partial aggregation

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* test: make partitioned vec_avg case distinguish weighted averages

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
dennis zhuang
2026-09-24 08:45:40 +00:00
committed by GitHub
parent 822042c7bb
commit 357c935804
5 changed files with 321 additions and 72 deletions
@@ -150,3 +150,48 @@ FROM (
UNION ALL
SELECT '[4.0, 5.0, 6.0, 10, -8, 100]' AS v
) ORDER BY v;
SELECT h, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v))
FROM (
SELECT 'a' AS h, '[1.0, 1.0]' AS v
UNION ALL
SELECT 'a' AS h, '[2.0, 2.0]' AS v
UNION ALL
SELECT 'b' AS h, '[3.0, 3.0]' AS v
) GROUP BY h ORDER BY h;
-- On partitioned tables the aggregates are split into partial state and merge.
-- Rows only land in two of the three partitions, the third one has an empty state.
-- The two non-empty partitions differ in row count and mean, so vec_avg must weight by count.
CREATE TABLE vector_aggr_partitioned (
ts TIMESTAMP TIME INDEX,
k INT,
g STRING,
v VECTOR(2),
PRIMARY KEY(k)
)
PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20);
INSERT INTO vector_aggr_partitioned VALUES
(1000, 1, 'a', '[1.0, 1.0]'),
(2000, 11, 'a', '[8.0, 8.0]'),
(3000, 2, 'a', '[3.0, 3.0]'),
(3500, 3, 'a', '[5.0, 5.0]'),
(4000, 12, 'b', '[4.0, 4.0]');
SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v))
FROM vector_aggr_partitioned;
SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v))
FROM vector_aggr_partitioned GROUP BY g ORDER BY g;
-- A NULL vector makes vec_sum and vec_product NULL, vec_avg skips it.
INSERT INTO vector_aggr_partitioned VALUES (5000, 21, 'b', NULL);
SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v))
FROM vector_aggr_partitioned;
SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v))
FROM vector_aggr_partitioned GROUP BY g ORDER BY g;
DROP TABLE vector_aggr_partitioned;