| # 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. |
| |
| from iceberg.api import Schema |
| from iceberg.api.expressions import Expressions |
| from iceberg.api.types import (BooleanType, |
| DateType, |
| DecimalType, |
| DoubleType, |
| FloatType, |
| IntegerType, |
| LongType, |
| NestedField, |
| StringType, |
| TimestampType) |
| from iceberg.core.filesystem import FileSystemInputFile, get_fs |
| from iceberg.parquet import ParquetReader |
| import pyarrow as pa |
| |
| |
| def test_basic_read(primitive_type_test_file, pyarrow_primitive_array, pyarrow_schema): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get()), |
| NestedField.optional(3, "str_col", StringType.get()), |
| NestedField.optional(4, "float_col", FloatType.get()), |
| NestedField.optional(5, "dbl_col", DoubleType.get()), |
| NestedField.optional(6, "decimal_col", DecimalType.of(9, 2)), |
| NestedField.optional(7, "big_decimal_col", DecimalType.of(19, 5)), |
| NestedField.optional(8, "huge_decimal_col", DecimalType.of(38, 9)), |
| NestedField.optional(9, "date_col", DateType.get()), |
| NestedField.optional(10, "ts_col", TimestampType.without_timezone()), |
| NestedField.optional(11, "ts_wtz_col", TimestampType.with_timezone()), |
| NestedField.optional(12, "bool_col", BooleanType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.always_true(), True) |
| |
| source_table = pa.table(pyarrow_primitive_array, schema=pyarrow_schema) |
| assert reader.read() == source_table |
| |
| |
| def test_projection(primitive_type_test_file, pyarrow_primitive_array, pyarrow_schema): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.always_true(), True) |
| |
| source_table = pa.table(pyarrow_primitive_array, schema=pyarrow_schema) |
| num_cols = source_table.num_columns |
| for i in range(1, num_cols - 1): |
| source_table = source_table.remove_column(num_cols - i) |
| |
| assert source_table == reader.read() |
| |
| |
| def test_column_rename(primitive_type_test_file): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get()), |
| NestedField.optional(3, "string_col", StringType.get()), |
| NestedField.optional(4, "float_col", FloatType.get()), |
| NestedField.optional(5, "dbl_col", DoubleType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.always_true(), True) |
| pyarrow_array = [pa.array([1, 2, 3, 4, 5], type=pa.int32()), |
| pa.array([1, 2, 3, None, 5], type=pa.int64()), |
| pa.array(['us', 'can', 'us', 'us', 'can'], type=pa.string()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float32()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float64())] |
| schema = pa.schema([pa.field("int_col", pa.int32(), False), |
| pa.field("bigint_col", pa.int64(), True), |
| pa.field("string_col", pa.string(), True), |
| pa.field("float_col", pa.float32(), True), |
| pa.field("dbl_col", pa.float64(), True)]) |
| |
| source_table = pa.table(pyarrow_array, schema=schema) |
| |
| target_table = reader.read() |
| assert source_table == target_table |
| |
| |
| def test_column_add(primitive_type_test_file): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get()), |
| NestedField.optional(3, "string_col", StringType.get()), |
| NestedField.optional(4, "float_col", FloatType.get()), |
| NestedField.optional(5, "dbl_col", DoubleType.get()), |
| NestedField.optional(13, "int_col2", IntegerType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.always_true(), True) |
| pyarrow_array = [pa.array([1, 2, 3, 4, 5], type=pa.int32()), |
| pa.array([1, 2, 3, None, 5], type=pa.int64()), |
| pa.array(['us', 'can', 'us', 'us', 'can'], type=pa.string()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float32()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float64()), |
| pa.array([None, None, None, None, None], type=pa.int32())] |
| source_table = pa.table(pyarrow_array, schema=pa.schema([pa.field("int_col", pa.int32(), nullable=False), |
| pa.field("bigint_col", pa.int64(), nullable=True), |
| pa.field("string_col", pa.string(), nullable=True), |
| pa.field("float_col", pa.float32(), nullable=True), |
| pa.field("dbl_col", pa.float64(), nullable=True), |
| pa.field("int_col2", pa.int32(), nullable=True), |
| ])) |
| |
| target_table = reader.read() |
| assert source_table == target_table |
| |
| |
| def test_decimal_column_add(primitive_type_test_file): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get()), |
| NestedField.optional(4, "float_col", FloatType.get()), |
| NestedField.optional(5, "dbl_col", DoubleType.get()), |
| NestedField.optional(13, "new_dec_col", DecimalType.of(38, 9)) |
| ]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.always_true(), True) |
| pyarrow_array = [pa.array([1, 2, 3, 4, 5], type=pa.int32()), |
| pa.array([1, 2, 3, None, 5], type=pa.int64()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float32()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float64()), |
| pa.array([None, None, None, None, None], type=pa.decimal128(38, 9))] |
| |
| source_table = pa.table(pyarrow_array, schema=pa.schema([pa.field("int_col", pa.int32(), nullable=False), |
| pa.field("bigint_col", pa.int64(), nullable=True), |
| pa.field("float_col", pa.float32(), nullable=True), |
| pa.field("dbl_col", pa.float64(), nullable=True), |
| pa.field("new_dec_col", pa.decimal128(38, 9), nullable=True) |
| ])) |
| |
| target_table = reader.read() |
| assert source_table == target_table |
| |
| |
| def test_column_reorder(primitive_type_test_file): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get()), |
| NestedField.optional(4, "float_col", FloatType.get()), |
| NestedField.optional(5, "dbl_col", DoubleType.get()), |
| NestedField.optional(3, "string_col", StringType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.always_true(), True) |
| pyarrow_array = [pa.array([1, 2, 3, 4, 5], type=pa.int32()), |
| pa.array([1, 2, 3, None, 5], type=pa.int64()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float32()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float64()), |
| pa.array(['us', 'can', 'us', 'us', 'can'], type=pa.string())] |
| source_table = pa.table(pyarrow_array, schema=pa.schema([pa.field("int_col", pa.int32(), nullable=False), |
| pa.field("bigint_col", pa.int64(), nullable=True), |
| pa.field("float_col", pa.float32(), nullable=True), |
| pa.field("dbl_col", pa.float64(), nullable=True), |
| pa.field("string_col", pa.string(), nullable=True) |
| ])) |
| |
| target_table = reader.read() |
| assert source_table == target_table |
| |
| |
| def test_column_upcast(primitive_type_test_file): |
| expected_schema = Schema([NestedField.required(1, "int_col", LongType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.always_true(), True) |
| pyarrow_array = [pa.array([1, 2, 3, 4, 5], type=pa.int32())] |
| source_table = pa.table(pyarrow_array, schema=pa.schema([pa.field("int_col", pa.int64(), nullable=False)])) |
| |
| target_table = reader.read() |
| assert source_table == target_table |
| |
| |
| def test_filter(primitive_type_test_file): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get()), |
| NestedField.optional(4, "float_col", FloatType.get()), |
| NestedField.optional(5, "dbl_col", DoubleType.get()), |
| NestedField.optional(3, "string_col", StringType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.equal("string_col", "us"), True) |
| pyarrow_array = [pa.array([1, 3, 4], type=pa.int32()), |
| pa.array([1, 3, None], type=pa.int64()), |
| pa.array([1.0, 3.0, 4.0], type=pa.float32()), |
| pa.array([1.0, 3.0, 4.0], type=pa.float64()), |
| pa.array(['us', 'us', 'us'], type=pa.string())] |
| source_table = pa.table(pyarrow_array, schema=pa.schema([pa.field("int_col", pa.int32(), nullable=False), |
| pa.field("bigint_col", pa.int64(), nullable=True), |
| pa.field("float_col", pa.float32(), nullable=True), |
| pa.field("dbl_col", pa.float64(), nullable=True), |
| pa.field("string_col", pa.string(), nullable=True) |
| ])) |
| |
| target_table = reader.read() |
| assert source_table == target_table |
| |
| |
| def test_compound_filter(primitive_type_test_file): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get()), |
| NestedField.optional(4, "float_col", FloatType.get()), |
| NestedField.optional(5, "dbl_col", DoubleType.get()), |
| NestedField.optional(3, "string_col", StringType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.and_(Expressions.equal("string_col", "us"), |
| Expressions.equal("int_col", 1)), |
| True) |
| pyarrow_array = [pa.array([1], type=pa.int32()), |
| pa.array([1], type=pa.int64()), |
| pa.array([1.0], type=pa.float32()), |
| pa.array([1.0], type=pa.float64()), |
| pa.array(['us'], type=pa.string())] |
| |
| source_table = pa.table(pyarrow_array, schema=pa.schema([pa.field("int_col", pa.int32(), nullable=False), |
| pa.field("bigint_col", pa.int64(), nullable=True), |
| pa.field("float_col", pa.float32(), nullable=True), |
| pa.field("dbl_col", pa.float64(), nullable=True), |
| pa.field("string_col", pa.string(), nullable=True) |
| ])) |
| |
| target_table = reader.read() |
| assert source_table == target_table |
| |
| |
| def test_schema_evolution_filter(primitive_type_test_file): |
| expected_schema = Schema([NestedField.required(1, "int_col", IntegerType.get()), |
| NestedField.optional(2, "bigint_col", LongType.get()), |
| NestedField.optional(16, "other_new_col", LongType.get()), |
| NestedField.optional(4, "float_col", FloatType.get()), |
| NestedField.optional(5, "dbl_col", DoubleType.get()), |
| NestedField.optional(3, "string_col", StringType.get()), |
| NestedField.optional(15, "new_col", StringType.get())]) |
| |
| input_file = FileSystemInputFile(get_fs(primitive_type_test_file, conf={}), primitive_type_test_file, {}) |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.not_null("new_col"), True) |
| |
| schema = pa.schema([pa.field("int_col", pa.int32(), nullable=False), |
| pa.field("bigint_col", pa.int64(), nullable=True), |
| pa.field("other_new_col", pa.int64(), nullable=True), |
| pa.field("float_col", pa.float32(), nullable=True), |
| pa.field("dbl_col", pa.float64(), nullable=True), |
| pa.field("string_col", pa.string(), nullable=True), |
| pa.field("new_col", pa.string(), nullable=True)]) |
| |
| pyarrow_not_null_array = [pa.array([], type=pa.int32()), |
| pa.array([], type=pa.int64()), |
| pa.array([], type=pa.int32()), |
| pa.array([], type=pa.float32()), |
| pa.array([], type=pa.float64()), |
| pa.array([], type=pa.string()), |
| pa.array([], type=pa.string())] |
| |
| not_null_table = pa.table(pyarrow_not_null_array, schema=schema) |
| pyarrow_null_array = [pa.array([1, 2, 3, 4, 5], type=pa.int32()), |
| pa.array([1, 2, 3, None, 5], type=pa.int64()), |
| pa.array([None, None, None, None, None], type=pa.int64()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float32()), |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float64()), |
| pa.array(['us', 'can', 'us', 'us', 'can'], type=pa.string()), |
| pa.array([None, None, None, None, None], type=pa.string())] |
| null_table = pa.table(pyarrow_null_array, schema=schema) |
| |
| target_table = reader.read() |
| assert not_null_table == target_table |
| |
| reader = ParquetReader(input_file, expected_schema, {}, Expressions.is_null("new_col"), True) |
| target_table = reader.read() |
| assert null_table == target_table |
