You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2022/04/25 13:11:15 UTC

[arrow-datafusion] branch master updated: Support scalar values for func Array (#2332)

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

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


The following commit(s) were added to refs/heads/master by this push:
     new 27799d58a Support scalar values for func Array (#2332)
27799d58a is described below

commit 27799d58ac38aa8376313f5a04bf0b45799b8657
Author: Yang Jiang <37...@users.noreply.github.com>
AuthorDate: Mon Apr 25 21:11:10 2022 +0800

    Support scalar values for func Array (#2332)
    
    * Support scalar values for func Array
    
    * Apply suggestions from code review
    
    Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
    
    * fix fmt
    
    Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
 datafusion/common/src/scalar.rs                   | 17 +++++++++++++++++
 datafusion/core/tests/sql/functions.rs            | 11 +++++++++++
 datafusion/physical-expr/src/array_expressions.rs | 20 +++++++-------------
 3 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs
index e6963db1b..03a59ff6d 100644
--- a/datafusion/common/src/scalar.rs
+++ b/datafusion/common/src/scalar.rs
@@ -1391,6 +1391,23 @@ impl ScalarValue {
                 }
                 Self::Struct(Some(Box::new(field_values)), Box::new(fields.clone()))
             }
+            DataType::FixedSizeList(nested_type, _len) => {
+                let list_array =
+                    array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
+                let value = match list_array.is_null(index) {
+                    true => None,
+                    false => {
+                        let nested_array = list_array.value(index);
+                        let scalar_vec = (0..nested_array.len())
+                            .map(|i| ScalarValue::try_from_array(&nested_array, i))
+                            .collect::<Result<Vec<_>>>()?;
+                        Some(scalar_vec)
+                    }
+                };
+                let value = value.map(Box::new);
+                let data_type = Box::new(nested_type.data_type().clone());
+                ScalarValue::List(value, data_type)
+            }
             other => {
                 return Err(DataFusionError::NotImplemented(format!(
                     "Can't create a scalar from array of type \"{:?}\"",
diff --git a/datafusion/core/tests/sql/functions.rs b/datafusion/core/tests/sql/functions.rs
index 06ddea09d..ee2da9c94 100644
--- a/datafusion/core/tests/sql/functions.rs
+++ b/datafusion/core/tests/sql/functions.rs
@@ -143,6 +143,17 @@ async fn query_array() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn query_array_scalar() -> Result<()> {
+    let ctx = SessionContext::new();
+
+    let sql = "SELECT array(1, 2, 3);";
+    let actual = execute(&ctx, sql).await;
+    let expected = vec![vec!["[1, 2, 3]"]];
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
 #[tokio::test]
 async fn query_count_distinct() -> Result<()> {
     let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs
index cf220149f..45ec881ba 100644
--- a/datafusion/physical-expr/src/array_expressions.rs
+++ b/datafusion/physical-expr/src/array_expressions.rs
@@ -59,7 +59,7 @@ macro_rules! array {
     }};
 }
 
-fn array_array(args: &[&dyn Array]) -> Result<ArrayRef> {
+fn array_array(args: &[ArrayRef]) -> Result<ArrayRef> {
     // do not accept 0 arguments.
     if args.is_empty() {
         return Err(DataFusionError::Internal(
@@ -90,18 +90,12 @@ fn array_array(args: &[&dyn Array]) -> Result<ArrayRef> {
 
 /// put values in an array.
 pub fn array(values: &[ColumnarValue]) -> Result<ColumnarValue> {
-    let arrays: Vec<&dyn Array> = values
+    let arrays: Vec<ArrayRef> = values
         .iter()
-        .map(|value| {
-            if let ColumnarValue::Array(value) = value {
-                Ok(value.as_ref())
-            } else {
-                Err(DataFusionError::NotImplemented(
-                    "Array is not implemented for scalar values.".to_string(),
-                ))
-            }
+        .map(|x| match x {
+            ColumnarValue::Array(array) => array.clone(),
+            ColumnarValue::Scalar(scalar) => scalar.to_array().clone(),
         })
-        .collect::<Result<_>>()?;
-
-    Ok(ColumnarValue::Array(array_array(&arrays)?))
+        .collect();
+    Ok(ColumnarValue::Array(array_array(arrays.as_slice())?))
 }