diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 7fd0c85b77..74a93b91b3 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -2574,6 +2574,7 @@ def data_file_statistics_from_parquet_metadata( parquet_metadata: pq.FileMetaData, stats_columns: dict[int, StatisticsCollector], parquet_column_mapping: dict[str, int], + nan_value_counts: dict[int, int] | None = None, ) -> DataFileStatistics: """ Compute and return DataFileStatistics that includes the following. @@ -2591,13 +2592,15 @@ def data_file_statistics_from_parquet_metadata( stats_columns (Dict[int, StatisticsCollector]): The statistics gathering plan. It is required to set the mode for column metrics collection parquet_column_mapping (Dict[str, int]): The mapping of the parquet file name to the field ID + nan_value_counts (Optional[Dict[int, int]]): The NaN counts collected from the written data, + as Parquet statistics do not track NaN values """ column_sizes: dict[int, int] = {} value_counts: dict[int, int] = {} split_offsets: list[int] = [] null_value_counts: dict[int, int] = {} - nan_value_counts: dict[int, int] = {} + nan_value_counts = dict(nan_value_counts) if nan_value_counts else {} col_aggs = {} @@ -2689,6 +2692,40 @@ def data_file_statistics_from_parquet_metadata( ) +def _nan_value_counts(table: pa.Table, schema: Schema, stats_columns: dict[int, StatisticsCollector]) -> dict[int, int]: + """Count the NaN values of the float and double columns that collect metrics.""" + nan_field_ids = { + field_id + for field_id, stats_col in stats_columns.items() + if isinstance(stats_col.iceberg_type, (FloatType, DoubleType)) and stats_col.mode != MetricsMode(MetricModeTypes.NONE) + } + nan_value_counts: dict[int, int] = {} + if not nan_field_ids: + return nan_value_counts + + def count(field_id: int, field_type: IcebergType, array: pa.Array) -> None: + if isinstance(field_type, StructType): + for nested_field in field_type.fields: + if array.type.get_field_index(nested_field.name) >= 0: + # struct_field also applies the validity of the parent struct + count(nested_field.field_id, nested_field.field_type, pc.struct_field(array, nested_field.name)) + elif isinstance(field_type, ListType): + count(field_type.element_id, field_type.element_type, array.flatten()) + elif isinstance(field_type, MapType): + entries = array.cast(pa.list_(pa.struct([array.type.key_field, array.type.item_field]))).flatten() + count(field_type.key_id, field_type.key_type, entries.field(0)) + count(field_type.value_id, field_type.value_type, entries.field(1)) + elif field_id in nan_field_ids and pa.types.is_floating(array.type): + nan_value_counts[field_id] = nan_value_counts.get(field_id, 0) + (pc.sum(pc.is_nan(array)).as_py() or 0) + + for field in schema.fields: + if table.schema.get_field_index(field.name) >= 0: + for chunk in table.column(field.name).chunks: + count(field.field_id, field.field_type, chunk) + + return nan_value_counts + + class ParquetFormatWriter(FileFormatWriter): """Writes Arrow tables to a Parquet file.""" @@ -2698,6 +2735,8 @@ def __init__(self, output_file: OutputFile, file_schema: Schema, properties: Pro self._properties = properties self._writer: pq.ParquetWriter | None = None self._fos: OutputStream | None = None + self._stats_columns = compute_statistics_plan(file_schema, properties) + self._nan_value_counts: dict[int, int] = {} self._parquet_writer_kwargs = _get_parquet_writer_kwargs(properties) self._row_group_size = property_as_int( properties=properties, @@ -2720,6 +2759,9 @@ def write(self, table: pa.Table) -> None: raise self._fos = fos self._writer.write(table, row_group_size=self._row_group_size) + # Parquet statistics do not track NaN values, so count them while writing + for field_id, nan_count in _nan_value_counts(table, self._file_schema, self._stats_columns).items(): + self._nan_value_counts[field_id] = self._nan_value_counts.get(field_id, 0) + nan_count def close(self) -> DataFileStatistics: if self._result is not None: @@ -2730,8 +2772,9 @@ def close(self) -> DataFileStatistics: self._writer.close() self._result = data_file_statistics_from_parquet_metadata( parquet_metadata=self._writer.writer.metadata, - stats_columns=compute_statistics_plan(self._file_schema, self._properties), + stats_columns=self._stats_columns, parquet_column_mapping=parquet_path_to_id_mapping(self._file_schema), + nan_value_counts=self._nan_value_counts, ) return self._result diff --git a/tests/catalog/test_catalog_behaviors.py b/tests/catalog/test_catalog_behaviors.py index 4310ed1eb6..b10598c32d 100644 --- a/tests/catalog/test_catalog_behaviors.py +++ b/tests/catalog/test_catalog_behaviors.py @@ -1345,6 +1345,27 @@ def test_append_nan_to_identity_partitioned_table(catalog: Catalog) -> None: assert sorted(tbl.scan(row_filter="value is nan").to_arrow()["id"].to_pylist()) == [2, 4] +def test_delete_keeps_nan_rows_that_do_not_match_the_filter(catalog: Catalog) -> None: + catalog.create_namespace("default") + identifier = f"default.delete_keeps_nan_rows_{catalog.name}" + iceberg_schema = Schema( + NestedField(1, "id", IntegerType(), required=False), + NestedField(2, "value", DoubleType(), required=False), + ) + tbl = catalog.create_table(identifier=identifier, schema=iceberg_schema) + tbl.append( + pa.Table.from_pydict( + {"id": [1, 2], "value": [1.0, float("nan")]}, + schema=schema_to_pyarrow(iceberg_schema), + ) + ) + + # NaN < 5.0 is false, so the NaN row must survive and the file cannot be dropped as a whole + tbl.delete("value < 5.0") + + assert tbl.scan().to_arrow()["id"].to_pylist() == [2] + + def test_record_batch_reader_consumed_exactly_once(catalog: Catalog) -> None: """The streaming path must consume the underlying generator exactly once. A regression that drained the reader twice (e.g. an extra .schema access diff --git a/tests/integration/test_writes/test_writes.py b/tests/integration/test_writes/test_writes.py index 30fdd76ab7..9e73f7a5f6 100644 --- a/tests/integration/test_writes/test_writes.py +++ b/tests/integration/test_writes/test_writes.py @@ -1687,7 +1687,7 @@ def test_merge_manifests_file_content(session_catalog: Catalog, arrow_table_with (11, b"\x01"), (12, b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"), ] - assert tbl_a_data_file["nan_value_counts"] == [] + assert tbl_a_data_file["nan_value_counts"] == [(6, 0), (7, 0)] assert tbl_a_data_file["null_value_counts"] == [ (1, 1), (2, 1), diff --git a/tests/io/test_format_writers.py b/tests/io/test_format_writers.py index 86ae65156b..db3ce36564 100644 --- a/tests/io/test_format_writers.py +++ b/tests/io/test_format_writers.py @@ -24,10 +24,10 @@ import pytest from pyiceberg.io.fileformat import FileFormatFactory, FileFormatModel -from pyiceberg.io.pyarrow import PyArrowFileIO +from pyiceberg.io.pyarrow import PyArrowFileIO, schema_to_pyarrow from pyiceberg.manifest import FileFormat from pyiceberg.schema import Schema -from pyiceberg.types import LongType, NestedField +from pyiceberg.types import DoubleType, FloatType, ListType, LongType, MapType, NestedField, StringType, StructType @pytest.fixture(params=FileFormatFactory.available_formats(), ids=lambda f: f.name.lower()) @@ -73,6 +73,57 @@ def test_null_handling(format_model: FileFormatModel, table_schema_simple: Schem assert stats.null_value_counts.get(1) == 1 +def test_nan_value_counts(format_model: FileFormatModel, tmp_path: Path) -> None: + """Float and double columns produce nan_value_counts, including nested and empty counts.""" + nan = float("nan") + schema = Schema( + NestedField(1, "id", LongType(), required=False), + NestedField(2, "f", FloatType(), required=False), + NestedField(3, "d", DoubleType(), required=False), + NestedField(4, "no_nans", DoubleType(), required=False), + NestedField(5, "s", StructType(NestedField(6, "x", DoubleType(), required=False)), required=False), + NestedField(7, "l", ListType(element_id=8, element_type=DoubleType(), element_required=False), required=False), + NestedField( + 9, + "m", + MapType(key_id=10, key_type=StringType(), value_id=11, value_type=DoubleType(), value_required=False), + required=False, + ), + ) + table = pa.Table.from_pylist( + [ + {"id": 1, "f": nan, "d": nan, "no_nans": 1.0, "s": {"x": nan}, "l": [nan, 1.0], "m": {"a": nan}}, + {"id": 2, "f": 1.0, "d": nan, "no_nans": None, "s": None, "l": None, "m": None}, + {"id": 3, "f": None, "d": 2.0, "no_nans": 2.0, "s": {"x": 1.0}, "l": [nan, nan], "m": {"b": 1.0, "c": nan}}, + ], + schema=schema_to_pyarrow(schema), + ) + file_path = str(tmp_path / f"test.{format_model.file_extension()}") + writer = format_model.create_writer(PyArrowFileIO().new_output(file_path), schema, {}) + writer.write(table) + writer.write(table.slice(0, 1)) + stats = writer.close() + + assert stats.nan_value_counts == {2: 2, 3: 3, 4: 0, 6: 2, 8: 4, 11: 3} + + +def test_nan_value_counts_respect_metrics_mode(format_model: FileFormatModel, tmp_path: Path) -> None: + """Columns with metrics mode none do not get a NaN count.""" + schema = Schema( + NestedField(1, "a", DoubleType(), required=False), + NestedField(2, "b", DoubleType(), required=False), + ) + table = pa.Table.from_pydict({"a": [float("nan")], "b": [float("nan")]}, schema=schema_to_pyarrow(schema)) + file_path = str(tmp_path / f"test.{format_model.file_extension()}") + writer = format_model.create_writer( + PyArrowFileIO().new_output(file_path), schema, {"write.metadata.metrics.column.b": "none"} + ) + writer.write(table) + stats = writer.close() + + assert stats.nan_value_counts == {1: 1} + + def test_context_manager_caches_result( format_model: FileFormatModel, table_schema_simple: Schema, arrow_table_simple: pa.Table, tmp_path: Path ) -> None: