feat: Add support for writing bloom filters

pyiceberg/io/pyarrow.py

@@ -181,7 +181,12 @@
 from pyiceberg.utils.config import Config
 from pyiceberg.utils.datetime import millis_to_datetime
 from pyiceberg.utils.decimal import unscaled_to_decimal
-from pyiceberg.utils.properties import get_first_property_value, property_as_bool, property_as_int
+from pyiceberg.utils.properties import (
+    get_first_property_value,
+    property_as_bool,
+    property_as_float,
+    property_as_int,
+)
 from pyiceberg.utils.singleton import Singleton
 from pyiceberg.utils.truncate import truncate_upper_bound_binary_string, truncate_upper_bound_text_string
 
@@ -2474,6 +2479,107 @@ def parquet_path_to_id_mapping(
     return result
 
 
+@dataclass(frozen=True)
+class BloomFilterOptions:
+    parquet_path: str
+    ndv: int | None
+    fpp: float | None
+
+
+class BloomFilterOptionsCollector(PreOrderSchemaVisitor[list[BloomFilterOptions]]):
+    _field_id: int = 0
+    _schema: Schema
+    _properties: dict[str, str]
+
+    def __init__(self, schema: Schema, properties: dict[str, str], id_to_parquet_path_mapping: dict[int, str]):
+        self._schema = schema
+        self._properties = properties
+        self._id_to_parquet_path_mapping = id_to_parquet_path_mapping
+
+    def schema(
+        self, schema: Schema, struct_result: Callable[[], builtins.list[BloomFilterOptions]]
+    ) -> builtins.list[BloomFilterOptions]:
+        return struct_result()
+
+    def struct(
+        self, struct: StructType, field_results: builtins.list[Callable[[], builtins.list[BloomFilterOptions]]]
+    ) -> builtins.list[BloomFilterOptions]:
+        return list(itertools.chain(*[result() for result in field_results]))
+
+    def field(
+        self, field: NestedField, field_result: Callable[[], builtins.list[BloomFilterOptions]]
+    ) -> builtins.list[BloomFilterOptions]:
+        self._field_id = field.field_id
+        return field_result()
+
+    def list(
+        self, list_type: ListType, element_result: Callable[[], builtins.list[BloomFilterOptions]]
+    ) -> builtins.list[BloomFilterOptions]:
+        self._field_id = list_type.element_id
+        return element_result()
+
+    def map(
+        self,
+        map_type: MapType,
+        key_result: Callable[[], builtins.list[BloomFilterOptions]],
+        value_result: Callable[[], builtins.list[BloomFilterOptions]],
+    ) -> builtins.list[BloomFilterOptions]:
+        self._field_id = map_type.key_id
+        k = key_result()
+        self._field_id = map_type.value_id
+        v = value_result()
+        return k + v
+
+    def primitive(self, primitive: PrimitiveType) -> builtins.list[BloomFilterOptions]:
+        from pyiceberg.table import TableProperties
+
+        column_name = self._schema.find_column_name(self._field_id)
+        if column_name is None:
+            return []
+
+        parquet_path = self._id_to_parquet_path_mapping.get(self._field_id)
+        if parquet_path is None:
+            return []
+
+        bloom_filter_enabled = property_as_bool(
+            self._properties, f"{TableProperties.PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX}.{column_name}", False
+        )
+        if not bloom_filter_enabled:
+            return []
+
+        bloom_filter_fpp = property_as_float(
+            self._properties, f"{TableProperties.PARQUET_BLOOM_FILTER_COLUMN_FPP_PREFIX}.{column_name}", None
+        )
+        bloom_filter_ndv = property_as_int(
+            self._properties, f"{TableProperties.PARQUET_BLOOM_FILTER_COLUMN_NDV_PREFIX}.{column_name}", None
+        )
+
+        return [BloomFilterOptions(parquet_path=parquet_path, ndv=bloom_filter_ndv, fpp=bloom_filter_fpp)]
+
+
+def get_bloom_filter_options(
+    schema: Schema,
+    table_properties: dict[str, str],
+    id_to_parquet_path: dict[int, str],
+) -> dict[str, dict[str, Any]]:
+    """
+    Get the bloom filter options from the table properties.
+
+    Args:
+        schema (pyiceberg.schema.Schema): The current table schema.
+        table_properties (dict[str, str]): The table properties.
+        id_to_parquet_path (Dict[int, str]): The mapping of the field ID to the parquet file name.
+    """
+    bloom_filter_options = pre_order_visit(schema, BloomFilterOptionsCollector(schema, table_properties, id_to_parquet_path))
+    result: dict[str, dict[str, Any]] = {}
+    for bf_opts in bloom_filter_options:
+        result[bf_opts.parquet_path] = {
+            **({"ndv": bf_opts.ndv} if bf_opts.ndv is not None else {}),
+            **({"fpp": bf_opts.fpp} if bf_opts.fpp is not None else {}),
+        }
+    return result
+
+
 def data_file_statistics_from_parquet_metadata(
     parquet_metadata: pq.FileMetaData,
     stats_columns: dict[int, StatisticsCollector],
@@ -2596,7 +2702,6 @@ def data_file_statistics_from_parquet_metadata(
 def write_file(io: FileIO, table_metadata: TableMetadata, tasks: Iterator[WriteTask]) -> Iterator[DataFile]:
     from pyiceberg.table import DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE, TableProperties
 
-    parquet_writer_kwargs = _get_parquet_writer_kwargs(table_metadata.properties)
     row_group_size = property_as_int(
         properties=table_metadata.properties,
         property_name=TableProperties.PARQUET_ROW_GROUP_LIMIT,
@@ -2613,6 +2718,9 @@ def write_parquet(task: WriteTask) -> DataFile:
         else:
             file_schema = table_schema
 
+        parquet_column_mapping = parquet_path_to_id_mapping(file_schema)
+        parquet_writer_kwargs = _get_parquet_writer_kwargs(table_metadata.properties, file_schema, parquet_column_mapping)
+
         downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False
         batches = [
             _to_requested_schema(
@@ -2638,7 +2746,7 @@ def write_parquet(task: WriteTask) -> DataFile:
         statistics = data_file_statistics_from_parquet_metadata(
             parquet_metadata=writer.writer.metadata,
             stats_columns=compute_statistics_plan(file_schema, table_metadata.properties),
-            parquet_column_mapping=parquet_path_to_id_mapping(file_schema),
+            parquet_column_mapping=parquet_column_mapping,
         )
         data_file = DataFile.from_args(
             content=DataFileContent.DATA,
@@ -2757,14 +2865,37 @@ def parquet_file_to_data_file(io: FileIO, table_metadata: TableMetadata, file_pa
 PYARROW_UNCOMPRESSED_CODEC = "none"
 
 
-def _get_parquet_writer_kwargs(table_properties: Properties) -> dict[str, Any]:
+def _get_parquet_writer_kwargs(
+    table_properties: Properties, file_schema: Schema, parquet_column_mapping: dict[str, int]
+) -> dict[str, Any]:
     from pyiceberg.table import TableProperties
 
-    for key_pattern in [
+    unsupported_key_patterns = [
         TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES,
         TableProperties.PARQUET_BLOOM_FILTER_MAX_BYTES,
-        f"{TableProperties.PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX}.*",
-    ]:
+    ]
+
+    from packaging import version
+
+    MIN_PYARROW_VERSION_SUPPORTING_BLOOM_FILTER_WRITES = "24.0.0"
+    pyarrow_supports_bloom_filter_writes = version.parse(pyarrow.__version__) >= version.parse(
+        MIN_PYARROW_VERSION_SUPPORTING_BLOOM_FILTER_WRITES
+    )
+
+    if not pyarrow_supports_bloom_filter_writes:
+        bloom_filter_key_patterns = [
+            f"{TableProperties.PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX}.*",
+            f"{TableProperties.PARQUET_BLOOM_FILTER_COLUMN_FPP_PREFIX}.*",
+            f"{TableProperties.PARQUET_BLOOM_FILTER_COLUMN_NDV_PREFIX}.*",
+        ]
+        if any(fnmatch.filter(table_properties, key_pattern) for key_pattern in bloom_filter_key_patterns):
+            warnings.warn(
+                f"Parquet writer option(s) for bloom filters require pyarrow {MIN_PYARROW_VERSION_SUPPORTING_BLOOM_FILTER_WRITES}"
+                " or higher",
+                stacklevel=2,
+            )
+
+    for key_pattern in unsupported_key_patterns:
         if unsupported_keys := fnmatch.filter(table_properties, key_pattern):
             warnings.warn(f"Parquet writer option(s) {unsupported_keys} not implemented", stacklevel=2)
 
@@ -2777,6 +2908,11 @@ def _get_parquet_writer_kwargs(table_properties: Properties) -> dict[str, Any]:
     if compression_codec == ICEBERG_UNCOMPRESSED_CODEC:
         compression_codec = PYARROW_UNCOMPRESSED_CODEC
 
+    bloom_filter_options = {}
+    if pyarrow_supports_bloom_filter_writes:
+        id_to_parquet_path = {field_id: parquet_path for parquet_path, field_id in parquet_column_mapping.items()}
+        bloom_filter_options = get_bloom_filter_options(file_schema, table_properties, id_to_parquet_path)
+
     return {
         "compression": compression_codec,
         "compression_level": compression_level,
@@ -2795,6 +2931,7 @@ def _get_parquet_writer_kwargs(table_properties: Properties) -> dict[str, Any]:
             property_name=TableProperties.PARQUET_PAGE_ROW_LIMIT,
             default=TableProperties.PARQUET_PAGE_ROW_LIMIT_DEFAULT,
         ),
+        **({"bloom_filter_options": bloom_filter_options} if bloom_filter_options else {}),
     }
 
 

pyiceberg/table/__init__.py

@@ -147,6 +147,10 @@ class TableProperties:
 
     PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX = "write.parquet.bloom-filter-enabled.column"
 
+    PARQUET_BLOOM_FILTER_COLUMN_FPP_PREFIX = "write.parquet.bloom-filter-fpp.column"
+
+    PARQUET_BLOOM_FILTER_COLUMN_NDV_PREFIX = "write.parquet.bloom-filter-ndv.column"
+
     WRITE_TARGET_FILE_SIZE_BYTES = "write.target-file-size-bytes"
     WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT = 512 * 1024 * 1024  # 512 MB
 

tests/integration/test_writes/test_writes.py

@@ -30,11 +30,13 @@
 import fastavro
 import pandas as pd
 import pandas.testing
+import pyarrow
 import pyarrow as pa
 import pyarrow.compute as pc
 import pyarrow.parquet as pq
 import pytest
 import pytz
+from packaging import version
 from pyarrow.fs import S3FileSystem
 from pydantic_core import ValidationError
 from pyspark.sql import SparkSession
@@ -68,6 +70,11 @@
 from pyiceberg.view.metadata import SQLViewRepresentation, ViewVersion
 from utils import TABLE_SCHEMA, _create_table
 
+skip_if_bloom_filter_not_supported = pytest.mark.skipif(
+    version.parse(pyarrow.__version__) < version.parse("24.0.0"),
+    reason="Requires pyarrow version >= 24.0.0",
+)
+
 
 @pytest.fixture(scope="session", autouse=True)
 def table_v1_with_null(session_catalog: Catalog, arrow_table_with_null: pa.Table) -> None:
@@ -695,7 +702,6 @@ def test_write_parquet_other_properties(
     "properties",
     [
         {"write.parquet.row-group-size-bytes": "42"},
-        {"write.parquet.bloom-filter-enabled.column.bool": "42"},
         {"write.parquet.bloom-filter-max-bytes": "42"},
     ],
 )
@@ -712,6 +718,49 @@ def test_write_parquet_unsupported_properties(
         tbl.append(arrow_table_with_null)
 
 
+@pytest.mark.integration
+@skip_if_bloom_filter_not_supported
+@pytest.mark.parametrize("format_version", [1, 2])
+def test_write_parquet_bloom_filter_properties(
+    spark: SparkSession, session_catalog: Catalog, arrow_table_with_null: pa.Table, format_version: int
+) -> None:
+    identifier = "default.write_parquet_bloom_filter_properties"
+
+    tbl = _create_table(
+        session_catalog,
+        identifier,
+        {
+            "format-version": format_version,
+            "write.parquet.bloom-filter-enabled.column.string": "true",
+            "write.parquet.bloom-filter-fpp.column.string": "0.1",
+            "write.parquet.bloom-filter-ndv.column.string": "100",
+        },
+        [arrow_table_with_null],
+    )
+
+    data_file_paths = [task.file.file_path for task in tbl.scan().plan_files()]
+
+    fs = S3FileSystem(
+        endpoint_override=session_catalog.properties["s3.endpoint"],
+        access_key=session_catalog.properties["s3.access-key-id"],
+        secret_key=session_catalog.properties["s3.secret-access-key"],
+    )
+    for data_file_path in data_file_paths:
+        uri = urlparse(data_file_path)
+        with fs.open_input_file(f"{uri.netloc}{uri.path}") as f:
+            parquet_metadata = pq.read_metadata(f)
+            for row_group_idx in range(parquet_metadata.num_row_groups):
+                row_group = parquet_metadata.row_group(row_group_idx)
+                for column_idx in range(parquet_metadata.num_columns):
+                    column = row_group.column(column_idx)
+                    if column.path_in_schema == "string":
+                        assert column.bloom_filter_offset is not None
+                        assert column.bloom_filter_length is not None
+                    else:
+                        assert column.bloom_filter_offset is None
+                        assert column.bloom_filter_length is None
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("format_version", [1, 2])
 def test_spark_writes_orc_pyiceberg_reads(spark: SparkSession, session_catalog: Catalog, format_version: int) -> None:

tests/table/test_metadata.py

@@ -26,6 +26,7 @@
 import pytest
 
 from pyiceberg.exceptions import ValidationError
+from pyiceberg.io.pyarrow import get_bloom_filter_options, parquet_path_to_id_mapping
 from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.serializers import FromByteStream
@@ -876,3 +877,39 @@ def test_new_table_metadata_format_v2_with_v3_schema_fails(field_type: Primitive
             location="s3://some_v1_location/",
             properties={"format-version": "2"},
         )
+
+
+def test_get_bloom_filter_options() -> None:
+    schema = Schema(
+        NestedField(field_id=10, name="foo", field_type=StringType(), required=False),
+        NestedField(field_id=22, name="bar", field_type=IntegerType(), required=False),
+        NestedField(field_id=33, name="baz", field_type=BooleanType(), required=False),
+        NestedField(
+            field_id=34,
+            name="qux",
+            field_type=StructType(
+                NestedField(field_id=35, name="quux", field_type=StringType(), required=False),
+                NestedField(field_id=36, name="quuux", field_type=IntegerType(), required=False),
+            ),
+            required=False,
+        ),
+    )
+
+    table_properties = {
+        "write.parquet.bloom-filter-enabled.column.foo": "true",
+        "write.parquet.bloom-filter-fpp.column.foo": "0.01",
+        "write.parquet.bloom-filter-ndv.column.foo": "1000",
+        "write.parquet.bloom-filter-enabled.column.bar": "false",
+        "write.parquet.bloom-filter-fpp.column.bar": "0.02",
+        "write.parquet.bloom-filter-ndv.column.bar": "2000",
+        "write.parquet.bloom-filter-enabled.column.qux.quux": "true",
+        "write.parquet.bloom-filter-fpp.column.qux.quux": "0.03",
+        "write.parquet.bloom-filter-ndv.column.qux.quux": "3000",
+    }
+
+    id_to_parquet_path = {field_id: parquet_path for parquet_path, field_id in parquet_path_to_id_mapping(schema).items()}
+    bloom_filter_options = get_bloom_filter_options(schema, table_properties, id_to_parquet_path)
+    assert bloom_filter_options == {
+        "foo": {"fpp": 0.01, "ndv": 1000},
+        "qux.quux": {"fpp": 0.03, "ndv": 3000},
+    }