# -*- coding: utf-8 -*-
from fastparquet.parquet_thrift import (
    ConvertedType,
    FieldRepetitionType,
    SchemaElement,
    Type,
)
from fastparquet.schema import SchemaHelper

from pyathena.fastparquet.util import to_column_info


def test_to_column_info():
    schema = SchemaHelper(
        [
            SchemaElement(
                type=None,
                type_length=None,
                repetition_type=None,
                name="hive_schema",
                num_children=16,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.BOOLEAN,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_boolean",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_tinyint",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_smallint",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_int",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT64,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_bigint",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.FLOAT,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_float",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.DOUBLE,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_double",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.BYTE_ARRAY,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_string",
                num_children=None,
                converted_type=ConvertedType.UTF8,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.BYTE_ARRAY,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_varchar",
                num_children=None,
                converted_type=ConvertedType.UTF8,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT96,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_timestamp",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_date",
                num_children=None,
                converted_type=ConvertedType.DATE,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.BYTE_ARRAY,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_binary",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=None,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_array",
                num_children=1,
                converted_type=ConvertedType.LIST,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=None,
                type_length=None,
                repetition_type=FieldRepetitionType.REPEATED,
                name="bag",
                num_children=1,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="array_element",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=None,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_map",
                num_children=1,
                converted_type=ConvertedType.MAP,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=None,
                type_length=None,
                repetition_type=FieldRepetitionType.REPEATED,
                name="key_value",
                num_children=2,
                converted_type=ConvertedType.MAP_KEY_VALUE,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.REQUIRED,
                name="key",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="value",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=None,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_struct",
                num_children=2,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="a",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.INT32,
                type_length=None,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="b",
                num_children=None,
                converted_type=None,
                scale=None,
                precision=None,
                field_id=None,
            ),
            SchemaElement(
                type=Type.FIXED_LEN_BYTE_ARRAY,
                type_length=5,
                repetition_type=FieldRepetitionType.OPTIONAL,
                name="col_decimal",
                num_children=None,
                converted_type=ConvertedType.DECIMAL,
                scale=1,
                precision=10,
                field_id=None,
            ),
        ]
    )
    assert to_column_info(schema) == (
        {
            "Name": "col_boolean",
            "Nullable": "NULLABLE",
            "Precision": 0,
            "Scale": 0,
            "Type": "boolean",
        },
        {
            "Name": "col_tinyint",
            "Nullable": "NULLABLE",
            "Precision": 10,
            "Scale": 0,
            "Type": "integer",
        },
        {
            "Name": "col_smallint",
            "Nullable": "NULLABLE",
            "Precision": 10,
            "Scale": 0,
            "Type": "integer",
        },
        {
            "Name": "col_int",
            "Nullable": "NULLABLE",
            "Precision": 10,
            "Scale": 0,
            "Type": "integer",
        },
        {
            "Name": "col_bigint",
            "Nullable": "NULLABLE",
            "Precision": 19,
            "Scale": 0,
            "Type": "bigint",
        },
        {
            "Name": "col_float",
            "Nullable": "NULLABLE",
            "Precision": 17,
            "Scale": 0,
            "Type": "float",
        },
        {
            "Name": "col_double",
            "Nullable": "NULLABLE",
            "Precision": 17,
            "Scale": 0,
            "Type": "double",
        },
        {
            "Name": "col_string",
            "Nullable": "NULLABLE",
            "Precision": 2147483647,
            "Scale": 0,
            "Type": "varchar",
        },
        {
            "Name": "col_varchar",
            "Nullable": "NULLABLE",
            "Precision": 2147483647,
            "Scale": 0,
            "Type": "varchar",
        },
        {
            "Name": "col_timestamp",
            "Nullable": "NULLABLE",
            "Precision": 3,
            "Scale": 0,
            "Type": "timestamp",
        },
        {
            "Name": "col_date",
            "Nullable": "NULLABLE",
            "Precision": 0,
            "Scale": 0,
            "Type": "date",
        },
        {
            "Name": "col_binary",
            "Nullable": "NULLABLE",
            "Precision": 1073741824,
            "Scale": 0,
            "Type": "varbinary",
        },
        {
            "Name": "col_array",
            "Nullable": "NULLABLE",
            "Precision": 0,
            "Scale": 0,
            "Type": "array",
        },
        {
            "Name": "col_map",
            "Nullable": "NULLABLE",
            "Precision": 0,
            "Scale": 0,
            "Type": "map",
        },
        {
            "Name": "col_decimal",
            "Nullable": "NULLABLE",
            "Precision": 10,
            "Scale": 1,
            "Type": "decimal",
        },
        {
            "Name": "col_struct.a",
            "Nullable": "NULLABLE",
            "Precision": 10,
            "Scale": 0,
            "Type": "integer",
        },
        {
            "Name": "col_struct.b",
            "Nullable": "NULLABLE",
            "Precision": 10,
            "Scale": 0,
            "Type": "integer",
        },
    )
