arrow_convert 0.12.1

use std::sync::Arc;

use arrow::datatypes::*;
use arrow_convert::{
    field::{with_list_element_metadata, with_list_element_name, DEFAULT_FIELD_NAME},
    serialize::TryIntoArrow,
    ArrowField, ArrowSerialize,
};
use pretty_assertions::assert_eq;

#[test]
fn test_schema_types() {
    #[derive(Debug, ArrowField)]
    #[allow(dead_code)]
    struct Root {
        name: Option<String>,
        is_deleted: bool,
        a1: Option<f64>,
        a2: i64,
        // binary
        a3: Option<Vec<u8>>,
        // date32
        a4: chrono::NaiveDate,
        // timestamp(ns, None)
        a5: chrono::NaiveDateTime,
        // timestamp(ns, None)
        a6: Option<chrono::NaiveDateTime>,
        // i128(precision, scale)
        #[arrow_field(type = "arrow_convert::field::I128<32, 32>")]
        a7: i128,
        // array of date times
        date_time_list: Vec<chrono::NaiveDateTime>,
        // optional list array of optional strings
        nullable_list: Option<Vec<Option<String>>>,
        // optional list array of required strings
        required_list: Vec<Option<String>>,
        // // custom type
        // custom: CustomType,
        // // custom optional type
        // nullable_custom: Option<CustomType>,
        // // vec custom type
        // custom_list: Vec<CustomType>,
        // nested struct
        child: Child,
        // int 32 array
        int32_array: Vec<i32>,
        // large binary
        #[arrow_field(type = "arrow_convert::field::LargeBinary")]
        large_binary: Vec<u8>,
        // fixed size binary
        #[arrow_field(type = "arrow_convert::field::FixedSizeBinary<3>")]
        fixed_size_binary: Vec<u8>,
        // large string
        #[arrow_field(type = "arrow_convert::field::LargeString")]
        large_string: String,
        // fixed chars
        area: [u8; 6],
        // string ref
        area_name: &'static str,
        // large vec
        #[arrow_field(type = "arrow_convert::field::LargeVec<i64>")]
        large_vec: Vec<i64>,
        // fixed size vec
        #[arrow_field(type = "arrow_convert::field::FixedSizeVec<i64, 3>")]
        fixed_size_vec: Vec<i64>,
    }

    #[derive(Debug, ArrowField)]
    #[allow(dead_code)]
    struct Child {
        a1: i64,
        a2: String,
        // nested struct array
        child_array: Vec<ChildChild>,
    }

    #[derive(Debug, ArrowField)]
    #[allow(dead_code)]
    pub struct ChildChild {
        a1: i32,
        bool_array: Vec<bool>,
        int64_array: Vec<i64>,
    }

    // enable Vec<CustomType>
    // arrow_convert::arrow_enable_vec_for_type!(CustomType);

    // #[derive(Debug)]
    // /// A newtype around a u64
    // pub struct CustomType(u64);

    // impl arrow_convert::field::ArrowField for CustomType {
    //     type Type = Self;

    //     fn data_type() -> arrow::datatypes::DataType {
    //         arrow::datatypes::DataType::Extension(
    //             "custom".to_string(),
    //             Arc::new(arrow::datatypes::DataType::UInt64),
    //             None,
    //         )
    //     }
    // }

    // impl arrow_convert::serialize::ArrowSerialize for CustomType {
    //     type ArrayBuilderType = arrow::array::UInt64Builder;

    //     #[inline]
    //     fn new_array() -> Self::ArrayBuilderType {
    //         unimplemented!();
    //     }

    //     #[inline]
    //     fn arrow_serialize(
    //         _v: &Self,
    //         _array: &mut Self::ArrayBuilderType,
    //     ) -> arrow::error::Result<()> {
    //         unimplemented!();
    //     }
    // }

    // impl arrow_convert::deserialize::ArrowDeserialize for CustomType {
    //     type ArrayType = arrow::array::PrimitiveArray<UInt64Type>;

    //     #[inline]
    //     fn arrow_deserialize(_v: Option<u64>) -> Option<Self> {
    //         unimplemented!();
    //     }
    // }

    assert_eq!(
        <Root as arrow_convert::field::ArrowField>::data_type(),
        DataType::Struct(Fields::from(vec![
            Field::new("name", DataType::Utf8, true),
            Field::new("is_deleted", DataType::Boolean, false),
            Field::new("a1", DataType::Float64, true),
            Field::new("a2", DataType::Int64, false),
            Field::new("a3", DataType::Binary, true),
            Field::new("a4", DataType::Date32, false),
            Field::new("a5", DataType::Timestamp(TimeUnit::Nanosecond, None), false),
            Field::new("a6", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
            Field::new("a7", DataType::Decimal128(32, 32), false),
            Field::new(
                "date_time_list",
                DataType::List(Arc::new(Field::new(
                    DEFAULT_FIELD_NAME,
                    DataType::Timestamp(TimeUnit::Nanosecond, None),
                    false
                ))),
                false
            ),
            Field::new(
                "nullable_list",
                DataType::List(Arc::new(Field::new(DEFAULT_FIELD_NAME, DataType::Utf8, true))),
                true
            ),
            Field::new(
                "required_list",
                DataType::List(Arc::new(Field::new(DEFAULT_FIELD_NAME, DataType::Utf8, true))),
                false
            ),
            // Field::new(
            //     "custom",
            //     DataType::Extension("custom".to_string(), Arc::new(DataType::UInt64), None),
            //     false
            // ),
            // Field::new(
            //     "nullable_custom",
            //     DataType::Extension("custom".to_string(), Arc::new(DataType::UInt64), None),
            //     true
            // ),
            // Field::new(
            //     "custom_list",
            //     DataType::List(Arc::new(Field::new(
            //         DEFAULT_FIELD_NAME,
            //         DataType::Extension("custom".to_string(), Arc::new(DataType::UInt64), None),
            //         false
            //     ))),
            //     false
            // ),
            Field::new(
                "child",
                DataType::Struct(Fields::from(vec![
                    Field::new("a1", DataType::Int64, false),
                    Field::new("a2", DataType::Utf8, false),
                    Field::new(
                        "child_array",
                        DataType::List(Arc::new(Field::new(
                            DEFAULT_FIELD_NAME,
                            DataType::Struct(Fields::from(vec![
                                Field::new("a1", DataType::Int32, false),
                                Field::new(
                                    "bool_array",
                                    DataType::List(Arc::new(Field::new(
                                        DEFAULT_FIELD_NAME,
                                        DataType::Boolean,
                                        false
                                    ))),
                                    false
                                ),
                                Field::new(
                                    "int64_array",
                                    DataType::List(Arc::new(Field::new(DEFAULT_FIELD_NAME, DataType::Int64, false))),
                                    false
                                ),
                            ])),
                            false
                        ))),
                        false
                    )
                ])),
                false
            ),
            Field::new(
                "int32_array",
                DataType::List(Arc::new(Field::new(DEFAULT_FIELD_NAME, DataType::Int32, false))),
                false
            ),
            Field::new("large_binary", DataType::LargeBinary, false),
            Field::new("fixed_size_binary", DataType::FixedSizeBinary(3), false),
            Field::new("large_string", DataType::LargeUtf8, false),
            Field::new("area", DataType::FixedSizeBinary(6), false),
            Field::new("area_name", DataType::Utf8, false),
            Field::new(
                "large_vec",
                DataType::LargeList(Arc::new(Field::new(DEFAULT_FIELD_NAME, DataType::Int64, false))),
                false
            ),
            Field::new(
                "fixed_size_vec",
                DataType::FixedSizeList(
                    Arc::new(Field::new(DEFAULT_FIELD_NAME, DataType::Int64, false)),
                    3
                ),
                false
            ),
        ]))
    );
}

#[test]
fn test_large_string_schema() {
    use arrow_convert::field::LargeString;

    assert_eq!(
        <LargeString as arrow_convert::field::ArrowField>::data_type(),
        DataType::LargeUtf8
    );
    assert!(!<LargeString as arrow_convert::field::ArrowField>::is_nullable());
    assert!(<Option<LargeString> as arrow_convert::field::ArrowField>::is_nullable());

    assert_eq!(
        <Vec<LargeString> as arrow_convert::field::ArrowField>::data_type(),
        DataType::List(Arc::new(Field::new(
            DEFAULT_FIELD_NAME,
            DataType::LargeUtf8,
            false
        )))
    );
}

#[test]
fn test_field_name_override_with_rename_all() {
    #[derive(Debug, ArrowField)]
    #[allow(dead_code)]
    #[arrow_field(rename_all = "camelCase")]
    struct Root {
        plain_field: i32,
        #[arrow_field(name = "custom_name")]
        renamed_field: i32,
        r#type: i32,
    }

    let DataType::Struct(fields) = <Root as arrow_convert::field::ArrowField>::data_type() else {
        panic!("expected struct datatype");
    };

    let names: Vec<_> = fields.iter().map(|field| field.name().to_string()).collect();
    assert_eq!(names, vec!["plainField", "custom_name", "type"]);
}

#[test]
fn test_rename_all_composes_with_type_name_and_skip() {
    #[derive(Debug, ArrowField)]
    #[allow(dead_code)]
    #[allow(non_snake_case)]
    #[arrow_field(rename_all = "snake_case")]
    struct Root {
        plainField: i32,
        #[arrow_field(type = "arrow_convert::field::LargeString")]
        optionalLabel: Option<String>,
        #[arrow_field(type = "arrow_convert::field::LargeVec<i64>", name = "custom_list")]
        ignoredNameByRenameAll: Vec<i64>,
        #[arrow_field(skip)]
        shouldSkip: i32,
    }

    let DataType::Struct(fields) = <Root as arrow_convert::field::ArrowField>::data_type() else {
        panic!("expected struct datatype");
    };

    assert_eq!(fields.len(), 3);

    assert_eq!(fields[0].name(), "plain_field");
    assert_eq!(fields[0].data_type(), &DataType::Int32);
    assert!(!fields[0].is_nullable());

    assert_eq!(fields[1].name(), "optional_label");
    assert_eq!(fields[1].data_type(), &DataType::LargeUtf8);
    assert!(!fields[1].is_nullable());

    assert_eq!(fields[2].name(), "custom_list");
    assert_eq!(
        fields[2].data_type(),
        &DataType::LargeList(Arc::new(Field::new(DEFAULT_FIELD_NAME, DataType::Int64, false)))
    );
    assert!(!fields[2].is_nullable());
}

#[test]
fn test_list_element_name_container_and_field_override() {
    #[derive(Debug, ArrowField)]
    #[allow(dead_code)]
    #[arrow_field(list_element_name = "entry")]
    struct Root {
        numbers: Vec<i32>,
        #[arrow_field(list_element_name = "node")]
        labels: Vec<String>,
        scalar: i64,
    }

    let DataType::Struct(fields) = <Root as arrow_convert::field::ArrowField>::data_type() else {
        panic!("expected struct datatype");
    };

    assert_eq!(
        fields[0].data_type(),
        &DataType::List(Arc::new(Field::new("entry", DataType::Int32, false)))
    );
    assert_eq!(
        fields[1].data_type(),
        &DataType::List(Arc::new(Field::new("node", DataType::Utf8, false)))
    );
    assert_eq!(fields[2].data_type(), &DataType::Int64);
}

#[test]
fn test_metadata_support_for_parquet_field_id_keys() {
    #[derive(Debug, ArrowField)]
    #[allow(dead_code)]
    #[arrow_field(list_element_metadata(scope = "container", PARQUET::field_id = "101"))]
    struct Root {
        #[arrow_field(
            metadata(role = "top", PARQUET::field_id = "7"),
            list_element_metadata(scope = "field", level = "1", PARQUET::field_id = "9")
        )]
        bids: Vec<i64>,
        asks: Vec<i64>,
    }

    let DataType::Struct(fields) = <Root as arrow_convert::field::ArrowField>::data_type() else {
        panic!("expected struct datatype");
    };

    let bids = &fields[0];
    assert_eq!(bids.metadata().get("role"), Some(&"top".to_string()));
    assert_eq!(bids.metadata().get("PARQUET:field_id"), Some(&"7".to_string()));
    let DataType::List(bids_element) = bids.data_type() else {
        panic!("expected list datatype for bids");
    };
    assert_eq!(bids_element.metadata().get("scope"), Some(&"field".to_string()));
    assert_eq!(bids_element.metadata().get("level"), Some(&"1".to_string()));
    assert_eq!(
        bids_element.metadata().get("PARQUET:field_id"),
        Some(&"9".to_string())
    );

    let asks = &fields[1];
    assert!(asks.metadata().get("role").is_none());
    let DataType::List(asks_element) = asks.data_type() else {
        panic!("expected list datatype for asks");
    };
    assert_eq!(
        asks_element.metadata().get("scope"),
        Some(&"container".to_string())
    );
    assert_eq!(
        asks_element.metadata().get("PARQUET:field_id"),
        Some(&"101".to_string())
    );
}

#[test]
fn test_list_element_metadata_field_override_wins_with_duplicate_container_keys() {
    #[derive(Debug, ArrowField)]
    #[allow(dead_code)]
    #[arrow_field(
        list_element_metadata(scope = "container_a"),
        list_element_metadata(scope = "container_b", keep = "container")
    )]
    struct Root {
        #[arrow_field(list_element_metadata(scope = "field"))]
        levels: Vec<i64>,
    }

    let DataType::Struct(fields) = <Root as arrow_convert::field::ArrowField>::data_type() else {
        panic!("expected struct datatype");
    };
    let DataType::List(levels_element) = fields[0].data_type() else {
        panic!("expected list datatype");
    };

    assert_eq!(
        levels_element.metadata().get("scope"),
        Some(&"field".to_string())
    );
    assert_eq!(
        levels_element.metadata().get("keep"),
        Some(&"container".to_string())
    );
}

#[test]
fn test_with_list_element_helpers_for_large_and_fixed_size_lists() {
    let large_field = Field::new(
        "values",
        DataType::LargeList(Arc::new(Field::new("_item", DataType::Int64, false))),
        false,
    );
    let large_named = with_list_element_name(large_field, Some("element"));
    let large_named = with_list_element_metadata(large_named, vec![("scope".to_string(), "book".to_string())]);
    let DataType::LargeList(large_element) = large_named.data_type() else {
        panic!("expected LargeList");
    };
    assert_eq!(large_element.name(), "element");
    assert_eq!(large_element.metadata().get("scope"), Some(&"book".to_string()));

    let fixed_field = Field::new(
        "values",
        DataType::FixedSizeList(Arc::new(Field::new("_item", DataType::Int64, false)), 3),
        false,
    );
    let fixed_named = with_list_element_name(fixed_field, Some("level"));
    let fixed_named = with_list_element_metadata(fixed_named, vec![("kind".to_string(), "depth".to_string())]);
    let DataType::FixedSizeList(fixed_element, size) = fixed_named.data_type() else {
        panic!("expected FixedSizeList");
    };
    assert_eq!(*size, 3);
    assert_eq!(fixed_element.name(), "level");
    assert_eq!(fixed_element.metadata().get("kind"), Some(&"depth".to_string()));
}

#[test]
fn test_serialize_respects_list_element_name_and_metadata() {
    #[derive(Debug, ArrowField, ArrowSerialize)]
    #[allow(dead_code)]
    #[arrow_field(list_element_metadata(scope = "container"))]
    struct Root {
        #[arrow_field(list_element_name = "level", list_element_metadata(PARQUET::field_id = "9"))]
        bids: Vec<i64>,
    }

    let rows = vec![Root { bids: vec![1, 2, 3] }];
    let array: arrow::array::ArrayRef = rows.try_into_arrow().expect("serialization should succeed");
    let struct_array = array
        .as_any()
        .downcast_ref::<arrow::array::StructArray>()
        .expect("expected StructArray");

    let fields = struct_array.fields();
    let bids = &fields[0];
    let DataType::List(element) = bids.data_type() else {
        panic!("expected list datatype");
    };

    assert_eq!(element.name(), "level");
    assert_eq!(element.metadata().get("scope"), Some(&"container".to_string()));
    assert_eq!(
        element.metadata().get("PARQUET:field_id"),
        Some(&"9".to_string())
    );
}