Jefffrey commented on code in PR #10991:
URL: https://github.com/apache/arrow-rs/pull/10991#discussion_r4165863652
##########
arrow-select/src/nullif.rs:
##########
@@ -42,19 +42,39 @@ use arrow_schema::{ArrowError, DataType};
/// assert_eq!(nulled.as_primitive(), &Int32Array::from(vec![None, None,
Some(1), Some(9)]));
/// ```
pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result<ArrayRef,
ArrowError> {
- let left_data = left.to_data();
-
- if left_data.len() != right.len() {
+ if left.len() != right.len() {
return Err(ArrowError::ComputeError(
"Cannot perform comparison operation on arrays of different
length".to_string(),
));
}
- let len = left_data.len();
+ let len = left.len();
- if len == 0 || left_data.data_type() == &DataType::Null {
- return Ok(make_array(left_data));
+ if len == 0 || left.data_type() == &DataType::Null {
+ return Ok(make_array(left.to_data()));
}
+ // Compute right_values & right_bitmap. True bits are positions that should
+ // become null; nulls in `right` do not introduce nulls.
+ let should_null = match right.nulls() {
+ Some(nulls) => right.values() & nulls.inner(),
+ None => right.values().clone(),
+ };
+
+ match left.data_type() {
Review Comment:
we should leave a comment here explaining why these types take a different
route
##########
arrow-select/src/nullif.rs:
##########
@@ -163,6 +195,125 @@ mod tests {
);
}
+ #[test]
+ fn test_nullif_run_end_encoded() {
+ let values = Int32Array::from(vec![10, 20]);
+ let ree = RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1,
2]), &values).unwrap();
+ let mask = BooleanArray::from(vec![Some(false), Some(true)]);
+
+ let result = nullif(&ree, &mask).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<RunArray<Int16Type>>()
+ .unwrap();
Review Comment:
```suggestion
let result = result.as_run::<Int16Type>();
```
##########
arrow-select/src/nullif.rs:
##########
@@ -163,6 +195,125 @@ mod tests {
);
}
+ #[test]
+ fn test_nullif_run_end_encoded() {
+ let values = Int32Array::from(vec![10, 20]);
+ let ree = RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1,
2]), &values).unwrap();
+ let mask = BooleanArray::from(vec![Some(false), Some(true)]);
+
+ let result = nullif(&ree, &mask).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<RunArray<Int16Type>>()
+ .unwrap();
+ assert_eq!(result.logical_null_count(), 1);
+ assert_eq!(
+ result
+ .values()
+ .as_primitive::<Int32Type>()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some(10), None]
+ );
+ }
+
+ #[test]
+ fn test_nullif_run_end_encoded_mask_nulls() {
+ let values = Int32Array::from(vec![10, 20, 30]);
+ let ree =
+ RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1, 2, 3]),
&values).unwrap();
+ let mask = BooleanArray::from(vec![Some(true), Some(false), None,
Some(true)]);
+ let mask = mask.slice(1, 3); // Some(false), None, Some(true)
+ let mask = mask.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let result = nullif(&ree, mask).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<RunArray<Int16Type>>()
+ .unwrap();
Review Comment:
```suggestion
let result = result.as_run::<Int16Type>();
```
##########
arrow-select/src/nullif.rs:
##########
@@ -163,6 +195,125 @@ mod tests {
);
}
+ #[test]
+ fn test_nullif_run_end_encoded() {
+ let values = Int32Array::from(vec![10, 20]);
+ let ree = RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1,
2]), &values).unwrap();
+ let mask = BooleanArray::from(vec![Some(false), Some(true)]);
+
+ let result = nullif(&ree, &mask).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<RunArray<Int16Type>>()
+ .unwrap();
+ assert_eq!(result.logical_null_count(), 1);
+ assert_eq!(
+ result
+ .values()
+ .as_primitive::<Int32Type>()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some(10), None]
+ );
+ }
+
+ #[test]
+ fn test_nullif_run_end_encoded_mask_nulls() {
+ let values = Int32Array::from(vec![10, 20, 30]);
+ let ree =
+ RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1, 2, 3]),
&values).unwrap();
+ let mask = BooleanArray::from(vec![Some(true), Some(false), None,
Some(true)]);
+ let mask = mask.slice(1, 3); // Some(false), None, Some(true)
+ let mask = mask.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let result = nullif(&ree, mask).unwrap();
Review Comment:
```suggestion
let result = nullif(&ree, &mask).unwrap();
```
##########
arrow-select/src/nullif.rs:
##########
@@ -163,6 +195,125 @@ mod tests {
);
}
+ #[test]
+ fn test_nullif_run_end_encoded() {
+ let values = Int32Array::from(vec![10, 20]);
+ let ree = RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1,
2]), &values).unwrap();
+ let mask = BooleanArray::from(vec![Some(false), Some(true)]);
+
+ let result = nullif(&ree, &mask).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<RunArray<Int16Type>>()
+ .unwrap();
+ assert_eq!(result.logical_null_count(), 1);
+ assert_eq!(
+ result
+ .values()
+ .as_primitive::<Int32Type>()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some(10), None]
+ );
+ }
+
+ #[test]
+ fn test_nullif_run_end_encoded_mask_nulls() {
+ let values = Int32Array::from(vec![10, 20, 30]);
+ let ree =
+ RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1, 2, 3]),
&values).unwrap();
+ let mask = BooleanArray::from(vec![Some(true), Some(false), None,
Some(true)]);
+ let mask = mask.slice(1, 3); // Some(false), None, Some(true)
+ let mask = mask.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let result = nullif(&ree, mask).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<RunArray<Int16Type>>()
+ .unwrap();
+ assert_eq!(result.logical_null_count(), 1);
+ assert_eq!(
+ result
+ .values()
+ .as_primitive::<Int32Type>()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some(10), Some(20), None]
+ );
+ }
+
+ #[test]
+ fn test_nullif_run_end_encoded_non_nullable_values() {
+ let ree = unsafe {
+ RunArray::<Int16Type>::new_unchecked(
+ DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int16, false)),
+ Arc::new(Field::new("values", DataType::Int32, false)),
+ ),
+ RunEndBuffer::new(vec![1_i16, 2].into(), 0, 2),
+ Arc::new(Int32Array::from(vec![10, 20])),
+ )
+ };
+ let mask = BooleanArray::from(vec![Some(false), Some(true)]);
+
+ let error = nullif(&ree, &mask).unwrap_err();
+ assert_eq!(
+ error.to_string(),
+ "Compute error: Cannot introduce nulls into a RunEndEncoded array
with a non-nullable values field"
+ );
+ }
+
+ #[test]
+ fn test_nullif_union() {
+ let fields =
+ UnionFields::try_new(vec![0], vec![Field::new("i",
DataType::Int32, true)]).unwrap();
+ let union = UnionArray::try_new(
+ fields,
+ ScalarBuffer::from(vec![0_i8, 0]),
+ Some(ScalarBuffer::from(vec![0_i32, 1])),
+ vec![Arc::new(Int32Array::from(vec![10, 20]))],
+ )
+ .unwrap();
+ let mask = BooleanArray::from(vec![Some(false), Some(true)]);
+
+ let result = nullif(&union, &mask).unwrap();
+ let result = result.as_any().downcast_ref::<UnionArray>().unwrap();
Review Comment:
```suggestion
let result = result.as_union();
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
