Skip to content

Commit b210b54

Browse files
TCeasonTCeason
committed
support array_sort array_reverse_sort
1 parent f58d32b commit b210b54

File tree

5 files changed

+253
-0
lines changed

5 files changed

+253
-0
lines changed

‎src/query/functions/src/scalars/array.rs

+73
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,10 @@ use common_expression::vectorize_with_builder_1_arg;
4444
use common_expression::vectorize_with_builder_2_arg;
4545
use common_expression::vectorize_with_builder_3_arg;
4646
use common_expression::with_number_mapped_type;
47+
use common_expression::BlockEntry;
4748
use common_expression::Column;
4849
use common_expression::ColumnBuilder;
50+
use common_expression::DataBlock;
4951
use common_expression::Domain;
5052
use common_expression::EvalContext;
5153
use common_expression::Function;
@@ -55,6 +57,7 @@ use common_expression::FunctionRegistry;
5557
use common_expression::FunctionSignature;
5658
use common_expression::Scalar;
5759
use common_expression::ScalarRef;
60+
use common_expression::SortColumnDescription;
5861
use common_expression::Value;
5962
use common_expression::ValueRef;
6063
use common_hashtable::HashtableKeyable;
@@ -134,6 +137,76 @@ pub fn register(registry: &mut FunctionRegistry) {
134137
}))
135138
});
136139

140+
registry.register_passthrough_nullable_1_arg::<EmptyArrayType, EmptyArrayType, _, _>(
141+
"array_sort",
142+
FunctionProperty::default(),
143+
|_| FunctionDomain::Full,
144+
vectorize_1_arg::<EmptyArrayType, EmptyArrayType>(|arr, _| arr),
145+
);
146+
147+
registry.register_passthrough_nullable_1_arg::<ArrayType<GenericType<0>>, ArrayType<GenericType<0>>, _, _>(
148+
"array_sort",
149+
FunctionProperty::default(),
150+
|_| FunctionDomain::Full,
151+
vectorize_1_arg::<ArrayType<GenericType<0>>, ArrayType<GenericType<0>>>(|arr, _| {
152+
let data_type = arr.data_type();
153+
let mut builder = ColumnBuilder::with_capacity(&data_type.clone(), arr.len());
154+
arr.iter().for_each(|val| { builder.push(val) });
155+
let sort_desc = vec![SortColumnDescription {
156+
offset: 0,
157+
asc: true,
158+
nulls_first: false,
159+
}];
160+
let columns = vec![builder.build()];
161+
let len = columns.get(0).map_or(1, |c| c.len());
162+
let columns = columns
163+
.iter()
164+
.map(|col| BlockEntry {
165+
data_type: col.data_type(),
166+
value: Value::Column(col.clone()),
167+
})
168+
.collect();
169+
let sort_block = DataBlock::sort(&DataBlock::new(columns, len), &sort_desc, None).unwrap();
170+
sort_block.columns()[0].value.convert_to_full_column(&data_type, 1)
171+
},
172+
),
173+
);
174+
175+
registry.register_passthrough_nullable_1_arg::<EmptyArrayType, EmptyArrayType, _, _>(
176+
"array_reverse_sort",
177+
FunctionProperty::default(),
178+
|_| FunctionDomain::Full,
179+
vectorize_1_arg::<EmptyArrayType, EmptyArrayType>(|arr, _| arr),
180+
);
181+
182+
registry.register_passthrough_nullable_1_arg::<ArrayType<GenericType<0>>, ArrayType<GenericType<0>>, _, _>(
183+
"array_reverse_sort",
184+
FunctionProperty::default(),
185+
|_| FunctionDomain::Full,
186+
vectorize_1_arg::<ArrayType<GenericType<0>>, ArrayType<GenericType<0>>>(|arr, _| {
187+
let data_type = arr.data_type();
188+
let mut builder = ColumnBuilder::with_capacity(&data_type.clone(), arr.len());
189+
arr.iter().for_each(|val| { builder.push(val) });
190+
let sort_desc = vec![SortColumnDescription {
191+
offset: 0,
192+
asc: false,
193+
nulls_first: false,
194+
}];
195+
let columns = vec![builder.build()];
196+
let len = columns.get(0).map_or(1, |c| c.len());
197+
let columns = columns
198+
.iter()
199+
.map(|col| BlockEntry {
200+
data_type: col.data_type(),
201+
value: Value::Column(col.clone()),
202+
})
203+
.collect();
204+
let sort_block = DataBlock::sort(&DataBlock::new(columns, len), &sort_desc, None).unwrap();
205+
sort_block.columns()[0].value.convert_to_full_column(&data_type, 1)
206+
},
207+
),
208+
);
209+
137210
registry.register_1_arg::<EmptyArrayType, NumberType<u8>, _, _>(
138211
"length",
139212
FunctionProperty::default(),

‎src/query/functions/tests/it/scalars/array.rs

+36
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ fn test_array() {
4343
test_array_count(file);
4444
test_array_max(file);
4545
test_array_min(file);
46+
test_array_sort(file);
47+
test_array_reverse_sort(file);
4648
}
4749

4850
fn test_create(file: &mut impl Write) {
@@ -436,3 +438,37 @@ fn test_array_min(file: &mut impl Write) {
436438
),
437439
]);
438440
}
441+
442+
fn test_array_sort(file: &mut impl Write) {
443+
run_ast(file, "array_sort([])", &[]);
444+
run_ast(file, "array_sort(NULL)", &[]);
445+
run_ast(file, "array_sort([8, 20, 1, 2, 3, 4, 5, 6, 7])", &[]);
446+
run_ast(file, "array_sort([9.32, 0, 1.2, 3.4, 5.6, 7.8])", &[]);
447+
run_ast(file, "array_sort(['x', 0, 1.2, 3.4, 5.6, 7.8])", &[]);
448+
run_ast(file, "array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL])", &[]);
449+
run_ast(file, "array_sort(['z', 'a', NULL, 'v', 'd', NULL])", &[]);
450+
}
451+
452+
fn test_array_reverse_sort(file: &mut impl Write) {
453+
run_ast(file, "array_reverse_sort([])", &[]);
454+
run_ast(file, "array_reverse_sort(NULL)", &[]);
455+
run_ast(file, "array_reverse_sort([8, 20, 1, 2, 3, 4, 5, 6, 7])", &[
456+
]);
457+
run_ast(
458+
file,
459+
"array_reverse_sort([9.32, 0, 1.2, 3.4, 5.6, 7.8])",
460+
&[],
461+
);
462+
run_ast(file, "array_reverse_sort(['x', 0, 1.2, 3.4, 5.6, 7.8])", &[
463+
]);
464+
run_ast(
465+
file,
466+
"array_reverse_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL])",
467+
&[],
468+
);
469+
run_ast(
470+
file,
471+
"array_reverse_sort(['z', 'a', NULL, 'v', 'd', NULL])",
472+
&[],
473+
);
474+
}

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

+126
Original file line numberDiff line numberDiff line change
@@ -1526,3 +1526,129 @@ evaluation (internal):
15261526
+--------+-------------------------------------------------------------------------+
15271527

15281528

1529+
ast : array_sort([])
1530+
raw expr : array_sort(array())
1531+
checked expr : array_sort<Array(Nothing)>(array<>())
1532+
optimized expr : [] :: Array(Nothing)
1533+
output type : Array(Nothing)
1534+
output domain : []
1535+
output : []
1536+
1537+
1538+
ast : array_sort(NULL)
1539+
raw expr : array_sort(NULL)
1540+
checked expr : array_sort<Array(Nothing) NULL>(CAST(NULL AS Array(Nothing) NULL))
1541+
optimized expr : NULL
1542+
output type : Array(Nothing) NULL
1543+
output domain : {NULL}
1544+
output : NULL
1545+
1546+
1547+
ast : array_sort([8, 20, 1, 2, 3, 4, 5, 6, 7])
1548+
raw expr : array_sort(array(8_u8, 20_u8, 1_u8, 2_u8, 3_u8, 4_u8, 5_u8, 6_u8, 7_u8))
1549+
checked expr : array_sort<T0=UInt8><Array(T0)>(array<T0=UInt8><T0, T0, T0, T0, T0, T0, T0, T0, T0>(8_u8, 20_u8, 1_u8, 2_u8, 3_u8, 4_u8, 5_u8, 6_u8, 7_u8))
1550+
optimized expr : [1, 2, 3, 4, 5, 6, 7, 8, 20]
1551+
output type : Array(UInt8)
1552+
output domain : [{1..=20}]
1553+
output : [1, 2, 3, 4, 5, 6, 7, 8, 20]
1554+
1555+
1556+
ast : array_sort([9.32, 0, 1.2, 3.4, 5.6, 7.8])
1557+
raw expr : array_sort(array(9.32_f64, 0_u8, 1.2_f64, 3.4_f64, 5.6_f64, 7.8_f64))
1558+
checked expr : array_sort<T0=Float64><Array(T0)>(array<T0=Float64><T0, T0, T0, T0, T0, T0>(9.32_f64, to_float64<UInt8>(0_u8), 1.2_f64, 3.4_f64, 5.6_f64, 7.8_f64))
1559+
optimized expr : [0, 1.2, 3.4, 5.6, 7.8, 9.32]
1560+
output type : Array(Float64)
1561+
output domain : [{0..=9.32}]
1562+
output : [0, 1.2, 3.4, 5.6, 7.8, 9.32]
1563+
1564+
1565+
ast : array_sort(['x', 0, 1.2, 3.4, 5.6, 7.8])
1566+
raw expr : array_sort(array("x", 0_u8, 1.2_f64, 3.4_f64, 5.6_f64, 7.8_f64))
1567+
checked expr : array_sort<T0=Variant><Array(T0)>(array<T0=Variant><T0, T0, T0, T0, T0, T0>(to_variant<T0=String><T0>("x"), to_variant<T0=UInt8><T0>(0_u8), to_variant<T0=Float64><T0>(1.2_f64), to_variant<T0=Float64><T0>(3.4_f64), to_variant<T0=Float64><T0>(5.6_f64), to_variant<T0=Float64><T0>(7.8_f64)))
1568+
optimized expr : [0, 1.2, 3.4, 5.6, 7.8, "x"]
1569+
output type : Array(Variant)
1570+
output domain : [Undefined]
1571+
output : [0, 1.2, 3.4, 5.6, 7.8, "x"]
1572+
1573+
1574+
ast : array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL])
1575+
raw expr : array_sort(array(1.2_f64, NULL, 3.4_f64, 5.6_f64, "2.2", NULL))
1576+
checked expr : array_sort<T0=Variant NULL><Array(T0)>(array<T0=Variant NULL><T0, T0, T0, T0, T0, T0>(CAST(1.2_f64 AS Variant NULL), CAST(NULL AS Variant NULL), CAST(3.4_f64 AS Variant NULL), CAST(5.6_f64 AS Variant NULL), CAST("2.2" AS Variant NULL), CAST(NULL AS Variant NULL)))
1577+
optimized expr : [1.2, 3.4, 5.6, "2.2", NULL, NULL]
1578+
output type : Array(Variant NULL)
1579+
output domain : [Undefined ∪ {NULL}]
1580+
output : [1.2, 3.4, 5.6, "2.2", NULL, NULL]
1581+
1582+
1583+
ast : array_sort(['z', 'a', NULL, 'v', 'd', NULL])
1584+
raw expr : array_sort(array("z", "a", NULL, "v", "d", NULL))
1585+
checked expr : array_sort<T0=String NULL><Array(T0)>(array<T0=String NULL><T0, T0, T0, T0, T0, T0>(CAST("z" AS String NULL), CAST("a" AS String NULL), CAST(NULL AS String NULL), CAST("v" AS String NULL), CAST("d" AS String NULL), CAST(NULL AS String NULL)))
1586+
optimized expr : ["a", "d", "v", "z", NULL, NULL]
1587+
output type : Array(String NULL)
1588+
output domain : [{""..="z"} ∪ {NULL}]
1589+
output : ["a", "d", "v", "z", NULL, NULL]
1590+
1591+
1592+
ast : array_reverse_sort([])
1593+
raw expr : array_reverse_sort(array())
1594+
checked expr : array_reverse_sort<Array(Nothing)>(array<>())
1595+
optimized expr : [] :: Array(Nothing)
1596+
output type : Array(Nothing)
1597+
output domain : []
1598+
output : []
1599+
1600+
1601+
ast : array_reverse_sort(NULL)
1602+
raw expr : array_reverse_sort(NULL)
1603+
checked expr : array_reverse_sort<Array(Nothing) NULL>(CAST(NULL AS Array(Nothing) NULL))
1604+
optimized expr : NULL
1605+
output type : Array(Nothing) NULL
1606+
output domain : {NULL}
1607+
output : NULL
1608+
1609+
1610+
ast : array_reverse_sort([8, 20, 1, 2, 3, 4, 5, 6, 7])
1611+
raw expr : array_reverse_sort(array(8_u8, 20_u8, 1_u8, 2_u8, 3_u8, 4_u8, 5_u8, 6_u8, 7_u8))
1612+
checked expr : array_reverse_sort<T0=UInt8><Array(T0)>(array<T0=UInt8><T0, T0, T0, T0, T0, T0, T0, T0, T0>(8_u8, 20_u8, 1_u8, 2_u8, 3_u8, 4_u8, 5_u8, 6_u8, 7_u8))
1613+
optimized expr : [20, 8, 7, 6, 5, 4, 3, 2, 1]
1614+
output type : Array(UInt8)
1615+
output domain : [{1..=20}]
1616+
output : [20, 8, 7, 6, 5, 4, 3, 2, 1]
1617+
1618+
1619+
ast : array_reverse_sort([9.32, 0, 1.2, 3.4, 5.6, 7.8])
1620+
raw expr : array_reverse_sort(array(9.32_f64, 0_u8, 1.2_f64, 3.4_f64, 5.6_f64, 7.8_f64))
1621+
checked expr : array_reverse_sort<T0=Float64><Array(T0)>(array<T0=Float64><T0, T0, T0, T0, T0, T0>(9.32_f64, to_float64<UInt8>(0_u8), 1.2_f64, 3.4_f64, 5.6_f64, 7.8_f64))
1622+
optimized expr : [9.32, 7.8, 5.6, 3.4, 1.2, 0]
1623+
output type : Array(Float64)
1624+
output domain : [{0..=9.32}]
1625+
output : [9.32, 7.8, 5.6, 3.4, 1.2, 0]
1626+
1627+
1628+
ast : array_reverse_sort(['x', 0, 1.2, 3.4, 5.6, 7.8])
1629+
raw expr : array_reverse_sort(array("x", 0_u8, 1.2_f64, 3.4_f64, 5.6_f64, 7.8_f64))
1630+
checked expr : array_reverse_sort<T0=Variant><Array(T0)>(array<T0=Variant><T0, T0, T0, T0, T0, T0>(to_variant<T0=String><T0>("x"), to_variant<T0=UInt8><T0>(0_u8), to_variant<T0=Float64><T0>(1.2_f64), to_variant<T0=Float64><T0>(3.4_f64), to_variant<T0=Float64><T0>(5.6_f64), to_variant<T0=Float64><T0>(7.8_f64)))
1631+
optimized expr : ["x", 7.8, 5.6, 3.4, 1.2, 0]
1632+
output type : Array(Variant)
1633+
output domain : [Undefined]
1634+
output : ["x", 7.8, 5.6, 3.4, 1.2, 0]
1635+
1636+
1637+
ast : array_reverse_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL])
1638+
raw expr : array_reverse_sort(array(1.2_f64, NULL, 3.4_f64, 5.6_f64, "2.2", NULL))
1639+
checked expr : array_reverse_sort<T0=Variant NULL><Array(T0)>(array<T0=Variant NULL><T0, T0, T0, T0, T0, T0>(CAST(1.2_f64 AS Variant NULL), CAST(NULL AS Variant NULL), CAST(3.4_f64 AS Variant NULL), CAST(5.6_f64 AS Variant NULL), CAST("2.2" AS Variant NULL), CAST(NULL AS Variant NULL)))
1640+
optimized expr : ["2.2", 5.6, 3.4, 1.2, NULL, NULL]
1641+
output type : Array(Variant NULL)
1642+
output domain : [Undefined ∪ {NULL}]
1643+
output : ["2.2", 5.6, 3.4, 1.2, NULL, NULL]
1644+
1645+
1646+
ast : array_reverse_sort(['z', 'a', NULL, 'v', 'd', NULL])
1647+
raw expr : array_reverse_sort(array("z", "a", NULL, "v", "d", NULL))
1648+
checked expr : array_reverse_sort<T0=String NULL><Array(T0)>(array<T0=String NULL><T0, T0, T0, T0, T0, T0>(CAST("z" AS String NULL), CAST("a" AS String NULL), CAST(NULL AS String NULL), CAST("v" AS String NULL), CAST("d" AS String NULL), CAST(NULL AS String NULL)))
1649+
optimized expr : ["z", "v", "d", "a", NULL, NULL]
1650+
output type : Array(String NULL)
1651+
output domain : [{""..="z"} ∪ {NULL}]
1652+
output : ["z", "v", "d", "a", NULL, NULL]
1653+
1654+

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

+8
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,14 @@ array_remove_last(Array(Nothing)) :: Array(Nothing)
6060
array_remove_last(Array(Nothing) NULL) :: Array(Nothing) NULL
6161
array_remove_last(Array(T0)) :: Array(T0)
6262
array_remove_last(Array(T0) NULL) :: Array(T0) NULL
63+
array_reverse_sort(Array(Nothing)) :: Array(Nothing)
64+
array_reverse_sort(Array(Nothing) NULL) :: Array(Nothing) NULL
65+
array_reverse_sort(Array(T0)) :: Array(T0)
66+
array_reverse_sort(Array(T0) NULL) :: Array(T0) NULL
67+
array_sort(Array(Nothing)) :: Array(Nothing)
68+
array_sort(Array(Nothing) NULL) :: Array(Nothing) NULL
69+
array_sort(Array(T0)) :: Array(T0)
70+
array_sort(Array(T0) NULL) :: Array(T0) NULL
6371
array_unique(Array(Nothing)) :: UInt64
6472
array_unique(Array(Nothing) NULL) :: UInt64 NULL
6573
array_unique(Array(T0)) :: UInt64

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

+10
Original file line numberDiff line numberDiff line change
@@ -101,5 +101,15 @@ select array_min(col1), array_min(col2), array_min(col3) from t
101101
----
102102
1 x 2022-02-02
103103

104+
query TTTTT
105+
select array_sort(col1),array_sort(col2),array_sort(col3),array_sort(col4),array_sort(col5) from t
106+
----
107+
[1,2,3,3] ['x','x','y','z'] ['2022-02-02'] ['2023-01-01 02:00:01.000000'] [[],[NULL],[1,2]]
108+
109+
query TTTTT
110+
select array_reverse_sort(col1),array_reverse_sort(col2),array_reverse_sort(col3),array_reverse_sort(col4),array_reverse_sort(col5) from t
111+
----
112+
[3,3,2,1] ['z','y','x','x'] ['2022-02-02'] ['2023-01-01 02:00:01.000000'] [[1,2],[NULL],[]]
113+
104114
statement ok
105115
DROP DATABASE array_func_test

0 commit comments

Comments
 (0)