native/spark-expr/src/array_funcs/size.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use arrow::array::builder::Int32Builder;
use arrow::array::{Array, ArrayRef, GenericListArray, Int32Array, OffsetSizeTrait};
use arrow::datatypes::{DataType, Field};
use datafusion::common::{exec_err, DataFusionError, Result as DataFusionResult, ScalarValue};
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
};
use std::sync::Arc;

/// Spark size() function that returns the size of arrays or maps.
/// Returns -1 for null inputs (Spark behavior differs from standard SQL).
pub fn spark_size(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
    if args.len() != 1 {
        return exec_err!("size function takes exactly one argument");
    }

    match &args[0] {
        ColumnarValue::Array(array) => {
            let result = spark_size_array(array)?;
            Ok(ColumnarValue::Array(result))
        }
        ColumnarValue::Scalar(scalar) => {
            let result = spark_size_scalar(scalar)?;
            Ok(ColumnarValue::Scalar(result))
        }
    }
}

#[derive(Debug, Hash, Eq, PartialEq)]
pub struct SparkSizeFunc {
    signature: Signature,
}

impl Default for SparkSizeFunc {
    fn default() -> Self {
        Self::new()
    }
}

impl SparkSizeFunc {
    pub fn new() -> Self {
        use DataType::*;
        Self {
            signature: Signature::uniform(
                1,
                vec![
                    List(Arc::new(Field::new("item", Null, true))),
                    LargeList(Arc::new(Field::new("item", Null, true))),
                    FixedSizeList(Arc::new(Field::new("item", Null, true)), -1),
                    Map(Arc::new(Field::new("entries", Null, true)), false),
                ],
                Volatility::Immutable,
            ),
        }
    }
}

impl ScalarUDFImpl for SparkSizeFunc {
    fn name(&self) -> &str {
        "size"
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult<DataType> {
        Ok(DataType::Int32)
    }

    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
        spark_size(&args.args)
    }
}

fn spark_size_array(array: &ArrayRef) -> Result<ArrayRef, DataFusionError> {
    let mut builder = Int32Array::builder(array.len());

    match array.data_type() {
        DataType::List(_) => {
            let list_array = array
                .as_any()
                .downcast_ref::<arrow::array::ListArray>()
                .ok_or_else(|| DataFusionError::Internal("Expected ListArray".to_string()))?;
            append_list_sizes(&mut builder, list_array);
        }
        DataType::LargeList(_) => {
            let list_array = array
                .as_any()
                .downcast_ref::<arrow::array::LargeListArray>()
                .ok_or_else(|| DataFusionError::Internal("Expected LargeListArray".to_string()))?;
            append_list_sizes(&mut builder, list_array);
        }
        DataType::FixedSizeList(_, size) => {
            let fixed_list_array = array
                .as_any()
                .downcast_ref::<arrow::array::FixedSizeListArray>()
                .ok_or_else(|| {
                    DataFusionError::Internal("Expected FixedSizeListArray".to_string())
                })?;

            for i in 0..fixed_list_array.len() {
                if fixed_list_array.is_null(i) {
                    builder.append_value(-1); // Spark behavior: return -1 for null
                } else {
                    builder.append_value(*size);
                }
            }
        }
        DataType::Map(_, _) => {
            let map_array = array
                .as_any()
                .downcast_ref::<arrow::array::MapArray>()
                .ok_or_else(|| DataFusionError::Internal("Expected MapArray".to_string()))?;

            for i in 0..map_array.len() {
                if map_array.is_null(i) {
                    builder.append_value(-1); // Spark behavior: return -1 for null
                } else {
                    let map_len = map_array.value_length(i);
                    builder.append_value(map_len);
                }
            }
        }
        _ => {
            return exec_err!(
                "size function only supports arrays and maps, got: {:?}",
                array.data_type()
            );
        }
    }

    Ok(Arc::new(builder.finish()))
}

/// Append the element count of each list row to `builder`, using `-1` for null
/// rows (Spark's behavior). `value_length` reads the row's element count from the
/// offset buffer, avoiding the per-row allocation that `value(i).len()` would incur
/// from materializing a sliced array.
fn append_list_sizes<O: OffsetSizeTrait>(
    builder: &mut Int32Builder,
    list_array: &GenericListArray<O>,
) {
    for i in 0..list_array.len() {
        if list_array.is_null(i) {
            builder.append_value(-1); // Spark behavior: return -1 for null
        } else {
            builder.append_value(list_array.value_length(i).as_usize() as i32);
        }
    }
}

fn spark_size_scalar(scalar: &ScalarValue) -> Result<ScalarValue, DataFusionError> {
    match scalar {
        ScalarValue::List(array) => {
            // ScalarValue::List contains a ListArray with exactly one row.
            // We need the length of that row's contents, not the row count.
            if array.is_null(0) {
                Ok(ScalarValue::Int32(Some(-1))) // Spark behavior: return -1 for null
            } else {
                let len = array.value(0).len() as i32;
                Ok(ScalarValue::Int32(Some(len)))
            }
        }
        ScalarValue::LargeList(array) => {
            if array.is_null(0) {
                Ok(ScalarValue::Int32(Some(-1)))
            } else {
                let len = array.value(0).len() as i32;
                Ok(ScalarValue::Int32(Some(len)))
            }
        }
        ScalarValue::FixedSizeList(array) => {
            if array.is_null(0) {
                Ok(ScalarValue::Int32(Some(-1)))
            } else {
                let len = array.value(0).len() as i32;
                Ok(ScalarValue::Int32(Some(len)))
            }
        }
        ScalarValue::Map(array) => {
            if array.is_null(0) {
                Ok(ScalarValue::Int32(Some(-1)))
            } else {
                let len = array.value_length(0);
                Ok(ScalarValue::Int32(Some(len)))
            }
        }
        ScalarValue::Null => {
            Ok(ScalarValue::Int32(Some(-1))) // Spark behavior: return -1 for null
        }
        _ => {
            exec_err!(
                "size function only supports arrays and maps, got: {:?}",
                scalar
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::{Int32Array, ListArray, NullBufferBuilder};
    use arrow::datatypes::{DataType, Field};
    use std::sync::Arc;

    #[test]
    fn test_spark_size_array() {
        // Create test data: [[1, 2, 3], [4, 5], null, []]
        let value_data = Int32Array::from(vec![1, 2, 3, 4, 5]);
        let value_offsets = arrow::buffer::OffsetBuffer::new(vec![0, 3, 5, 5, 5].into());
        let field = Arc::new(Field::new("item", DataType::Int32, true));

        let mut null_buffer = NullBufferBuilder::new(4);
        null_buffer.append(true); // [1, 2, 3] - not null
        null_buffer.append(true); // [4, 5] - not null
        null_buffer.append(false); // null
        null_buffer.append(true); // [] - not null but empty

        let list_array = ListArray::try_new(
            field,
            value_offsets,
            Arc::new(value_data),
            null_buffer.finish(),
        )
        .unwrap();

        let array_ref: ArrayRef = Arc::new(list_array);
        let result = spark_size_array(&array_ref).unwrap();
        let result = result.as_any().downcast_ref::<Int32Array>().unwrap();

        // Expected: [3, 2, -1, 0]
        assert_eq!(result.value(0), 3); // [1, 2, 3] has 3 elements
        assert_eq!(result.value(1), 2); // [4, 5] has 2 elements
        assert_eq!(result.value(2), -1); // null returns -1
        assert_eq!(result.value(3), 0); // [] has 0 elements
    }

    #[test]
    fn test_spark_size_scalar() {
        // Test non-null list with 3 elements
        let values = Int32Array::from(vec![1, 2, 3]);
        let field = Arc::new(Field::new("item", DataType::Int32, true));
        let offsets = arrow::buffer::OffsetBuffer::new(vec![0, 3].into());
        let list_array = ListArray::try_new(field, offsets, Arc::new(values), None).unwrap();
        let scalar = ScalarValue::List(Arc::new(list_array));
        let result = spark_size_scalar(&scalar).unwrap();
        assert_eq!(result, ScalarValue::Int32(Some(3))); // The array [1,2,3] has 3 elements

        // Test empty list
        let empty_values = Int32Array::from(vec![] as Vec<i32>);
        let field = Arc::new(Field::new("item", DataType::Int32, true));
        let offsets = arrow::buffer::OffsetBuffer::new(vec![0, 0].into());
        let empty_list_array =
            ListArray::try_new(field, offsets, Arc::new(empty_values), None).unwrap();
        let scalar = ScalarValue::List(Arc::new(empty_list_array));
        let result = spark_size_scalar(&scalar).unwrap();
        assert_eq!(result, ScalarValue::Int32(Some(0))); // Empty array has 0 elements

        // Test null handling
        let scalar = ScalarValue::Null;
        let result = spark_size_scalar(&scalar).unwrap();
        assert_eq!(result, ScalarValue::Int32(Some(-1)));
    }

    #[test]
    fn test_spark_size_map_array() {
        use arrow::array::{Int32Array, MapArray, StringArray};

        // Create test data: [{"key1": 1, "key2": 2}, {"key3": 3}, {}, null]
        let keys = StringArray::from(vec![Some("key1"), Some("key2"), Some("key3")]);
        let values = Int32Array::from(vec![Some(1), Some(2), Some(3)]);

        // Create entry offsets: [0, 2, 3, 3, 3] representing:
        // - Map 1: entries 0-1 (2 key-value pairs)
        // - Map 2: entry 2 (1 key-value pair)
        // - Map 3: entries 3-2 (0 key-value pairs, empty map)
        // - Map 4: null (handled by null buffer)
        let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0i32, 2, 3, 3, 3].into());

        let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
        let value_field = Arc::new(Field::new("value", DataType::Int32, true));

        let entries = arrow::array::StructArray::new(
            arrow::datatypes::Fields::from(vec![key_field, value_field]),
            vec![Arc::new(keys), Arc::new(values)],
            None,
        );

        // Create null buffer for the map array (fourth map is null)
        let mut null_buffer = NullBufferBuilder::new(4);
        null_buffer.append(true); // Map with 2 entries - not null
        null_buffer.append(true); // Map with 1 entry - not null
        null_buffer.append(true); // Empty map - not null
        null_buffer.append(false); // null map

        let map_field = Arc::new(Field::new(
            "entries",
            DataType::Struct(arrow::datatypes::Fields::from(vec![
                Field::new("key", DataType::Utf8, false),
                Field::new("value", DataType::Int32, true),
            ])),
            false,
        ));

        let map_array = MapArray::try_new(
            map_field,
            entry_offsets,
            entries,
            null_buffer.finish(),
            false,
        )
        .unwrap();

        let array_ref: ArrayRef = Arc::new(map_array);
        let result = spark_size_array(&array_ref).unwrap();
        let result = result.as_any().downcast_ref::<Int32Array>().unwrap();

        // Expected: [2, 1, 0, -1]
        assert_eq!(result.value(0), 2); // Map with 2 key-value pairs
        assert_eq!(result.value(1), 1); // Map with 1 key-value pair
        assert_eq!(result.value(2), 0); // empty map has 0 pairs
        assert_eq!(result.value(3), -1); // null map returns -1
    }

    #[test]
    fn test_spark_size_scalar_map() {
        use arrow::array::{Int32Array, MapArray, StringArray};

        // Test non-null map with 2 entries
        let keys = StringArray::from(vec![Some("a"), Some("b")]);
        let values = Int32Array::from(vec![Some(1), Some(2)]);
        let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0i32, 2].into());

        let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
        let value_field = Arc::new(Field::new("value", DataType::Int32, true));

        let entries = arrow::array::StructArray::new(
            arrow::datatypes::Fields::from(vec![key_field, value_field]),
            vec![Arc::new(keys), Arc::new(values)],
            None,
        );

        let map_field = Arc::new(Field::new(
            "entries",
            DataType::Struct(arrow::datatypes::Fields::from(vec![
                Field::new("key", DataType::Utf8, false),
                Field::new("value", DataType::Int32, true),
            ])),
            false,
        ));

        let map_array = MapArray::try_new(map_field, entry_offsets, entries, None, false).unwrap();
        let scalar = ScalarValue::Map(Arc::new(map_array));
        let result = spark_size_scalar(&scalar).unwrap();
        assert_eq!(result, ScalarValue::Int32(Some(2))); // Map with 2 entries
    }

    #[test]
    fn test_spark_size_scalar_null_map() {
        use arrow::array::{Int32Array, MapArray, StringArray};

        // Test null map - Spark returns -1 for null map input
        let keys = StringArray::from(vec![Some("a")]);
        let values = Int32Array::from(vec![Some(1)]);
        let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0i32, 1].into());

        let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
        let value_field = Arc::new(Field::new("value", DataType::Int32, true));

        let entries = arrow::array::StructArray::new(
            arrow::datatypes::Fields::from(vec![key_field, value_field]),
            vec![Arc::new(keys), Arc::new(values)],
            None,
        );

        let map_field = Arc::new(Field::new(
            "entries",
            DataType::Struct(arrow::datatypes::Fields::from(vec![
                Field::new("key", DataType::Utf8, false),
                Field::new("value", DataType::Int32, true),
            ])),
            false,
        ));

        // Mark the single map entry as null
        let mut null_buffer = NullBufferBuilder::new(1);
        null_buffer.append(false); // null map

        let map_array = MapArray::try_new(
            map_field,
            entry_offsets,
            entries,
            null_buffer.finish(),
            false,
        )
        .unwrap();
        let scalar = ScalarValue::Map(Arc::new(map_array));
        let result = spark_size_scalar(&scalar).unwrap();
        assert_eq!(result, ScalarValue::Int32(Some(-1))); // null map returns -1
    }

    #[test]
    fn test_spark_size_fixed_size_list_array() {
        use arrow::array::FixedSizeListArray;

        // Create test data: fixed-size arrays of size 3
        // [[1, 2, 3], [4, 5, 6], null]
        let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 0, 0, 0]); // Last 3 values are for the null entry
        let list_size = 3;

        let mut null_buffer = NullBufferBuilder::new(3);
        null_buffer.append(true); // [1, 2, 3] - not null
        null_buffer.append(true); // [4, 5, 6] - not null
        null_buffer.append(false); // null

        let list_field = Arc::new(Field::new("item", DataType::Int32, true));

        let fixed_list_array = FixedSizeListArray::new(
            list_field,
            list_size,
            Arc::new(values),
            null_buffer.finish(),
        );

        let array_ref: ArrayRef = Arc::new(fixed_list_array);
        let result = spark_size_array(&array_ref).unwrap();
        let result = result.as_any().downcast_ref::<Int32Array>().unwrap();

        // Expected: [3, 3, -1]
        assert_eq!(result.value(0), 3); // Fixed-size list always has size 3
        assert_eq!(result.value(1), 3); // Fixed-size list always has size 3
        assert_eq!(result.value(2), -1); // null returns -1
    }

    #[test]
    fn test_spark_size_large_list_array() {
        use arrow::array::LargeListArray;

        // Create test data: [[1, 2, 3, 4], [5], null, []]
        let value_data = Int32Array::from(vec![1, 2, 3, 4, 5]);
        let value_offsets = arrow::buffer::OffsetBuffer::new(vec![0i64, 4, 5, 5, 5].into());
        let field = Arc::new(Field::new("item", DataType::Int32, true));

        let mut null_buffer = NullBufferBuilder::new(4);
        null_buffer.append(true); // [1, 2, 3, 4] - not null
        null_buffer.append(true); // [5] - not null
        null_buffer.append(false); // null
        null_buffer.append(true); // [] - not null but empty

        let large_list_array = LargeListArray::try_new(
            field,
            value_offsets,
            Arc::new(value_data),
            null_buffer.finish(),
        )
        .unwrap();

        let array_ref: ArrayRef = Arc::new(large_list_array);
        let result = spark_size_array(&array_ref).unwrap();
        let result = result.as_any().downcast_ref::<Int32Array>().unwrap();

        // Expected: [4, 1, -1, 0]
        assert_eq!(result.value(0), 4); // [1, 2, 3, 4] has 4 elements
        assert_eq!(result.value(1), 1); // [5] has 1 element
        assert_eq!(result.value(2), -1); // null returns -1
        assert_eq!(result.value(3), 0); // [] has 0 elements
    }
}