Dandandan commented on a change in pull request #813:
URL: https://github.com/apache/arrow-datafusion/pull/813#discussion_r682060135
##########
File path: datafusion/src/physical_plan/expressions/in_list.rs
##########
@@ -99,6 +124,104 @@ macro_rules! make_contains {
}};
}
+macro_rules! make_contains_primitive {
+ ($ARRAY:expr, $LIST_VALUES:expr, $NEGATED:expr, $SCALAR_VALUE:ident,
$ARRAY_TYPE:ident) => {{
+ let array = $ARRAY.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
+
+ let mut contains_null = false;
+ let values = $LIST_VALUES
+ .iter()
+ .flat_map(|expr| match expr {
+ ColumnarValue::Scalar(s) => match s {
+ ScalarValue::$SCALAR_VALUE(Some(v)) => Some(*v),
+ ScalarValue::$SCALAR_VALUE(None) => {
+ contains_null = true;
+ None
+ }
+ ScalarValue::Utf8(None) => {
+ contains_null = true;
+ None
+ }
+ datatype => unimplemented!("Unexpected type {} for
InList", datatype),
+ },
+ ColumnarValue::Array(_) => {
+ unimplemented!("InList does not yet support nested
columns.")
+ }
+ })
+ .collect::<Vec<_>>();
+
+ if $NEGATED {
+ if contains_null {
+ Ok(ColumnarValue::Array(Arc::new(
+ array
+ .iter()
+ .map(|x| match x.map(|v| !values.contains(&v)) {
+ Some(true) => None,
+ x => x,
+ })
+ .collect::<BooleanArray>(),
+ )))
+ } else {
+ Ok(ColumnarValue::Array(Arc::new(
+ values_not_in_list_primitive(array, &values)?,
+ )))
+ }
+ } else {
+ if contains_null {
+ Ok(ColumnarValue::Array(Arc::new(
+ array
+ .iter()
+ .map(|x| match x.map(|v| values.contains(&v)) {
+ Some(false) => None,
+ x => x,
+ })
+ .collect::<BooleanArray>(),
+ )))
+ } else {
+ Ok(ColumnarValue::Array(Arc::new(values_in_list_primitive(
+ array, &values,
+ )?)))
+ }
+ }
+ }};
+}
+
+fn values_in_list_primitive<T: ArrowPrimitiveType>(
Review comment:
I could for clarity. Note that the `values` slice can not contain nulls
based on the type, but the `array` can have nulls (and should work for that
case).
--
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]
