Function test_align_array_dimensions

datafusion/functions-nested/src/utils.rs:288–333

()

    /// Only test internal functions, array-related sql functions will be tested in sqllogictest `array.slt`
    #[test]
    fn test_align_array_dimensions() {
        let array1d_1: ArrayRef =
            Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
                Some(vec![Some(1), Some(2), Some(3)]),
                Some(vec![Some(4), Some(5)]),
            ]));
        let array1d_2: ArrayRef =
            Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
                Some(vec![Some(6), Some(7), Some(8)]),
            ]));

        let array2d_1: ArrayRef = Arc::new(
            SingleRowListArrayBuilder::new(Arc::clone(&array1d_1)).build_list_array(),
        );
        let array2d_2 = Arc::new(
            SingleRowListArrayBuilder::new(Arc::clone(&array1d_2)).build_list_array(),
        );

        let res = align_array_dimensions::<i32>(vec![
            array1d_1.to_owned(),
            array2d_2.to_owned(),
        ])
        .unwrap();

        let expected = as_list_array(&array2d_1).unwrap();
        let expected_dim = datafusion_common::utils::list_ndims(array2d_1.data_type());
        assert_ne!(as_list_array(&res[0]).unwrap(), expected);
        assert_eq!(
            datafusion_common::utils::list_ndims(res[0].data_type()),
            expected_dim
        );

        let array3d_1: ArrayRef =
            Arc::new(SingleRowListArrayBuilder::new(array2d_1).build_list_array());
        let array3d_2: ArrayRef =
            Arc::new(SingleRowListArrayBuilder::new(array2d_2).build_list_array());
        let res = align_array_dimensions::<i32>(vec![array1d_1, array3d_2]).unwrap();

        let expected = as_list_array(&array3d_1).unwrap();
        let expected_dim = datafusion_common::utils::list_ndims(array3d_1.data_type());
        assert_ne!(as_list_array(&res[0]).unwrap(), expected);
        assert_eq!(
            datafusion_common::utils::list_ndims(res[0].data_type()),
            expected_dim
        );
    }
}