Skip to content

Commit 0fde339

Browse files
committed
SubList function (#5529)
* SubList function #5529 Signed-off-by: Sai charan raj Gudala <s.charancherry22@gmail.com>
1 parent f136f48 commit 0fde339

4 files changed

Lines changed: 241 additions & 1 deletion

File tree

data-prepper-expression/src/main/antlr/DataPrepperExpression.g4

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,7 @@ fragment
167167
FunctionArg
168168
: JsonPointer
169169
| String
170+
| Integer
170171
;
171172

172173
variableIdentifier
@@ -187,7 +188,7 @@ literal
187188

188189
Integer
189190
: ZERO
190-
| NonZeroDigit Digit*
191+
| '-'? NonZeroDigit Digit*
191192
;
192193

193194
Float

data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/ParseTreeCoercionService.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,11 @@ public Object coercePrimaryTerminalNode(final TerminalNode node, final Event eve
5858
}
5959
argList.add(trimmedArg);
6060
} else {
61+
try {
62+
argList.add(Integer.parseInt(trimmedArg));
63+
continue;
64+
} catch (final Exception e) {
65+
}
6166
throw new RuntimeException("Unsupported type passed as function argument");
6267
}
6368
}
Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package org.opensearch.dataprepper.expression;
2+
3+
import java.util.ArrayList;
4+
import java.util.List;
5+
import java.util.function.Function;
6+
7+
import org.opensearch.dataprepper.model.event.Event;
8+
9+
import javax.inject.Named;
10+
11+
@Named
12+
public class SubListExpressionFunction implements ExpressionFunction {
13+
14+
@Override
15+
public String getFunctionName() {
16+
return "subList";
17+
}
18+
19+
@Override
20+
public Object evaluate(List<Object> args, Event event, Function<Object, Object> convertLiteralType) {
21+
if (args.size() != 3) {
22+
throw new IllegalArgumentException("subList() takes 3 arguments");
23+
}
24+
if (!(args.get(0) instanceof String)) {
25+
throw new IllegalArgumentException("subList() takes 1st argument as string type");
26+
}
27+
int startIndex, endIndex;
28+
try {
29+
if (args.get(1) instanceof Integer) {
30+
startIndex = (Integer) args.get(1);
31+
} else {
32+
String str = (String) args.get(1);
33+
startIndex = Integer.parseInt(str.substring(1, str.length() - 1));
34+
}
35+
if (args.get(2) instanceof Integer) {
36+
endIndex = (Integer) args.get(2);
37+
} else {
38+
String str = (String) args.get(2);
39+
endIndex = Integer.parseInt(str.substring(1, str.length() - 1));
40+
}
41+
} catch (NumberFormatException | ClassCastException e) {
42+
throw new IllegalArgumentException("subList() takes 2nd and 3rd arguments as integers");
43+
}
44+
45+
String key = (String)args.get(0);
46+
final Object value = event.get(key, Object.class);
47+
if (value == null) return null;
48+
if (!(value instanceof List)) {
49+
throw new RuntimeException(key + " is not of list type");
50+
}
51+
List<?> sourceList = (List<?>)value;
52+
53+
if (endIndex == -1) endIndex = sourceList.size();
54+
55+
if (startIndex < 0 || startIndex >= sourceList.size()) {
56+
throw new RuntimeException("subList() start index should be between 0 and list length (inclusive)");
57+
}
58+
59+
if (endIndex < 0 || endIndex > sourceList.size()) {
60+
throw new RuntimeException("subList() end index should be between 0 and list length or -1 for list length (exclusive)");
61+
}
62+
if (startIndex > endIndex) {
63+
throw new RuntimeException("subList() start index should be less than or equal to end index");
64+
}
65+
return new ArrayList<>(sourceList.subList(startIndex, endIndex));
66+
}
67+
68+
}
Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
package org.opensearch.dataprepper.expression;
2+
3+
import org.junit.jupiter.api.Test;
4+
import org.opensearch.dataprepper.model.event.Event;
5+
import org.opensearch.dataprepper.model.event.JacksonEvent;
6+
7+
import static org.junit.jupiter.api.Assertions.assertEquals;
8+
import static org.mockito.Mockito.mock;
9+
import static org.hamcrest.CoreMatchers.equalTo;
10+
import static org.hamcrest.MatcherAssert.assertThat;
11+
import static org.junit.jupiter.api.Assertions.assertThrows;
12+
13+
import java.util.List;
14+
import java.util.Map;
15+
import java.util.function.Function;
16+
17+
public class SubListExpressionFunctionTest {
18+
private SubListExpressionFunction subListExpressionFunction;
19+
private Event testEvent;
20+
private Function<Object, Object> testFunction;
21+
22+
private Event createTestEvent(final Object data) {
23+
return JacksonEvent.builder().withEventType("event").withData(data).build();
24+
}
25+
26+
27+
public SubListExpressionFunction createObjectUnderTest() {
28+
testFunction = mock(Function.class);
29+
return new SubListExpressionFunction();
30+
}
31+
32+
@Test
33+
void testFunctionName() {
34+
subListExpressionFunction = createObjectUnderTest();
35+
assertEquals("subList", subListExpressionFunction.getFunctionName());
36+
}
37+
38+
@Test
39+
void testWithValidArguments() {
40+
subListExpressionFunction = createObjectUnderTest();
41+
List<Integer> testList = List.of(1, 2, 3, 4);
42+
testEvent = createTestEvent(Map.of("list", testList));
43+
assertThat(subListExpressionFunction.evaluate(List.of("/list", "\"1\"", "\"3\""), testEvent, testFunction), equalTo(List.of(2, 3)));
44+
}
45+
46+
@Test
47+
void testWithValidArgumentsCase2() {
48+
subListExpressionFunction = createObjectUnderTest();
49+
List<Integer> testList = List.of(1, 2, 3, 4);
50+
testEvent = createTestEvent(Map.of("list", testList));
51+
assertThat(subListExpressionFunction.evaluate(List.of("/list", "\"1\"", "\"3\""), testEvent, testFunction), equalTo(List.of(2, 3)));
52+
}
53+
54+
@Test
55+
void testWithValidArgumentsCase3() {
56+
subListExpressionFunction = createObjectUnderTest();
57+
List<Integer> testList = List.of(1, 2, 3, 4);
58+
testEvent = createTestEvent(Map.of("main", Map.of("list", testList)));
59+
assertThat(subListExpressionFunction.evaluate(List.of("/main/list", "\"1\"", "\"3\""), testEvent, testFunction), equalTo(List.of(2, 3)));
60+
}
61+
62+
@Test
63+
void testWithValidArgumentsCase4() {
64+
subListExpressionFunction = createObjectUnderTest();
65+
List<Integer> testList = List.of(1, 2, 3, 4);
66+
testEvent = createTestEvent(Map.of("main", Map.of("list", testList)));
67+
assertThat(subListExpressionFunction.evaluate(List.of("/main/list", "\"1\"", "\"-1\""), testEvent, testFunction), equalTo(List.of(2, 3, 4)));
68+
}
69+
70+
@Test
71+
void testWithValidArgumentsCase5() {
72+
subListExpressionFunction = createObjectUnderTest();
73+
List<Integer> testList = List.of(1, 2, 3, 4);
74+
testEvent = createTestEvent(Map.of("list", testList));
75+
assertThat(subListExpressionFunction.evaluate(List.of("/list", 1, 3), testEvent, testFunction), equalTo(List.of(2, 3)));
76+
}
77+
78+
@Test
79+
void testWithOutOfBoundArgumentsCase1() {
80+
subListExpressionFunction = createObjectUnderTest();
81+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
82+
testEvent = createTestEvent(Map.of("list", testList));
83+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"-1\"", "\"4\""), testEvent, testFunction));
84+
assertEquals("subList() start index should be between 0 and list length (inclusive)", exception.getMessage());
85+
}
86+
87+
@Test
88+
void testWithOutOfBoundArgumentsCase2() {
89+
subListExpressionFunction = createObjectUnderTest();
90+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
91+
testEvent = createTestEvent(Map.of("list", testList));
92+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"10\"", "\"4\""), testEvent, testFunction));
93+
assertEquals("subList() start index should be between 0 and list length (inclusive)", exception.getMessage());
94+
}
95+
96+
@Test
97+
void testWithOutOfBoundArgumentsCase3() {
98+
subListExpressionFunction = createObjectUnderTest();
99+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
100+
testEvent = createTestEvent(Map.of("list", testList));
101+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"4\"", "\"2\""), testEvent, testFunction));
102+
assertEquals("subList() start index should be less than or equal to end index", exception.getMessage());
103+
}
104+
105+
@Test
106+
void testWithOutOfBoundArgumentsCase4() {
107+
subListExpressionFunction = createObjectUnderTest();
108+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
109+
testEvent = createTestEvent(Map.of("list", testList));
110+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"1\"", "\"10\""), testEvent, testFunction));
111+
assertEquals("subList() end index should be between 0 and list length or -1 for list length (exclusive)", exception.getMessage());
112+
}
113+
114+
@Test
115+
void testWithOutOfBoundArgumentsCase5() {
116+
subListExpressionFunction = createObjectUnderTest();
117+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
118+
testEvent = createTestEvent(Map.of("list", testList));
119+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"1\"", "\"-2\""), testEvent, testFunction));
120+
assertEquals("subList() end index should be between 0 and list length or -1 for list length (exclusive)", exception.getMessage());
121+
}
122+
123+
@Test
124+
void testWithInvalidArguments() {
125+
subListExpressionFunction = createObjectUnderTest();
126+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
127+
testEvent = createTestEvent(Map.of("list", testList));
128+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"5\"", "\"2\""), testEvent, testFunction));
129+
assertEquals("subList() start index should be less than or equal to end index", exception.getMessage());
130+
}
131+
132+
@Test
133+
void testWithInvalidArgumentsCase2() {
134+
subListExpressionFunction = createObjectUnderTest();
135+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
136+
testEvent = createTestEvent(Map.of("list", testList));
137+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"five\"", "\"two\""), testEvent, testFunction));
138+
assertEquals("subList() takes 2nd and 3rd arguments as integers", exception.getMessage());
139+
}
140+
141+
@Test
142+
void testWithInvalidArgumentsCase3() {
143+
subListExpressionFunction = createObjectUnderTest();
144+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
145+
testEvent = createTestEvent(Map.of("list", testList));
146+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"0\""), testEvent, testFunction));
147+
assertEquals("subList() takes 3 arguments", exception.getMessage());
148+
}
149+
150+
@Test
151+
void testWithInvalidArgumentsCase4() {
152+
subListExpressionFunction = createObjectUnderTest();
153+
List<Integer> testList = List.of(1, 2, 3, 4, 5, 6);
154+
testEvent = createTestEvent(Map.of("list", testList));
155+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of(1, "\"0\"", "\"2\""), testEvent, testFunction));
156+
assertEquals("subList() takes 1st argument as string type", exception.getMessage());
157+
}
158+
159+
@Test
160+
void testWithInvalidArgumentsCase5() {
161+
subListExpressionFunction = createObjectUnderTest();
162+
testEvent = createTestEvent(Map.of("list", "testList"));
163+
Exception exception = assertThrows(RuntimeException.class, () -> subListExpressionFunction.evaluate(List.of("/list", "\"0\"", "\"2\""), testEvent, testFunction));
164+
assertEquals("/list is not of list type", exception.getMessage());
165+
}
166+
}

0 commit comments

Comments
 (0)