floe-python 0.8.0__py3-none-any.whl

floe/iceberg/unique.py

@@ -0,0 +1,61 @@
+"""Unique keys already in an Iceberg table, for append (Rust `seed_unique_tracker_for_append`)."""
+
+from collections.abc import Mapping
+
+import polars as pl
+from pyiceberg.table import Table
+
+from floe.checks.unique import (
+    UniqueTracker,
+    _float_bits,  # pyright: ignore[reportPrivateUsage] -- the tracker's own float key
+)
+from floe.iceberg.catalog import run_error
+
+
+def seed_from_table(
+    tracker: UniqueTracker,
+    table: Table,
+    batch: pl.DataFrame,
+    stored: Mapping[str, str],
+) -> None:
+    """Seed `tracker` with the keys of `table` that `batch` also holds; `stored` maps runtime to table names.
+
+    A target key the batch lacks is never hit, so the flags are Rust's while memory follows the
+    batch: the table is streamed, key columns only. The join matches -0.0 to 0.0 and NaN to any
+    NaN; the tracker then compares float bits, so those extra keys are never hit.
+    A key with a column the table lacks (`add_columns` adds it on write) is skipped: historical
+    rows hold NULL there, which never duplicates.
+    """
+    columns = set(table.schema().column_names)
+    for key in tracker.keys:
+        names = [stored.get(c, c) for c in key]
+        if not columns.issuperset(names):
+            continue
+        # positional names (k0, b0, ...): no stored column name can collide with a temporary
+        ids = [f"k{i}" for i in range(len(key))]
+        distinct = (
+            batch.select(pl.col(c).alias(k) for c, k in zip(key, ids, strict=True))
+            .drop_nulls()
+            .unique()
+        )
+        if not distinct.height:  # all-NULL keys are never checked; nothing to look up
+            continue
+        # distinct by float bits, as the tracker keys them: `unique` on a float merges -0.0 with 0.0
+        floats = [k for c, k in zip(key, ids, strict=True) if batch.schema[c].is_float()]  # fmt: skip
+        bits = {f"b{k}": pl.col(k).map_batches(_float_bits, return_dtype=pl.UInt64, is_elementwise=True) for k in floats}  # fmt: skip
+        subset = [k for k in ids if k not in floats] + list(bits)
+        with run_error("iceberg unique seed scan failed"):
+            scan = pl.scan_iceberg(table).select(pl.col(n).alias(k) for n, k in zip(names, ids, strict=True))  # fmt: skip
+            # the table's dtypes can differ (unsigned widened, datetimes in us); a key the table type
+            # cannot hold casts to NULL, which never matches
+            lookup = distinct.cast(scan.collect_schema(), strict=False).lazy()
+            hits = (
+                scan.join(lookup, on=ids, how="semi")
+                .with_columns(**bits)
+                # distinct keys only: memory follows the batch, not the history
+                .unique(subset)
+                .drop(*bits)
+                .collect(engine="streaming")
+            )
+        back = pl.Schema({k: batch.schema[c] for c, k in zip(key, ids, strict=True)})
+        tracker.seed(hits.cast(back).rename(dict(zip(ids, key, strict=True))))

floe/iceberg/write.py

@@ -0,0 +1,270 @@
+"""Overwrite and append of an Arrow stream in one catalog commit, one snapshot (#597's `single` path)."""
+
+# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportUnknownParameterType=false, reportMissingTypeStubs=false
+# ^ pyarrow ships no type stubs
+
+import itertools
+import uuid
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass, replace
+from datetime import datetime
+from pathlib import PurePosixPath
+from typing import Literal
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from pyiceberg.exceptions import CommitFailedException
+from pyiceberg.io.pyarrow import (
+    _check_pyarrow_schema_compatible,  # pyright: ignore[reportPrivateUsage] -- same check `Transaction.append` makes
+    _dataframe_to_data_files,  # pyright: ignore[reportPrivateUsage] -- feeds one snapshot producer from many chunks
+)
+from pyiceberg.manifest import DataFile, DataFileContent, ManifestEntryStatus
+from pyiceberg.table import Table, TableProperties
+from pyiceberg.table.snapshots import ancestors_of
+
+from floe.errors import FloeRunError
+from floe.iceberg.catalog import run_error
+from floe.state import (
+    _backoff,  # pyright: ignore[reportPrivateUsage] -- one jittered backoff for every CAS retry
+)
+
+WriteMode = Literal["overwrite", "append"]
+
+RUN_ID = "floe.run_id"
+INPUT_HASH = "floe.input_hash"
+
+# `metrics::DEFAULT_SMALL_FILE_THRESHOLD_BYTES`
+SMALL_FILE_THRESHOLD_BYTES = 16 * 1024 * 1024
+_BYTES_PER_MIB = 1024 * 1024
+_REPORTED_FILE_PATHS = 50
+CHUNK_ROWS = 1_000_000
+COMMIT_ATTEMPTS = 5
+
+
+@dataclass(frozen=True)
+class IcebergWriteResult:
+    files_written: int
+    snapshot_id: int | None
+    metadata_version: int | None
+    file_paths: list[str]
+    total_bytes_written: int
+    avg_file_size_mb: float | None
+    small_files_count: int
+    already_committed: bool = (
+        False  # a conflict showed another run committed this batch (#545)
+    )
+
+
+def _chunks(reader: pa.RecordBatchReader, chunk_rows: int) -> Iterator[pa.Table]:
+    pending: list[pa.RecordBatch] = []
+    rows = 0
+    for batch in reader:
+        offset = 0
+        while offset < batch.num_rows:
+            part = batch.slice(offset, chunk_rows - rows)
+            offset += part.num_rows
+            pending.append(part)
+            rows += part.num_rows
+            if rows >= chunk_rows:
+                yield pa.Table.from_batches(pending, reader.schema)
+                pending, rows = [], 0
+    if rows:
+        yield pa.Table.from_batches(pending, reader.schema)
+
+
+def _metadata_version(location: str | None) -> int | None:
+    prefix = PurePosixPath(location or "").name.split("-", 1)[0]
+    return int(prefix) if prefix.isdigit() else None
+
+
+def write_table(
+    table: Table,
+    reader: pa.RecordBatchReader,
+    mode: WriteMode,
+    run_id: str,
+    input_hash: str,
+    small_file_threshold_bytes: int = SMALL_FILE_THRESHOLD_BYTES,
+    chunk_rows: int = CHUNK_ROWS,
+    unique_keys: Sequence[Sequence[str]] = (),
+) -> IcebergWriteResult:
+    """Data files are written chunk by chunk (partitioned tables split each chunk), then committed once.
+
+    Overwrite replaces every existing file in the same snapshot. An append of no rows commits nothing.
+    On a conflict an append checks `unique_keys` (table column names) against the files added since the base
+    snapshot and fails closed, committing nothing, when a rival already holds a key (#549).
+    """
+    with run_error("iceberg data write failed"):
+        _check_pyarrow_schema_compatible(table.schema(), reader.schema)
+        counter = itertools.count()
+        write_uuid = uuid.uuid4()
+        data_files: list[DataFile] = [
+            data_file
+            for chunk in _chunks(reader, chunk_rows)
+            for data_file in _dataframe_to_data_files(
+                table.metadata, chunk, table.io, write_uuid, counter
+            )
+        ]
+    base = (
+        table.metadata.current_snapshot_id
+    )  # only commits after it can be a rival's copy of this batch
+    for attempt in range(COMMIT_ATTEMPTS):
+        if attempt:  # a conflict: the data files stay, only the transaction is rebuilt on the fresh table
+            _backoff(attempt - 1)
+            with run_error("rest iceberg table refresh failed"):
+                table.refresh()
+                seen = mode == "append" and find_committed(
+                    table, input_hash, after=base
+                )
+            if seen:
+                return replace(write_result(table, [], small_file_threshold_bytes), already_committed=True)  # fmt: skip
+            if mode == "append":
+                _revalidate_unique(table, base, data_files, unique_keys)
+        replaced: list[DataFile] = []
+        if mode == "overwrite":
+            with run_error("iceberg overwrite scan failed"):
+                replaced = [task.file for task in table.scan().plan_files()]
+        if not (data_files or replaced):
+            break
+        try:
+            with run_error("rest iceberg commit failed"):
+                # PyIceberg retries a 409 itself (refresh, rebuild) and would skip the guard; in memory only, never sent
+                table.metadata = table.metadata.model_copy(update={"properties": {**table.properties, TableProperties.COMMIT_NUM_RETRIES: "0"}})  # fmt: skip
+                tx = table.transaction()
+                update = tx.update_snapshot(
+                    snapshot_properties={RUN_ID: run_id, INPUT_HASH: input_hash}
+                )
+                with (
+                    update.overwrite()
+                    if mode == "overwrite"
+                    else update.fast_append() as snap
+                ):
+                    for old in replaced:
+                        snap.delete_data_file(old)
+                    for data_file in data_files:
+                        snap.append_data_file(data_file)
+                table = tx.commit_transaction()
+        except FloeRunError as e:
+            if not isinstance(e.__cause__, CommitFailedException):
+                raise
+            continue
+        break
+    else:
+        raise FloeRunError(f"rest iceberg commit failed: conflicted {COMMIT_ATTEMPTS} times in a row")  # fmt: skip
+    return write_result(table, data_files, small_file_threshold_bytes)
+
+
+def added_data_files(table: Table, since: int | None) -> list[DataFile]:
+    """Data files added by the snapshots after `since` (a DuckDB MERGE commits two)."""
+    files: list[DataFile] = []
+    snapshot = table.current_snapshot()
+    while snapshot is not None and snapshot.snapshot_id != since:
+        files += [
+            entry.data_file
+            for manifest in snapshot.manifests(table.io)
+            if manifest.added_snapshot_id == snapshot.snapshot_id
+            for entry in manifest.fetch_manifest_entry(table.io)
+            if entry.status == ManifestEntryStatus.ADDED
+            and entry.data_file.content == DataFileContent.DATA
+        ]
+        parent = snapshot.parent_snapshot_id
+        snapshot = None if parent is None else table.snapshot_by_id(parent)
+    return files
+
+
+def _keys(
+    table: Table, file: DataFile, key: Sequence[str]
+) -> Iterator[tuple[object, ...]]:
+    """The non-NULL key tuples of a data file; a file without a key column (evolved later) holds none."""
+    with table.io.new_input(file.file_path).open() as f:
+        schema = pq.read_schema(f)
+        if not set(key) <= set(schema.names):
+            return
+        f.seek(0)
+        rows = pq.read_table(f, columns=list(key))
+    columns = [
+        # the tracker's float key: bit patterns, so NaN matches NaN and -0.0 differs from 0.0
+        col.combine_chunks().cast(pa.float64()).view(pa.uint64())
+        if pa.types.is_floating(col.type)
+        else col
+        for col in (rows[c] for c in key)
+    ]
+    yield from (k for k in zip(*(c.to_pylist() for c in columns)) if None not in k)
+
+
+def _revalidate_unique(
+    table: Table,
+    base: int | None,
+    data_files: list[DataFile],
+    unique_keys: Sequence[Sequence[str]],
+) -> None:
+    """Raise when a key of `data_files` is in a file a rival added since `base`.
+
+    ponytail: the rivals' keys sit in memory (cost follows the concurrent writes, not the history);
+    row-level deletes are not applied, so a deleted rival row still fails closed.
+    """
+    with run_error("iceberg unique revalidation failed"):
+        live = {task.file.file_path for task in table.scan().plan_files()}
+        rival = [f for f in added_data_files(table, base) if f.file_path in live]
+        clash = next(
+            (
+                key
+                for key in unique_keys
+                if (held := {k for file in rival for k in _keys(table, file, key)})
+                and any(
+                    k in held for file in data_files for k in _keys(table, file, key)
+                )
+            ),
+            None,
+        )
+    if clash:
+        raise FloeRunError(f"unique key ({', '.join(clash)}) already committed by a concurrent run; nothing was committed, rerun to revalidate")  # fmt: skip
+
+
+def write_result(
+    table: Table,
+    data_files: list[DataFile],
+    small_file_threshold_bytes: int = SMALL_FILE_THRESHOLD_BYTES,
+) -> IcebergWriteResult:
+    """Rust's report fields for the files a commit added."""
+    sizes = [f.file_size_in_bytes for f in data_files]
+    snapshot = table.current_snapshot()
+    return IcebergWriteResult(
+        files_written=len(data_files),
+        snapshot_id=None if snapshot is None else snapshot.snapshot_id,
+        metadata_version=_metadata_version(table.metadata_location),
+        file_paths=[PurePosixPath(f.file_path).name for f in data_files][
+            :_REPORTED_FILE_PATHS
+        ],
+        total_bytes_written=sum(sizes),
+        avg_file_size_mb=sum(sizes) / len(sizes) / _BYTES_PER_MIB if sizes else None,
+        small_files_count=sum(s < small_file_threshold_bytes for s in sizes),
+    )
+
+
+def find_committed(
+    table: Table,
+    input_hash: str,
+    since: datetime | None = None,
+    after: int | None = None,
+) -> bool | None:
+    """Whether a snapshot ancestor carries `floe.input_hash` (#545); the metadata is in memory, no manifest I/O.
+
+    The walk stops at snapshots older than `since`, the claim's `acquired_at`: a batch can't commit before its
+    claim. None when the history back to `since` was expired, so the answer is unknown. `after` is a snapshot id
+    that ends the walk: only the commits made since it count.
+    """
+    stop = None if since is None else since.timestamp() * 1000
+    oldest = None
+    for snapshot in ancestors_of(table.current_snapshot(), table.metadata):
+        if snapshot.snapshot_id == after or (
+            stop is not None and snapshot.timestamp_ms < stop
+        ):
+            return False
+        if (
+            snapshot.summary is not None
+            and snapshot.summary.get(INPUT_HASH) == input_hash
+        ):
+            return True
+        oldest = snapshot
+    expired = oldest is not None and oldest.parent_snapshot_id is not None
+    return None if stop is not None and expired else False

floe/lineage/events.py

@@ -0,0 +1,163 @@
+"""OpenLineage `RunEvent` bodies, built as Rust builds them (Rust: `lineage/mod.rs`)."""
+
+import os
+from collections.abc import Sequence
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from floe import __version__
+from floe.config.project import LineageConfig
+from floe.config.types import EntityConfig
+
+DEFAULT_PRODUCER = f"https://github.com/OpenLakeForge/floe/releases/tag/v{__version__}"
+RUN_EVENT_SCHEMA = "https://openlineage.io/spec/1-0-5/OpenLineage.json#/$defs/RunEvent"
+DQ_SCHEMA = "https://openlineage.io/spec/facets/1-0-2/DataQualityMetricsOutputDatasetFacet.json"  # fmt: skip
+CLOUD_PREFIXES = ("s3://", "gs://", "gcs://", "az://", "abfss://", "abfs://")
+
+Datasets = list[dict[str, Any]]
+
+
+def split_storage_uri(uri: str) -> tuple[str, str]:
+    """(namespace, name): a cloud URI splits at its authority; anything else is a `file` path."""
+    for prefix in CLOUD_PREFIXES:
+        if uri.startswith(prefix):
+            slash = uri.find("/", len(prefix))
+            return (uri, "") if slash < 0 else (uri[:slash], uri[slash + 1 :])
+    return "file", uri
+
+
+def resolve_api_key(configured: str | None, ambient: str | None) -> str | None:
+    """An explicit key wins; else `OPENLINEAGE_API_KEY`. An unexpanded `{{...}}` counts as absent."""
+    key = (configured or "").strip()
+    if key and not (key.startswith("{{") and key.endswith("}}")):
+        return key
+    return (ambient or "").strip() or None
+
+
+def event_type(status: str) -> str:
+    return "FAIL" if status in ("failed", "aborted") else "COMPLETE"
+
+
+def iso(ms: int) -> str:
+    """RFC 3339 in UTC with trailing fractional zeros trimmed, as Rust's `time` crate prints it."""
+    base = datetime.fromtimestamp(ms // 1000, UTC).strftime("%Y-%m-%dT%H:%M:%S")
+    frac = f"{ms % 1000:03d}".rstrip("0")
+    return f"{base}.{frac}Z" if frac else f"{base}Z"
+
+
+def _dataset(namespace: str, name: str) -> dict[str, Any]:
+    return {"namespace": namespace, "name": name}
+
+
+class EventBuilder:
+    """Job identity and per-entity datasets, fixed when the run starts."""
+
+    def __init__(
+        self, config: LineageConfig, entities: Sequence[EntityConfig], config_path: str
+    ) -> None:
+        self.namespace = config.namespace
+        self.producer = config.producer or DEFAULT_PRODUCER
+        self.run_job_name = config.job_name or Path(config_path).stem or "floe-run"
+        ds = config.dataset_namespace
+        dataset_ns = ds if ds and ds.strip() else config.namespace
+        self.entities = {e.id: e for e in entities}
+        self.uris: dict[str, tuple[dict[str, Any], ...]] = {}
+        for e in entities:
+            iceberg = e.sink.accepted.iceberg
+            if iceberg:
+                ns = e.name if iceberg.namespace is None else iceberg.namespace
+                table = e.name if iceberg.table is None else iceberg.table
+                accepted = _dataset(dataset_ns, f"{ns}.{table}")
+            else:
+                accepted = _dataset(*split_storage_uri(e.sink.accepted.path))
+            rejected = e.sink.rejected
+            self.uris[e.id] = (
+                _dataset(*split_storage_uri(e.source.path)),
+                accepted,
+                *([_dataset(*split_storage_uri(rejected.path))] if rejected else []),
+            )
+
+    def body(
+        self,
+        event_type: str,
+        ts_ms: int,
+        run_id: str,
+        job: str,
+        inputs: Datasets | None = None,
+        outputs: Datasets | None = None,
+    ) -> dict[str, Any]:
+        parent = self.parent()
+        return {
+            "eventType": event_type,
+            "eventTime": iso(ts_ms),
+            "run": {"runId": run_id, "facets": {"parent": parent} if parent else {}},
+            "job": {"namespace": self.namespace, "name": job, "facets": {}},
+            "inputs": inputs or [],
+            "outputs": outputs or [],
+            "producer": self.producer,
+            "schemaURL": RUN_EVENT_SCHEMA,
+        }
+
+    def _facet(self, schema: str, **fields: Any) -> dict[str, Any]:
+        return fields | {"_producer": self.producer, "_schemaURL": schema}
+
+    def datasets(self, name: str, s: dict[str, Any]) -> tuple[Datasets, Datasets]:
+        """Inputs and outputs of an `entity_finished` event; none for an unknown entity."""
+        if name not in self.uris:
+            return [], []
+        source, accepted, *rejected = self.uris[name]
+        columns = self.entities[name].schema_.columns
+        facets: dict[str, Any] = {
+            "schema": self._facet(
+                "https://openlineage.io/spec/facets/1-1-1/SchemaDatasetFacet.json",
+                fields=[{"name": c.name, "type": c.column_type} for c in columns],
+            ),
+            "dataQualityMetrics": self._facet(
+                DQ_SCHEMA,
+                rowCount=s["accepted"],
+                validCount=s["accepted"],
+                invalidCount=0,
+            ),
+            "floeQualityRun": self._facet(
+                "https://github.com/OpenLakeForge/floe/schemas/FloeQualityRunFacet.json",
+                entity=name,
+                rejectionRate=s["rejected"] / s["rows"] if s["rows"] > 0 else 0.0,
+                **{k: s[k] for k in ("files", "rows", "warnings", "errors")},
+            ),
+        }
+        if columns:
+            facets["columnLineage"] = self._facet(
+                "https://openlineage.io/spec/facets/1-1-1/ColumnLineageDatasetFacet.json",
+                fields={
+                    c.name: {"inputFields": [source | {"field": c.source or c.name}]}
+                    for c in columns
+                },
+            )
+        outputs = [accepted | {"facets": facets}]
+        for r in rejected:
+            dq = self._facet(
+                DQ_SCHEMA,
+                rowCount=s["rejected"],
+                validCount=0,
+                invalidCount=s["rejected"],
+            )
+            outputs.append(r | {"facets": {"dataQualityMetrics": dq}})
+        return [source], outputs
+
+    def parent(self) -> dict[str, Any] | None:
+        """The Airflow, else Dagster, parent run read from the environment."""
+        env = os.environ
+        if (run_id := env.get("AIRFLOW_CTX_DAG_RUN_ID")) is not None:
+            dag = env.get("AIRFLOW_CTX_DAG_ID", "")
+            task = env.get("AIRFLOW_CTX_TASK_ID", "")
+            job = f"{dag}.{task}" if task else dag
+        elif (run_id := env.get("DAGSTER_RUN_ID")) is not None:
+            job = env.get("DAGSTER_JOB_NAME", "")
+        else:
+            return None
+        return self._facet(
+            "https://openlineage.io/spec/facets/1-0-0/ParentRunFacet.json",
+            run={"runId": run_id},
+            job={"namespace": self.namespace, "name": job},
+        )