perf(avro): cache Avro-to-Iceberg schema conversion

pyiceberg/avro/file.py

@@ -25,6 +25,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from enum import Enum
+from functools import lru_cache
 from types import TracebackType
 from typing import (
     Generic,
@@ -68,6 +69,11 @@
 _SCHEMA_KEY = "avro.schema"
 
 
+@lru_cache
+def _cached_avro_to_iceberg(avro_schema_string: str) -> Schema:
+    return AvroSchemaConversion().avro_to_iceberg(json.loads(avro_schema_string))
+
+
 class AvroFileHeader(Record):
     @property
     def magic(self) -> bytes:
@@ -97,9 +103,7 @@ def compression_codec(self) -> type[Codec] | None:
 
     def get_schema(self) -> Schema:
         if _SCHEMA_KEY in self.meta:
-            avro_schema_string = self.meta[_SCHEMA_KEY]
-            avro_schema = json.loads(avro_schema_string)
-            return AvroSchemaConversion().avro_to_iceberg(avro_schema)
+            return _cached_avro_to_iceberg(self.meta[_SCHEMA_KEY])
         else:
             raise ValueError("No schema found in Avro file headers")
 

tests/avro/test_file.py

@@ -15,6 +15,7 @@
 #  specific language governing permissions and limitations
 #  under the License.
 import inspect
+import json
 from _decimal import Decimal
 from datetime import datetime
 from enum import Enum
@@ -87,6 +88,38 @@ def test_missing_schema() -> None:
     assert "No schema found in Avro file headers" in str(exc_info.value)
 
 
+def test_get_schema_is_cached_per_schema_string() -> None:
+    avro_schema_string = json.dumps(
+        {
+            "type": "record",
+            "name": "manifest_entry",
+            "fields": [
+                {"name": "status", "type": "int", "field-id": 0},
+                {"name": "snapshot_id", "type": ["null", "long"], "default": None, "field-id": 1},
+            ],
+        }
+    )
+    expected = AvroSchemaConversion().avro_to_iceberg(json.loads(avro_schema_string))
+
+    header = AvroFileHeader(bytes(0), {"avro.schema": avro_schema_string}, bytes(16))
+    other_header = AvroFileHeader(bytes(0), {"avro.schema": avro_schema_string}, bytes(16))
+
+    assert header.get_schema() == expected
+    # Identical schema strings resolve to the same cached object; different strings do not.
+    assert header.get_schema() is other_header.get_schema()
+
+    different_schema_string = json.dumps(
+        {
+            "type": "record",
+            "name": "manifest_entry",
+            "fields": [{"name": "status", "type": "int", "field-id": 0}],
+        }
+    )
+    different_header = AvroFileHeader(bytes(0), {"avro.schema": different_schema_string}, bytes(16))
+    assert different_header.get_schema() is not header.get_schema()
+    assert different_header.get_schema() == AvroSchemaConversion().avro_to_iceberg(json.loads(different_schema_string))
+
+
 # helper function to serialize our objects to dicts to enable
 # direct comparison with the dicts returned by fastavro
 def todict(obj: Any) -> Any:

tests/benchmark/test_avro_schema_cache_benchmark.py

@@ -0,0 +1,76 @@
+# 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.
+"""Benchmark for scan planning over a table with many manifests.
+
+Every manifest embeds an identical Avro schema string, so converting it to an
+Iceberg schema on every manifest open is redundant. With that conversion cached,
+only the unique schemas are converted rather than one per manifest.
+
+Run with: uv run pytest tests/benchmark/test_avro_schema_cache_benchmark.py -v -s -m benchmark
+"""
+
+from __future__ import annotations
+
+import statistics
+import timeit
+
+import pyarrow as pa
+import pytest
+
+from pyiceberg.catalog.memory import InMemoryCatalog
+
+
+@pytest.fixture
+def memory_catalog(tmp_path_factory: pytest.TempPathFactory) -> InMemoryCatalog:
+    warehouse_path = str(tmp_path_factory.mktemp("warehouse"))
+    catalog = InMemoryCatalog("memory_test", warehouse=f"file://{warehouse_path}")
+    catalog.create_namespace("default")
+    return catalog
+
+
+@pytest.mark.benchmark
+def test_scan_planning_many_manifests(memory_catalog: InMemoryCatalog) -> None:
+    """Time `scan().plan_files()` on a table with many manifests.
+
+    Each append creates a new manifest, so this exercises the per-manifest cost
+    of scan planning, dominated by the Avro-to-Iceberg schema conversion run once
+    per manifest open.
+    """
+    num_appends = 150
+    data = pa.table(
+        {
+            "id": pa.array(range(1000), pa.int64()),
+            "val": pa.array([f"v{i % 50}" for i in range(1000)]),
+        }
+    )
+    table = memory_catalog.create_table("default.scan_bench", schema=data.schema)
+    for _ in range(num_appends):
+        table.append(data)
+
+    table = memory_catalog.load_table("default.scan_bench")
+    # Warm up, and sanity check that each append contributes one data file to plan.
+    assert len(list(table.scan().plan_files())) == num_appends
+
+    num_runs = 10
+    runs = []
+    for _ in range(num_runs):
+        start_time = timeit.default_timer()
+        list(table.scan().plan_files())
+        runs.append(timeit.default_timer() - start_time)
+
+    print(f"\n--- scan planning over {num_appends} manifests ---")
+    print(f"median: {statistics.median(runs) * 1000:.1f} ms  min: {min(runs) * 1000:.1f} ms  ({num_runs} runs)")