Skip to content

Commit 925298b

Browse files
gertjanalebyhr
authored andcommitted
Add array functions array_first and array_last
1 parent fcb1e28 commit 925298b

File tree

7 files changed

+250
-0
lines changed

7 files changed

+250
-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
@@ -100,9 +100,11 @@
100100
import io.trino.operator.scalar.ArrayElementAtFunction;
101101
import io.trino.operator.scalar.ArrayExceptFunction;
102102
import io.trino.operator.scalar.ArrayFilterFunction;
103+
import io.trino.operator.scalar.ArrayFirstFunction;
103104
import io.trino.operator.scalar.ArrayHistogramFunction;
104105
import io.trino.operator.scalar.ArrayIntersectFunction;
105106
import io.trino.operator.scalar.ArrayJoin;
107+
import io.trino.operator.scalar.ArrayLastFunction;
106108
import io.trino.operator.scalar.ArrayMaxFunction;
107109
import io.trino.operator.scalar.ArrayMinFunction;
108110
import io.trino.operator.scalar.ArrayNgramsFunction;
@@ -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(ArrayFirstFunction.class)
498+
.scalar(ArrayLastFunction.class)
495499
.scalar(ArraySortFunction.class)
496500
.scalar(ArraySortComparatorFunction.class)
497501
.scalar(ArrayShuffleFunction.class)
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("array_first")
33+
@Description("Get the first element of an array")
34+
public final class ArrayFirstFunction
35+
{
36+
private ArrayFirstFunction() {}
37+
38+
@TypeParameter("E")
39+
@SqlNullable
40+
@SqlType("E")
41+
public static Object arrayFirst(
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+
}
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("array_last")
33+
@Description("Get the last element of an array")
34+
public final class ArrayLastFunction
35+
{
36+
private ArrayLastFunction() {}
37+
38+
@TypeParameter("E")
39+
@SqlNullable
40+
@SqlType("E")
41+
public static Object arrayLast(
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: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2161,6 +2161,140 @@ public void testElementAt()
21612161
.isEqualTo(decimal("2.22222222222222222", createDecimalType(18, 17)));
21622162
}
21632163

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

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,12 @@ Returns an array of the elements in the union of `x` and `y`, without duplicates
9898
Returns an array of elements in `x` but not in `y`, without duplicates.
9999
:::
100100

101+
:::{function} array_first(array(E)) -> E
102+
Returns the first element of an `array`.
103+
If the array is empty, the function returns `NULL`, whereas
104+
the subscript operator would fail in such a case.
105+
:::
106+
101107
:::{function} array_histogram(x) -> map<K, bigint>
102108
Returns a map where the keys are the unique elements in the input array
103109
`x` and the values are the number of times that each element appears in
@@ -127,6 +133,12 @@ Null elements are omitted in the result.
127133
Concatenates the elements of the given array using the delimiter and an optional string to replace nulls.
128134
:::
129135

136+
:::{function} array_last(array(E)) -> E
137+
Returns the last element of an `array`.
138+
If the array is empty, the function returns `NULL`, whereas
139+
the subscript operator would fail in such a case.
140+
:::
141+
130142
:::{function} array_max(x) -> x
131143
Returns the maximum value of input array.
132144
:::

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,11 @@ For more details, see {doc}`array`
5656
- {func}`any_match`
5757
- {func}`array_distinct`
5858
- {func}`array_except`
59+
- {func}`array_first`
5960
- {func}`array_histogram`
6061
- {func}`array_intersect`
6162
- {func}`array_join`
63+
- {func}`array_last`
6264
- {func}`array_max`
6365
- {func}`array_min`
6466
- {func}`array_position`

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,11 @@
3636
- {func}`array_agg`
3737
- {func}`array_distinct`
3838
- {func}`array_except`
39+
- {func}`array_first`
3940
- {func}`array_histogram`
4041
- {func}`array_intersect`
4142
- {func}`array_join`
43+
- {func}`array_last`
4244
- {func}`array_max`
4345
- {func}`array_min`
4446
- {func}`array_position`

0 commit comments

Comments
 (0)