This is an automated email from the ASF dual-hosted git repository.
wjones127 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 1e6ff64151 feat: support eq_array and to_array_of_size for FSL (#8225)
1e6ff64151 is described below
commit 1e6ff64151edd4d16362b05598885cd02ef04453
Author: Will Jones <[email protected]>
AuthorDate: Thu Nov 16 11:24:40 2023 -0800
feat: support eq_array and to_array_of_size for FSL (#8225)
---
datafusion/common/src/scalar.rs | 63 ++++++++++++++++++++++++++++++++++-------
1 file changed, 53 insertions(+), 10 deletions(-)
diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs
index 211ac13e19..e8dac2a7f4 100644
--- a/datafusion/common/src/scalar.rs
+++ b/datafusion/common/src/scalar.rs
@@ -1725,7 +1725,7 @@ impl ScalarValue {
///
/// Errors if `self` is
/// - a decimal that fails be converted to a decimal array of size
- /// - a `Fixedsizelist` that is not supported yet
+ /// - a `Fixedsizelist` that fails to be concatenated into an array of size
/// - a `List` that fails to be concatenated into an array of size
/// - a `Dictionary` that fails be converted to a dictionary array of size
pub fn to_array_of_size(&self, size: usize) -> Result<ArrayRef> {
@@ -1846,10 +1846,7 @@ impl ScalarValue {
.collect::<LargeBinaryArray>(),
),
},
- ScalarValue::FixedSizeList(..) => {
- return _not_impl_err!("FixedSizeList is not supported yet")
- }
- ScalarValue::List(arr) => {
+ ScalarValue::List(arr) | ScalarValue::FixedSizeList(arr) => {
let arrays = std::iter::repeat(arr.as_ref())
.take(size)
.collect::<Vec<_>>();
@@ -2324,8 +2321,6 @@ impl ScalarValue {
///
/// Errors if
/// - it fails to downcast `array` to the data type of `self`
- /// - `self` is a `Fixedsizelist`
- /// - `self` is a `List`
/// - `self` is a `Struct`
///
/// # Panics
@@ -2398,10 +2393,10 @@ impl ScalarValue {
ScalarValue::LargeBinary(val) => {
eq_array_primitive!(array, index, LargeBinaryArray, val)?
}
- ScalarValue::FixedSizeList(..) => {
- return _not_impl_err!("FixedSizeList is not supported yet")
+ ScalarValue::List(arr) | ScalarValue::FixedSizeList(arr) => {
+ let right = array.slice(index, 1);
+ arr == &right
}
- ScalarValue::List(_) => return _not_impl_err!("List is not
supported yet"),
ScalarValue::Date32(val) => {
eq_array_primitive!(array, index, Date32Array, val)?
}
@@ -3103,6 +3098,27 @@ mod tests {
assert_eq!(&arr, actual_list_arr);
}
+ #[test]
+ fn test_to_array_of_size_for_fsl() {
+ let values = Int32Array::from_iter([Some(1), None, Some(2)]);
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let arr = FixedSizeListArray::new(field.clone(), 3, Arc::new(values),
None);
+ let sv = ScalarValue::FixedSizeList(Arc::new(arr));
+ let actual_arr = sv
+ .to_array_of_size(2)
+ .expect("Failed to convert to array of size");
+
+ let expected_values =
+ Int32Array::from_iter([Some(1), None, Some(2), Some(1), None,
Some(2)]);
+ let expected_arr =
+ FixedSizeListArray::new(field, 3, Arc::new(expected_values), None);
+
+ assert_eq!(
+ &expected_arr,
+ as_fixed_size_list_array(actual_arr.as_ref()).unwrap()
+ );
+ }
+
#[test]
fn test_list_to_array_string() {
let scalars = vec![
@@ -3181,6 +3197,33 @@ mod tests {
assert_eq!(result, &expected);
}
+ #[test]
+ fn test_list_scalar_eq_to_array() {
+ let list_array: ArrayRef =
+ Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ Some(vec![Some(0), Some(1), Some(2)]),
+ None,
+ Some(vec![None, Some(5)]),
+ ]));
+
+ let fsl_array: ArrayRef =
+ Arc::new(FixedSizeListArray::from_iter_primitive::<Int32Type, _,
_>(
+ vec![
+ Some(vec![Some(0), Some(1), Some(2)]),
+ None,
+ Some(vec![Some(3), None, Some(5)]),
+ ],
+ 3,
+ ));
+
+ for arr in [list_array, fsl_array] {
+ for i in 0..arr.len() {
+ let scalar = ScalarValue::List(arr.slice(i, 1));
+ assert!(scalar.eq_array(&arr, i).unwrap());
+ }
+ }
+ }
+
#[test]
fn scalar_add_trait_test() -> Result<()> {
let float_value = ScalarValue::Float64(Some(123.));
