Skip to content

Commit 63ab393

Browse files
gstvggstvg
committed
impl null arg execution
1 parent 6f051c5 commit 63ab393

8 files changed

Lines changed: 42 additions & 5 deletions

File tree

‎datafusion/functions-nested/src/lambda_utils.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ pub(crate) fn coerce_single_list_arg(
6565
DataType::List(Arc::clone(field))
6666
}
6767
DataType::LargeListView(field) => DataType::LargeList(Arc::clone(field)),
68-
DataType::Null => DataType::Null,
68+
DataType::Null => DataType::new_list(DataType::Null, true),
6969
_ => return plan_err!("{name} expected a list as first argument, got {list}"),
7070
};
7171

‎datafusion/functions-nested/src/map_extract.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,11 @@ impl ScalarUDFImpl for MapExtract {
105105

106106
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
107107
let [map_type, _] = take_function_args(self.name(), arg_types)?;
108+
109+
if map_type.is_null() {
110+
return Ok(DataType::Null);
111+
}
112+
108113
let map_fields = get_map_entry_field(map_type)?;
109114
Ok(DataType::List(Arc::new(Field::new_list_field(
110115
map_fields.last().unwrap().data_type().clone(),
@@ -189,6 +194,7 @@ fn map_extract_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
189194

190195
let map_array = match map_arg.data_type() {
191196
DataType::Map(_, _) => as_map_array(&map_arg)?,
197+
DataType::Null => return Ok(Arc::clone(map_arg)),
192198
_ => return exec_err!("The first argument in map_extract must be a map"),
193199
};
194200

‎datafusion/spark/src/function/array/repeat.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -74,13 +74,14 @@ impl ScalarUDFImpl for SparkArrayRepeat {
7474

7575
// Coerce the second argument to Int64/UInt64 if it's a numeric type
7676
let second = match second_type {
77-
DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
78-
DataType::Int64
79-
}
77+
DataType::Int8
78+
| DataType::Int16
79+
| DataType::Int32
80+
| DataType::Int64
81+
| DataType::Null => DataType::Int64,
8082
DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => {
8183
DataType::UInt64
8284
}
83-
DataType::Null => DataType::Null,
8485
_ => return exec_err!("count must be an integer type"),
8586
};
8687

‎datafusion/sqllogictest/test_files/array/array_any_match.slt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,12 @@ SELECT list_any_match([1, 2, 3], x -> x > 2);
103103
----
104104
true
105105

106+
# null arg
107+
query B
108+
SELECT array_any_match(NULL, x -> x > 2);
109+
----
110+
NULL
111+
106112
statement ok
107113
drop table t;
108114

‎datafusion/sqllogictest/test_files/array/array_filter.slt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,12 @@ SELECT array_transform(array_filter(list, v -> v > 1), v -> v * 3) FROM with_nul
204204
[6]
205205
NULL
206206

207+
# null arg
208+
query ?
209+
SELECT array_filter(NULL, x -> x > 2);
210+
----
211+
NULL
212+
207213
statement ok
208214
drop table t;
209215

‎datafusion/sqllogictest/test_files/array/array_transform.slt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,12 @@ physical_plan
393393
02)--ProjectionExec: expr=[text@0 as text, list@1 as list, number@2 as number, CASE WHEN number@2 > 30 THEN array_transform(make_array(make_array(list@1)), (list) -> array_transform(list@3, (list) -> array_transform(list@4, (v) -> number@2 + v@5 + array_element(list@4, 1)))) ELSE array_transform(make_array(make_array(list@1)), (list) -> array_transform(list@3, (list) -> array_transform(list@4, (v) -> number@2 + array_element(list@4, 1)))) END as CASE WHEN t.number > Int64(30) THEN array_transform(make_array(make_array(t.list)),(list) -> array_transform(list,(list) -> array_transform(list,(v) -> t.number + v + list[Int64(1)]))) ELSE array_transform(make_array(make_array(t.list)),(list) -> array_transform(list,(list) -> array_transform(list,(v) -> t.number + list[Int64(1)]))) END]
394394
03)----DataSourceExec: partitions=1, partition_sizes=[1]
395395

396+
# null arg
397+
query ?
398+
SELECT array_transform(NULL, x -> x * 2);
399+
----
400+
NULL
401+
396402
query error
397403
select array_transform();
398404
----

‎datafusion/sqllogictest/test_files/map.slt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -642,6 +642,12 @@ select map_extract(MAP {1: 1, 2: 2, 3:3}, '1'), map_extract(MAP {1: 1, 2: 2, 3:3
642642
----
643643
[1] [1] [1] [NULL] [1]
644644

645+
# null arg
646+
query ?
647+
select map_extract(NULL, 'a');
648+
----
649+
NULL
650+
645651
# map_extract with columns
646652
query ???
647653
select map_extract(column1, 1), map_extract(column1, 5), map_extract(column1, 7) from map_array_table_1;

‎datafusion/sqllogictest/test_files/spark/string/repeat.slt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,3 +25,9 @@
2525
## PySpark 3.5.5 Result: {'repeat(123, 2)': '123123', 'typeof(repeat(123, 2))': 'string', 'typeof(123)': 'string', 'typeof(2)': 'int'}
2626
#query
2727
#SELECT repeat('123'::string, 2::int);
28+
29+
# null count
30+
query T
31+
select repeat('a', NULL);
32+
----
33+
NULL

0 commit comments

Comments
 (0)