1use arrow::array::{
21 Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, AsArray, BooleanArray,
22 BooleanBufferBuilder, Datum, MAX_INLINE_VIEW_LEN, PrimitiveArray, Scalar,
23 StringArrayType, StringViewArray,
24};
25use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer};
26use arrow::datatypes::DataType;
27use arrow::downcast_primitive_array;
28use arrow::row::{RowConverter, Rows, SortField};
29use datafusion_common::cast::{as_fixed_size_list_array, as_generic_list_array};
30use datafusion_common::utils::string_utils::string_array_to_vec;
31use datafusion_common::utils::take_function_args;
32use datafusion_common::{DataFusionError, Result, ScalarValue, exec_err};
33use datafusion_expr::expr::ScalarFunction;
34use datafusion_expr::simplify::ExprSimplifyResult;
35use datafusion_expr::{
36 ColumnarValue, Documentation, Expr, ScalarFunctionArgs, ScalarUDFImpl, Signature,
37 Volatility, in_list,
38};
39use datafusion_macros::user_doc;
40use datafusion_physical_expr_common::datum::compare_with_eq;
41use itertools::Itertools;
42
43use crate::make_array::make_array_udf;
44use crate::utils::make_scalar_function;
45
46use hashbrown::HashSet;
47use std::ops::Range;
48use std::sync::Arc;
49
50make_udf_expr_and_func!(ArrayHas,
52 array_has,
53 haystack_array element, "returns true, if the element appears in the first array, otherwise false.", array_has_udf );
57make_udf_expr_and_func!(ArrayHasAll,
58 array_has_all,
59 haystack_array needle_array, "returns true if each element of the second array appears in the first array; otherwise, it returns false.", array_has_all_udf );
63make_udf_expr_and_func!(ArrayHasAny,
64 array_has_any,
65 first_array second_array, "returns true if at least one element of the second array appears in the first array; otherwise, it returns false.", array_has_any_udf );
69
70#[user_doc(
71 doc_section(label = "Array Functions"),
72 description = "Returns true if the array contains the element.",
73 syntax_example = "array_has(array, element)",
74 sql_example = r#"```sql
75> select array_has([1, 2, 3], 2);
76+-----------------------------+
77| array_has(List([1,2,3]), 2) |
78+-----------------------------+
79| true |
80+-----------------------------+
81```"#,
82 argument(
83 name = "array",
84 description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
85 ),
86 argument(
87 name = "element",
88 description = "Scalar or Array expression. Can be a constant, column, or function, and any combination of array operators."
89 )
90)]
91#[derive(Debug, PartialEq, Eq, Hash)]
92pub struct ArrayHas {
93 signature: Signature,
94 aliases: Vec<String>,
95}
96
97impl Default for ArrayHas {
98 fn default() -> Self {
99 Self::new()
100 }
101}
102
103impl ArrayHas {
104 pub fn new() -> Self {
105 Self {
106 signature: Signature::array_and_element(Volatility::Immutable),
107 aliases: vec![
108 String::from("list_has"),
109 String::from("array_contains"),
110 String::from("list_contains"),
111 ],
112 }
113 }
114}
115
116impl ScalarUDFImpl for ArrayHas {
117 fn name(&self) -> &str {
118 "array_has"
119 }
120
121 fn signature(&self) -> &Signature {
122 &self.signature
123 }
124
125 fn return_type(&self, _: &[DataType]) -> Result<DataType> {
126 Ok(DataType::Boolean)
127 }
128
129 fn simplify(
130 &self,
131 mut args: Vec<Expr>,
132 _info: &datafusion_expr::simplify::SimplifyContext,
133 ) -> Result<ExprSimplifyResult> {
134 let [haystack, needle] = take_function_args(self.name(), &mut args)?;
135
136 match haystack {
139 Expr::Literal(scalar, _) if scalar.is_null() => {
140 return Ok(ExprSimplifyResult::Simplified(Expr::Literal(
141 ScalarValue::Boolean(None),
142 None,
143 )));
144 }
145 Expr::Literal(
146 scalar @ ScalarValue::List(_) | scalar @ ScalarValue::LargeList(_),
148 _,
149 ) => {
150 if let Ok(scalar_values) =
151 ScalarValue::convert_array_to_scalar_vec(&scalar.to_array()?)
152 {
153 assert_eq!(scalar_values.len(), 1);
154 let list = scalar_values
155 .into_iter()
156 .flatten()
157 .flatten()
158 .map(|v| Expr::Literal(v, None))
159 .collect();
160
161 return Ok(ExprSimplifyResult::Simplified(in_list(
162 std::mem::take(needle),
163 list,
164 false,
165 )));
166 }
167 }
168 Expr::ScalarFunction(ScalarFunction { func, args })
169 if func == &make_array_udf() =>
170 {
171 return Ok(ExprSimplifyResult::Simplified(in_list(
173 std::mem::take(needle),
174 std::mem::take(args),
175 false,
176 )));
177 }
178 _ => {}
179 };
180 Ok(ExprSimplifyResult::Original(args))
181 }
182
183 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
184 let [first_arg, second_arg] = take_function_args(self.name(), &args.args)?;
185 if first_arg.data_type().is_null() {
186 return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
189 }
190
191 match &second_arg {
192 ColumnarValue::Array(array_needle) => {
193 let haystack = first_arg.to_array(array_needle.len())?;
195 let array = array_has_inner_for_array(&haystack, array_needle)?;
196 Ok(ColumnarValue::Array(array))
197 }
198 ColumnarValue::Scalar(scalar_needle) => {
199 if scalar_needle.is_null() {
202 return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
203 }
204
205 let haystack = first_arg.to_array(1)?;
207 let needle = scalar_needle.to_array_of_size(1)?;
208 let needle = Scalar::new(needle);
209 let array = array_has_inner_for_scalar(&haystack, &needle)?;
210 if let ColumnarValue::Scalar(_) = &first_arg {
211 let scalar_value = ScalarValue::try_from_array(&array, 0)?;
213 Ok(ColumnarValue::Scalar(scalar_value))
214 } else {
215 Ok(ColumnarValue::Array(array))
216 }
217 }
218 }
219 }
220
221 fn aliases(&self) -> &[String] {
222 &self.aliases
223 }
224
225 fn documentation(&self) -> Option<&Documentation> {
226 self.doc()
227 }
228}
229
230fn array_has_inner_for_scalar(
231 haystack: &ArrayRef,
232 needle: &dyn Datum,
233) -> Result<ArrayRef> {
234 let haystack = haystack.as_ref().try_into()?;
235 array_has_dispatch_for_scalar(haystack, needle)
236}
237
238fn array_has_inner_for_array(haystack: &ArrayRef, needle: &ArrayRef) -> Result<ArrayRef> {
239 let haystack = haystack.as_ref().try_into()?;
240 array_has_dispatch_for_array(haystack, needle)
241}
242
243#[derive(Copy, Clone)]
244enum ArrayWrapper<'a> {
245 FixedSizeList(&'a arrow::array::FixedSizeListArray),
246 List(&'a arrow::array::GenericListArray<i32>),
247 LargeList(&'a arrow::array::GenericListArray<i64>),
248}
249
250impl<'a> TryFrom<&'a dyn Array> for ArrayWrapper<'a> {
251 type Error = DataFusionError;
252
253 fn try_from(
254 value: &'a dyn Array,
255 ) -> std::result::Result<ArrayWrapper<'a>, Self::Error> {
256 match value.data_type() {
257 DataType::List(_) => {
258 Ok(ArrayWrapper::List(as_generic_list_array::<i32>(value)?))
259 }
260 DataType::LargeList(_) => Ok(ArrayWrapper::LargeList(
261 as_generic_list_array::<i64>(value)?,
262 )),
263 DataType::FixedSizeList(_, _) => Ok(ArrayWrapper::FixedSizeList(
264 as_fixed_size_list_array(value)?,
265 )),
266 _ => exec_err!("array_has does not support type '{}'.", value.data_type()),
267 }
268 }
269}
270
271impl<'a> ArrayWrapper<'a> {
272 fn len(&self) -> usize {
273 match self {
274 ArrayWrapper::FixedSizeList(arr) => arr.len(),
275 ArrayWrapper::List(arr) => arr.len(),
276 ArrayWrapper::LargeList(arr) => arr.len(),
277 }
278 }
279
280 fn iter(&self) -> Box<dyn Iterator<Item = Option<ArrayRef>> + 'a> {
281 match self {
282 ArrayWrapper::FixedSizeList(arr) => Box::new(arr.iter()),
283 ArrayWrapper::List(arr) => Box::new(arr.iter()),
284 ArrayWrapper::LargeList(arr) => Box::new(arr.iter()),
285 }
286 }
287
288 fn values(&self) -> &ArrayRef {
289 match self {
290 ArrayWrapper::FixedSizeList(arr) => arr.values(),
291 ArrayWrapper::List(arr) => arr.values(),
292 ArrayWrapper::LargeList(arr) => arr.values(),
293 }
294 }
295
296 fn value_type(&self) -> DataType {
297 match self {
298 ArrayWrapper::FixedSizeList(arr) => arr.value_type(),
299 ArrayWrapper::List(arr) => arr.value_type(),
300 ArrayWrapper::LargeList(arr) => arr.value_type(),
301 }
302 }
303
304 fn offsets(&self) -> Box<dyn Iterator<Item = usize> + 'a> {
305 match self {
306 ArrayWrapper::FixedSizeList(arr) => {
307 let value_length = arr.value_length() as usize;
308 Box::new((0..=arr.len()).map(move |i| i * value_length))
309 }
310 ArrayWrapper::List(arr) => {
311 Box::new(arr.offsets().iter().map(|o| (*o) as usize))
312 }
313 ArrayWrapper::LargeList(arr) => {
314 Box::new(arr.offsets().iter().map(|o| (*o) as usize))
315 }
316 }
317 }
318
319 fn nulls(&self) -> Option<&NullBuffer> {
320 match self {
321 ArrayWrapper::FixedSizeList(arr) => arr.nulls(),
322 ArrayWrapper::List(arr) => arr.nulls(),
323 ArrayWrapper::LargeList(arr) => arr.nulls(),
324 }
325 }
326}
327
328fn array_has_dispatch_for_array<'a>(
334 haystack: ArrayWrapper<'a>,
335 needle: &ArrayRef,
336) -> Result<ArrayRef> {
337 let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls());
338 let needle = needle.as_ref();
339
340 let raw = OffsetBuffer::new(
343 haystack
344 .offsets()
345 .map(|o| o as i64)
346 .collect::<Vec<_>>()
347 .into(),
348 );
349 let first = raw[0];
350 let visible_values = haystack
351 .values()
352 .slice(first as usize, (raw[raw.len() - 1] - first) as usize);
353 let visible_values = visible_values.as_ref();
354 let offsets: Vec<usize> = raw.subtract(first).iter().map(|&o| o as usize).collect();
355
356 let fast_path = if visible_values.data_type() != needle.data_type() {
359 None
360 } else {
361 downcast_primitive_array! {
362 visible_values => {
363 let num_rows = offsets.len() - 1;
368 if num_rows > 0
369 && offsets[num_rows] / num_rows > NULL_FAST_PATH_MAX_LEN
370 && visible_values.null_count() > 0
371 {
372 None
373 } else {
374 Some(array_has_array_primitive(
375 visible_values, needle, &offsets,
376 combined_nulls.as_ref(),
377 ))
378 }
379 },
380 DataType::Utf8 => Some(array_has_array_string(
381 visible_values.as_string::<i32>(),
382 needle.as_string::<i32>(),
383 &offsets,
384 combined_nulls.as_ref(),
385 )),
386 DataType::LargeUtf8 => Some(array_has_array_string(
387 visible_values.as_string::<i64>(),
388 needle.as_string::<i64>(),
389 &offsets,
390 combined_nulls.as_ref(),
391 )),
392 DataType::Utf8View => Some(array_has_array_string_view(
393 visible_values.as_string_view(),
394 needle.as_string_view(),
395 &offsets,
396 combined_nulls.as_ref(),
397 )),
398 _ => None,
399 }
400 };
401
402 if let Some(values) = fast_path {
403 return Ok(Arc::new(BooleanArray::new(values, combined_nulls)));
404 }
405
406 let mut result = BooleanBufferBuilder::new(haystack.len());
408 for (i, arr) in haystack.iter().enumerate() {
409 if combined_nulls.as_ref().is_some_and(|n| n.is_null(i)) {
410 result.append(false);
411 continue;
412 }
413 let arr = arr.unwrap();
414 let is_nested = arr.data_type().is_nested();
415 let needle_row = Scalar::new(needle.slice(i, 1));
416 let eq_array = compare_with_eq(&arr, &needle_row, is_nested)?;
417 result.append(eq_array.has_true());
418 }
419
420 Ok(Arc::new(BooleanArray::new(result.finish(), combined_nulls)))
421}
422
423const NULL_FAST_PATH_MAX_LEN: usize = 512;
426
427fn array_has_array_primitive<T: ArrowPrimitiveType>(
434 values: &PrimitiveArray<T>,
435 needle: &dyn Array,
436 offsets: &[usize],
437 combined_nulls: Option<&NullBuffer>,
438) -> BooleanBuffer
439where
440 T::Native: ArrowNativeTypeOp,
441{
442 let needle = needle.as_primitive::<T>();
443 let num_rows = offsets.len() - 1;
444 let value_slice = values.values();
445 let needle_slice = needle.values();
446
447 let Some(element_nulls) = values.nulls() else {
448 return BooleanBuffer::collect_bool(num_rows, |i| {
449 if combined_nulls.is_some_and(|n| n.is_null(i)) {
450 return false;
451 }
452 let needle_val = needle_slice[i];
454 let start = offsets[i];
455 let end = offsets[i + 1];
456 value_slice[start..end]
457 .iter()
458 .fold(false, |acc, &v| acc | v.is_eq(needle_val))
459 });
460 };
461
462 let mut result = BooleanBufferBuilder::new(num_rows);
464 let mut needle_expanded: Vec<T::Native> = Vec::new();
465 for chunk_start in (0..num_rows).step_by(ROW_CONVERSION_CHUNK_SIZE) {
466 let chunk_end = (chunk_start + ROW_CONVERSION_CHUNK_SIZE).min(num_rows);
467 let elem_start = offsets[chunk_start];
468 let elem_end = offsets[chunk_end];
469
470 needle_expanded.clear();
473 for i in chunk_start..chunk_end {
474 needle_expanded.extend(std::iter::repeat_n(
475 needle_slice[i],
476 offsets[i + 1] - offsets[i],
477 ));
478 }
479 let chunk_values = &value_slice[elem_start..elem_end];
480 let eq_bits = BooleanBuffer::collect_bool(chunk_values.len(), |k| {
481 chunk_values[k].is_eq(needle_expanded[k])
482 });
483 let matched = &eq_bits
484 & &element_nulls
485 .inner()
486 .slice(elem_start, elem_end - elem_start);
487
488 for i in chunk_start..chunk_end {
489 if combined_nulls.is_some_and(|n| n.is_null(i)) {
490 result.append(false);
491 continue;
492 }
493 let start = offsets[i] - elem_start;
494 let end = offsets[i + 1] - elem_start;
495 result.append(matched.slice(start, end - start).has_true());
496 }
497 }
498 result.finish()
499}
500
501fn array_has_array_string<'a, S: StringArrayType<'a> + Copy>(
503 values: S,
504 needle: S,
505 offsets: &[usize],
506 combined_nulls: Option<&NullBuffer>,
507) -> BooleanBuffer {
508 let num_rows = offsets.len() - 1;
509 BooleanBuffer::collect_bool(num_rows, |i| {
510 if combined_nulls.is_some_and(|n| n.is_null(i)) {
511 return false;
512 }
513 let needle_val = needle.value(i);
515 let start = offsets[i];
516 let end = offsets[i + 1];
517 (start..end).any(|k| values.value(k) == needle_val && !values.is_null(k))
520 })
521}
522
523fn array_has_array_string_view(
529 values: &StringViewArray,
530 needle: &StringViewArray,
531 offsets: &[usize],
532 combined_nulls: Option<&NullBuffer>,
533) -> BooleanBuffer {
534 let num_rows = offsets.len() - 1;
535 let value_views = values.views();
536 let needle_views = needle.views();
537 BooleanBuffer::collect_bool(num_rows, |i| {
538 if combined_nulls.is_some_and(|n| n.is_null(i)) {
539 return false;
540 }
541 let needle_view = needle_views[i];
543 let needle_inline = (needle_view as u32) <= MAX_INLINE_VIEW_LEN;
545 let needle_lo = needle_view as u64;
546 let needle_val = needle.value(i);
547 let start = offsets[i];
548 let end = offsets[i + 1];
549 (start..end).any(|k| {
550 let v = value_views[k];
551 let matched = if needle_inline {
552 v == needle_view
554 } else {
555 (v as u64) == needle_lo && values.value(k) == needle_val
557 };
558 matched && !values.is_null(k)
559 })
560 })
561}
562
563fn array_has_dispatch_for_scalar(
564 haystack: ArrayWrapper<'_>,
565 needle: &dyn Datum,
566) -> Result<ArrayRef> {
567 if haystack.len() == 0 {
570 return Ok(Arc::new(BooleanArray::new(
571 BooleanBuffer::new_unset(haystack.len()),
572 None,
573 )));
574 }
575
576 let offsets: Vec<usize> = haystack.offsets().collect();
579 let first_offset = offsets[0];
580 let visible_values = haystack
581 .values()
582 .slice(first_offset, offsets[offsets.len() - 1] - first_offset);
583
584 let is_nested = visible_values.data_type().is_nested();
585 let eq_array = compare_with_eq(&visible_values, needle, is_nested)?;
586
587 let eq_bits = match eq_array.nulls() {
595 Some(nulls) => eq_array.values() & nulls.inner(),
596 None => eq_array.values().clone(),
597 };
598
599 let validity = match &haystack {
600 ArrayWrapper::FixedSizeList(arr) => arr.nulls(),
601 ArrayWrapper::List(arr) => arr.nulls(),
602 ArrayWrapper::LargeList(arr) => arr.nulls(),
603 };
604 let mut matches = eq_bits.set_indices().peekable();
605 let mut result = BooleanBufferBuilder::new(haystack.len());
606 result.append_n(haystack.len(), false);
607
608 for (i, window) in offsets.windows(2).enumerate() {
611 let end = window[1] - first_offset;
612
613 let has_match = matches.peek().is_some_and(|&p| p < end);
614
615 while matches.peek().is_some_and(|&p| p < end) {
617 matches.next();
618 }
619
620 if has_match && validity.is_none_or(|v| v.is_valid(i)) {
621 result.set_bit(i, true);
622 }
623 }
624
625 Ok(Arc::new(BooleanArray::new(
628 result.finish(),
629 validity.cloned(),
630 )))
631}
632
633fn array_has_all_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
634 array_has_all_and_any_inner(args, ComparisonType::All)
635}
636
637const ROW_CONVERSION_CHUNK_SIZE: usize = 512;
642
643fn general_array_has_for_all_and_any<'a>(
645 haystack: ArrayWrapper<'a>,
646 needle: ArrayWrapper<'a>,
647 comparison_type: ComparisonType,
648) -> Result<ArrayRef> {
649 let num_rows = haystack.len();
650 let converter = RowConverter::new(vec![SortField::new(haystack.value_type())])?;
651
652 let h_offsets: Vec<usize> = haystack.offsets().collect();
653 let n_offsets: Vec<usize> = needle.offsets().collect();
654
655 let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls());
656 let mut result = BooleanBufferBuilder::new(num_rows);
657
658 for chunk_start in (0..num_rows).step_by(ROW_CONVERSION_CHUNK_SIZE) {
659 let chunk_end = (chunk_start + ROW_CONVERSION_CHUNK_SIZE).min(num_rows);
660
661 let h_elem_start = h_offsets[chunk_start];
664 let h_elem_end = h_offsets[chunk_end];
665 let n_elem_start = n_offsets[chunk_start];
666 let n_elem_end = n_offsets[chunk_end];
667
668 let h_vals = haystack
669 .values()
670 .slice(h_elem_start, h_elem_end - h_elem_start);
671 let n_vals = needle
672 .values()
673 .slice(n_elem_start, n_elem_end - n_elem_start);
674
675 let chunk_h_rows = converter.convert_columns(&[h_vals])?;
676 let chunk_n_rows = converter.convert_columns(&[n_vals])?;
677
678 for i in chunk_start..chunk_end {
679 if combined_nulls.as_ref().is_some_and(|n| n.is_null(i)) {
680 result.append(false);
681 continue;
682 }
683 result.append(general_array_has_all_and_any_kernel(
684 &chunk_h_rows,
685 (h_offsets[i] - h_elem_start)..(h_offsets[i + 1] - h_elem_start),
686 &chunk_n_rows,
687 (n_offsets[i] - n_elem_start)..(n_offsets[i + 1] - n_elem_start),
688 comparison_type,
689 ));
690 }
691 }
692
693 Ok(Arc::new(BooleanArray::new(result.finish(), combined_nulls)))
694}
695
696fn array_has_all_and_any_string_internal<'a>(
698 haystack: ArrayWrapper<'a>,
699 needle: ArrayWrapper<'a>,
700 comparison_type: ComparisonType,
701) -> Result<ArrayRef> {
702 let num_rows = haystack.len();
703
704 let h_offsets: Vec<usize> = haystack.offsets().collect();
705 let n_offsets: Vec<usize> = needle.offsets().collect();
706
707 let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls());
708 let mut result = BooleanBufferBuilder::new(num_rows);
709
710 for chunk_start in (0..num_rows).step_by(ROW_CONVERSION_CHUNK_SIZE) {
711 let chunk_end = (chunk_start + ROW_CONVERSION_CHUNK_SIZE).min(num_rows);
712
713 let h_elem_start = h_offsets[chunk_start];
714 let h_elem_end = h_offsets[chunk_end];
715 let n_elem_start = n_offsets[chunk_start];
716 let n_elem_end = n_offsets[chunk_end];
717
718 let h_vals = haystack
719 .values()
720 .slice(h_elem_start, h_elem_end - h_elem_start);
721 let n_vals = needle
722 .values()
723 .slice(n_elem_start, n_elem_end - n_elem_start);
724
725 let chunk_h_strings = string_array_to_vec(h_vals.as_ref());
726 let chunk_n_strings = string_array_to_vec(n_vals.as_ref());
727
728 for i in chunk_start..chunk_end {
729 if combined_nulls.as_ref().is_some_and(|n| n.is_null(i)) {
730 result.append(false);
731 continue;
732 }
733 let h_start = h_offsets[i] - h_elem_start;
734 let h_end = h_offsets[i + 1] - h_elem_start;
735 let n_start = n_offsets[i] - n_elem_start;
736 let n_end = n_offsets[i + 1] - n_elem_start;
737 result.append(array_has_string_kernel(
738 &chunk_h_strings[h_start..h_end],
739 &chunk_n_strings[n_start..n_end],
740 comparison_type,
741 ));
742 }
743 }
744
745 Ok(Arc::new(BooleanArray::new(result.finish(), combined_nulls)))
746}
747
748fn array_has_all_and_any_dispatch<'a>(
749 haystack: ArrayWrapper<'a>,
750 needle: ArrayWrapper<'a>,
751 comparison_type: ComparisonType,
752) -> Result<ArrayRef> {
753 if needle.values().is_empty() {
754 let buffer = match comparison_type {
755 ComparisonType::All => BooleanBuffer::new_set(haystack.len()),
756 ComparisonType::Any => BooleanBuffer::new_unset(haystack.len()),
757 };
758 Ok(Arc::new(BooleanArray::from(buffer)))
759 } else {
760 match needle.value_type() {
761 DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
762 array_has_all_and_any_string_internal(haystack, needle, comparison_type)
763 }
764 _ => general_array_has_for_all_and_any(haystack, needle, comparison_type),
765 }
766 }
767}
768
769fn array_has_all_and_any_inner(
770 args: &[ArrayRef],
771 comparison_type: ComparisonType,
772) -> Result<ArrayRef> {
773 let haystack: ArrayWrapper = args[0].as_ref().try_into()?;
774 let needle: ArrayWrapper = args[1].as_ref().try_into()?;
775 array_has_all_and_any_dispatch(haystack, needle, comparison_type)
776}
777
778fn array_has_any_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
779 array_has_all_and_any_inner(args, ComparisonType::Any)
780}
781
782fn array_has_any_with_scalar(
784 columnar_arg: &ColumnarValue,
785 scalar_arg: &ScalarValue,
786) -> Result<ColumnarValue> {
787 if scalar_arg.is_null() {
788 return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
789 }
790
791 let scalar_array = scalar_arg.to_array_of_size(1)?;
793 let scalar_list: ArrayWrapper = scalar_array.as_ref().try_into()?;
794 let offsets: Vec<usize> = scalar_list.offsets().collect();
795 let scalar_values = scalar_list
796 .values()
797 .slice(offsets[0], offsets[1] - offsets[0]);
798
799 if scalar_values.is_empty() {
801 return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))));
802 }
803
804 match scalar_values.data_type() {
805 DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
806 array_has_any_with_scalar_string(columnar_arg, &scalar_values)
807 }
808 _ => array_has_any_with_scalar_general(columnar_arg, &scalar_values),
809 }
810}
811
812const SCALAR_SMALL_THRESHOLD: usize = 8;
817
818fn array_has_any_with_scalar_string(
820 columnar_arg: &ColumnarValue,
821 scalar_values: &ArrayRef,
822) -> Result<ColumnarValue> {
823 let (col_arr, is_scalar_output) = match columnar_arg {
824 ColumnarValue::Array(arr) => (Arc::clone(arr), false),
825 ColumnarValue::Scalar(s) => (s.to_array_of_size(1)?, true),
826 };
827
828 let col_list: ArrayWrapper = col_arr.as_ref().try_into()?;
829 let col_values = col_list.values();
830 let col_offsets: Vec<usize> = col_list.offsets().collect();
831 let col_nulls = col_list.nulls();
832
833 let scalar_lookup = ScalarStringLookup::new(scalar_values);
834 let has_null_scalar = scalar_values.null_count() > 0;
835
836 let result = match col_values.data_type() {
837 DataType::Utf8 => array_has_any_string_inner(
838 col_values.as_string::<i32>(),
839 &col_offsets,
840 col_nulls,
841 has_null_scalar,
842 &scalar_lookup,
843 ),
844 DataType::LargeUtf8 => array_has_any_string_inner(
845 col_values.as_string::<i64>(),
846 &col_offsets,
847 col_nulls,
848 has_null_scalar,
849 &scalar_lookup,
850 ),
851 DataType::Utf8View => array_has_any_string_inner(
852 col_values.as_string_view(),
853 &col_offsets,
854 col_nulls,
855 has_null_scalar,
856 &scalar_lookup,
857 ),
858 _ => unreachable!("array_has_any_with_scalar_string called with non-string type"),
859 };
860
861 if is_scalar_output {
862 Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
863 &result, 0,
864 )?))
865 } else {
866 Ok(ColumnarValue::Array(result))
867 }
868}
869
870enum ScalarStringLookup<'a> {
872 Set(HashSet<&'a str>),
874 List(Vec<Option<&'a str>>),
876}
877
878impl<'a> ScalarStringLookup<'a> {
879 fn new(scalar_values: &'a ArrayRef) -> Self {
880 let strings = string_array_to_vec(scalar_values.as_ref());
881 if strings.len() > SCALAR_SMALL_THRESHOLD {
882 ScalarStringLookup::Set(strings.into_iter().flatten().collect())
883 } else {
884 ScalarStringLookup::List(strings)
885 }
886 }
887
888 fn contains(&self, value: &str) -> bool {
889 match self {
890 ScalarStringLookup::Set(set) => set.contains(value),
891 ScalarStringLookup::List(list) => list.contains(&Some(value)),
892 }
893 }
894}
895
896fn array_has_any_string_inner<'a, C: StringArrayType<'a> + Copy>(
899 col_strings: C,
900 col_offsets: &[usize],
901 col_nulls: Option<&NullBuffer>,
902 has_null_scalar: bool,
903 scalar_lookup: &ScalarStringLookup<'_>,
904) -> ArrayRef {
905 let num_rows = col_offsets.len() - 1;
906 let mut result = BooleanBufferBuilder::new(num_rows);
907
908 for i in 0..num_rows {
909 if col_nulls.is_some_and(|v| v.is_null(i)) {
910 result.append(false);
911 continue;
912 }
913 let start = col_offsets[i];
914 let end = col_offsets[i + 1];
915 let found = (start..end).any(|j| {
916 if col_strings.is_null(j) {
917 has_null_scalar
918 } else {
919 scalar_lookup.contains(col_strings.value(j))
920 }
921 });
922 result.append(found);
923 }
924
925 Arc::new(BooleanArray::new(result.finish(), col_nulls.cloned()))
926}
927
928fn array_has_any_with_scalar_general(
931 columnar_arg: &ColumnarValue,
932 scalar_values: &ArrayRef,
933) -> Result<ColumnarValue> {
934 let converter =
935 RowConverter::new(vec![SortField::new(scalar_values.data_type().clone())])?;
936 let scalar_rows = converter.convert_columns(&[Arc::clone(scalar_values)])?;
937
938 let (col_arr, is_scalar_output) = match columnar_arg {
939 ColumnarValue::Array(arr) => (Arc::clone(arr), false),
940 ColumnarValue::Scalar(s) => (s.to_array_of_size(1)?, true),
941 };
942
943 let col_list: ArrayWrapper = col_arr.as_ref().try_into()?;
944 let col_rows = converter.convert_columns(&[Arc::clone(col_list.values())])?;
945 let col_offsets: Vec<usize> = col_list.offsets().collect();
946 let col_nulls = col_list.nulls();
947
948 let mut result = BooleanBufferBuilder::new(col_list.len());
949 let num_scalar = scalar_rows.num_rows();
950
951 if num_scalar > SCALAR_SMALL_THRESHOLD {
952 let scalar_set: HashSet<Box<[u8]>> = (0..num_scalar)
954 .map(|i| Box::from(scalar_rows.row(i).as_ref()))
955 .collect();
956
957 for i in 0..col_list.len() {
958 if col_nulls.is_some_and(|v| v.is_null(i)) {
959 result.append(false);
960 continue;
961 }
962 let start = col_offsets[i];
963 let end = col_offsets[i + 1];
964 let found =
965 (start..end).any(|j| scalar_set.contains(col_rows.row(j).as_ref()));
966 result.append(found);
967 }
968 } else {
969 for i in 0..col_list.len() {
971 if col_nulls.is_some_and(|v| v.is_null(i)) {
972 result.append(false);
973 continue;
974 }
975 let start = col_offsets[i];
976 let end = col_offsets[i + 1];
977 let found = (start..end)
978 .any(|j| (0..num_scalar).any(|k| col_rows.row(j) == scalar_rows.row(k)));
979 result.append(found);
980 }
981 }
982
983 let output: ArrayRef =
984 Arc::new(BooleanArray::new(result.finish(), col_nulls.cloned()));
985
986 if is_scalar_output {
987 Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
988 &output, 0,
989 )?))
990 } else {
991 Ok(ColumnarValue::Array(output))
992 }
993}
994
995#[user_doc(
996 doc_section(label = "Array Functions"),
997 description = "Returns true if all elements of sub_array exist in array.",
998 syntax_example = "array_has_all(array, sub_array)",
999 sql_example = r#"```sql
1000> select array_has_all([1, 2, 3, 4], [2, 3]);
1001+---------------------------------------------+
1002| array_has_all(List([1,2,3,4]), List([2,3])) |
1003+---------------------------------------------+
1004| true |
1005+---------------------------------------------+
1006```"#,
1007 argument(
1008 name = "array",
1009 description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
1010 ),
1011 argument(
1012 name = "sub_array",
1013 description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
1014 )
1015)]
1016#[derive(Debug, PartialEq, Eq, Hash)]
1017pub struct ArrayHasAll {
1018 signature: Signature,
1019 aliases: Vec<String>,
1020}
1021
1022impl Default for ArrayHasAll {
1023 fn default() -> Self {
1024 Self::new()
1025 }
1026}
1027
1028impl ArrayHasAll {
1029 pub fn new() -> Self {
1030 Self {
1031 signature: Signature::arrays(2, None, Volatility::Immutable),
1032 aliases: vec![String::from("list_has_all")],
1033 }
1034 }
1035}
1036
1037impl ScalarUDFImpl for ArrayHasAll {
1038 fn name(&self) -> &str {
1039 "array_has_all"
1040 }
1041
1042 fn signature(&self) -> &Signature {
1043 &self.signature
1044 }
1045
1046 fn return_type(&self, _: &[DataType]) -> Result<DataType> {
1047 Ok(DataType::Boolean)
1048 }
1049
1050 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
1051 make_scalar_function(array_has_all_inner)(&args.args)
1052 }
1053
1054 fn aliases(&self) -> &[String] {
1055 &self.aliases
1056 }
1057
1058 fn documentation(&self) -> Option<&Documentation> {
1059 self.doc()
1060 }
1061}
1062
1063#[user_doc(
1064 doc_section(label = "Array Functions"),
1065 description = "Returns true if the arrays have any elements in common.",
1066 syntax_example = "array_has_any(array1, array2)",
1067 sql_example = r#"```sql
1068> select array_has_any([1, 2, 3], [3, 4]);
1069+------------------------------------------+
1070| array_has_any(List([1,2,3]), List([3,4])) |
1071+------------------------------------------+
1072| true |
1073+------------------------------------------+
1074```"#,
1075 argument(
1076 name = "array1",
1077 description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
1078 ),
1079 argument(
1080 name = "array2",
1081 description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
1082 )
1083)]
1084#[derive(Debug, PartialEq, Eq, Hash)]
1085pub struct ArrayHasAny {
1086 signature: Signature,
1087 aliases: Vec<String>,
1088}
1089
1090impl Default for ArrayHasAny {
1091 fn default() -> Self {
1092 Self::new()
1093 }
1094}
1095
1096impl ArrayHasAny {
1097 pub fn new() -> Self {
1098 Self {
1099 signature: Signature::arrays(2, None, Volatility::Immutable),
1100 aliases: vec![String::from("list_has_any"), String::from("arrays_overlap")],
1101 }
1102 }
1103}
1104
1105impl ScalarUDFImpl for ArrayHasAny {
1106 fn name(&self) -> &str {
1107 "array_has_any"
1108 }
1109
1110 fn signature(&self) -> &Signature {
1111 &self.signature
1112 }
1113
1114 fn return_type(&self, _: &[DataType]) -> Result<DataType> {
1115 Ok(DataType::Boolean)
1116 }
1117
1118 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
1119 let [first_arg, second_arg] = take_function_args(self.name(), &args.args)?;
1120
1121 match (&first_arg, &second_arg) {
1123 (cv, ColumnarValue::Scalar(scalar)) | (ColumnarValue::Scalar(scalar), cv) => {
1124 array_has_any_with_scalar(cv, scalar)
1125 }
1126 _ => make_scalar_function(array_has_any_inner)(&args.args),
1127 }
1128 }
1129
1130 fn aliases(&self) -> &[String] {
1131 &self.aliases
1132 }
1133
1134 fn documentation(&self) -> Option<&Documentation> {
1135 self.doc()
1136 }
1137}
1138
1139#[derive(Debug, PartialEq, Clone, Copy)]
1141enum ComparisonType {
1142 All,
1144 Any,
1146}
1147
1148fn array_has_string_kernel(
1149 haystack: &[Option<&str>],
1150 needle: &[Option<&str>],
1151 comparison_type: ComparisonType,
1152) -> bool {
1153 match comparison_type {
1154 ComparisonType::All => needle
1155 .iter()
1156 .dedup()
1157 .all(|x| haystack.iter().dedup().any(|y| y == x)),
1158 ComparisonType::Any => needle
1159 .iter()
1160 .dedup()
1161 .any(|x| haystack.iter().dedup().any(|y| y == x)),
1162 }
1163}
1164
1165fn general_array_has_all_and_any_kernel(
1166 haystack_rows: &Rows,
1167 h_range: Range<usize>,
1168 needle_rows: &Rows,
1169 mut n_range: Range<usize>,
1170 comparison_type: ComparisonType,
1171) -> bool {
1172 let h_start = h_range.start;
1173 let h_end = h_range.end;
1174
1175 match comparison_type {
1176 ComparisonType::All => n_range.all(|ni| {
1177 let needle_row = needle_rows.row(ni);
1178 (h_start..h_end).any(|hi| haystack_rows.row(hi) == needle_row)
1179 }),
1180 ComparisonType::Any => n_range.any(|ni| {
1181 let needle_row = needle_rows.row(ni);
1182 (h_start..h_end).any(|hi| haystack_rows.row(hi) == needle_row)
1183 }),
1184 }
1185}
1186
1187#[cfg(test)]
1188mod tests {
1189 use std::sync::Arc;
1190
1191 use arrow::datatypes::Int32Type;
1192 use arrow::{
1193 array::{
1194 Array, ArrayRef, AsArray, FixedSizeListArray, Int32Array, ListArray,
1195 create_array,
1196 },
1197 buffer::OffsetBuffer,
1198 datatypes::{DataType, Field},
1199 };
1200 use datafusion_common::{
1201 DataFusionError, ScalarValue, config::ConfigOptions,
1202 utils::SingleRowListArrayBuilder,
1203 };
1204 use datafusion_expr::simplify::SimplifyContext;
1205 use datafusion_expr::{
1206 ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDFImpl, col, lit,
1207 simplify::ExprSimplifyResult,
1208 };
1209
1210 use crate::expr_fn::make_array;
1211
1212 use super::{ArrayHas, ArrayHasAll, ArrayHasAny};
1213
1214 #[test]
1215 fn test_simplify_array_has_to_in_list() {
1216 let haystack = lit(SingleRowListArrayBuilder::new(create_array!(
1217 Int32,
1218 [1, 2, 3]
1219 ))
1220 .build_list_scalar());
1221 let needle = col("c");
1222
1223 let context = SimplifyContext::default();
1224
1225 let Ok(ExprSimplifyResult::Simplified(Expr::InList(in_list))) =
1226 ArrayHas::new().simplify(vec![haystack, needle.clone()], &context)
1227 else {
1228 panic!("Expected simplified expression");
1229 };
1230
1231 assert_eq!(
1232 in_list,
1233 datafusion_expr::expr::InList {
1234 expr: Box::new(needle),
1235 list: vec![lit(1), lit(2), lit(3)],
1236 negated: false,
1237 }
1238 );
1239 }
1240
1241 #[test]
1242 fn test_simplify_array_has_with_make_array_to_in_list() {
1243 let haystack = make_array(vec![lit(1), lit(2), lit(3)]);
1244 let needle = col("c");
1245
1246 let context = SimplifyContext::default();
1247
1248 let Ok(ExprSimplifyResult::Simplified(Expr::InList(in_list))) =
1249 ArrayHas::new().simplify(vec![haystack, needle.clone()], &context)
1250 else {
1251 panic!("Expected simplified expression");
1252 };
1253
1254 assert_eq!(
1255 in_list,
1256 datafusion_expr::expr::InList {
1257 expr: Box::new(needle),
1258 list: vec![lit(1), lit(2), lit(3)],
1259 negated: false,
1260 }
1261 );
1262 }
1263
1264 #[test]
1265 fn test_simplify_array_has_with_null_to_null() {
1266 let haystack = Expr::Literal(ScalarValue::Null, None);
1267 let needle = col("c");
1268
1269 let context = SimplifyContext::default();
1270 let Ok(ExprSimplifyResult::Simplified(simplified)) =
1271 ArrayHas::new().simplify(vec![haystack, needle], &context)
1272 else {
1273 panic!("Expected simplified expression");
1274 };
1275
1276 assert_eq!(simplified, Expr::Literal(ScalarValue::Boolean(None), None));
1277 }
1278
1279 #[test]
1280 fn test_simplify_array_has_with_null_list_to_null() {
1281 let haystack =
1282 ListArray::from_iter_primitive::<Int32Type, [Option<i32>; 0], _>([None]);
1283 let haystack = Expr::Literal(ScalarValue::List(Arc::new(haystack)), None);
1284 let needle = col("c");
1285
1286 let context = SimplifyContext::default();
1287 let Ok(ExprSimplifyResult::Simplified(simplified)) =
1288 ArrayHas::new().simplify(vec![haystack, needle], &context)
1289 else {
1290 panic!("Expected simplified expression");
1291 };
1292
1293 assert_eq!(simplified, Expr::Literal(ScalarValue::Boolean(None), None));
1294 }
1295
1296 #[test]
1297 fn test_array_has_complex_list_not_simplified() {
1298 let haystack = col("c1");
1299 let needle = col("c2");
1300
1301 let context = SimplifyContext::default();
1302
1303 let Ok(ExprSimplifyResult::Original(args)) =
1304 ArrayHas::new().simplify(vec![haystack, needle.clone()], &context)
1305 else {
1306 panic!("Expected simplified expression");
1307 };
1308
1309 assert_eq!(args, vec![col("c1"), col("c2")],);
1310 }
1311
1312 #[test]
1313 fn test_array_has_list_empty_child() -> Result<(), DataFusionError> {
1314 let haystack_field = Arc::new(Field::new_list(
1315 "haystack",
1316 Field::new_list("", Field::new("", DataType::Int32, true), true),
1317 true,
1318 ));
1319
1320 let needle_field = Arc::new(Field::new("needle", DataType::Int32, true));
1321 let return_field = Arc::new(Field::new("return", DataType::Boolean, true));
1322 let haystack = ListArray::new(
1323 Field::new_list_field(DataType::Int32, true).into(),
1324 OffsetBuffer::new(vec![0, 0].into()),
1325 Arc::new(Int32Array::from(Vec::<i32>::new())) as ArrayRef,
1326 Some(vec![false].into()),
1327 );
1328
1329 let haystack = ColumnarValue::Array(Arc::new(haystack));
1330 let needle = ColumnarValue::Scalar(ScalarValue::Int32(Some(1)));
1331 let result = ArrayHas::new().invoke_with_args(ScalarFunctionArgs {
1332 args: vec![haystack, needle],
1333 arg_fields: vec![haystack_field, needle_field],
1334 number_rows: 1,
1335 return_field,
1336 config_options: Arc::new(ConfigOptions::default()),
1337 })?;
1338
1339 let output = result.into_array(1)?;
1340 let output = output.as_boolean();
1341 assert_eq!(output.len(), 1);
1342 assert!(output.is_null(0));
1343
1344 Ok(())
1345 }
1346
1347 #[test]
1348 fn test_array_has_sliced_list() -> Result<(), DataFusionError> {
1349 let list = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1351 Some(vec![Some(10), Some(20)]),
1352 Some(vec![Some(30), Some(40)]),
1353 Some(vec![Some(50), Some(60)]),
1354 Some(vec![Some(70), Some(80)]),
1355 ]);
1356 let sliced = list.slice(1, 2);
1357 let haystack_field =
1358 Arc::new(Field::new("haystack", sliced.data_type().clone(), true));
1359 let needle_field = Arc::new(Field::new("needle", DataType::Int32, true));
1360 let return_field = Arc::new(Field::new("return", DataType::Boolean, true));
1361
1362 let invoke = |needle: i32| -> Result<ArrayRef, DataFusionError> {
1365 ArrayHas::new()
1366 .invoke_with_args(ScalarFunctionArgs {
1367 args: vec![
1368 ColumnarValue::Array(Arc::new(sliced.clone())),
1369 ColumnarValue::Scalar(ScalarValue::Int32(Some(needle))),
1370 ],
1371 arg_fields: vec![
1372 Arc::clone(&haystack_field),
1373 Arc::clone(&needle_field),
1374 ],
1375 number_rows: 2,
1376 return_field: Arc::clone(&return_field),
1377 config_options: Arc::new(ConfigOptions::default()),
1378 })?
1379 .into_array(2)
1380 };
1381
1382 let output = invoke(10)?.as_boolean().clone();
1383 assert!(!output.value(0));
1384 assert!(!output.value(1));
1385
1386 let output = invoke(70)?.as_boolean().clone();
1387 assert!(!output.value(0));
1388 assert!(!output.value(1));
1389
1390 Ok(())
1391 }
1392
1393 #[test]
1394 fn test_array_has_list_null_haystack() -> Result<(), DataFusionError> {
1395 let haystack_field = Arc::new(Field::new("haystack", DataType::Null, true));
1396 let needle_field = Arc::new(Field::new("needle", DataType::Int32, true));
1397 let return_field = Arc::new(Field::new("return", DataType::Boolean, true));
1398 let haystack =
1399 ListArray::from_iter_primitive::<Int32Type, [Option<i32>; 0], _>([
1400 None, None, None,
1401 ]);
1402
1403 let haystack = ColumnarValue::Array(Arc::new(haystack));
1404 let needle = ColumnarValue::Scalar(ScalarValue::Int32(Some(1)));
1405 let result = ArrayHas::new().invoke_with_args(ScalarFunctionArgs {
1406 args: vec![haystack, needle],
1407 arg_fields: vec![haystack_field, needle_field],
1408 number_rows: 1,
1409 return_field,
1410 config_options: Arc::new(ConfigOptions::default()),
1411 })?;
1412
1413 let output = result.into_array(1)?;
1414 let output = output.as_boolean();
1415 assert_eq!(output.len(), 3);
1416 for i in 0..3 {
1417 assert!(output.is_null(i));
1418 }
1419
1420 Ok(())
1421 }
1422
1423 fn invoke_and_assert(
1426 udf: &dyn ScalarUDFImpl,
1427 haystack: &ArrayRef,
1428 needle: ArrayRef,
1429 expected: &[Option<bool>],
1430 ) {
1431 let num_rows = haystack.len();
1432 let list_type = haystack.data_type();
1433 let result = udf
1434 .invoke_with_args(ScalarFunctionArgs {
1435 args: vec![
1436 ColumnarValue::Array(Arc::clone(haystack)),
1437 ColumnarValue::Array(needle),
1438 ],
1439 arg_fields: vec![
1440 Arc::new(Field::new("haystack", list_type.clone(), false)),
1441 Arc::new(Field::new("needle", list_type.clone(), false)),
1442 ],
1443 number_rows: num_rows,
1444 return_field: Arc::new(Field::new("return", DataType::Boolean, true)),
1445 config_options: Arc::new(ConfigOptions::default()),
1446 })
1447 .unwrap();
1448 let output = result.into_array(num_rows).unwrap();
1449 assert_eq!(output.as_boolean().iter().collect::<Vec<_>>(), expected);
1450 }
1451
1452 #[test]
1453 fn test_sliced_list_offsets() {
1454 let field: Arc<Field> = Arc::new(Field::new("item", DataType::Int32, false));
1460 let full_values = Arc::new(Int32Array::from(vec![1, 2, 11, 12, 21, 22, 31, 32]));
1461 let full_offsets = OffsetBuffer::new(vec![0, 2, 4, 6, 8].into());
1462 let full = ListArray::new(Arc::clone(&field), full_offsets, full_values, None);
1463 let sliced_haystack: ArrayRef = Arc::new(full.slice(1, 2));
1464
1465 let needle_all: ArrayRef = Arc::new(ListArray::new(
1467 Arc::clone(&field),
1468 OffsetBuffer::new(vec![0, 1, 2].into()),
1469 Arc::new(Int32Array::from(vec![11, 21])),
1470 None,
1471 ));
1472 invoke_and_assert(
1473 &ArrayHasAll::new(),
1474 &sliced_haystack,
1475 needle_all,
1476 &[Some(true), Some(true)],
1477 );
1478
1479 let needle_any: ArrayRef = Arc::new(ListArray::new(
1481 field,
1482 OffsetBuffer::new(vec![0, 2, 4].into()),
1483 Arc::new(Int32Array::from(vec![99, 11, 99, 21])),
1484 None,
1485 ));
1486 invoke_and_assert(
1487 &ArrayHasAny::new(),
1488 &sliced_haystack,
1489 needle_any,
1490 &[Some(true), Some(true)],
1491 );
1492 }
1493
1494 #[test]
1495 fn test_sliced_fixed_size_list_offsets() {
1496 let field = Arc::new(Field::new("item", DataType::Int32, false));
1498 let full_values = Arc::new(Int32Array::from(vec![1, 2, 11, 12, 21, 22, 31, 32]));
1499 let full = FixedSizeListArray::new(Arc::clone(&field), 2, full_values, None);
1500 let sliced_haystack: ArrayRef = Arc::new(full.slice(1, 2));
1501
1502 let needle_all: ArrayRef = Arc::new(FixedSizeListArray::new(
1504 Arc::clone(&field),
1505 2,
1506 Arc::new(Int32Array::from(vec![11, 12, 21, 22])),
1507 None,
1508 ));
1509 invoke_and_assert(
1510 &ArrayHasAll::new(),
1511 &sliced_haystack,
1512 needle_all,
1513 &[Some(true), Some(true)],
1514 );
1515
1516 let needle_any: ArrayRef = Arc::new(FixedSizeListArray::new(
1518 field,
1519 2,
1520 Arc::new(Int32Array::from(vec![99, 12, 99, 22])),
1521 None,
1522 ));
1523 invoke_and_assert(
1524 &ArrayHasAny::new(),
1525 &sliced_haystack,
1526 needle_any,
1527 &[Some(true), Some(true)],
1528 );
1529 }
1530
1531 fn invoke_array_has_array(haystack: ArrayRef, needle: ArrayRef) -> ArrayRef {
1534 let num_rows = haystack.len();
1535 let haystack_type = haystack.data_type().clone();
1536 let needle_type = needle.data_type().clone();
1537 ArrayHas::new()
1538 .invoke_with_args(ScalarFunctionArgs {
1539 args: vec![ColumnarValue::Array(haystack), ColumnarValue::Array(needle)],
1540 arg_fields: vec![
1541 Arc::new(Field::new("haystack", haystack_type, false)),
1542 Arc::new(Field::new("needle", needle_type, false)),
1543 ],
1544 number_rows: num_rows,
1545 return_field: Arc::new(Field::new("return", DataType::Boolean, true)),
1546 config_options: Arc::new(ConfigOptions::default()),
1547 })
1548 .unwrap()
1549 .into_array(num_rows)
1550 .unwrap()
1551 }
1552
1553 #[test]
1554 fn test_array_has_array_needle_sliced() {
1555 let full = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1561 Some(vec![Some(1), Some(2)]),
1562 Some(vec![Some(10), Some(20), Some(30)]), Some(vec![Some(40)]), Some(vec![Some(50), Some(60)]), Some(vec![Some(70)]),
1566 ]);
1567 let sliced_haystack: ArrayRef = Arc::new(full.slice(1, 3));
1568 let sliced_needle: ArrayRef =
1569 Arc::new(Int32Array::from(vec![999, 20, 41, 60, 999]).slice(1, 3));
1570 let result = invoke_array_has_array(sliced_haystack, sliced_needle);
1571 assert_eq!(
1572 result.as_boolean().iter().collect::<Vec<_>>(),
1573 vec![Some(true), Some(false), Some(true)]
1574 );
1575
1576 let field = Arc::new(Field::new("item", DataType::Int32, true));
1579 let fsl_values = Arc::new(Int32Array::from(vec![1, 2, 11, 12, 21, 22, 31, 32]));
1580 let fsl: ArrayRef =
1581 Arc::new(FixedSizeListArray::new(field, 2, fsl_values, None).slice(1, 2));
1582 let needle: ArrayRef = Arc::new(Int32Array::from(vec![11, 99]));
1583 let result = invoke_array_has_array(fsl, needle);
1584 assert_eq!(
1585 result.as_boolean().iter().collect::<Vec<_>>(),
1586 vec![Some(true), Some(false)]
1587 );
1588 }
1589}