risingwave_expr_impl/table_function/
generate_subscripts.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use auto_enums::auto_enum;
16use risingwave_common::array::ListRef;
17use risingwave_common::types::ScalarRefImpl;
18use risingwave_expr::function;
19
20/// ```slt
21/// query I
22/// SELECT generate_subscripts(ARRAY['foo', 'bar', null], -1, false) AS s;
23/// ----
24///
25/// query I
26/// SELECT generate_subscripts(ARRAY['foo', 'bar', null], 0, false) AS s;
27/// ----
28///
29/// query I
30/// SELECT generate_subscripts(ARRAY['foo', 'bar', null], 1, true) AS s;
31/// ----
32/// 3
33/// 2
34/// 1
35///
36/// query I
37/// SELECT generate_subscripts(ARRAY['foo', 'bar', null], 1, false) AS s;
38/// ----
39/// 1
40/// 2
41/// 3
42///
43/// query I
44/// SELECT generate_subscripts(ARRAY[ARRAY['foo'], ARRAY['bar'], ARRAY[null]], 0, false) AS s;
45/// ----
46///
47/// query I
48/// SELECT generate_subscripts(ARRAY[ARRAY['foo'], ARRAY['bar'], ARRAY[null]], 1, false) AS s;
49/// ----
50/// 1
51/// 2
52/// 3
53///
54/// query I
55/// SELECT generate_subscripts(ARRAY[ARRAY['foo'], ARRAY['bar'], ARRAY[null]], 2, false) AS s;
56/// ----
57/// 1
58/// ```
59#[function("generate_subscripts(anyarray, int4, boolean) -> setof int4")]
60fn generate_subscripts_reverse(
61    array: ListRef<'_>,
62    dim: i32,
63    reverse: bool,
64) -> impl Iterator<Item = i32> {
65    generate_subscripts_iterator(array, dim, reverse)
66}
67
68/// ```slt
69/// query I
70/// SELECT generate_subscripts(ARRAY['foo', 'bar', null], -1, false) AS s;
71/// ----
72///
73/// query I
74/// SELECT generate_subscripts(ARRAY['foo', 'bar', null], 0, false) AS s;
75/// ----
76///
77/// query I
78/// SELECT generate_subscripts(ARRAY['foo', 'bar', null], 1) AS s;
79/// ----
80/// 1
81/// 2
82/// 3
83///
84/// query I
85/// SELECT generate_subscripts(ARRAY['foo', 'bar', null], 1) AS s;
86/// ----
87/// 1
88/// 2
89/// 3
90///
91/// query I
92/// SELECT generate_subscripts(ARRAY[ARRAY['foo'], ARRAY['bar'], ARRAY[null]], 0) AS s;
93/// ----
94///
95/// query I
96/// SELECT generate_subscripts(ARRAY[ARRAY['foo'], ARRAY['bar'], ARRAY[null]], 1) AS s;
97/// ----
98/// 1
99/// 2
100/// 3
101///
102/// query I
103/// SELECT generate_subscripts(ARRAY[ARRAY['foo'], ARRAY['bar'], ARRAY[null]], 2) AS s;
104/// ----
105/// 1
106/// ```
107#[function("generate_subscripts(anyarray, int4) -> setof int4")]
108fn generate_subscripts(array: ListRef<'_>, dim: i32) -> impl Iterator<Item = i32> {
109    generate_subscripts_iterator(array, dim, false)
110}
111
112#[auto_enum(Iterator)]
113fn generate_subscripts_iterator(
114    array: ListRef<'_>,
115    dim: i32,
116    reverse: bool,
117) -> impl Iterator<Item = i32> {
118    let (cur, end) = generate_subscripts_inner(array, dim);
119
120    if reverse {
121        return (cur..end).rev();
122    } else {
123        return cur..end;
124    }
125}
126
127fn generate_subscripts_inner(array: ListRef<'_>, dim: i32) -> (i32, i32) {
128    let nothing = (0, 0);
129    match dim {
130        ..=0 => nothing,
131        1 => (1, array.len() as i32 + 1),
132        // Although RW's array can be zig-zag, we just look at the first element.
133        2.. => match array.get(0) {
134            Some(Some(ScalarRefImpl::List(list))) => generate_subscripts_inner(list, dim - 1),
135            _ => nothing,
136        },
137    }
138}
risingwave_expr_impl/table_function/generate_subscripts.rs

risingwave_expr_impl/table_function/
generate_subscripts.rs
