diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index d616b01e7750..6f3f80e26571 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -378,8 +378,17 @@ fn filter_array(values: &dyn Array, predicate: &FilterPredicate) -> Result { + if predicate.filter.len() != values.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Filter predicate of length {} is not equal than target array of length {}", + predicate.filter.len(), + values.len() + ))); + } + + // Safety: We have checked that the predicate and values have the same length downcast_run_array!{ - values => Ok(Arc::new(filter_run_end_array(values, predicate)?)), + values => Ok(Arc::new(unsafe { filter_run_end_array(values, predicate)? })), t => unimplemented!("Filter not supported for RunEndEncoded type {:?}", t) } } @@ -422,7 +431,10 @@ fn filter_array(values: &dyn Array, predicate: &FilterPredicate) -> Result( +/// +/// # Safety +/// The caller must ensure that the `pred` and `re_arr` are the same length. +unsafe fn filter_run_end_array( re_arr: &RunArray, pred: &FilterPredicate, ) -> Result, ArrowError> @@ -1280,6 +1292,18 @@ mod tests { assert_eq!(0, actual.len()); } + #[test] + fn test_filter_run_end_encoding_array_safety() { + let run_ends = Int64Array::from(vec![2, 3, 8, 10]); + let values = Int64Array::from(vec![7, -2, 9, -8]); + let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray"); + let b = BooleanArray::from(vec![ + true, false, false, false, false, false, false, false, false, + ]); + let c = filter(&a, &b); + assert!(c.is_err()); + } + #[test] fn test_filter_dictionary_array() { let values = [Some("hello"), None, Some("world"), Some("!")];