Skip to content

Commit f1c7624

Browse files
b41shb41sh
authored
fix(functions): fix unnest with null argument (#10764)
* fix(functions): fix unnest with null argument * fix
1 parent 363a831 commit f1c7624

File tree

4 files changed

+222
-18
lines changed

4 files changed

+222
-18
lines changed

‎src/query/functions/src/srfs/mod.rs‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use common_expression::types::array::ArrayColumnBuilder;
1616
use common_expression::types::AnyType;
1717
use common_expression::types::DataType;
1818
use common_expression::Column;
19+
use common_expression::ColumnBuilder;
1920
use common_expression::Function;
2021
use common_expression::FunctionEval;
2122
use common_expression::FunctionKind;
@@ -39,6 +40,10 @@ pub fn register(registry: &mut FunctionRegistry) {
3940
ArrayColumnBuilder::<AnyType>::repeat(&col, num_rows).build()
4041
}
4142
Value::Column(Column::Array(col)) => *col,
43+
Value::Column(Column::Nullable(box nullable_column)) => match nullable_column.column {
44+
Column::Array(col) => *col,
45+
_ => unreachable!(),
46+
},
4247
_ => unreachable!(),
4348
};
4449
debug_assert_eq!(unnest_array.len(), num_rows);
@@ -63,6 +68,18 @@ pub fn register(registry: &mut FunctionRegistry) {
6368
))))),
6469
))))),
6570
))))];
71+
registry.register_function(Function {
72+
signature: FunctionSignature {
73+
name: "unnest".to_string(),
74+
args_type: vec![args_type[0].wrap_nullable()],
75+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
76+
DataType::Generic(0),
77+
))]),
78+
},
79+
eval: FunctionEval::SRF {
80+
eval: Box::new(unnest_impl),
81+
},
82+
});
6683
registry.register_function(Function {
6784
signature: FunctionSignature {
6885
name: "unnest".to_string(),
@@ -88,6 +105,18 @@ pub fn register(registry: &mut FunctionRegistry) {
88105
))))),
89106
))))),
90107
))))];
108+
registry.register_function(Function {
109+
signature: FunctionSignature {
110+
name: "unnest".to_string(),
111+
args_type: vec![args_type[0].wrap_nullable()],
112+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
113+
DataType::Generic(0),
114+
))]),
115+
},
116+
eval: FunctionEval::SRF {
117+
eval: Box::new(unnest_impl),
118+
},
119+
});
91120
registry.register_function(Function {
92121
signature: FunctionSignature {
93122
name: "unnest".to_string(),
@@ -111,6 +140,18 @@ pub fn register(registry: &mut FunctionRegistry) {
111140
))))),
112141
))))),
113142
))))];
143+
registry.register_function(Function {
144+
signature: FunctionSignature {
145+
name: "unnest".to_string(),
146+
args_type: vec![args_type[0].wrap_nullable()],
147+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
148+
DataType::Generic(0),
149+
))]),
150+
},
151+
eval: FunctionEval::SRF {
152+
eval: Box::new(unnest_impl),
153+
},
154+
});
114155
registry.register_function(Function {
115156
signature: FunctionSignature {
116157
name: "unnest".to_string(),
@@ -134,6 +175,18 @@ pub fn register(registry: &mut FunctionRegistry) {
134175
))))),
135176
))))),
136177
))))];
178+
registry.register_function(Function {
179+
signature: FunctionSignature {
180+
name: "unnest".to_string(),
181+
args_type: vec![args_type[0].wrap_nullable()],
182+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
183+
DataType::Generic(0),
184+
))]),
185+
},
186+
eval: FunctionEval::SRF {
187+
eval: Box::new(unnest_impl),
188+
},
189+
});
137190
registry.register_function(Function {
138191
signature: FunctionSignature {
139192
name: "unnest".to_string(),
@@ -157,6 +210,18 @@ pub fn register(registry: &mut FunctionRegistry) {
157210
))))),
158211
))))),
159212
))))];
213+
registry.register_function(Function {
214+
signature: FunctionSignature {
215+
name: "unnest".to_string(),
216+
args_type: vec![args_type[0].wrap_nullable()],
217+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
218+
DataType::Generic(0),
219+
))]),
220+
},
221+
eval: FunctionEval::SRF {
222+
eval: Box::new(unnest_impl),
223+
},
224+
});
160225
registry.register_function(Function {
161226
signature: FunctionSignature {
162227
name: "unnest".to_string(),
@@ -178,6 +243,18 @@ pub fn register(registry: &mut FunctionRegistry) {
178243
Box::new(DataType::Nullable(Box::new(DataType::Generic(0)))),
179244
))))),
180245
))))];
246+
registry.register_function(Function {
247+
signature: FunctionSignature {
248+
name: "unnest".to_string(),
249+
args_type: vec![args_type[0].wrap_nullable()],
250+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
251+
DataType::Generic(0),
252+
))]),
253+
},
254+
eval: FunctionEval::SRF {
255+
eval: Box::new(unnest_impl),
256+
},
257+
});
181258
registry.register_function(Function {
182259
signature: FunctionSignature {
183260
name: "unnest".to_string(),
@@ -199,6 +276,18 @@ pub fn register(registry: &mut FunctionRegistry) {
199276
Box::new(DataType::Generic(0)),
200277
))))),
201278
))))];
279+
registry.register_function(Function {
280+
signature: FunctionSignature {
281+
name: "unnest".to_string(),
282+
args_type: vec![args_type[0].wrap_nullable()],
283+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
284+
DataType::Generic(0),
285+
))]),
286+
},
287+
eval: FunctionEval::SRF {
288+
eval: Box::new(unnest_impl),
289+
},
290+
});
202291
registry.register_function(Function {
203292
signature: FunctionSignature {
204293
name: "unnest".to_string(),
@@ -218,6 +307,18 @@ pub fn register(registry: &mut FunctionRegistry) {
218307
let args_type = vec![DataType::Array(Box::new(DataType::Array(Box::new(
219308
DataType::Array(Box::new(DataType::Nullable(Box::new(DataType::Generic(0))))),
220309
))))];
310+
registry.register_function(Function {
311+
signature: FunctionSignature {
312+
name: "unnest".to_string(),
313+
args_type: vec![args_type[0].wrap_nullable()],
314+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
315+
DataType::Generic(0),
316+
))]),
317+
},
318+
eval: FunctionEval::SRF {
319+
eval: Box::new(unnest_impl),
320+
},
321+
});
221322
registry.register_function(Function {
222323
signature: FunctionSignature {
223324
name: "unnest".to_string(),
@@ -237,6 +338,18 @@ pub fn register(registry: &mut FunctionRegistry) {
237338
let args_type = vec![DataType::Array(Box::new(DataType::Array(Box::new(
238339
DataType::Nullable(Box::new(DataType::Generic(0))),
239340
))))];
341+
registry.register_function(Function {
342+
signature: FunctionSignature {
343+
name: "unnest".to_string(),
344+
args_type: vec![args_type[0].wrap_nullable()],
345+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
346+
DataType::Generic(0),
347+
))]),
348+
},
349+
eval: FunctionEval::SRF {
350+
eval: Box::new(unnest_impl),
351+
},
352+
});
240353
registry.register_function(Function {
241354
signature: FunctionSignature {
242355
name: "unnest".to_string(),
@@ -256,6 +369,18 @@ pub fn register(registry: &mut FunctionRegistry) {
256369
let args_type = vec![DataType::Array(Box::new(DataType::Nullable(Box::new(
257370
DataType::Generic(0),
258371
))))];
372+
registry.register_function(Function {
373+
signature: FunctionSignature {
374+
name: "unnest".to_string(),
375+
args_type: vec![args_type[0].wrap_nullable()],
376+
return_type: DataType::Tuple(vec![DataType::Nullable(Box::new(
377+
DataType::Generic(0),
378+
))]),
379+
},
380+
eval: FunctionEval::SRF {
381+
eval: Box::new(unnest_impl),
382+
},
383+
});
259384
registry.register_function(Function {
260385
signature: FunctionSignature {
261386
name: "unnest".to_string(),
@@ -269,4 +394,26 @@ pub fn register(registry: &mut FunctionRegistry) {
269394
},
270395
});
271396
}
397+
398+
{
399+
// Unnest NULL
400+
let args_type = vec![DataType::Null];
401+
registry.register_function(Function {
402+
signature: FunctionSignature {
403+
name: "unnest".to_string(),
404+
args_type,
405+
return_type: DataType::Tuple(vec![DataType::Null]),
406+
},
407+
eval: FunctionEval::SRF {
408+
eval: Box::new(|_, num_rows| {
409+
let mut columns = Vec::with_capacity(num_rows);
410+
(0..num_rows).for_each(|_| {
411+
let column = ColumnBuilder::with_capacity(&DataType::Null, 0).build();
412+
columns.push((Value::Column(Column::Tuple(vec![column])), 0));
413+
});
414+
columns
415+
}),
416+
},
417+
});
418+
}
272419
}

‎src/query/functions/tests/it/scalars/testdata/function_list.txt‎

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3757,16 +3757,27 @@ Functions overloads:
37573757
0 typeof(T0) :: String
37583758
0 unhex(String) :: String
37593759
1 unhex(String NULL) :: String NULL
3760-
0 unnest(Array(Array(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL))))))))))) :: Tuple(T0 NULL,)
3761-
1 unnest(Array(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL)))))))))) :: Tuple(T0 NULL,)
3762-
2 unnest(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL))))))))) :: Tuple(T0 NULL,)
3763-
3 unnest(Array(Array(Array(Array(Array(Array(Array(T0 NULL)))))))) :: Tuple(T0 NULL,)
3764-
4 unnest(Array(Array(Array(Array(Array(Array(T0 NULL))))))) :: Tuple(T0 NULL,)
3765-
5 unnest(Array(Array(Array(Array(Array(T0 NULL)))))) :: Tuple(T0 NULL,)
3766-
6 unnest(Array(Array(Array(Array(T0 NULL))))) :: Tuple(T0 NULL,)
3767-
7 unnest(Array(Array(Array(T0 NULL)))) :: Tuple(T0 NULL,)
3768-
8 unnest(Array(Array(T0 NULL))) :: Tuple(T0 NULL,)
3769-
9 unnest(Array(T0 NULL)) :: Tuple(T0 NULL,)
3760+
0 unnest(Array(Array(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL)))))))))) NULL) :: Tuple(T0 NULL,)
3761+
1 unnest(Array(Array(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL))))))))))) :: Tuple(T0 NULL,)
3762+
2 unnest(Array(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL))))))))) NULL) :: Tuple(T0 NULL,)
3763+
3 unnest(Array(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL)))))))))) :: Tuple(T0 NULL,)
3764+
4 unnest(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL)))))))) NULL) :: Tuple(T0 NULL,)
3765+
5 unnest(Array(Array(Array(Array(Array(Array(Array(Array(T0 NULL))))))))) :: Tuple(T0 NULL,)
3766+
6 unnest(Array(Array(Array(Array(Array(Array(Array(T0 NULL))))))) NULL) :: Tuple(T0 NULL,)
3767+
7 unnest(Array(Array(Array(Array(Array(Array(Array(T0 NULL)))))))) :: Tuple(T0 NULL,)
3768+
8 unnest(Array(Array(Array(Array(Array(Array(T0 NULL)))))) NULL) :: Tuple(T0 NULL,)
3769+
9 unnest(Array(Array(Array(Array(Array(Array(T0 NULL))))))) :: Tuple(T0 NULL,)
3770+
10 unnest(Array(Array(Array(Array(Array(T0 NULL))))) NULL) :: Tuple(T0 NULL,)
3771+
11 unnest(Array(Array(Array(Array(Array(T0 NULL)))))) :: Tuple(T0 NULL,)
3772+
12 unnest(Array(Array(Array(Array(T0 NULL)))) NULL) :: Tuple(T0 NULL,)
3773+
13 unnest(Array(Array(Array(Array(T0 NULL))))) :: Tuple(T0 NULL,)
3774+
14 unnest(Array(Array(Array(T0 NULL))) NULL) :: Tuple(T0 NULL,)
3775+
15 unnest(Array(Array(Array(T0 NULL)))) :: Tuple(T0 NULL,)
3776+
16 unnest(Array(Array(T0 NULL)) NULL) :: Tuple(T0 NULL,)
3777+
17 unnest(Array(Array(T0 NULL))) :: Tuple(T0 NULL,)
3778+
18 unnest(Array(T0 NULL) NULL) :: Tuple(T0 NULL,)
3779+
19 unnest(Array(T0 NULL)) :: Tuple(T0 NULL,)
3780+
20 unnest(NULL) :: Tuple(NULL,)
37703781
0 upper(String) :: String
37713782
1 upper(String NULL) :: String NULL
37723783
0 xor(Boolean, Boolean) :: Boolean

‎src/query/sql/src/evaluator/block_operator.rs‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ use common_expression::Evaluator;
2626
use common_expression::Expr;
2727
use common_expression::FieldIndex;
2828
use common_expression::FunctionContext;
29+
use common_expression::ScalarRef;
2930
use common_expression::Value;
3031
use common_functions::BUILTIN_FUNCTIONS;
3132
use common_pipeline_core::processors::port::InputPort;
@@ -140,14 +141,28 @@ impl BlockOperator {
140141
// TODO(leiysky): this can be optimized by using a `zip` array function
141142
if repeat_times < max_num_rows {
142143
for field in fields {
143-
let nullable_column = field.as_nullable().unwrap();
144-
let mut column_builder = NullableColumnBuilder::from_column(
145-
(**nullable_column).clone(),
146-
);
147-
(0..(max_num_rows - repeat_times)).for_each(|_| {
148-
column_builder.push_null();
149-
});
150-
*field = Column::Nullable(Box::new(column_builder.build()));
144+
match field {
145+
Column::Null { .. } => {
146+
*field = ColumnBuilder::repeat(
147+
&ScalarRef::Null,
148+
max_num_rows,
149+
&DataType::Null,
150+
)
151+
.build();
152+
}
153+
Column::Nullable(box nullable_column) => {
154+
let mut column_builder =
155+
NullableColumnBuilder::from_column(
156+
(*nullable_column).clone(),
157+
);
158+
(0..(max_num_rows - repeat_times)).for_each(|_| {
159+
column_builder.push_null();
160+
});
161+
*field =
162+
Column::Nullable(Box::new(column_builder.build()));
163+
}
164+
_ => unreachable!(),
165+
}
151166
}
152167
}
153168
}

‎tests/sqllogictests/suites/query/02_function/02_0062_function_unnest‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@ query I
22
select unnest([]::array(int));
33
----
44

5+
query I
6+
select unnest(null);
7+
----
8+
59
query I
610
select unnest([1,2,3]);
711
----
@@ -252,6 +256,33 @@ select unnest(a), a from t;
252256
statement ok
253257
drop table t;
254258

259+
statement ok
260+
create table t (a array(int) null);
261+
262+
statement ok
263+
insert into t values ([1,2]), (null), ([3,4,5]);
264+
265+
query I
266+
select unnest(a) from t;
267+
----
268+
1
269+
2
270+
3
271+
4
272+
5
273+
274+
query IT
275+
select unnest(a), unnest(null) from t;
276+
----
277+
1 NULL
278+
2 NULL
279+
3 NULL
280+
4 NULL
281+
5 NULL
282+
283+
statement ok
284+
drop table t;
285+
255286
statement ok
256287
DROP DATABASE db_02_0062;
257288

0 commit comments

Comments
 (0)