From c26409c314b08e411170ee68404843da53251788 Mon Sep 17 00:00:00 2001 From: Janardhan Reddy Kasireddy Date: Wed, 19 Aug 2026 00:52:04 -0400 Subject: [PATCH] fix: enable nested field statistics with inferred metrics limit --- pyiceberg/io/pyarrow.py | 123 +++++++++++++++-- pyiceberg/table/__init__.py | 3 + tests/io/test_pyarrow_stats.py | 233 +++++++++++++++++++++++++++++++-- 3 files changed, 340 insertions(+), 19 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..2ecdaa6230 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -2336,6 +2336,80 @@ def match_metrics_mode(mode: str) -> MetricsMode: raise ValueError(f"Unsupported metrics mode: {mode}") +class _LimitFieldIds(PreOrderSchemaVisitor[set[int]]): + """Select fields that should receive inferred default metrics.""" + + def __init__(self, limit: int) -> None: + self._limit = limit + self._field_ids: set[int] = set() + + def _should_continue(self) -> bool: + return len(self._field_ids) < self._limit + + @staticmethod + def _metrics_eligible(field_type: IcebergType) -> bool: + return isinstance(field_type, PrimitiveType) + + def schema(self, schema: Schema, struct_result: Callable[[], set[int]]) -> set[int]: + struct_result() + return self._field_ids + + def struct( + self, + struct: StructType, + field_results: builtins.list[Callable[[], set[int]]], + ) -> set[int]: + # Prefer primitive fields at the current struct level before + # descending into nested structures. + for field in struct.fields: + if not self._should_continue(): + break + + if self._metrics_eligible(field.field_type): + self._field_ids.add(field.field_id) + + for result in field_results: + if not self._should_continue(): + break + result() + + return self._field_ids + + def field(self, field: NestedField, field_result: Callable[[], set[int]]) -> set[int]: + field_result() + return self._field_ids + + def list(self, list_type: ListType, element_result: Callable[[], set[int]]) -> set[int]: + if self._should_continue() and self._metrics_eligible(list_type.element_type): + self._field_ids.add(list_type.element_id) + + element_result() + return self._field_ids + + def map( + self, + map_type: MapType, + key_result: Callable[[], set[int]], + value_result: Callable[[], set[int]], + ) -> set[int]: + if self._should_continue() and self._metrics_eligible(map_type.key_type): + self._field_ids.add(map_type.key_id) + + if self._should_continue() and self._metrics_eligible(map_type.value_type): + self._field_ids.add(map_type.value_id) + + key_result() + value_result() + return self._field_ids + + def primitive(self, primitive: PrimitiveType) -> set[int]: + return self._field_ids + + +def _limit_field_ids(schema: Schema, limit: int) -> set[int]: + return pre_order_visit(schema, _LimitFieldIds(limit)) + + @dataclass(frozen=True) class StatisticsCollector: field_id: int @@ -2349,15 +2423,44 @@ class PyArrowStatisticsCollector(PreOrderSchemaVisitor[list[StatisticsCollector] _schema: Schema _properties: dict[str, str] _default_mode: str + _inferred_field_ids: set[int] | None def __init__(self, schema: Schema, properties: dict[str, str]): from pyiceberg.table import TableProperties self._schema = schema self._properties = properties - self._default_mode = self._properties.get( - TableProperties.DEFAULT_WRITE_METRICS_MODE, TableProperties.DEFAULT_WRITE_METRICS_MODE_DEFAULT - ) + + configured_default_mode = self._properties.get(TableProperties.DEFAULT_WRITE_METRICS_MODE) + + if configured_default_mode is not None: + # An explicitly configured default applies to all columns. + self._default_mode = configured_default_mode + self._inferred_field_ids = None + else: + self._default_mode = TableProperties.DEFAULT_WRITE_METRICS_MODE_DEFAULT + + max_inferred_columns = property_as_int( + self._properties, + TableProperties.METRICS_MAX_INFERRED_COLUMN_DEFAULTS, + TableProperties.METRICS_MAX_INFERRED_COLUMN_DEFAULTS_DEFAULT, + ) + + if max_inferred_columns is None: + max_inferred_columns = TableProperties.METRICS_MAX_INFERRED_COLUMN_DEFAULTS_DEFAULT + + if max_inferred_columns < 0: + logger.warning( + "Invalid value for %s (negative): %s, falling back to %s", + TableProperties.METRICS_MAX_INFERRED_COLUMN_DEFAULTS, + max_inferred_columns, + TableProperties.METRICS_MAX_INFERRED_COLUMN_DEFAULTS_DEFAULT, + ) + max_inferred_columns = TableProperties.METRICS_MAX_INFERRED_COLUMN_DEFAULTS_DEFAULT + + self._inferred_field_ids = ( + _limit_field_ids(schema, max_inferred_columns) if len(schema.field_ids) > max_inferred_columns else None + ) def schema( self, schema: Schema, struct_result: Callable[[], builtins.list[StatisticsCollector]] @@ -2402,6 +2505,9 @@ def primitive(self, primitive: PrimitiveType) -> builtins.list[StatisticsCollect metrics_mode = match_metrics_mode(self._default_mode) + if self._inferred_field_ids is not None and self._field_id not in self._inferred_field_ids: + metrics_mode = MetricsMode(MetricModeTypes.NONE) + col_mode = self._properties.get(f"{TableProperties.METRICS_MODE_COLUMN_CONF_PREFIX}.{column_name}") if col_mode: metrics_mode = match_metrics_mode(col_mode) @@ -2412,11 +2518,6 @@ def primitive(self, primitive: PrimitiveType) -> builtins.list[StatisticsCollect ): metrics_mode = MetricsMode(MetricModeTypes.FULL) - is_nested = column_name.find(".") >= 0 - - if is_nested and metrics_mode.type in [MetricModeTypes.TRUNCATE, MetricModeTypes.FULL]: - metrics_mode = MetricsMode(MetricModeTypes.COUNTS) - return [StatisticsCollector(field_id=self._field_id, iceberg_type=primitive, mode=metrics_mode, column_name=column_name)] @@ -2429,9 +2530,9 @@ def compute_statistics_plan( The resulting list is assumed to have the same length and same order as the columns in the pyarrow table. This allows the list to map from the column index to the Iceberg column ID. - For each element, the desired metrics collection that was provided by the user in the configuration - is computed and then adjusted according to the data type of the column. For nested columns the minimum - and maximum values are not computed. And truncation is only applied to text of binary strings. + For each element, the desired metrics collection configured by the user is + computed and adjusted according to the data type of the column. Truncation + is only applied to string and binary types. Args: table_properties (from pyiceberg.table.metadata.TableMetadata): The Iceberg table metadata properties. diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 3dffc2270c..1307950282 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -170,6 +170,9 @@ class TableProperties: DEFAULT_WRITE_METRICS_MODE = "write.metadata.metrics.default" DEFAULT_WRITE_METRICS_MODE_DEFAULT = "truncate(16)" + METRICS_MAX_INFERRED_COLUMN_DEFAULTS = "write.metadata.metrics.max-inferred-column-defaults" + METRICS_MAX_INFERRED_COLUMN_DEFAULTS_DEFAULT = 100 + METRICS_MODE_COLUMN_CONF_PREFIX = "write.metadata.metrics.column" WRITE_PARTITION_SUMMARY_LIMIT = "write.summary.partition-limit" diff --git a/tests/io/test_pyarrow_stats.py b/tests/io/test_pyarrow_stats.py index 0e628829eb..e31a38e36c 100644 --- a/tests/io/test_pyarrow_stats.py +++ b/tests/io/test_pyarrow_stats.py @@ -65,7 +65,9 @@ BooleanType, FloatType, IntegerType, + NestedField, StringType, + StructType, ) from pyiceberg.utils.datetime import date_to_days, datetime_to_micros, time_to_micros @@ -273,14 +275,29 @@ def test_bounds() -> None: ) datafile = DataFile.from_args(**statistics.to_serialized_dict()) - assert len(datafile.lower_bounds) == 2 + assert set(datafile.lower_bounds) == {1, 2, 6, 7, 8, 9, 10} + assert datafile.lower_bounds[1].decode() == "aaaaaaaaaaaaaaaa" assert datafile.lower_bounds[2] == STRUCT_FLOAT.pack(1.69) - assert len(datafile.upper_bounds) == 2 + # Nested list/map/struct primitive fields + assert datafile.lower_bounds[6] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[7] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[8] == STRUCT_INT64.pack(2) + assert datafile.lower_bounds[9] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[10] == STRUCT_FLOAT.pack(-1.34) + + assert set(datafile.upper_bounds) == {1, 2, 6, 7, 8, 9, 10} + assert datafile.upper_bounds[1].decode() == "zzzzzzzzzzzzzzz{" assert datafile.upper_bounds[2] == STRUCT_FLOAT.pack(100) + assert datafile.upper_bounds[6] == STRUCT_INT64.pack(9) + assert datafile.upper_bounds[7] == STRUCT_INT64.pack(5) + assert datafile.upper_bounds[8] == STRUCT_INT64.pack(6) + assert datafile.upper_bounds[9] == STRUCT_INT64.pack(54) + assert datafile.upper_bounds[10] == STRUCT_FLOAT.pack(0.2) + def test_metrics_mode_parsing() -> None: assert match_metrics_mode("none") == MetricsMode(MetricModeTypes.NONE) @@ -324,6 +341,94 @@ def test_metrics_mode_none() -> None: assert len(datafile.upper_bounds) == 0 +def test_metrics_mode_nested_primitive_fields() -> None: + schema = Schema( + NestedField( + field_id=1, + name="location", + field_type=StructType( + NestedField( + field_id=2, + name="latitude", + field_type=FloatType(), + ), + NestedField( + field_id=3, + name="longitude", + field_type=FloatType(), + ), + ), + ) + ) + + statistics_plan = compute_statistics_plan(schema, {}) + + assert statistics_plan[2].column_name == "location.latitude" + assert statistics_plan[3].column_name == "location.longitude" + + assert statistics_plan[2].mode == MetricsMode(MetricModeTypes.FULL) + assert statistics_plan[3].mode == MetricsMode(MetricModeTypes.FULL) + + +def test_metrics_mode_max_inferred_column_defaults() -> None: + schema = Schema( + NestedField(field_id=1, name="a", field_type=IntegerType()), + NestedField(field_id=2, name="b", field_type=IntegerType()), + NestedField(field_id=3, name="c", field_type=IntegerType()), + NestedField(field_id=4, name="d", field_type=IntegerType()), + ) + + statistics_plan = compute_statistics_plan( + schema, + { + "write.metadata.metrics.max-inferred-column-defaults": "2", + }, + ) + + assert statistics_plan[1].mode == MetricsMode(MetricModeTypes.FULL) + assert statistics_plan[2].mode == MetricsMode(MetricModeTypes.FULL) + + assert statistics_plan[3].mode == MetricsMode(MetricModeTypes.NONE) + assert statistics_plan[4].mode == MetricsMode(MetricModeTypes.NONE) + + +def test_metrics_mode_inferred_limit_prioritizes_top_level_fields() -> None: + schema = Schema( + NestedField( + field_id=1, + name="location", + field_type=StructType( + NestedField( + field_id=2, + name="latitude", + field_type=FloatType(), + ), + NestedField( + field_id=3, + name="longitude", + field_type=FloatType(), + ), + ), + ), + NestedField( + field_id=4, + name="id", + field_type=IntegerType(), + ), + ) + + statistics_plan = compute_statistics_plan( + schema, + { + "write.metadata.metrics.max-inferred-column-defaults": "1", + }, + ) + + assert statistics_plan[4].mode == MetricsMode(MetricModeTypes.FULL) + assert statistics_plan[2].mode == MetricsMode(MetricModeTypes.NONE) + assert statistics_plan[3].mode == MetricsMode(MetricModeTypes.NONE) + + def test_metrics_mode_counts() -> None: metadata, table_metadata = construct_test_table() @@ -359,14 +464,26 @@ def test_metrics_mode_full() -> None: assert len(datafile.null_value_counts) == 7 assert len(datafile.nan_value_counts) == 0 - assert len(datafile.lower_bounds) == 2 + assert set(datafile.lower_bounds) == {1, 2, 6, 7, 8, 9, 10} assert datafile.lower_bounds[1].decode() == "aaaaaaaaaaaaaaaaaaaa" assert datafile.lower_bounds[2] == STRUCT_FLOAT.pack(1.69) - assert len(datafile.upper_bounds) == 2 + assert datafile.lower_bounds[6] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[7] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[8] == STRUCT_INT64.pack(2) + assert datafile.lower_bounds[9] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[10] == STRUCT_FLOAT.pack(-1.34) + + assert set(datafile.upper_bounds) == {1, 2, 6, 7, 8, 9, 10} assert datafile.upper_bounds[1].decode() == "zzzzzzzzzzzzzzzzzzzz" assert datafile.upper_bounds[2] == STRUCT_FLOAT.pack(100) + assert datafile.upper_bounds[6] == STRUCT_INT64.pack(9) + assert datafile.upper_bounds[7] == STRUCT_INT64.pack(5) + assert datafile.upper_bounds[8] == STRUCT_INT64.pack(6) + assert datafile.upper_bounds[9] == STRUCT_INT64.pack(54) + assert datafile.upper_bounds[10] == STRUCT_FLOAT.pack(0.2) + def test_metrics_mode_non_default_trunc() -> None: metadata, table_metadata = construct_test_table() @@ -384,14 +501,28 @@ def test_metrics_mode_non_default_trunc() -> None: assert len(datafile.null_value_counts) == 7 assert len(datafile.nan_value_counts) == 0 - assert len(datafile.lower_bounds) == 2 + assert set(datafile.lower_bounds) == {1, 2, 6, 7, 8, 9, 10} + assert datafile.lower_bounds[1].decode() == "aa" assert datafile.lower_bounds[2] == STRUCT_FLOAT.pack(1.69) - assert len(datafile.upper_bounds) == 2 + assert datafile.lower_bounds[6] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[7] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[8] == STRUCT_INT64.pack(2) + assert datafile.lower_bounds[9] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[10] == STRUCT_FLOAT.pack(-1.34) + + assert set(datafile.upper_bounds) == {1, 2, 6, 7, 8, 9, 10} + assert datafile.upper_bounds[1].decode() == "z{" assert datafile.upper_bounds[2] == STRUCT_FLOAT.pack(100) + assert datafile.upper_bounds[6] == STRUCT_INT64.pack(9) + assert datafile.upper_bounds[7] == STRUCT_INT64.pack(5) + assert datafile.upper_bounds[8] == STRUCT_INT64.pack(6) + assert datafile.upper_bounds[9] == STRUCT_INT64.pack(54) + assert datafile.upper_bounds[10] == STRUCT_FLOAT.pack(0.2) + def test_column_metrics_mode() -> None: metadata, table_metadata = construct_test_table() @@ -410,15 +541,101 @@ def test_column_metrics_mode() -> None: assert len(datafile.null_value_counts) == 6 assert len(datafile.nan_value_counts) == 0 - assert len(datafile.lower_bounds) == 1 + assert set(datafile.lower_bounds) == {2, 6, 7, 8, 9, 10} + assert datafile.lower_bounds[2] == STRUCT_FLOAT.pack(1.69) + assert datafile.lower_bounds[6] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[7] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[8] == STRUCT_INT64.pack(2) + assert datafile.lower_bounds[9] == STRUCT_INT64.pack(1) + assert datafile.lower_bounds[10] == STRUCT_FLOAT.pack(-1.34) assert 1 not in datafile.lower_bounds - assert len(datafile.upper_bounds) == 1 + assert set(datafile.upper_bounds) == {2, 6, 7, 8, 9, 10} + assert datafile.upper_bounds[2] == STRUCT_FLOAT.pack(100) + assert datafile.upper_bounds[6] == STRUCT_INT64.pack(9) + assert datafile.upper_bounds[7] == STRUCT_INT64.pack(5) + assert datafile.upper_bounds[8] == STRUCT_INT64.pack(6) + assert datafile.upper_bounds[9] == STRUCT_INT64.pack(54) + assert datafile.upper_bounds[10] == STRUCT_FLOAT.pack(0.2) assert 1 not in datafile.upper_bounds +def test_metrics_mode_column_override_bypasses_inferred_limit() -> None: + schema = Schema( + NestedField(field_id=1, name="a", field_type=IntegerType()), + NestedField(field_id=2, name="b", field_type=IntegerType()), + NestedField(field_id=3, name="c", field_type=IntegerType()), + ) + + statistics_plan = compute_statistics_plan( + schema, + { + "write.metadata.metrics.max-inferred-column-defaults": "1", + "write.metadata.metrics.column.c": "full", + }, + ) + + assert statistics_plan[1].mode == MetricsMode(MetricModeTypes.FULL) + assert statistics_plan[2].mode == MetricsMode(MetricModeTypes.NONE) + assert statistics_plan[3].mode == MetricsMode(MetricModeTypes.FULL) + + +def test_explicit_default_metrics_mode_ignores_inferred_limit() -> None: + schema = Schema( + NestedField(field_id=1, name="a", field_type=IntegerType()), + NestedField(field_id=2, name="b", field_type=IntegerType()), + NestedField(field_id=3, name="c", field_type=IntegerType()), + ) + + statistics_plan = compute_statistics_plan( + schema, + { + "write.metadata.metrics.default": "full", + "write.metadata.metrics.max-inferred-column-defaults": "1", + }, + ) + + assert statistics_plan[1].mode == MetricsMode(MetricModeTypes.FULL) + assert statistics_plan[2].mode == MetricsMode(MetricModeTypes.FULL) + assert statistics_plan[3].mode == MetricsMode(MetricModeTypes.FULL) + + +def test_metrics_mode_zero_max_inferred_column_defaults() -> None: + schema = Schema( + NestedField(field_id=1, name="a", field_type=IntegerType()), + NestedField(field_id=2, name="b", field_type=IntegerType()), + ) + + statistics_plan = compute_statistics_plan( + schema, + { + "write.metadata.metrics.max-inferred-column-defaults": "0", + }, + ) + + assert statistics_plan[1].mode == MetricsMode(MetricModeTypes.NONE) + assert statistics_plan[2].mode == MetricsMode(MetricModeTypes.NONE) + + +def test_metrics_mode_negative_max_inferred_column_defaults_uses_default() -> None: + schema = Schema( + NestedField(field_id=1, name="a", field_type=IntegerType()), + NestedField(field_id=2, name="b", field_type=IntegerType()), + ) + + statistics_plan = compute_statistics_plan( + schema, + { + "write.metadata.metrics.max-inferred-column-defaults": "-1", + }, + ) + + assert statistics_plan[1].mode == MetricsMode(MetricModeTypes.FULL) + assert statistics_plan[2].mode == MetricsMode(MetricModeTypes.FULL) + + def construct_test_table_primitive_types() -> tuple[pq.FileMetaData, TableMetadataV1 | TableMetadataV2]: table_metadata = { "format-version": 2,