Skip to content

Commit 8ec63ae

Browse files
nielspardonnielspardon
committed
fix(extensions)!: change distribution option to enum arg for std_dev and variance
Signed-off-by: Niels Pardon <par@zurich.ibm.com>
1 parent fcd29c6 commit 8ec63ae

4 files changed

Lines changed: 143 additions & 10 deletions

File tree

‎extensions/functions_arithmetic.yaml‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1395,21 +1395,21 @@ aggregate_functions:
13951395
- args:
13961396
- name: x
13971397
value: fp32
1398+
- name: distribution
1399+
options: [ SAMPLE, POPULATION]
13981400
options:
13991401
rounding:
14001402
values: [ TIE_TO_EVEN, TIE_AWAY_FROM_ZERO, TRUNCATE, CEILING, FLOOR ]
1401-
distribution:
1402-
values: [ SAMPLE, POPULATION]
14031403
nullability: DECLARED_OUTPUT
14041404
return: fp32?
14051405
- args:
14061406
- name: x
14071407
value: fp64
1408+
- name: distribution
1409+
options: [ SAMPLE, POPULATION]
14081410
options:
14091411
rounding:
14101412
values: [ TIE_TO_EVEN, TIE_AWAY_FROM_ZERO, TRUNCATE, CEILING, FLOOR ]
1411-
distribution:
1412-
values: [ SAMPLE, POPULATION]
14131413
nullability: DECLARED_OUTPUT
14141414
return: fp64?
14151415
- name: "variance"
@@ -1418,21 +1418,21 @@ aggregate_functions:
14181418
- args:
14191419
- name: x
14201420
value: fp32
1421+
- name: distribution
1422+
options: [ SAMPLE, POPULATION]
14211423
options:
14221424
rounding:
14231425
values: [ TIE_TO_EVEN, TIE_AWAY_FROM_ZERO, TRUNCATE, CEILING, FLOOR ]
1424-
distribution:
1425-
values: [ SAMPLE, POPULATION]
14261426
nullability: DECLARED_OUTPUT
14271427
return: fp32?
14281428
- args:
14291429
- name: x
14301430
value: fp64
1431+
- name: distribution
1432+
options: [ SAMPLE, POPULATION]
14311433
options:
14321434
rounding:
14331435
values: [ TIE_TO_EVEN, TIE_AWAY_FROM_ZERO, TRUNCATE, CEILING, FLOOR ]
1434-
distribution:
1435-
values: [ SAMPLE, POPULATION]
14361436
nullability: DECLARED_OUTPUT
14371437
return: fp64?
14381438
- name: "corr"

‎tests/cases/arithmetic/std_dev.test‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
### SUBSTRAIT_AGGREGATE_TEST: v1.0
2+
### SUBSTRAIT_INCLUDE: '/extensions/functions_arithmetic.yaml'
3+
4+
# basic: Basic examples without any special cases
5+
((1.0, 2.0, 3.0, 4.0, 5.0)) std_dev(col0::fp32, SAMPLE::enum) = 1.5811388::fp32
6+
((1.0, 2.0, 3.0, 4.0, 5.0)) std_dev(col0::fp64, SAMPLE::enum) = 1.5811388300841898::fp64
7+
((1.0, 2.0, 3.0, 4.0, 5.0)) std_dev(col0::fp32, POPULATION::enum) = 1.4142135::fp32
8+
((1.0, 2.0, 3.0, 4.0, 5.0)) std_dev(col0::fp64, POPULATION::enum) = 1.4142135623730951::fp64
9+
10+
# uniform_values: Standard deviation of uniform values
11+
((5.0, 5.0, 5.0, 5.0)) std_dev(col0::fp32, SAMPLE::enum) = 0.0::fp32
12+
((5.0, 5.0, 5.0, 5.0)) std_dev(col0::fp64, POPULATION::enum) = 0.0::fp64
13+
14+
# single_value: Standard deviation with single value
15+
((42.0)) std_dev(col0::fp32, SAMPLE::enum) = Null::fp32
16+
((42.0)) std_dev(col0::fp64, POPULATION::enum) = 0.0::fp64
17+
18+
# negative_values: Standard deviation with negative values
19+
((-5.0, -3.0, -1.0, 1.0, 3.0, 5.0)) std_dev(col0::fp32, SAMPLE::enum) = 3.8944404::fp32
20+
((-5.0, -3.0, -1.0, 1.0, 3.0, 5.0)) std_dev(col0::fp64, SAMPLE::enum) = 3.8944404818493075::fp64
21+
((-10.0, -5.0, 0.0, 5.0, 10.0)) std_dev(col0::fp32, POPULATION::enum) = 7.0710678::fp32
22+
((-10.0, -5.0, 0.0, 5.0, 10.0)) std_dev(col0::fp64, POPULATION::enum) = 7.0710678118654755::fp64
23+
24+
# decimal_precision: Standard deviation with decimal values
25+
((1.5, 2.5, 3.5, 4.5, 5.5)) std_dev(col0::fp32, SAMPLE::enum) = 1.5811388::fp32
26+
((1.5, 2.5, 3.5, 4.5, 5.5)) std_dev(col0::fp64, SAMPLE::enum) = 1.5811388300841898::fp64
27+
((0.1, 0.2, 0.3, 0.4, 0.5)) std_dev(col0::fp64, POPULATION::enum) = 0.14142135623730953::fp64
28+
29+
# large_values: Standard deviation with large values
30+
((1000.0, 2000.0, 3000.0, 4000.0, 5000.0)) std_dev(col0::fp32, SAMPLE::enum) = 1581.1388::fp32
31+
((1000.0, 2000.0, 3000.0, 4000.0, 5000.0)) std_dev(col0::fp64, SAMPLE::enum) = 1581.1388300841898::fp64
32+
33+
# small_values: Standard deviation with small values
34+
((0.001, 0.002, 0.003, 0.004, 0.005)) std_dev(col0::fp64, SAMPLE::enum) = 0.0015811388300841896::fp64
35+
((0.001, 0.002, 0.003, 0.004, 0.005)) std_dev(col0::fp64, POPULATION::enum) = 0.0014142135623730951::fp64
36+
37+
# null_handling: Examples with null as input or output
38+
((Null, Null, Null)) std_dev(col0::fp32, SAMPLE::enum) = Null::fp32
39+
(()) std_dev(col0::fp32, SAMPLE::enum) = Null::fp32
40+
((1.0, Null, 3.0, Null, 5.0)) std_dev(col0::fp32, SAMPLE::enum) = 2.0::fp32
41+
((1.0, Null, 3.0, Null, 5.0)) std_dev(col0::fp64, POPULATION::enum) = 1.632993161855452::fp64
42+
43+
# rounding: Examples with different rounding modes
44+
((1.1, 2.2, 3.3, 4.4, 5.5)) std_dev(col0::fp32, SAMPLE::enum) [rounding:TIE_TO_EVEN] = 1.7406897::fp32
45+
((1.1, 2.2, 3.3, 4.4, 5.5)) std_dev(col0::fp64, SAMPLE::enum) [rounding:TRUNCATE] = 1.7406897166664838::fp64
46+
47+
# two_values: Standard deviation with two values
48+
((10.0, 20.0)) std_dev(col0::fp32, SAMPLE::enum) = 7.071068::fp32
49+
((10.0, 20.0)) std_dev(col0::fp64, POPULATION::enum) = 5.0::fp64
50+
51+
# mixed_range: Standard deviation with mixed range values
52+
((0.0, 100.0, 50.0, 25.0, 75.0)) std_dev(col0::fp32, SAMPLE::enum) = 41.010193::fp32
53+
((0.0, 100.0, 50.0, 25.0, 75.0)) std_dev(col0::fp64, POPULATION::enum) = 36.66060555964672::fp64

‎tests/cases/arithmetic/variance.test‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
### SUBSTRAIT_AGGREGATE_TEST: v1.0
2+
### SUBSTRAIT_INCLUDE: '/extensions/functions_arithmetic.yaml'
3+
4+
# basic: Basic examples without any special cases
5+
((1.0, 2.0, 3.0, 4.0, 5.0)) variance(col0::fp32, SAMPLE::enum) = 2.5::fp32
6+
((1.0, 2.0, 3.0, 4.0, 5.0)) variance(col0::fp64, SAMPLE::enum) = 2.5::fp64
7+
((1.0, 2.0, 3.0, 4.0, 5.0)) variance(col0::fp32, POPULATION::enum) = 2.0::fp32
8+
((1.0, 2.0, 3.0, 4.0, 5.0)) variance(col0::fp64, POPULATION::enum) = 2.0::fp64
9+
10+
# uniform_values: Variance of uniform values
11+
((5.0, 5.0, 5.0, 5.0)) variance(col0::fp32, SAMPLE::enum) = 0.0::fp32
12+
((5.0, 5.0, 5.0, 5.0)) variance(col0::fp64, POPULATION::enum) = 0.0::fp64
13+
14+
# single_value: Variance with single value
15+
((42.0)) variance(col0::fp32, SAMPLE::enum) = Null::fp32
16+
((42.0)) variance(col0::fp64, POPULATION::enum) = 0.0::fp64
17+
18+
# negative_values: Variance with negative values
19+
((-5.0, -3.0, -1.0, 1.0, 3.0, 5.0)) variance(col0::fp32, SAMPLE::enum) = 15.166667::fp32
20+
((-5.0, -3.0, -1.0, 1.0, 3.0, 5.0)) variance(col0::fp64, SAMPLE::enum) = 15.166666666666666::fp64
21+
((-10.0, -5.0, 0.0, 5.0, 10.0)) variance(col0::fp32, POPULATION::enum) = 50.0::fp32
22+
((-10.0, -5.0, 0.0, 5.0, 10.0)) variance(col0::fp64, POPULATION::enum) = 50.0::fp64
23+
24+
# decimal_precision: Variance with decimal values
25+
((1.5, 2.5, 3.5, 4.5, 5.5)) variance(col0::fp32, SAMPLE::enum) = 2.5::fp32
26+
((1.5, 2.5, 3.5, 4.5, 5.5)) variance(col0::fp64, SAMPLE::enum) = 2.5::fp64
27+
((0.1, 0.2, 0.3, 0.4, 0.5)) variance(col0::fp64, POPULATION::enum) = 0.020000000000000004::fp64
28+
29+
# large_values: Variance with large values
30+
((1000.0, 2000.0, 3000.0, 4000.0, 5000.0)) variance(col0::fp32, SAMPLE::enum) = 2500000.0::fp32
31+
((1000.0, 2000.0, 3000.0, 4000.0, 5000.0)) variance(col0::fp64, SAMPLE::enum) = 2500000.0::fp64
32+
33+
# small_values: Variance with small values
34+
((0.001, 0.002, 0.003, 0.004, 0.005)) variance(col0::fp64, SAMPLE::enum) = 0.0000025::fp64
35+
((0.001, 0.002, 0.003, 0.004, 0.005)) variance(col0::fp64, POPULATION::enum) = 0.000002::fp64
36+
37+
# null_handling: Examples with null as input or output
38+
((Null, Null, Null)) variance(col0::fp32, SAMPLE::enum) = Null::fp32
39+
(()) variance(col0::fp32, SAMPLE::enum) = Null::fp32
40+
((1.0, Null, 3.0, Null, 5.0)) variance(col0::fp32, SAMPLE::enum) = 4.0::fp32
41+
((1.0, Null, 3.0, Null, 5.0)) variance(col0::fp64, POPULATION::enum) = 2.666666666666667::fp64
42+
43+
# rounding: Examples with different rounding modes
44+
((1.1, 2.2, 3.3, 4.4, 5.5)) variance(col0::fp32, SAMPLE::enum) [rounding:TIE_TO_EVEN] = 3.03::fp32
45+
((1.1, 2.2, 3.3, 4.4, 5.5)) variance(col0::fp64, SAMPLE::enum) [rounding:TRUNCATE] = 3.0299999999999994::fp64
46+
47+
# two_values: Variance with two values
48+
((10.0, 20.0)) variance(col0::fp32, SAMPLE::enum) = 50.0::fp32
49+
((10.0, 20.0)) variance(col0::fp64, POPULATION::enum) = 25.0::fp64
50+
51+
# mixed_range: Variance with mixed range values
52+
((0.0, 100.0, 50.0, 25.0, 75.0)) variance(col0::fp32, SAMPLE::enum) = 1681.25::fp32
53+
((0.0, 100.0, 50.0, 25.0, 75.0)) variance(col0::fp64, POPULATION::enum) = 1345.0::fp64
54+
55+
# zero_mean: Variance with values around zero
56+
((-2.0, -1.0, 0.0, 1.0, 2.0)) variance(col0::fp32, SAMPLE::enum) = 2.5::fp32
57+
((-2.0, -1.0, 0.0, 1.0, 2.0)) variance(col0::fp64, POPULATION::enum) = 2.0::fp64
58+
59+
# three_values: Variance with three values
60+
((10.0, 20.0, 30.0)) variance(col0::fp32, SAMPLE::enum) = 100.0::fp32
61+
((10.0, 20.0, 30.0)) variance(col0::fp64, POPULATION::enum) = 66.66666666666667::fp64

‎tests/coverage/nodes.py‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,10 +56,29 @@ def is_return_type_error(self):
5656
return isinstance(self.result, SubstraitError)
5757

5858
def get_arg_types(self):
59-
return [arg.get_base_type() for arg in self.args]
59+
types = []
60+
for arg in self.args:
61+
if isinstance(arg, CaseLiteral):
62+
types.append(arg.get_base_type())
63+
elif isinstance(arg, AggregateArgument):
64+
# For aggregate arguments, use column_type if available, otherwise extract from scalar_value
65+
if arg.column_type:
66+
types.append(arg.column_type)
67+
elif arg.scalar_value:
68+
types.append(arg.scalar_value.get_base_type())
69+
return types
6070

6171
def get_signature(self):
62-
return f"{self.func_name}({', '.join([arg.type for arg in self.args])}) = {self.get_return_type()}"
72+
arg_types = []
73+
for arg in self.args:
74+
if isinstance(arg, CaseLiteral):
75+
arg_types.append(arg.type)
76+
elif isinstance(arg, AggregateArgument):
77+
if arg.column_type:
78+
arg_types.append(arg.column_type)
79+
elif arg.scalar_value:
80+
arg_types.append(arg.scalar_value.type)
81+
return f"{self.func_name}({', '.join(arg_types)}) = {self.get_return_type()}"
6382

6483

6584
@dataclass

0 commit comments

Comments
 (0)