From cd14ac484b8c6f4c79ee5410cfaf52a578303241 Mon Sep 17 00:00:00 2001 From: Vishnu Prakash Date: Wed, 15 Jul 2026 18:13:09 +0530 Subject: [PATCH] perf(avro): cache Avro-to-Iceberg schema conversion --- pyiceberg/avro/file.py | 10 ++- tests/avro/test_file.py | 33 ++++++++ .../test_avro_schema_cache_benchmark.py | 76 +++++++++++++++++++ 3 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 tests/benchmark/test_avro_schema_cache_benchmark.py diff --git a/pyiceberg/avro/file.py b/pyiceberg/avro/file.py index 7db92818fe..f124f6cbec 100644 --- a/pyiceberg/avro/file.py +++ b/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") diff --git a/tests/avro/test_file.py b/tests/avro/test_file.py index 137215ebc8..2d11377275 100644 --- a/tests/avro/test_file.py +++ b/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: diff --git a/tests/benchmark/test_avro_schema_cache_benchmark.py b/tests/benchmark/test_avro_schema_cache_benchmark.py new file mode 100644 index 0000000000..2a08c3492e --- /dev/null +++ b/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)")