Skip to content

Commit 8368e60

Browse files
vamsimanoharvamsimanohar
authored
eval sum, avg implementation (#3986)
Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>
1 parent e2375fe commit 8368e60

File tree

6 files changed

+596
-6
lines changed

6 files changed

+596
-6
lines changed

‎docs/user/ppl/functions/math.rst‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,84 @@ Example::
132132
+--------------+
133133

134134

135+
SUM
136+
---
137+
138+
Description
139+
>>>>>>>>>>>
140+
141+
Usage: sum(x, y, ...) calculates the sum of all provided arguments. This function accepts a variable number of arguments.
142+
143+
Note: This function is only available in the eval command context and is rewritten to arithmetic addition while query parsing.
144+
145+
Argument type: Variable number of INTEGER/LONG/FLOAT/DOUBLE arguments
146+
147+
Return type: Wider number type among all arguments
148+
149+
Example::
150+
151+
os> source=accounts | eval `SUM(1, 2, 3)` = SUM(1, 2, 3) | fields `SUM(1, 2, 3)`
152+
fetched rows / total rows = 4/4
153+
+--------------+
154+
| SUM(1, 2, 3) |
155+
|--------------|
156+
| 6 |
157+
| 6 |
158+
| 6 |
159+
| 6 |
160+
+--------------+
161+
162+
os> source=accounts | eval total = SUM(age, 10, 5) | fields age, total
163+
fetched rows / total rows = 4/4
164+
+-----+-------+
165+
| age | total |
166+
|-----+-------|
167+
| 32 | 47 |
168+
| 36 | 51 |
169+
| 28 | 43 |
170+
| 33 | 48 |
171+
+-----+-------+
172+
173+
174+
AVG
175+
---
176+
177+
Description
178+
>>>>>>>>>>>
179+
180+
Usage: avg(x, y, ...) calculates the average (arithmetic mean) of all provided arguments. This function accepts a variable number of arguments.
181+
182+
Note: This function is only available in the eval command context and is rewritten to arithmetic expression (sum / count) at query parsing time.
183+
184+
Argument type: Variable number of INTEGER/LONG/FLOAT/DOUBLE arguments
185+
186+
Return type: DOUBLE
187+
188+
Example::
189+
190+
os> source=accounts | eval `AVG(1, 2, 3)` = AVG(1, 2, 3) | fields `AVG(1, 2, 3)`
191+
fetched rows / total rows = 4/4
192+
+--------------+
193+
| AVG(1, 2, 3) |
194+
|--------------|
195+
| 2.0 |
196+
| 2.0 |
197+
| 2.0 |
198+
| 2.0 |
199+
+--------------+
200+
201+
os> source=accounts | eval average = AVG(age, 30) | fields age, average
202+
fetched rows / total rows = 4/4
203+
+-----+---------+
204+
| age | average |
205+
|-----+---------|
206+
| 32 | 31.0 |
207+
| 36 | 33.0 |
208+
| 28 | 29.0 |
209+
| 33 | 31.5 |
210+
+-----+---------+
211+
212+
135213
ACOS
136214
----
137215

‎integ-test/src/test/java/org/opensearch/sql/ppl/MathematicalFunctionIT.java‎

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -547,4 +547,225 @@ public void testRint() throws IOException {
547547
verifySchema(result, schema("f", null, "double"));
548548
verifySome(result.getJSONArray("datarows"), rows(Math.rint(1.7)));
549549
}
550+
551+
// SUM function tests for eval command
552+
@Test
553+
public void testEvalSumSingleInteger() throws IOException {
554+
JSONObject result =
555+
executeQuery(
556+
String.format("source=%s | eval f = sum(42) | fields f | head 5", TEST_INDEX_BANK));
557+
verifySchema(result, schema("f", null, "int"));
558+
verifyDataRows(result, rows(42), rows(42), rows(42), rows(42), rows(42));
559+
}
560+
561+
@Test
562+
public void testEvalSumMultipleIntegers() throws IOException {
563+
JSONObject result =
564+
executeQuery(
565+
String.format(
566+
"source=%s | eval f = sum(1, 2, 3) | fields f | head 5", TEST_INDEX_BANK));
567+
verifySchema(result, schema("f", null, "int"));
568+
verifyDataRows(result, rows(6), rows(6), rows(6), rows(6), rows(6));
569+
}
570+
571+
@Test
572+
public void testEvalSumMixedNumericTypes() throws IOException {
573+
JSONObject result =
574+
executeQuery(
575+
String.format("source=%s | eval f = sum(1, 2.5) | fields f | head 5", TEST_INDEX_BANK));
576+
verifySchema(result, schema("f", null, "double"));
577+
verifyDataRows(result, rows(3.5), rows(3.5), rows(3.5), rows(3.5), rows(3.5));
578+
}
579+
580+
@Test
581+
public void testEvalSumWithFields() throws IOException {
582+
JSONObject result =
583+
executeQuery(
584+
String.format(
585+
"source=%s | eval f = sum(age, 10) | fields f | head 7", TEST_INDEX_BANK));
586+
verifySchema(result, schema("f", null, "int"));
587+
verifyDataRows(result, rows(42), rows(46), rows(38), rows(43), rows(46), rows(49), rows(44));
588+
}
589+
590+
@Test
591+
public void testEvalSumMultipleNumericArguments() throws IOException {
592+
JSONObject result =
593+
executeQuery(
594+
String.format(
595+
"source=%s | eval f = sum(1, 2, 3, 4, 5) | fields f | head 5", TEST_INDEX_BANK));
596+
verifySchema(result, schema("f", null, "int"));
597+
verifyDataRows(result, rows(15), rows(15), rows(15), rows(15), rows(15));
598+
}
599+
600+
// AVG function tests for eval command
601+
@Test
602+
public void testEvalAvgSingleInteger() throws IOException {
603+
JSONObject result =
604+
executeQuery(
605+
String.format("source=%s | eval f = avg(42) | fields f | head 5", TEST_INDEX_BANK));
606+
verifySchema(result, schema("f", null, "double"));
607+
verifyDataRows(result, rows(42.0), rows(42.0), rows(42.0), rows(42.0), rows(42.0));
608+
}
609+
610+
@Test
611+
public void testEvalAvgMultipleIntegers() throws IOException {
612+
JSONObject result =
613+
executeQuery(
614+
String.format(
615+
"source=%s | eval f = avg(1, 2, 3) | fields f | head 5", TEST_INDEX_BANK));
616+
verifySchema(result, schema("f", null, "double"));
617+
verifyDataRows(result, rows(2.0), rows(2.0), rows(2.0), rows(2.0), rows(2.0));
618+
}
619+
620+
@Test
621+
public void testEvalAvgTwoIntegers() throws IOException {
622+
JSONObject result =
623+
executeQuery(
624+
String.format("source=%s | eval f = avg(1, 2) | fields f | head 5", TEST_INDEX_BANK));
625+
verifySchema(result, schema("f", null, "double"));
626+
verifyDataRows(result, rows(1.5), rows(1.5), rows(1.5), rows(1.5), rows(1.5));
627+
}
628+
629+
@Test
630+
public void testEvalAvgMixedNumericTypes() throws IOException {
631+
JSONObject result =
632+
executeQuery(
633+
String.format("source=%s | eval f = avg(1, 2.5) | fields f | head 5", TEST_INDEX_BANK));
634+
verifySchema(result, schema("f", null, "double"));
635+
verifyDataRows(result, rows(1.75), rows(1.75), rows(1.75), rows(1.75), rows(1.75));
636+
}
637+
638+
@Test
639+
public void testEvalAvgWithFields() throws IOException {
640+
JSONObject result =
641+
executeQuery(
642+
String.format(
643+
"source=%s | eval f = avg(age, 10) | fields f | head 7", TEST_INDEX_BANK));
644+
verifySchema(result, schema("f", null, "double"));
645+
verifyDataRows(
646+
result, rows(21.0), rows(23.0), rows(19.0), rows(21.5), rows(23.0), rows(24.5), rows(22.0));
647+
}
648+
649+
@Test
650+
public void testEvalAvgMultipleValues() throws IOException {
651+
JSONObject result =
652+
executeQuery(
653+
String.format("source=%s | eval f = avg(1, 4) | fields f | head 5", TEST_INDEX_BANK));
654+
verifySchema(result, schema("f", null, "double"));
655+
verifyDataRows(result, rows(2.5), rows(2.5), rows(2.5), rows(2.5), rows(2.5));
656+
}
657+
658+
@Test
659+
public void testEvalAvgFiveValues() throws IOException {
660+
JSONObject result =
661+
executeQuery(
662+
String.format(
663+
"source=%s | eval f = avg(1, 2, 3, 4, 5) | fields f | head 5", TEST_INDEX_BANK));
664+
verifySchema(result, schema("f", null, "double"));
665+
verifyDataRows(result, rows(3.0), rows(3.0), rows(3.0), rows(3.0), rows(3.0));
666+
}
667+
668+
// Combined sum and avg tests
669+
@Test
670+
public void testEvalSumAndAvgComparison() throws IOException {
671+
JSONObject result =
672+
executeQuery(
673+
String.format(
674+
"source=%s | eval sum_val = sum(10, 20, 30), avg_val = avg(10, 20, 30) | fields"
675+
+ " sum_val, avg_val | head 5",
676+
TEST_INDEX_BANK));
677+
verifySchema(result, schema("sum_val", null, "int"), schema("avg_val", null, "double"));
678+
verifyDataRows(
679+
result, rows(60, 20.0), rows(60, 20.0), rows(60, 20.0), rows(60, 20.0), rows(60, 20.0));
680+
}
681+
682+
@Test
683+
public void testEvalSumInWhereClause() throws IOException {
684+
JSONObject result =
685+
executeQuery(
686+
String.format(
687+
"source=%s | where sum(age, 10) > 40 | eval f = sum(age, 10) | fields f | head 6",
688+
TEST_INDEX_BANK));
689+
verifySchema(result, schema("f", null, "int"));
690+
// Should return rows where age + 10 > 40, so age > 30
691+
verifyDataRows(result, rows(42), rows(46), rows(43), rows(46), rows(49), rows(44));
692+
}
693+
694+
@Test
695+
public void testEvalAvgInWhereClause() throws IOException {
696+
JSONObject result =
697+
executeQuery(
698+
String.format(
699+
"source=%s | where avg(age, 10) > 20 | eval f = avg(age, 10) | fields f | head 6",
700+
TEST_INDEX_BANK));
701+
verifySchema(result, schema("f", null, "double"));
702+
// Should return rows where (age + 10) / 2 > 20, so age > 30
703+
verifyDataRows(result, rows(21.0), rows(23.0), rows(21.5), rows(23.0), rows(24.5), rows(22.0));
704+
}
705+
706+
@Test
707+
public void testEvalComplexExpression() throws IOException {
708+
JSONObject result =
709+
executeQuery(
710+
String.format(
711+
"source=%s | eval f = sum(age, 5) + avg(10, 20) | fields f | head 5",
712+
TEST_INDEX_BANK));
713+
verifySchema(result, schema("f", null, "double"));
714+
// sum(age, 5) + avg(10, 20) = (age + 5) + 15.0
715+
verifyDataRows(result, rows(52.0), rows(56.0), rows(48.0), rows(53.0), rows(56.0));
716+
}
717+
718+
@Test
719+
public void testEvalNestedSumAvg() throws IOException {
720+
// Note: This tests the arithmetic expression rewriting, not actual nested function calls
721+
JSONObject result =
722+
executeQuery(
723+
String.format(
724+
"source=%s | eval f = sum(avg(20, 30), 10) | fields f | head 5", TEST_INDEX_BANK));
725+
verifySchema(result, schema("f", null, "double"));
726+
// avg(20, 30) = 25.0, sum(25.0, 10) = 35.0
727+
verifyDataRows(result, rows(35.0), rows(35.0), rows(35.0), rows(35.0), rows(35.0));
728+
}
729+
730+
@Test
731+
public void testEvalSumWithMultipleFields() throws IOException {
732+
JSONObject result =
733+
executeQuery(
734+
String.format(
735+
"source=%s | eval f = sum(age, age, 10) | fields f | head 5", TEST_INDEX_BANK));
736+
verifySchema(result, schema("f", null, "int"));
737+
// sum(age, age, 10) = age + age + 10 = 2*age + 10
738+
verifyDataRows(result, rows(74), rows(82), rows(66), rows(76), rows(82));
739+
}
740+
741+
@Test
742+
public void testEvalAvgWithExpression() throws IOException {
743+
JSONObject result =
744+
executeQuery(
745+
String.format(
746+
"source=%s | eval f = avg(age * 2, 10) | fields f | head 5", TEST_INDEX_BANK));
747+
verifySchema(result, schema("f", null, "double"));
748+
// avg(age * 2, 10) = (age * 2 + 10) / 2 = age + 5
749+
verifyDataRows(result, rows(37.0), rows(41.0), rows(33.0), rows(38.0), rows(41.0));
750+
}
751+
752+
@Test
753+
public void testEvalSumWithNegativeNumbers() throws IOException {
754+
JSONObject result =
755+
executeQuery(
756+
String.format(
757+
"source=%s | eval f = sum(-5, 10, -3) | fields f | head 5", TEST_INDEX_BANK));
758+
verifySchema(result, schema("f", null, "int"));
759+
verifyDataRows(result, rows(2), rows(2), rows(2), rows(2), rows(2));
760+
}
761+
762+
@Test
763+
public void testEvalAvgWithNegativeNumbers() throws IOException {
764+
JSONObject result =
765+
executeQuery(
766+
String.format(
767+
"source=%s | eval f = avg(-10, 10) | fields f | head 5", TEST_INDEX_BANK));
768+
verifySchema(result, schema("f", null, "double"));
769+
verifyDataRows(result, rows(0.0), rows(0.0), rows(0.0), rows(0.0), rows(0.0));
770+
}
550771
}

‎ppl/src/main/antlr/OpenSearchPPLParser.g4‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -713,6 +713,8 @@ mathematicalFunctionName
713713
| TRUNCATE
714714
| RINT
715715
| SIGNUM
716+
| SUM
717+
| AVG
716718
| trigonometricFunctionName
717719
;
718720

0 commit comments

Comments
 (0)