iceberg-loader 0.0.4__py3-none-any.whl

iceberg_loader/__about__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: 2025-present Ivan Matveev <skioneim@gmail.com>
+#
+# SPDX-License-Identifier: MIT
+__version__ = '0.0.4'

iceberg_loader/__init__.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: 2025-present Ivan Matveev <skioneim@gmail.com>
+#
+# SPDX-License-Identifier: MIT
+
+from iceberg_loader.__about__ import __version__
+from iceberg_loader.iceberg_loader import (
+    IcebergLoader,
+    LoaderConfig,
+    load_batches_to_iceberg,
+    load_data_to_iceberg,
+    load_ipc_stream_to_iceberg,
+)
+from iceberg_loader.maintenance import expire_snapshots
+
+__all__ = [
+    'load_data_to_iceberg',
+    'load_batches_to_iceberg',
+    'load_ipc_stream_to_iceberg',
+    'expire_snapshots',
+    '__version__',
+    'IcebergLoader',
+    'LoaderConfig',
+]

iceberg_loader/arrow_utils.py

@@ -0,0 +1,151 @@
+import json
+import logging
+from collections.abc import Iterator
+from functools import partial
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+_json_dumps = partial(json.dumps, ensure_ascii=False, separators=(',', ':'))
+logger = logging.getLogger(__name__)
+
+
+def _get_memory_pool() -> pa.MemoryPool:
+    return pa.system_memory_pool()
+
+
+def create_arrow_table_from_data(data: list[dict[str, Any]]) -> pa.Table:
+    """
+    Creates a PyArrow Table from a list of dictionaries, handling complex types
+    by serializing them to JSON strings if necessary.
+
+    Args:
+        data: A list of dictionaries representing the rows of the table.
+
+    Returns:
+        A PyArrow Table.
+    """
+    if not data:
+        return pa.Table.from_arrays([], schema=pa.schema([]))
+
+    return _create_table_native(data)
+
+
+def _create_table_native(data: list[dict[str, Any]]) -> pa.Table:
+    """
+    Manually constructs Arrow arrays from a list of dicts.
+    Handles mixed types by falling back to string (JSON) serialization.
+    """
+    if not data:
+        return pa.Table.from_arrays([], schema=pa.schema([]))
+
+    # Single pass to collect all unique keys would be better, but
+    # for simplicity and memory safety we just iterate.
+    # Optimization: using set comprehension is slightly faster than loop.
+    all_keys = set().union(*(d.keys() for d in data))
+
+    arrays = []
+    fields = []
+
+    # Pre-allocate pool
+    pool = _get_memory_pool()
+
+    for key in all_keys:
+        column_values: list[str | None] = []
+        for item in data:
+            value = item.get(key)
+            if isinstance(value, dict | list):
+                column_values.append(_json_dumps(value))
+            else:
+                column_values.append(str(value) if value is not None else None)
+
+        try:
+            array = pa.array(column_values, memory_pool=pool)
+            # Use nullable=True by default for robustness
+            field = pa.field(key, array.type, nullable=True)
+        except (pa.ArrowInvalid, pa.ArrowTypeError):
+            # Fallback to string for mixed types
+            str_values = [str(v) if v is not None else None for v in column_values]
+            array = pa.array(str_values, type=pa.string(), memory_pool=pool)
+            field = pa.field(key, pa.string(), nullable=True)
+
+        # Handle all-null columns
+        if pa.types.is_null(array.type):
+            array = pa.nulls(len(column_values), type=pa.string(), memory_pool=pool)
+            field = pa.field(key, pa.string(), nullable=True)
+
+        arrays.append(array)
+        fields.append(field)
+
+    return pa.Table.from_arrays(arrays, schema=pa.schema(fields))
+
+
+def convert_column_type(column: pa.Array, target_type: pa.DataType, column_name: str | None = None) -> pa.Array:
+    """Casts a single column to the target type, handling errors by nulling."""
+    if column.type == target_type:
+        return column
+
+    try:
+        return pc.cast(column, target_type)
+    except (ValueError, TypeError, pa.ArrowInvalid):
+        logger.warning('Cast failed for column %s to %s, filling nulls', column_name or '<unknown>', target_type)
+        return pa.nulls(len(column), type=target_type, memory_pool=_get_memory_pool())
+
+
+def convert_table_types(table: pa.Table, target_schema: pa.Schema) -> pa.Table:
+    """Convert PyArrow Table to match target schema, casting types and adding missing columns."""
+    if table.schema.equals(target_schema):
+        return table
+
+    return _convert_table_types_internal(table, target_schema)
+
+
+def _convert_table_types_internal(table: pa.Table, target_schema: pa.Schema) -> pa.Table:
+    source_fields = {field.name: field for field in table.schema}
+    new_arrays, new_fields = [], []
+
+    for field in target_schema:
+        if field.name in source_fields:
+            source_field = source_fields[field.name]
+            column = table[field.name]
+
+            if source_field.type.equals(field.type):
+                new_arrays.append(column)
+                new_fields.append(field)
+            else:
+                new_array = convert_column_type(column, field.type, field.name)
+                new_arrays.append(new_array)
+                new_fields.append(
+                    pa.field(field.name, new_array.type, nullable=field.nullable, metadata=field.metadata)
+                )
+        else:
+            # Add missing column as nulls
+            null_array = pa.nulls(len(table), type=field.type, memory_pool=_get_memory_pool())
+            new_arrays.append(null_array)
+            new_fields.append(field)
+
+    return pa.Table.from_arrays(new_arrays, schema=pa.schema(new_fields))
+
+
+def create_record_batches_from_dicts(
+    data_iterator: Iterator[dict[str, Any]], batch_size: int = 10000
+) -> Iterator[pa.RecordBatch]:
+    """Convert iterator of dicts to PyArrow RecordBatches with specified batch size."""
+    batch = []
+
+    for item in data_iterator:
+        batch.append(item)
+
+        if len(batch) >= batch_size:
+            table = create_arrow_table_from_data(batch)
+            batches = table.to_batches(max_chunksize=batch_size)
+            if batches:
+                yield batches[0]
+            batch = []
+
+    if batch:
+        table = create_arrow_table_from_data(batch)
+        batches = table.to_batches(max_chunksize=len(batch))
+        if batches:
+            yield batches[0]

iceberg_loader/iceberg_loader.py

@@ -0,0 +1,350 @@
+import logging
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Any, BinaryIO, Literal
+
+import pyarrow as pa
+from pyiceberg.catalog import Catalog
+
+from iceberg_loader.arrow_utils import convert_table_types
+from iceberg_loader.schema import SchemaManager
+from iceberg_loader.settings import TABLE_PROPERTIES
+from iceberg_loader.strategies import get_write_strategy
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LoaderConfig:
+    """Defaults for IcebergLoader operations."""
+
+    write_mode: Literal['overwrite', 'append', 'upsert'] = 'overwrite'
+    partition_col: str | None = None
+    replace_filter: str | None = None
+    schema_evolution: bool = False
+    table_properties: dict[str, Any] | None = None
+    commit_interval: int = 0
+    join_cols: list[str] | None = None
+
+
+class IcebergLoader:
+    """
+    Facade for loading data into Iceberg tables.
+    Orchestrates SchemaManager and WriteStrategy to handle complex ingestion scenarios.
+    """
+
+    def __init__(
+        self,
+        catalog: Catalog,
+        table_properties: dict[str, Any] | None = None,
+        default_config: LoaderConfig | None = None,
+    ):
+        self.catalog = catalog
+        self.table_properties = TABLE_PROPERTIES.copy()
+        if table_properties:
+            self.table_properties.update(table_properties)
+
+        self.schema_manager = SchemaManager(self.catalog, self.table_properties)
+        self.default_config = default_config or LoaderConfig()
+
+    def load_data(
+        self,
+        table_data: pa.Table,
+        table_identifier: tuple[str, str],
+        write_mode: Literal['overwrite', 'append', 'upsert'] | None = None,
+        partition_col: str | None = None,
+        replace_filter: str | None = None,
+        schema_evolution: bool | None = None,
+        table_properties: dict[str, Any] | None = None,
+        join_cols: list[str] | None = None,
+        config: LoaderConfig | None = None,
+    ) -> dict[str, Any]:
+        """
+        Load PyArrow Table into Iceberg table.
+        Delegates to load_data_batches for consistency.
+        """
+        batches = table_data.to_batches()
+        return self.load_data_batches(
+            batch_iterator=iter(batches),
+            table_identifier=table_identifier,
+            write_mode=write_mode,
+            partition_col=partition_col,
+            replace_filter=replace_filter,
+            schema_evolution=schema_evolution,
+            table_properties=table_properties,
+            join_cols=join_cols,
+            config=config,
+        )
+
+    def load_ipc_stream(
+        self,
+        stream_source: str | BinaryIO | pa.NativeFile,
+        table_identifier: tuple[str, str],
+        write_mode: Literal['overwrite', 'append', 'upsert'] | None = None,
+        partition_col: str | None = None,
+        replace_filter: str | None = None,
+        schema_evolution: bool | None = None,
+        commit_interval: int | None = None,
+        join_cols: list[str] | None = None,
+        table_properties: dict[str, Any] | None = None,
+        config: LoaderConfig | None = None,
+    ) -> dict[str, Any]:
+        """Loads data from an Apache Arrow IPC stream source."""
+        with pa.ipc.open_stream(stream_source) as reader:
+            return self.load_data_batches(
+                batch_iterator=reader,
+                table_identifier=table_identifier,
+                write_mode=write_mode,
+                partition_col=partition_col,
+                replace_filter=replace_filter,
+                schema_evolution=schema_evolution,
+                commit_interval=commit_interval,
+                join_cols=join_cols,
+                table_properties=table_properties,
+                config=config,
+            )
+
+    def load_data_batches(
+        self,
+        batch_iterator: Iterator[pa.RecordBatch] | pa.RecordBatchReader,
+        table_identifier: tuple[str, str],
+        write_mode: Literal['overwrite', 'append', 'upsert'] | None = None,
+        partition_col: str | None = None,
+        replace_filter: str | None = None,
+        schema_evolution: bool | None = None,
+        commit_interval: int | None = None,
+        join_cols: list[str] | None = None,
+        table_properties: dict[str, Any] | None = None,
+        config: LoaderConfig | None = None,
+    ) -> dict[str, Any]:
+        """
+        Main orchestration method.
+        Iterates over batches, manages buffers, and delegates writing.
+        """
+        total_rows = 0
+        batches_processed = 0
+
+        # Buffer
+        pending_batches: list[pa.RecordBatch] = []
+
+        resolved = self._resolve_config(
+            config=config,
+            write_mode=write_mode,
+            partition_col=partition_col,
+            replace_filter=replace_filter,
+            schema_evolution=schema_evolution,
+            commit_interval=commit_interval,
+            join_cols=join_cols,
+            table_properties=table_properties,
+        )
+        self.table_properties = resolved['table_properties']
+        strategy = get_write_strategy(resolved['write_mode'], resolved['replace_filter'], resolved['join_cols'])
+        effective_schema_evolution = resolved['schema_evolution']
+
+        # State tracking
+        table = None
+        is_first_write = True
+        new_table_created = False
+
+        def process_buffer(batches: list[pa.RecordBatch]) -> None:
+            nonlocal table, new_table_created, is_first_write, total_rows
+
+            if not batches:
+                return
+
+            combined_table = None
+
+            if effective_schema_evolution:
+                # Try fast path
+                try:
+                    combined_table = pa.Table.from_batches(batches)
+                except pa.ArrowInvalid:
+                    # Mixed schemas in buffer
+                    logger.info('Mixed schemas in batch buffer. Normalizing...')
+
+                    if table is None:
+                        table = self.schema_manager.ensure_table_exists(
+                            table_identifier, batches[0].schema, resolved['partition_col']
+                        )
+                        if table.current_snapshot() is None:
+                            new_table_created = True
+
+                    # Evolve schema for all batches
+                    for b in batches:
+                        self.schema_manager.evolve_schema_if_needed(table, b.schema)
+
+                    target_arrow_schema = self.schema_manager.get_arrow_schema(table)
+
+                    # Cast all batches
+                    normalized_tables = []
+                    for b in batches:
+                        t = pa.Table.from_batches([b])
+                        t_casted = convert_table_types(t, target_arrow_schema)
+                        normalized_tables.append(t_casted)
+
+                    combined_table = pa.concat_tables(normalized_tables)
+
+            if combined_table is None:
+                combined_table = pa.Table.from_batches(batches)
+
+            # 1. Ensure Table Exists
+            if table is None:
+                table = self.schema_manager.ensure_table_exists(
+                    table_identifier, combined_table.schema, resolved['partition_col']
+                )
+                if table.current_snapshot() is None:
+                    new_table_created = True
+
+            # 2. Schema Evolution
+            if effective_schema_evolution:
+                self.schema_manager.evolve_schema_if_needed(table, combined_table.schema)
+
+            # 3. Type Conversion
+            target_schema = self.schema_manager.get_arrow_schema(table)
+            combined_table = convert_table_types(combined_table, target_schema)
+
+            # The strategy now handles transaction management internally
+            strategy.write(table, combined_table, is_first_write)
+
+            is_first_write = False
+            total_rows += len(combined_table)
+
+        # Main Loop
+        for batch in batch_iterator:
+            pending_batches.append(batch)
+            batches_processed += 1
+
+            limit = 1 if resolved['commit_interval'] <= 1 else resolved['commit_interval']
+
+            if len(pending_batches) >= limit:
+                process_buffer(pending_batches)
+                pending_batches = []
+
+        # Flush remainder
+        if pending_batches:
+            process_buffer(pending_batches)
+
+        return {
+            'rows_loaded': total_rows,
+            'write_mode': resolved['write_mode'],
+            'partition_col': resolved['partition_col'] if resolved['partition_col'] else 'none',
+            'table_location': table.location() if table else 'none',
+            'snapshot_id': table.current_snapshot().snapshot_id if table and table.current_snapshot() else 'none',
+            'batches_processed': batches_processed,
+            'new_table_created': new_table_created,
+        }
+
+    def _resolve_config(
+        self,
+        config: LoaderConfig | None,
+        write_mode: Literal['overwrite', 'append', 'upsert'] | None,
+        partition_col: str | None,
+        replace_filter: str | None,
+        schema_evolution: bool | None,
+        commit_interval: int | None,
+        join_cols: list[str] | None,
+        table_properties: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        base = config or self.default_config
+        merged_table_props = self.table_properties.copy()
+        if base.table_properties:
+            merged_table_props.update(base.table_properties)
+        if table_properties:
+            merged_table_props.update(table_properties)
+        return {
+            'write_mode': write_mode or base.write_mode,
+            'partition_col': partition_col if partition_col is not None else base.partition_col,
+            'replace_filter': replace_filter if replace_filter is not None else base.replace_filter,
+            'schema_evolution': base.schema_evolution if schema_evolution is None else schema_evolution,
+            'commit_interval': base.commit_interval if commit_interval is None else commit_interval,
+            'join_cols': join_cols if join_cols is not None else base.join_cols,
+            'table_properties': merged_table_props,
+        }
+
+
+# Public API functions (thin wrappers)
+
+
+def load_data_to_iceberg(
+    table_data: pa.Table,
+    table_identifier: tuple[str, str],
+    catalog: Catalog,
+    write_mode: Literal['overwrite', 'append', 'upsert'] | None = None,
+    partition_col: str | None = None,
+    replace_filter: str | None = None,
+    schema_evolution: bool | None = None,
+    table_properties: dict[str, Any] | None = None,
+    join_cols: list[str] | None = None,
+    config: LoaderConfig | None = None,
+) -> dict[str, Any]:
+    """Public wrapper around IcebergLoader.load_data using an optional LoaderConfig."""
+    loader = IcebergLoader(catalog, table_properties, default_config=config)
+    return loader.load_data(
+        table_data,
+        table_identifier,
+        write_mode,
+        partition_col,
+        replace_filter,
+        schema_evolution,
+        table_properties,
+        join_cols,
+        config=config,
+    )
+
+
+def load_batches_to_iceberg(
+    batch_iterator: Iterator[pa.RecordBatch] | pa.RecordBatchReader,
+    table_identifier: tuple[str, str],
+    catalog: Catalog,
+    write_mode: Literal['overwrite', 'append', 'upsert'] | None = None,
+    partition_col: str | None = None,
+    replace_filter: str | None = None,
+    schema_evolution: bool | None = None,
+    table_properties: dict[str, Any] | None = None,
+    commit_interval: int | None = None,
+    join_cols: list[str] | None = None,
+    config: LoaderConfig | None = None,
+) -> dict[str, Any]:
+    """Public wrapper around IcebergLoader.load_data_batches using an optional LoaderConfig."""
+    loader = IcebergLoader(catalog, table_properties, default_config=config)
+    return loader.load_data_batches(
+        batch_iterator,
+        table_identifier,
+        write_mode,
+        partition_col,
+        replace_filter,
+        schema_evolution,
+        commit_interval,
+        join_cols,
+        table_properties,
+        config,
+    )
+
+
+def load_ipc_stream_to_iceberg(
+    stream_source: str | BinaryIO | pa.NativeFile,
+    table_identifier: tuple[str, str],
+    catalog: Catalog,
+    write_mode: Literal['overwrite', 'append', 'upsert'] | None = None,
+    partition_col: str | None = None,
+    replace_filter: str | None = None,
+    schema_evolution: bool | None = None,
+    table_properties: dict[str, Any] | None = None,
+    commit_interval: int | None = None,
+    join_cols: list[str] | None = None,
+    config: LoaderConfig | None = None,
+) -> dict[str, Any]:
+    """Public wrapper around IcebergLoader.load_ipc_stream using an optional LoaderConfig."""
+    loader = IcebergLoader(catalog, table_properties, default_config=config)
+    return loader.load_ipc_stream(
+        stream_source,
+        table_identifier,
+        write_mode,
+        partition_col,
+        replace_filter,
+        schema_evolution,
+        commit_interval,
+        join_cols,
+        table_properties,
+        config,
+    )

iceberg_loader/maintenance.py

@@ -0,0 +1,95 @@
+import logging
+from datetime import datetime
+from typing import Any
+
+try:
+    from pyiceberg.exceptions import CommitFailedException as IcebergError
+except ImportError:  # pragma: no cover - fallback for environments without pyiceberg
+
+    class IcebergError(Exception):  # type: ignore[no-redef]
+        """Fallback when pyiceberg is not installed."""
+
+
+logger = logging.getLogger(__name__)
+
+
+class SnapshotMaintenance:
+    """
+    Manages Iceberg table snapshot maintenance operations.
+
+    Provides functionality to expire old snapshots and prevent metadata bloat.
+    """
+
+    def expire_snapshots(self, table: Any, keep_last: int = 1, older_than_ms: int | None = None) -> None:
+        """
+        Expire old snapshots to prevent metadata issues.
+
+        Args:
+            table: Iceberg table instance.
+            keep_last: How many most recent snapshots to keep (default: 1).
+            older_than_ms: Optional timestamp in milliseconds; expire snapshots strictly older than this moment.
+                           If provided, it overrides keep_last logic.
+        """
+        try:
+            table.refresh()
+
+            snapshots = list(table.snapshots())
+            logger.info('Found %d snapshots', len(snapshots))
+
+            if len(snapshots) == 0:
+                logger.info('No snapshots found, nothing to fix')
+                return
+
+            sorted_snapshots = sorted(snapshots, key=lambda s: s.timestamp_ms)
+            current_snapshot = table.current_snapshot()
+
+            logger.info('Snapshot details (sorted by timestamp):')
+            for i, snap in enumerate(sorted_snapshots):
+                is_current = current_snapshot and snap.snapshot_id == current_snapshot.snapshot_id
+                marker = ' <-- CURRENT' if is_current else ''
+                logger.info('  %d. ID=%s, timestamp=%s%s', i + 1, snap.snapshot_id, snap.timestamp_ms, marker)
+
+            expire = table.maintenance.expire_snapshots()
+
+            # Determine cutoff strategy
+            if older_than_ms is not None:
+                cutoff_datetime = datetime.fromtimestamp(older_than_ms / 1000.0)
+                logger.info('Expiring snapshots older than %s (ms=%d)', cutoff_datetime, older_than_ms)
+                expire = expire.older_than(cutoff_datetime)
+            else:
+                if keep_last < 0:
+                    logger.info('keep_last < 0 specified, skipping expiration')
+                    return
+                if len(snapshots) <= keep_last:
+                    logger.info('Table has %d snapshots, keep_last=%d → nothing to expire', len(snapshots), keep_last)
+                    return
+                cutoff_snapshot = sorted_snapshots[-keep_last]
+                cutoff_datetime = datetime.fromtimestamp((cutoff_snapshot.timestamp_ms - 1) / 1000.0)
+                logger.info(
+                    'Expiring snapshots older than %s to keep last %d snapshot(s)',
+                    cutoff_datetime,
+                    keep_last,
+                )
+                expire = expire.older_than(cutoff_datetime)
+
+            before = len(snapshots)
+            expire.commit()
+            table.refresh()
+            remaining_snapshots = list(table.snapshots())
+            after = len(remaining_snapshots)
+            logger.info('Successfully expired snapshots: %d removed, %d remaining', before - after, after)
+
+        except (IcebergError, OSError, ValueError, RuntimeError) as e:
+            logger.warning('Failed to expire snapshots for table %s: %s', table.name(), e)
+
+
+def expire_snapshots(table: Any, keep_last: int = 1, older_than_ms: int | None = None) -> None:
+    """
+    Convenience function to expire snapshots without instantiating SnapshotMaintenance.
+
+    Args:
+        table: Iceberg table instance.
+        keep_last: How many most recent snapshots to keep (default: 1).
+        older_than_ms: Optional timestamp in milliseconds; expire snapshots strictly older than this moment.
+    """
+    SnapshotMaintenance().expire_snapshots(table, keep_last=keep_last, older_than_ms=older_than_ms)

iceberg_loader/py.typed

@@ -0,0 +1,2 @@
+
+