Skip to content

Commit 7e4201f

Browse files
gertjanalgertjanal
committed
Add array functions element_first and element_last
1 parent 7e0539a commit 7e4201f

File tree

7 files changed

+244
-0
lines changed

7 files changed

+244
-0
lines changed

‎core/trino-main/src/main/java/io/trino/metadata/SystemFunctionBundle.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,8 @@
9898
import io.trino.operator.scalar.ArrayContainsSequence;
9999
import io.trino.operator.scalar.ArrayDistinctFunction;
100100
import io.trino.operator.scalar.ArrayElementAtFunction;
101+
import io.trino.operator.scalar.ArrayElementFirstFunction;
102+
import io.trino.operator.scalar.ArrayElementLastFunction;
101103
import io.trino.operator.scalar.ArrayExceptFunction;
102104
import io.trino.operator.scalar.ArrayFilterFunction;
103105
import io.trino.operator.scalar.ArrayHistogramFunction;
@@ -492,6 +494,8 @@ public static FunctionBundle create(FeaturesConfig featuresConfig, TypeOperators
492494
.functions(IDENTITY_CAST, CAST_FROM_UNKNOWN)
493495
.scalar(ArrayRemoveFunction.class)
494496
.scalar(ArrayElementAtFunction.class)
497+
.scalar(ArrayElementFirstFunction.class)
498+
.scalar(ArrayElementLastFunction.class)
495499
.scalar(ArraySortFunction.class)
496500
.scalar(ArraySortComparatorFunction.class)
497501
.scalar(ArrayShuffleFunction.class)

‎core/trino-main/src/main/java/io/trino/operator/scalar/ArrayElementFirstFunction.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.operator.scalar;
15+
16+
import io.trino.spi.block.Block;
17+
import io.trino.spi.function.Convention;
18+
import io.trino.spi.function.Description;
19+
import io.trino.spi.function.OperatorDependency;
20+
import io.trino.spi.function.ScalarFunction;
21+
import io.trino.spi.function.SqlNullable;
22+
import io.trino.spi.function.SqlType;
23+
import io.trino.spi.function.TypeParameter;
24+
25+
import java.lang.invoke.MethodHandle;
26+
27+
import static io.trino.operator.scalar.ArrayElementAtFunction.elementAt;
28+
import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION_NOT_NULL;
29+
import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL;
30+
import static io.trino.spi.function.OperatorType.READ_VALUE;
31+
32+
@ScalarFunction("element_first")
33+
@Description("Get first element of array")
34+
public final class ArrayElementFirstFunction
35+
{
36+
private ArrayElementFirstFunction() {}
37+
38+
@TypeParameter("E")
39+
@SqlNullable
40+
@SqlType("E")
41+
public static Object elementFirst(
42+
@OperatorDependency(operator = READ_VALUE, argumentTypes = "E", convention = @Convention(arguments = BLOCK_POSITION_NOT_NULL, result = FAIL_ON_NULL)) MethodHandle readValue,
43+
@SqlType("array(E)") Block array)
44+
throws Throwable
45+
{
46+
return elementAt(readValue, array, 1);
47+
}
48+
}

‎core/trino-main/src/main/java/io/trino/operator/scalar/ArrayElementLastFunction.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.operator.scalar;
15+
16+
import io.trino.spi.block.Block;
17+
import io.trino.spi.function.Convention;
18+
import io.trino.spi.function.Description;
19+
import io.trino.spi.function.OperatorDependency;
20+
import io.trino.spi.function.ScalarFunction;
21+
import io.trino.spi.function.SqlNullable;
22+
import io.trino.spi.function.SqlType;
23+
import io.trino.spi.function.TypeParameter;
24+
25+
import java.lang.invoke.MethodHandle;
26+
27+
import static io.trino.operator.scalar.ArrayElementAtFunction.elementAt;
28+
import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION_NOT_NULL;
29+
import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL;
30+
import static io.trino.spi.function.OperatorType.READ_VALUE;
31+
32+
@ScalarFunction("element_last")
33+
@Description("Get last element of array")
34+
public final class ArrayElementLastFunction
35+
{
36+
private ArrayElementLastFunction() {}
37+
38+
@TypeParameter("E")
39+
@SqlNullable
40+
@SqlType("E")
41+
public static Object elementLast(
42+
@OperatorDependency(operator = READ_VALUE, argumentTypes = "E", convention = @Convention(arguments = BLOCK_POSITION_NOT_NULL, result = FAIL_ON_NULL)) MethodHandle readValue,
43+
@SqlType("array(E)") Block array)
44+
throws Throwable
45+
{
46+
return elementAt(readValue, array, -1);
47+
}
48+
}

‎core/trino-main/src/test/java/io/trino/type/TestArrayOperators.java‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2161,6 +2161,134 @@ public void testElementAt()
21612161
.isEqualTo(decimal("2.22222222222222222", createDecimalType(18, 17)));
21622162
}
21632163

2164+
@Test
2165+
public void testElementFirst()
2166+
{
2167+
assertThat(assertions.function("element_first", "ARRAY[]"))
2168+
.isNull(UNKNOWN);
2169+
2170+
assertThat(assertions.function("element_first", "ARRAY[NULL]"))
2171+
.isNull(UNKNOWN);
2172+
2173+
assertThat(assertions.function("element_first", "ARRAY[2, 1, 3]"))
2174+
.isEqualTo(2);
2175+
2176+
assertThat(assertions.function("element_first", "ARRAY[NULL, 1]"))
2177+
.isNull(INTEGER);
2178+
2179+
assertThat(assertions.function("element_first", "ARRAY[BIGINT '2', 1, 3]"))
2180+
.isEqualTo(2L);
2181+
2182+
assertThat(assertions.function("element_first", "ARRAY[1.0E0, 2.5E0, 3.5E0]"))
2183+
.isEqualTo(1.0);
2184+
2185+
assertThat(assertions.function("element_first", "ARRAY[ARRAY[1, 2], ARRAY[3]]"))
2186+
.hasType(new ArrayType(INTEGER))
2187+
.isEqualTo(ImmutableList.of(1, 2));
2188+
2189+
assertThat(assertions.function("element_first", "ARRAY[NULL, ARRAY[1, 2], ARRAY[3]]"))
2190+
.isNull(new ArrayType(INTEGER));
2191+
2192+
assertThat(assertions.function("element_first", "ARRAY['puppies', 'kittens']"))
2193+
.hasType(createVarcharType(7))
2194+
.isEqualTo("puppies");
2195+
2196+
assertThat(assertions.function("element_first", "ARRAY[NULL, 'puppies', 'kittens']"))
2197+
.isNull(createVarcharType(7));
2198+
2199+
assertThat(assertions.function("element_first", "ARRAY[TRUE, FALSE]"))
2200+
.isEqualTo(true);
2201+
2202+
assertThat(assertions.function("element_first", "ARRAY[NULL, FALSE]"))
2203+
.isNull(BOOLEAN);
2204+
2205+
assertThat(assertions.function("element_first", "ARRAY[TIMESTAMP '1970-01-01 00:00:01', TIMESTAMP '1973-07-08 22:00:01']"))
2206+
.hasType(createTimestampType(0))
2207+
.isEqualTo(sqlTimestampOf(0, 1970, 1, 1, 0, 0, 1, 0));
2208+
2209+
assertThat(assertions.function("element_first", "ARRAY[infinity()]"))
2210+
.isEqualTo(POSITIVE_INFINITY);
2211+
2212+
assertThat(assertions.function("element_first", "ARRAY[-infinity()]"))
2213+
.isEqualTo(NEGATIVE_INFINITY);
2214+
2215+
assertThat(assertions.function("element_first", "ARRAY[sqrt(-1)]"))
2216+
.isEqualTo(NaN);
2217+
2218+
assertThat(assertions.function("element_first", "ARRAY[2.1, 2.2, 2.3]"))
2219+
.isEqualTo(decimal("2.1", createDecimalType(2, 1)));
2220+
2221+
assertThat(assertions.function("element_first", "ARRAY[2.111111222111111114111, 2.22222222222222222, 2.222222222222223]"))
2222+
.isEqualTo(decimal("2.111111222111111114111", createDecimalType(22, 21)));
2223+
2224+
assertThat(assertions.function("element_first", "ARRAY[1.9, 2, 2.3]"))
2225+
.isEqualTo(decimal("0000000001.9", createDecimalType(11, 1)));
2226+
}
2227+
2228+
@Test
2229+
public void testElementLast()
2230+
{
2231+
assertThat(assertions.function("element_last", "ARRAY[]"))
2232+
.isNull(UNKNOWN);
2233+
2234+
assertThat(assertions.function("element_last", "ARRAY[NULL]"))
2235+
.isNull(UNKNOWN);
2236+
2237+
assertThat(assertions.function("element_last", "ARRAY[2, 1, 3]"))
2238+
.isEqualTo(3);
2239+
2240+
assertThat(assertions.function("element_last", "ARRAY[1, NULL]"))
2241+
.isNull(INTEGER);
2242+
2243+
assertThat(assertions.function("element_last", "ARRAY[BIGINT '2', 1, 3]"))
2244+
.isEqualTo(3L);
2245+
2246+
assertThat(assertions.function("element_last", "ARRAY[1.0E0, 2.5E0, 3.5E0]"))
2247+
.isEqualTo(3.5);
2248+
2249+
assertThat(assertions.function("element_last", "ARRAY[ARRAY[1, 2], ARRAY[3]]"))
2250+
.hasType(new ArrayType(INTEGER))
2251+
.isEqualTo(ImmutableList.of(3));
2252+
2253+
assertThat(assertions.function("element_last", "ARRAY[ARRAY[1, 2], ARRAY[3], NULL]"))
2254+
.isNull(new ArrayType(INTEGER));
2255+
2256+
assertThat(assertions.function("element_last", "ARRAY['puppies', 'kittens']"))
2257+
.hasType(createVarcharType(7))
2258+
.isEqualTo("kittens");
2259+
2260+
assertThat(assertions.function("element_last", "ARRAY['puppies', 'kittens', NULL]"))
2261+
.isNull(createVarcharType(7));
2262+
2263+
assertThat(assertions.function("element_last", "ARRAY[TRUE, FALSE]"))
2264+
.isEqualTo(false);
2265+
2266+
assertThat(assertions.function("element_last", "ARRAY[FALSE, NULL]"))
2267+
.isNull(BOOLEAN);
2268+
2269+
assertThat(assertions.function("element_last", "ARRAY[TIMESTAMP '1970-01-01 00:00:01', TIMESTAMP '1973-07-08 22:00:01']"))
2270+
.hasType(createTimestampType(0))
2271+
.isEqualTo(sqlTimestampOf(0, 1973, 7, 8, 22, 0, 1, 0));
2272+
2273+
assertThat(assertions.function("element_last", "ARRAY[infinity()]"))
2274+
.isEqualTo(POSITIVE_INFINITY);
2275+
2276+
assertThat(assertions.function("element_last", "ARRAY[-infinity()]"))
2277+
.isEqualTo(NEGATIVE_INFINITY);
2278+
2279+
assertThat(assertions.function("element_last", "ARRAY[sqrt(-1)]"))
2280+
.isEqualTo(NaN);
2281+
2282+
assertThat(assertions.function("element_last", "ARRAY[2.1, 2.2, 2.3]"))
2283+
.isEqualTo(decimal("2.3", createDecimalType(2, 1)));
2284+
2285+
assertThat(assertions.function("element_last", "ARRAY[2.111111222111111114111, 2.22222222222222222, 2.222222222222223]"))
2286+
.isEqualTo(decimal("2.222222222222223000000", createDecimalType(22, 21)));
2287+
2288+
assertThat(assertions.function("element_last", "ARRAY[1.9, 2, 2.3]"))
2289+
.isEqualTo(decimal("0000000002.3", createDecimalType(11, 1)));
2290+
}
2291+
21642292
@Test
21652293
public void testSort()
21662294
{

‎docs/src/main/sphinx/functions/array.md‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,18 @@ the subscript operator would fail in such a case.
252252
If `index` \< 0, `element_at` accesses elements from the last to the first.
253253
:::
254254

255+
:::{function} element_first(array(E)) -> E
256+
Returns the first element of `array`.
257+
If the array is empty, the function returns `NULL`, whereas
258+
the subscript operator would fail in such a case.
259+
:::
260+
261+
:::{function} element_last(array(E)) -> E
262+
Returns the last element of `array`.
263+
If the array is empty, the function returns `NULL`, whereas
264+
the subscript operator would fail in such a case.
265+
:::
266+
255267
:::{function} filter(array(T), function(T,boolean)) -> array(T)
256268
Constructs an array from those elements of `array` for which `function` returns true:
257269

‎docs/src/main/sphinx/functions/list-by-topic.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ For more details, see {doc}`array`
7070
- `concat()`
7171
- {func}`contains`
7272
- {func}`element_at`
73+
- {func}`element_first`
74+
- {func}`element_last`
7375
- {func}`filter`
7476
- {func}`flatten`
7577
- {func}`ngrams`

‎docs/src/main/sphinx/functions/list.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,8 @@
139139

140140
- {func}`e`
141141
- {func}`element_at`
142+
- {func}`element_first`
143+
- {func}`element_last`
142144
- {func}`empty_approx_set`
143145
- `evaluate_classifier_predictions`
144146
- {func}`every`

0 commit comments

Comments
 (0)