Skip to content

Commit 588d3c1

Browse files
b41shb41sh
authored
feat(query): Array Functions Supporting ARRAY and VARIANT Types (#18213)
feat(query): Add Array Functions Supporting ARRAY and VARIANT Types
1 parent aff7201 commit 588d3c1

File tree

12 files changed

+3523
-240
lines changed

12 files changed

+3523
-240
lines changed

‎src/meta/proto-conv/tests/it/main.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,4 +120,5 @@ mod v125_table_index;
120120
mod v126_iceberg_storage_catalog_option;
121121
mod v127_user_option_workload_group;
122122
mod v128_storage_network_config;
123+
mod v129_vector_datatype;
123124
mod v130_udf_imports_packages;

‎src/meta/proto-conv/tests/it/v129_vector_datatype.rs‎

Lines changed: 275 additions & 0 deletions
Large diffs are not rendered by default.

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

Lines changed: 377 additions & 53 deletions
Large diffs are not rendered by default.

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

Lines changed: 540 additions & 0 deletions
Large diffs are not rendered by default.

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

Lines changed: 195 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,11 @@ fn test_array() {
5353
test_array_skewness(file);
5454
test_array_sort(file);
5555
test_arrays_zip(file);
56+
test_array_compact(file);
57+
test_array_except(file);
58+
test_array_overlap(file);
59+
test_array_remove(file);
60+
test_array_reverse(file);
5661
}
5762

5863
fn test_create(file: &mut impl Write) {
@@ -81,6 +86,8 @@ fn test_length(file: &mut impl Write) {
8186
fn test_range(file: &mut impl Write) {
8287
run_ast(file, "range(10, 20)", &[]);
8388
run_ast(file, "range(10, 500000011)", &[]);
89+
run_ast(file, "range(10, 20, 3)", &[]);
90+
run_ast(file, "range(-1, -10, -2)", &[]);
8491
}
8592

8693
fn test_get(file: &mut impl Write) {
@@ -119,7 +126,7 @@ fn test_slice(file: &mut impl Write) {
119126
run_ast(file, "slice([0, 1, 2, 3], 1, 5)", &[]);
120127
run_ast(file, "slice(['a', 'b', 'c', 'd'], 0, 2)", &[]);
121128
run_ast(file, "slice(['a', 'b', 'c', 'd'], 1, 4)", &[]);
122-
run_ast(file, "slice(['a', 'b', 'c', 'd'], 2, 6)", &[]);
129+
run_ast(file, "slice(['a', 'b', 'c', 'd'], -3, -1)", &[]);
123130
run_ast(file, "slice([a, b, c], 1, 2)", &[
124131
("a", Int16Type::from_data(vec![0i16, 1, 2])),
125132
("b", Int16Type::from_data(vec![3i16, 4, 5])),
@@ -441,6 +448,11 @@ fn test_array_count(file: &mut impl Write) {
441448
UInt64Type::from_data_with_validity(vec![4u64, 6, 5, 0], vec![true, true, true, false]),
442449
),
443450
]);
451+
452+
// Test with variant type
453+
run_ast(file, "array_count(parse_json('[1, 2, 3, 4, 5]'))", &[]);
454+
run_ast(file, "array_count(parse_json('[1, 2, null, 4, 5]'))", &[]);
455+
run_ast(file, "array_count(parse_json('[1.2, 3.4, 5.6, 7.8]'))", &[]);
444456
}
445457

446458
fn test_array_max(file: &mut impl Write) {
@@ -477,6 +489,15 @@ fn test_array_max(file: &mut impl Write) {
477489
UInt64Type::from_data_with_validity(vec![4u64, 6, 5, 0], vec![true, true, true, false]),
478490
),
479491
]);
492+
493+
// Test with variant type
494+
run_ast(file, "array_max(parse_json('[1, 2, 3, 4, 5]'))", &[]);
495+
run_ast(file, "array_max(parse_json('[1, 2, null, 4, 5]'))", &[]);
496+
run_ast(
497+
file,
498+
"array_max(parse_json('[\"a\", \"b\", \"c\", \"d\"]'))",
499+
&[],
500+
);
480501
}
481502

482503
fn test_array_min(file: &mut impl Write) {
@@ -513,6 +534,15 @@ fn test_array_min(file: &mut impl Write) {
513534
UInt64Type::from_data_with_validity(vec![4u64, 6, 5, 0], vec![true, true, true, false]),
514535
),
515536
]);
537+
538+
// Test with variant type
539+
run_ast(file, "array_min(parse_json('[1, 2, 3, 4, 5]'))", &[]);
540+
run_ast(file, "array_min(parse_json('[1, 2, null, 4, 5]'))", &[]);
541+
run_ast(
542+
file,
543+
"array_min(parse_json('[\"a\", \"b\", \"c\", \"d\"]'))",
544+
&[],
545+
);
516546
}
517547

518548
fn test_array_any(file: &mut impl Write) {
@@ -547,6 +577,11 @@ fn test_array_any(file: &mut impl Write) {
547577
UInt64Type::from_data_with_validity(vec![4u64, 6, 5, 0], vec![true, true, true, false]),
548578
),
549579
]);
580+
581+
// Test with variant type
582+
run_ast(file, "array_any(parse_json('[1, 2, 3]'))", &[]);
583+
run_ast(file, "array_any(parse_json('[null, 2, 3]'))", &[]);
584+
run_ast(file, "array_any(parse_json('[\"a\", \"b\", \"c\"]'))", &[]);
550585
}
551586

552587
fn test_array_stddev_samp(file: &mut impl Write) {
@@ -745,3 +780,162 @@ fn test_arrays_zip(file: &mut impl Write) {
745780
run_ast(file, "arrays_zip([1,2,3], ['a','b','c'], 10)", &[]);
746781
run_ast(file, "arrays_zip([1,2,3], ['a','b'], 10)", &[]);
747782
}
783+
784+
fn test_array_compact(file: &mut impl Write) {
785+
// Test with arrays containing nulls
786+
run_ast(file, "array_compact([])", &[]);
787+
run_ast(file, "array_compact([1, NULL, 2, NULL, 3])", &[]);
788+
run_ast(file, "array_compact([NULL, NULL, NULL])", &[]);
789+
790+
// Test with array without nulls
791+
run_ast(file, "array_compact([1, 2, 3])", &[]);
792+
793+
// Test with column data
794+
run_ast(file, "array_compact([a, b, c])", &[
795+
(
796+
"a",
797+
Int16Type::from_data_with_validity(vec![1i16, 2, 3], vec![true, false, true]),
798+
),
799+
(
800+
"b",
801+
Int16Type::from_data_with_validity(vec![4i16, 5, 6], vec![false, true, false]),
802+
),
803+
(
804+
"c",
805+
Int16Type::from_data_with_validity(vec![7i16, 8, 9], vec![true, true, true]),
806+
),
807+
]);
808+
809+
// Test with variant type
810+
run_ast(
811+
file,
812+
"array_compact(parse_json('[1, null, 2, null, 3]'))",
813+
&[],
814+
);
815+
run_ast(file, "array_compact(parse_json('[null, null]'))", &[]);
816+
run_ast(file, "array_compact(parse_json('[1, 2, 3]'))", &[]);
817+
}
818+
819+
fn test_array_except(file: &mut impl Write) {
820+
// Test with simple arrays
821+
run_ast(file, "array_except([], [])", &[]);
822+
run_ast(file, "array_except([1, 2, 3], [2, 3, 4])", &[]);
823+
run_ast(file, "array_except([1, 2, 3], [4, 5, 6])", &[]);
824+
run_ast(file, "array_except([1, 2, 3], [1, 2, 3])", &[]);
825+
826+
// Test with nulls
827+
run_ast(file, "array_except([1, NULL, 2, 3], [2, 3])", &[]);
828+
run_ast(file, "array_except([1, 2, 3], [NULL, 2, 3])", &[]);
829+
run_ast(file, "array_except([NULL, NULL], [NULL])", &[]);
830+
831+
// Test with column data
832+
run_ast(file, "array_except([a, b], [c, d])", &[
833+
("a", Int16Type::from_data(vec![1i16, 2, 3])),
834+
("b", Int16Type::from_data(vec![4i16, 5, 6])),
835+
("c", Int16Type::from_data(vec![4i16, 5, 7])),
836+
("d", Int16Type::from_data(vec![8i16, 9, 10])),
837+
]);
838+
839+
// Test with variant type
840+
run_ast(
841+
file,
842+
"array_except(parse_json('[1, 2, 3]'), parse_json('[2, 3, 4]'))",
843+
&[],
844+
);
845+
run_ast(
846+
file,
847+
"array_except(parse_json('[1, 2, 3]'), parse_json('[4, 5, 6]'))",
848+
&[],
849+
);
850+
}
851+
852+
fn test_array_overlap(file: &mut impl Write) {
853+
// Test with simple arrays
854+
run_ast(file, "array_overlap([], [])", &[]);
855+
run_ast(file, "array_overlap([1, 2, 3], [2, 3, 4])", &[]);
856+
run_ast(file, "array_overlap([1, 2, 3], [4, 5, 6])", &[]);
857+
run_ast(file, "array_overlap([1, 2, 3], [1, 2, 3])", &[]);
858+
859+
// Test with nulls
860+
run_ast(file, "array_overlap([1, NULL, 2, 3], [2, 3])", &[]);
861+
run_ast(file, "array_overlap([1, 2, 3], [NULL, 2, 3])", &[]);
862+
run_ast(file, "array_overlap([NULL, NULL], [NULL])", &[]);
863+
864+
// Test with column data
865+
run_ast(file, "array_overlap([a, b], [c, d])", &[
866+
("a", Int16Type::from_data(vec![1i16, 2, 3])),
867+
("b", Int16Type::from_data(vec![4i16, 5, 6])),
868+
("c", Int16Type::from_data(vec![4i16, 5, 7])),
869+
("d", Int16Type::from_data(vec![8i16, 9, 10])),
870+
]);
871+
872+
// Test with variant type
873+
run_ast(
874+
file,
875+
"array_overlap(parse_json('[1, 2, 3]'), parse_json('[2, 3, 4]'))",
876+
&[],
877+
);
878+
run_ast(
879+
file,
880+
"array_overlap(parse_json('[1, 2, 3]'), parse_json('[4, 5, 6]'))",
881+
&[],
882+
);
883+
}
884+
885+
fn test_array_remove(file: &mut impl Write) {
886+
// Test with simple arrays
887+
run_ast(file, "array_remove([], 1)", &[]);
888+
run_ast(file, "array_remove([1, 2, 3, 2], 2)", &[]);
889+
run_ast(file, "array_remove([1, 2, 3], 4)", &[]);
890+
891+
// Test with strings
892+
run_ast(file, "array_remove(['a', 'b', 'c', 'b'], 'b')", &[]);
893+
894+
// Test with nulls
895+
run_ast(file, "array_remove([1, NULL, 2, 3], NULL)", &[]);
896+
run_ast(file, "array_remove([1, NULL, 2, 3], 2)", &[]);
897+
898+
// Test with column data
899+
run_ast(file, "array_remove([a, b, c], d)", &[
900+
("a", Int16Type::from_data(vec![1i16, 2, 3])),
901+
("b", Int16Type::from_data(vec![4i16, 5, 6])),
902+
("c", Int16Type::from_data(vec![7i16, 8, 9])),
903+
("d", Int16Type::from_data(vec![4i16, 5, 7])),
904+
]);
905+
906+
// Test with variant type
907+
run_ast(file, "array_remove(parse_json('[1, 2, 3, 2]'), 2)", &[]);
908+
run_ast(
909+
file,
910+
"array_remove(parse_json('[\"a\", \"b\", \"c\", \"b\"]'), 'b')",
911+
&[],
912+
);
913+
}
914+
915+
fn test_array_reverse(file: &mut impl Write) {
916+
// Test with simple arrays
917+
run_ast(file, "array_reverse([])", &[]);
918+
run_ast(file, "array_reverse([1, 2, 3])", &[]);
919+
run_ast(file, "array_reverse([1])", &[]);
920+
921+
// Test with strings
922+
run_ast(file, "array_reverse(['a', 'b', 'c'])", &[]);
923+
924+
// Test with nulls
925+
run_ast(file, "array_reverse([1, NULL, 2, 3])", &[]);
926+
927+
// Test with column data
928+
run_ast(file, "array_reverse([a, b, c])", &[
929+
("a", Int16Type::from_data(vec![1i16, 2, 3])),
930+
("b", Int16Type::from_data(vec![4i16, 5, 6])),
931+
("c", Int16Type::from_data(vec![7i16, 8, 9])),
932+
]);
933+
934+
// Test with variant type
935+
run_ast(file, "array_reverse(parse_json('[1, 2, 3]'))", &[]);
936+
run_ast(
937+
file,
938+
"array_reverse(parse_json('[\"a\", \"b\", \"c\"]'))",
939+
&[],
940+
);
941+
}

0 commit comments

Comments
 (0)