(arrow-datafusion) 02/03: Fix columnar value to array

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch alamb/array_crate
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git

commit 2d6512aa567ef1f0b9dbf7db55a704684858298c
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Fri Feb 2 09:49:41 2024 -0500

    Fix columnar value to array
---
 datafusion/core/src/execution/context/mod.rs |  1 -
 datafusion/expr/src/columnar_value.rs        | 40 +++++++++++-----------------
 2 files changed, 16 insertions(+), 25 deletions(-)

diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs
index abeffa6354..3661f6973c 100644
--- a/datafusion/core/src/execution/context/mod.rs
+++ b/datafusion/core/src/execution/context/mod.rs
@@ -1353,7 +1353,6 @@ impl SessionState {
         datafusion_functions_array::register_all(&mut new_self)
             .expect("can not register array expressions");
 
-
         new_self
     }
     /// Returns new [`SessionState`] using the provided
diff --git a/datafusion/expr/src/columnar_value.rs b/datafusion/expr/src/columnar_value.rs
index 95369c2242..5f2481ebc6 100644
--- a/datafusion/expr/src/columnar_value.rs
+++ b/datafusion/expr/src/columnar_value.rs
@@ -20,7 +20,7 @@
 use arrow::array::ArrayRef;
 use arrow::array::NullArray;
 use arrow::datatypes::DataType;
-use datafusion_common::{DataFusionError, Result, ScalarValue};
+use datafusion_common::{internal_err, DataFusionError, Result, ScalarValue};
 use std::sync::Arc;
 
 /// Represents the result of evaluating an expression: either a single
@@ -86,33 +86,25 @@ impl ColumnarValue {
             return Ok(vec![]);
         }
 
-        let len = args
-            .iter()
-            .fold(Option::<usize>::None, |acc, arg| match arg {
-                ColumnarValue::Scalar(_) if acc.is_none() => Some(1),
-                ColumnarValue::Scalar(_) => {
-                    if let Some(1) = acc {
-                        acc
+        let mut array_len = None;
+        for arg in args {
+            array_len = match (arg, array_len) {
+                (ColumnarValue::Array(a), None) => Some(a.len()),
+                (ColumnarValue::Array(a), Some(array_len)) => {
+                    if array_len == a.len() {
+                        Some(array_len)
                     } else {
-                        None
+                        return internal_err!(
+                            "Arguments has mixed length. Expected length: {array_len}, found length: {}",a.len()
+                        );
                     }
                 }
-                ColumnarValue::Array(a) => {
-                    if let Some(l) = acc {
-                        if l == a.len() {
-                            acc
-                        } else {
-                            None
-                        }
-                    } else {
-                        Some(a.len())
-                    }
-                }
-            });
+                (ColumnarValue::Scalar(_), array_len) => array_len,
+            }
+        }
 
-        let inferred_length = len.ok_or(DataFusionError::Internal(
-            "Arguments has mixed length".to_string(),
-        ))?;
+        // If array_len is none, it means there are only scalars, so make a 1 element array
+        let inferred_length = array_len.unwrap_or(1);
 
         let args = args
             .iter()