diff --git a/README.md b/README.md index 635d99de..0fa2ed9b 100755 --- a/README.md +++ b/README.md @@ -108,8 +108,9 @@ shop = Shop( ) ``` -Your own destination is two methods, sync or async. `PartitionedDestination` and -`DatabaseDestination` handle partition scoping for you. +Your own destination is two methods, sync or async, each handling one partition, or the +unpartitioned whole. Windows are split and gathered for you; `DatabaseDestination` adds the +delete-then-insert dance for tables. ```python import json @@ -120,14 +121,17 @@ from pathlib import Path class JSONLDestination(il.Destination): base_path: str = "" - def write(self, context: il.IOContext, data) -> None: - path = Path(self.base_path) / context.asset.dataset / f"{context.asset.table}.jsonl" + def _path(self, context: il.IOContext, partition: il.Partition | None) -> Path: + base = Path(self.base_path) / context.asset.dataset / context.asset.table + return base / ("data.jsonl" if partition is None else f"{partition.id}.jsonl") + + def write_partition(self, context: il.IOContext, partition: il.Partition | None, data) -> None: + path = self._path(context, partition) path.parent.mkdir(parents=True, exist_ok=True) path.write_text("\n".join(json.dumps(row, default=str) for row in data)) - def read(self, context: il.IOContext): - path = Path(self.base_path) / context.asset.dataset / f"{context.asset.table}.jsonl" - return [json.loads(line) for line in path.read_text().splitlines()] + def read_partition(self, context: il.IOContext, partition: il.Partition | None): + return [json.loads(line) for line in self._path(context, partition).read_text().splitlines()] ``` ### Running diff --git a/docs/extending/representations.md b/docs/extending/representations.md index 04715818..c4b1f893 100644 --- a/docs/extending/representations.md +++ b/docs/extending/representations.md @@ -65,7 +65,7 @@ rows, whose record coercion rejects non-tabular data with a clear error. - Conform resolves the conformer through `Representation.of(result)`. - `Partition.slice()` and `TimePartition.slice()` filter through the representation, so window writes split correctly for any table type. -- `PartitionedDestination` and `DatabaseDestination` convert through `to_records` and +- `Destination` and `DatabaseDestination` convert through `to_records` and `from_records`. - `DatabaseDestination.read_representation` and `@destination(read_representation=...)` name the representation reads should materialize into. diff --git a/docs/guide/destinations.md b/docs/guide/destinations.md index 11bdf73a..76e913cf 100644 --- a/docs/guide/destinations.md +++ b/docs/guide/destinations.md @@ -88,14 +88,15 @@ Every `read()` and `write()` receives an immutable `IOContext`: | Field | Meaning | |-------|---------| | `asset` | The asset being read or written. `asset.table`, `asset.dataset`, `asset.partitioning` name the storage location. | -| `partition_or_window` | The scope of this call, or `None` for an unpartitioned asset. | +| `partition_or_window` | The partition or window of this call, or `None` for an unpartitioned asset. | | `schema` | The effective schema of the data: the declared one, or the one inferred during conform. `None` when none could be resolved. | | `metadata` | Run metadata (run id, backfill id). | ## Custom destinations -Subclass `il.Destination`, or decorate a plain class with `@il.destination`, and implement -`read()` and `write()`. Both may be sync or `async def`: +A destination stores data one **partition** at a time, `None` standing for the whole of an +unpartitioned asset. Subclass `il.Destination`, or decorate a plain class with `@il.destination`, +and implement the two partition hooks. Both may be sync or `async def`: ```py import json @@ -108,19 +109,28 @@ import interloper as il class JSONDestination(il.Destination): base_path: str = "" - def write(self, context: il.IOContext, data: Any) -> None: - path = Path(self.base_path) / context.asset.dataset / f"{context.asset.table}.json" + def _path(self, context: il.IOContext, partition: il.Partition | None) -> Path: + base = Path(self.base_path) / (context.asset.dataset or "") / context.asset.table + if partition is None: + return base / "data.json" + return base / f"{context.asset.partitioning.column}={partition.id}" / "data.json" + + def write_partition(self, context: il.IOContext, partition: il.Partition | None, data: Any) -> None: + path = self._path(context, partition) path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(data, default=str)) - def read(self, context: il.IOContext) -> Any: - path = Path(self.base_path) / context.asset.dataset / f"{context.asset.table}.json" - return json.loads(path.read_text()) + def read_partition(self, context: il.IOContext, partition: il.Partition | None) -> Any: + return json.loads(self._path(context, partition).read_text()) ``` -This form writes one file per asset and ignores `context.partition_or_window`: a partition -rewrite replaces every partition and a window lands in one file. It suits unpartitioned assets; -partitioned ones belong on `il.PartitionedDestination` below. +`write()` and `read()` are the base class's: a window write is split into one `write_partition` call +per partition, slicing the data through its [representation](../extending/representations.md) +on the partition column; a window read returns one result per partition, newest first. A +destination written this way is partition-correct by construction, and `CSVDestination`, +`FileDestination`, `MemoryDestination` and `GCSDestination` are all built exactly like this. The +partitions a context covers are on the context itself, `context.partitions` and `context.slices(data)`, +for a backend that needs them. The decorator accepts the class's public ClassVars and field defaults, plus `relations=`; see the [decorators reference](../reference/decorators.md). A destination's own connection is a relation, @@ -130,32 +140,15 @@ declared as an annotation or through `relations=`; see Override `partition_row_counts(context)` to report rows per partition; `asset.partition_row_counts()` and coverage tooling call it. -### Partition-aware destinations - -`il.PartitionedDestination` implements the partition dispatch once. Subclasses implement two -scope hooks and are partition-correct by construction: - -```py -class JSONDestination(il.PartitionedDestination): - base_path: str = "" - - def _write_scope(self, context, partition, data) -> None: - # partition is None for the unpartitioned whole - ... - - def _read_scope(self, context, partition): - ... -``` - -A window write is split into one `_write_scope` call per partition, slicing the data through -its [representation](../extending/representations.md) on the partition column. A window read -returns one result per partition, newest first. `CSVDestination` and `MemoryDestination` are -built this way. +A backend whose storage is not per partition may override `write()` or `read()` +wholesale. `DatabaseDestination` below does that for writes: rows carry the partition column, so +a window clears every partition it covers and inserts the whole batch once. ### Database destinations `DatabaseDestination` (imported from `interloper.destination`, together with -`WriteDisposition`) targets stores addressed by table and schema. Reads and writes reduce to a +`WriteDisposition`) targets stores addressed by table and schema. Its `write_partition` clears one +partition and inserts; its `write` batches a window into one insert. Reads and writes reduce to a small set of row operations: | Hook | Called for | @@ -176,7 +169,7 @@ representation directly (a DataFrame into a Parquet load job) using `context.sch Behaviour the base class owns: - **Write disposition**: `write_disposition = WriteDisposition.REPLACE` (default) deletes the - matching scope before inserting; `APPEND` never deletes. A class attribute, not a field. + matching partition before inserting; `APPEND` never deletes. A class attribute, not a field. - **Time partitions are scoped by bounds**, not by equality, because rows of a monthly partition carry daily dates. - **Read representation**: rows are materialized into the representation named by diff --git a/packages/interloper-core/src/interloper/__init__.py b/packages/interloper-core/src/interloper/__init__.py index 275a4421..a63ad888 100644 --- a/packages/interloper-core/src/interloper/__init__.py +++ b/packages/interloper-core/src/interloper/__init__.py @@ -26,7 +26,6 @@ FileDestination, IOContext, MemoryDestination, - PartitionedDestination, destination, ) from interloper.events import Event, EventBus, EventType @@ -137,7 +136,6 @@ "Partition", "PartitionConfig", "PartitionWindow", - "PartitionedDestination", "RESTClient", "RangePaginator", "RefreshTokenOAuthConnection", diff --git a/packages/interloper-core/src/interloper/destination/__init__.py b/packages/interloper-core/src/interloper/destination/__init__.py index 83cb3c56..1710fbcc 100644 --- a/packages/interloper-core/src/interloper/destination/__init__.py +++ b/packages/interloper-core/src/interloper/destination/__init__.py @@ -7,7 +7,6 @@ from interloper.destination.decorator import destination from interloper.destination.file import FileDestination from interloper.destination.memory import MemoryDestination -from interloper.destination.partitioned import PartitionedDestination __all__ = [ "CSVDestination", @@ -17,7 +16,6 @@ "FileDestination", "IOContext", "MemoryDestination", - "PartitionedDestination", "WriteDisposition", "destination", ] diff --git a/packages/interloper-core/src/interloper/destination/base.py b/packages/interloper-core/src/interloper/destination/base.py index ebdec41f..f3a579e7 100644 --- a/packages/interloper-core/src/interloper/destination/base.py +++ b/packages/interloper-core/src/interloper/destination/base.py @@ -2,11 +2,11 @@ from __future__ import annotations -from abc import abstractmethod from typing import Any, ClassVar from interloper.component import Component, ComponentDefinition from interloper.destination.context import IOContext +from interloper.partitioning.base import Partition from interloper.utils.text import to_label @@ -20,24 +20,33 @@ class DestinationDefinition(ComponentDefinition): class Destination(Component): - """A component that reads and writes asset data. - - Subclass and implement ``read()`` and ``write()``. They may be written - as plain sync methods (the common case: most warehouse/file clients are - sync) or as ``async def`` for native async I/O (e.g. asyncpg, aiofiles). - The engine is async-native: it awaits async implementations directly and - offloads sync ones to a worker thread, so a destination never blocks the - event loop either way. An annotation naming a component class declares a + """A component that reads and writes asset data, one partition at a time. + + A destination stores data per **partition**, ``None`` standing for the + whole of an unpartitioned asset. Subclass and implement + :meth:`write_partition` and :meth:`read_partition` for a single one; + :meth:`write` and :meth:`read` own the rest, splitting a window write into + one call per partition and gathering a window read into one result per + partition, so a destination is partition-correct by construction. A + backend that does not store per partition (a database that clears each + partition and inserts a window in one batch) overrides :meth:`write` or + :meth:`read` instead. + + The hooks may be plain sync methods (the common case: most warehouse and + file clients are sync) or ``async def`` for native async I/O. The engine + is async-native: it awaits async implementations directly and offloads + sync ones to a worker thread, so a destination never blocks the event + loop either way. An annotation naming a component class declares a relation, which the destination resolves by name:: - class PostgresDestination(Destination): - connection: PostgresConnection + class JSONDestination(Destination): + connection: BucketConnection - def read(self, context: IOContext) -> Any: - return query_table(self.connection.connection_string, context.table) + def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None: + self.connection.put(self._path(context, partition), json.dumps(data, default=str)) - def write(self, context: IOContext, data: Any) -> None: - insert_into(self.connection.connection_string, context.table, data) + def read_partition(self, context: IOContext, partition: Partition | None) -> Any: + return json.loads(self.connection.get(self._path(context, partition))) """ tags: ClassVar[list[str]] = [] @@ -68,25 +77,63 @@ def definition(cls) -> DestinationDefinition: relations=dict(cls.relations), ) - @abstractmethod - def read(self, context: IOContext) -> Any: - """Read data from this destination. + def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None: + """Store *data* for one partition. Args: - context: Destination context with asset, partition, and metadata. + context: IO context carrying the target asset and the effective schema. + partition: The partition being stored, or ``None`` for the + unpartitioned whole. + data: The partition's slice of the data to store. + + Raises: + NotImplementedError: Every destination implements this, unless it + overrides :meth:`write` for storage that is not per partition. + """ + raise NotImplementedError(f"{type(self).__name__} must implement write_partition()") - Returns: - The data read from the destination. + def read_partition(self, context: IOContext, partition: Partition | None) -> Any: + """Load one partition. + + Args: + context: IO context carrying the target asset and the effective schema. + partition: The partition to load, or ``None`` for the unpartitioned + whole. + + Raises: + NotImplementedError: Every destination implements this, unless it + overrides :meth:`read` for storage that is not per partition. """ + raise NotImplementedError(f"{type(self).__name__} must implement read_partition()") - @abstractmethod def write(self, context: IOContext, data: Any) -> None: - """Write data to this destination. + """Write data, one partition at a time. + + A window is split into one :meth:`write_partition` call per partition, + each receiving its slice of the data; a single partition or the + unpartitioned whole is one call receiving the data as is. Args: - context: Destination context with asset, partition, and metadata. - data: The data to write. + context: IO context carrying the target asset, the partition or window, + and the effective schema. + data: The data to write, in its native representation. + """ + for partition, chunk in context.slices(data): + self.write_partition(context, partition, chunk) + + def read(self, context: IOContext) -> Any: + """Read data for the context's partition, or window. + + Args: + context: IO context carrying the target asset, the partition or window, + and the effective schema. + + Returns: + The partition's data; a window returns one result per partition, in + window order. """ + results = [self.read_partition(context, partition) for partition in context.partitions] + return results if context.window else results[0] def partition_row_counts(self, context: IOContext) -> dict[str, int]: """Return row counts grouped by the asset's partition column. diff --git a/packages/interloper-core/src/interloper/destination/context.py b/packages/interloper-core/src/interloper/destination/context.py index f8b181a8..0f745f0f 100644 --- a/packages/interloper-core/src/interloper/destination/context.py +++ b/packages/interloper-core/src/interloper/destination/context.py @@ -5,6 +5,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any +from interloper.errors import ConfigError from interloper.partitioning.base import Partition, PartitionWindow if TYPE_CHECKING: @@ -16,16 +17,69 @@ class IOContext: """Immutable context passed to :meth:`Destination.read` and :meth:`Destination.write`. - Carries the target asset, optional partition scope, and arbitrary metadata + Carries the target asset, optional partition or window, and arbitrary metadata so that destination implementations can resolve the correct storage location. - ``schema`` is the *effective* schema of the data being written or read — + ``schema`` is the *effective* schema of the data being written or read: the asset's declared schema when set, otherwise the schema inferred during - conform. Destinations use it for DDL, typed load jobs, and restoring - types on read. ``None`` when no schema could be resolved. + conform. Destinations use it for DDL, typed load jobs, and restoring + types on read. ``None`` when no schema could be resolved. + + A destination stores data per **partition**, ``None`` standing for the + unpartitioned whole; :attr:`partitions` and :meth:`slices` spell that out + for the three shapes ``partition_or_window`` can take, so no destination + branches on them itself. """ asset: Asset partition_or_window: Partition | PartitionWindow | None = None metadata: dict[str, Any] = field(default_factory=dict) schema: type[Schema] | None = None + + @property + def window(self) -> bool: + """Whether this context spans several partitions. + + Returns: + True for a partition window, False for one partition or the + unpartitioned whole. + """ + return isinstance(self.partition_or_window, PartitionWindow) + + @property + def partitions(self) -> list[Partition | None]: + """The partitions this context covers, in window order. + + Returns: + ``[None]`` for the unpartitioned whole, ``[partition]`` for one + partition, one entry per partition for a window. + """ + target = self.partition_or_window + if isinstance(target, PartitionWindow): + return list(target) + return [target] + + def slices(self, data: Any) -> list[tuple[Partition | None, Any]]: + """Pair each partition with its slice of *data*. + + A window's data is split on the asset's partition column through the + data's representation; a single partition, or the whole, receives the + data as is. Data whose representation is not recognised cannot be + split and is handed to every partition as is. + + Args: + data: The data being written, in its native representation. + + Returns: + ``(partition, slice)`` pairs, one per partition, in :attr:`partitions` order. + + Raises: + ConfigError: If a window is written for an asset that declares no + partitioning, since nothing says which column to slice on. + """ + if not self.window: + return [(partition, data) for partition in self.partitions] + if self.asset.partitioning is None: + raise ConfigError(f"Asset '{self.asset.key}' is not partitioned and cannot be written over a window") + column = self.asset.partitioning.column + return [(partition, partition.slice(data, column)) for partition in self.partitions if partition is not None] diff --git a/packages/interloper-core/src/interloper/destination/csv.py b/packages/interloper-core/src/interloper/destination/csv.py index da9145a5..f6d33a69 100644 --- a/packages/interloper-core/src/interloper/destination/csv.py +++ b/packages/interloper-core/src/interloper/destination/csv.py @@ -6,23 +6,22 @@ from pathlib import Path from typing import Any +from interloper.destination.base import Destination from interloper.destination.context import IOContext from interloper.destination.decorator import destination -from interloper.destination.partitioned import PartitionedDestination from interloper.errors import DataNotFoundError from interloper.partitioning.base import Partition from interloper.representation import Representation @destination(name="CSV") -class CSVDestination(PartitionedDestination): +class CSVDestination(Destination): """Destination that reads and writes CSV files on the local filesystem. Data is stored under ``{base_path}/{dataset}/{table}/data.csv`` (or ``{base_path}/{table}/data.csv`` when no dataset is set). Partitioned assets add a ``{column}={id}`` subdirectory; the partition - dispatch (including window splitting) comes from - :class:`PartitionedDestination`. + dispatch (including window splitting) is :class:`Destination`'s. Data is viewed as ``list[dict]`` records through its registered representation on write (each dict is a row; the keys of the first dict @@ -43,18 +42,18 @@ def _asset_path(self, context: IOContext) -> Path: """ return Path(self.base_path) / (context.asset.dataset or "") / context.asset.table - def _scope_path(self, context: IOContext, partition: Partition | None) -> Path: - """Return the data file path for a scope. + def _partition_path(self, context: IOContext, partition: Partition | None) -> Path: + """Return the data file path for a partition. Args: context: IO context whose asset supplies the dataset, table, and partition column. - partition: The scope's partition, or ``None`` for the unpartitioned + partition: The partition, or ``None`` for the unpartitioned whole. Returns: ``.../data.csv``, inside a ``{column}={id}`` subdirectory for - partition scopes. + partitions. """ base = self._asset_path(context) if partition is None: @@ -62,21 +61,21 @@ def _scope_path(self, context: IOContext, partition: Partition | None) -> Path: assert context.asset.partitioning return base / f"{context.asset.partitioning.column}={partition.id}" / "data.csv" - def _write_scope(self, context: IOContext, partition: Partition | None, data: Any) -> None: - """Write one scope's data as CSV (converted to records through its representation). + def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None: + """Write one partition's data as CSV (converted to records through its representation). Args: context: IO context whose asset supplies the dataset, table, and partition column. partition: The partition being written, or ``None`` for the unpartitioned whole. - data: The scope's slice of the data, in its native representation. + data: The partition's slice of the data, in its native representation. """ rows = Representation.of(data).to_records(data) - self._write_csv(self._scope_path(context, partition), rows) + self._write_csv(self._partition_path(context, partition), rows) - def _read_scope(self, context: IOContext, partition: Partition | None) -> list[dict[str, Any]]: - """Read one scope's CSV file. + def read_partition(self, context: IOContext, partition: Partition | None) -> list[dict[str, Any]]: + """Read one partition's CSV file. Args: context: IO context whose asset supplies the dataset, table, and @@ -87,7 +86,7 @@ def _read_scope(self, context: IOContext, partition: Partition | None) -> list[d Returns: Rows as a list of dicts (typed via the context schema when set). """ - return self._read_csv(self._scope_path(context, partition), context) + return self._read_csv(self._partition_path(context, partition), context) def _write_csv(self, file_path: Path, data: list[dict[str, Any]]) -> None: """Write a list of row dicts to a CSV file. diff --git a/packages/interloper-core/src/interloper/destination/database.py b/packages/interloper-core/src/interloper/destination/database.py index c755932f..94eaea77 100644 --- a/packages/interloper-core/src/interloper/destination/database.py +++ b/packages/interloper-core/src/interloper/destination/database.py @@ -9,10 +9,11 @@ from enum import Enum from typing import Any, ClassVar +from interloper.destination.base import Destination from interloper.destination.context import IOContext -from interloper.destination.partitioned import PartitionedDestination +from interloper.errors import ConfigError from interloper.normalizer import MaterializationStrategy -from interloper.partitioning.base import Partition, PartitionWindow +from interloper.partitioning.base import Partition from interloper.partitioning.time import TimePartition from interloper.representation import REPRESENTATIONS, Representation from interloper.resource.fields import SelectField @@ -32,7 +33,7 @@ class WriteDisposition(str, Enum): APPEND = "append" -class DatabaseDestination(PartitionedDestination): +class DatabaseDestination(Destination): """Abstract base class for database-backed destination implementations. Provides the partition-aware write/read dispatch logic that is common to any @@ -199,11 +200,10 @@ def partition_row_counts(self, context: IOContext) -> dict[str, int]: context: IO context whose asset supplies the table, dataset, and partition column. """ - assert context.asset.partitioning is not None return self._count_by_partition( context.asset.table, context.asset.dataset or None, - context.asset.partitioning.column, + self._partition_column(context), ) # -- Data conversion ------------------------------------------------------- @@ -241,62 +241,101 @@ def _insert_data(self, table: str, schema: str | None, data: Any, context: IOCon self._insert(table, schema, rows) def write(self, context: IOContext, data: Any) -> None: - """Write data to the database table. + """Write data to the database table, a window as one batch. - Overrides the :class:`PartitionedDestination` template: database - backends delete per scope but insert the whole batch once (rows - carry the partition column), instead of storing per-partition - slices. With ``REPLACE``, matching rows are deleted before - inserting; with ``APPEND``, rows are inserted without any prior - deletion. + Rows carry the partition column, so a database need not store per + partition: a window clears every partition it covers and inserts the + whole batch once, instead of one :meth:`write_partition` per partition + (one load job rather than one per day). A single partition, or the + whole, is one :meth:`write_partition`. The write-time schema strategy + is applied to the data once, before either path. Args: - context: IO context carrying the target asset, the partition scope, + context: IO context carrying the target asset, the partition or window, and the effective schema. data: The data to write, in its native representation. """ - table = context.asset.table - schema = context.asset.dataset or None - if is_empty(data): return - data = self._apply_materialization_strategy(data, context) + self._warn_missing_partition_column(data, context) + if not context.window: + self.write_partition(context, context.partitions[0], data) + return + table, schema = context.asset.table, context.asset.dataset or None + with self._transaction(): + self._clear(table, schema, context, context.partitions) + self._insert_data(table, schema, data, context) - if context.partition_or_window is not None and context.asset.partitioning is not None: - col = context.asset.partitioning.column - columns = Representation.of(data).columns(data) - if columns and col not in columns: - warnings.warn( - f"Partition column '{col}' not found in data for asset " - f"'{context.asset.key}'. Columns present: {sorted(columns)}. " - f"Downstream reads by partition will fail.", - UserWarning, - stacklevel=2, - ) + def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None: + """Replace one partition's rows: clear it, then insert the data. - replacing = self.write_disposition is WriteDisposition.REPLACE + With ``REPLACE`` the partition's rows are deleted first; with ``APPEND`` + the data is inserted on top of whatever is there. + Args: + context: IO context carrying the target asset and the effective schema. + partition: The partition being stored, or ``None`` for the whole table. + data: The data to store, in its native representation. + """ + table, schema = context.asset.table, context.asset.dataset or None with self._transaction(): - if context.partition_or_window is None: - if replacing: - self._delete_all(table, schema) + self._clear(table, schema, context, [partition]) + self._insert_data(table, schema, data, context) - elif isinstance(context.partition_or_window, PartitionWindow): - assert context.asset.partitioning - col = context.asset.partitioning.column - if replacing: - for partition in context.partition_or_window: - self._delete_scope(table, schema, col, partition) + def _clear(self, table: str, schema: str | None, context: IOContext, partitions: list[Partition | None]) -> None: + """Delete the rows of the given partitions when the write disposition replaces. + Args: + table: Target table name. + schema: Database schema holding the table, or ``None`` for the backend default. + context: IO context whose asset supplies the partition column. + partitions: The partitions to clear; ``None`` clears the whole table. + """ + if self.write_disposition is not WriteDisposition.REPLACE: + return + for partition in partitions: + if partition is None: + self._delete_all(table, schema) else: - assert isinstance(context.partition_or_window, Partition) - assert context.asset.partitioning - col = context.asset.partitioning.column - if replacing: - self._delete_scope(table, schema, col, context.partition_or_window) + self._clear_partition(table, schema, self._partition_column(context), partition) - self._insert_data(table, schema, data, context) + def _warn_missing_partition_column(self, data: Any, context: IOContext) -> None: + """Warn when partitioned data lacks its partition column, since reads by partition would find nothing. + + Args: + data: The data about to be written. + context: IO context whose asset supplies the partition column. + """ + if context.partition_or_window is None or context.asset.partitioning is None: + return + column = context.asset.partitioning.column + columns = Representation.of(data).columns(data) + if columns and column not in columns: + warnings.warn( + f"Partition column '{column}' not found in data for asset " + f"'{context.asset.key}'. Columns present: {sorted(columns)}. " + f"Downstream reads by partition will fail.", + UserWarning, + stacklevel=3, + ) + + @staticmethod + def _partition_column(context: IOContext) -> str: + """The asset's partition column, which a partitioned write or read guarantees exists. + + Args: + context: IO context whose asset is partitioned. + + Returns: + The partition column name. + + Raises: + ConfigError: If the asset declares no partitioning. + """ + if context.asset.partitioning is None: + raise ConfigError(f"Asset '{context.asset.key}' is not partitioned") + return context.asset.partitioning.column def _apply_materialization_strategy(self, data: Any, context: IOContext) -> Any: """Enforce this backend's write-time schema strategy. @@ -326,7 +365,7 @@ def _apply_materialization_strategy(self, data: Any, context: IOContext) -> Any: conformer.validate(data, context.schema, strict=True) return data - def _delete_scope(self, table: str, schema: str | None, column: str, partition: Partition) -> None: + def _clear_partition(self, table: str, schema: str | None, column: str, partition: Partition) -> None: """Delete one partition's rows: by bounds for a time partition, by id otherwise. A time partition's rows may carry values anywhere inside the period @@ -345,8 +384,8 @@ def _delete_scope(self, table: str, schema: str | None, column: str, partition: else: self._delete_partition(table, schema, column, partition.id) - def _read_scope(self, context: IOContext, partition: Partition | None) -> Any: - """Load one scope from the database table. + def read_partition(self, context: IOContext, partition: Partition | None) -> Any: + """Load one partition from the database table. Args: context: IO context whose asset supplies the table, dataset, and @@ -354,14 +393,13 @@ def _read_scope(self, context: IOContext, partition: Partition | None) -> Any: partition: The partition to load, or ``None`` for the whole table. Returns: - The scope's rows, materialized into the read representation. + The partition's rows, materialized into the read representation. """ table = context.asset.table schema = context.asset.dataset or None if partition is None: return self._from_rows(self._select_all(table, schema)) - assert context.asset.partitioning - column = context.asset.partitioning.column + column = self._partition_column(context) if isinstance(partition, TimePartition): start, end = partition.bounds return self._from_rows(self._select_partition_range(table, schema, column, start, end)) diff --git a/packages/interloper-core/src/interloper/destination/file.py b/packages/interloper-core/src/interloper/destination/file.py index dcecaa05..28961148 100644 --- a/packages/interloper-core/src/interloper/destination/file.py +++ b/packages/interloper-core/src/interloper/destination/file.py @@ -6,22 +6,21 @@ from pathlib import Path from typing import Any +from interloper.destination.base import Destination from interloper.destination.context import IOContext from interloper.destination.decorator import destination -from interloper.destination.partitioned import PartitionedDestination from interloper.errors import DataNotFoundError from interloper.partitioning.base import Partition @destination(name="File") -class FileDestination(PartitionedDestination): +class FileDestination(Destination): """Destination that reads and writes pickle files on the local filesystem. Data is stored under ``{base_path}/{dataset}/{table}/data.pkl`` (or ``{base_path}/{table}/data.pkl`` when no dataset is set). Partitioned assets add a ``{column}={id}`` subdirectory; the partition - dispatch (including window splitting) comes from - :class:`PartitionedDestination`. + dispatch (including window splitting) is :class:`Destination`'s. Unlike :class:`~interloper.destination.csv.CSVDestination` this stores whatever the asset returned, tabular or not — so a window write of a @@ -43,18 +42,18 @@ def _asset_path(self, context: IOContext) -> Path: """ return Path(self.base_path) / (context.asset.dataset or "") / context.asset.table - def _scope_path(self, context: IOContext, partition: Partition | None) -> Path: - """Return the data file path for a scope. + def _partition_path(self, context: IOContext, partition: Partition | None) -> Path: + """Return the data file path for a partition. Args: context: IO context whose asset supplies the dataset, table, and partition column. - partition: The scope's partition, or ``None`` for the unpartitioned + partition: The partition, or ``None`` for the unpartitioned whole. Returns: ``.../data.pkl``, inside a ``{column}={id}`` subdirectory for - partition scopes. + partitions. """ base = self._asset_path(context) if partition is None: @@ -62,23 +61,23 @@ def _scope_path(self, context: IOContext, partition: Partition | None) -> Path: assert context.asset.partitioning return base / f"{context.asset.partitioning.column}={partition.id}" / "data.pkl" - def _write_scope(self, context: IOContext, partition: Partition | None, data: Any) -> None: - """Pickle one scope's data to its file. + def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None: + """Pickle one partition's data to its file. Args: context: IO context whose asset supplies the dataset, table, and partition column. partition: The partition being written, or ``None`` for the unpartitioned whole. - data: The scope's slice of the data, stored as-is. + data: The partition's slice of the data, stored as-is. """ - path = self._scope_path(context, partition) + path = self._partition_path(context, partition) path.parent.mkdir(parents=True, exist_ok=True) with path.open("wb") as f: pickle.dump(data, f) - def _read_scope(self, context: IOContext, partition: Partition | None) -> Any: - """Unpickle one scope's file. + def read_partition(self, context: IOContext, partition: Partition | None) -> Any: + """Unpickle one partition's file. Args: context: IO context whose asset supplies the dataset, table, and @@ -90,9 +89,9 @@ def _read_scope(self, context: IOContext, partition: Partition | None) -> Any: The deserialized data. Raises: - DataNotFoundError: If the scope's data file does not exist. + DataNotFoundError: If the partition's data file does not exist. """ - path = self._scope_path(context, partition) + path = self._partition_path(context, partition) if not path.exists(): raise DataNotFoundError(f"No data file for '{context.asset}': {path}") with path.open("rb") as f: @@ -102,7 +101,7 @@ def partition_row_counts(self, context: IOContext) -> dict[str, int]: """Return row counts grouped by partition by scanning pickle files on disk. Counts items (``len(data)`` for lists, ``1`` otherwise), since a - pickled scope need not be tabular. + pickled partition need not be tabular. Args: context: IO context whose asset supplies the dataset, table, and diff --git a/packages/interloper-core/src/interloper/destination/memory.py b/packages/interloper-core/src/interloper/destination/memory.py index d01fa42d..be2fad26 100644 --- a/packages/interloper-core/src/interloper/destination/memory.py +++ b/packages/interloper-core/src/interloper/destination/memory.py @@ -4,39 +4,39 @@ from typing import Any, ClassVar +from interloper.destination.base import Destination from interloper.destination.context import IOContext from interloper.destination.decorator import destination -from interloper.destination.partitioned import PartitionedDestination from interloper.errors import DataNotFoundError from interloper.partitioning.base import Partition, PartitionConfig @destination(name="Memory") -class MemoryDestination(PartitionedDestination): +class MemoryDestination(Destination): """Destination that stores data in a class-level dict keyed by ``{dataset}/{table}/{partition}``. All instances share a single ``_storage`` dict so data written by one asset is visible to others. The partition dispatch (including window - splitting) comes from :class:`PartitionedDestination`. Call - :meth:`clear` between test runs. + splitting) is :class:`Destination`'s. Call :meth:`clear` between test + runs. """ _storage: ClassVar[dict[str, Any]] = {} - def _write_scope(self, context: IOContext, partition: Partition | None, data: Any) -> None: - """Store one scope's data under its path-style key. + def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None: + """Store one partition's data under its path-style key. Args: context: IO context whose asset supplies the table, dataset, and partitioning. partition: The partition being stored, or ``None`` for the unpartitioned whole. - data: The scope's data, stored as-is. + data: The partition's data, stored as-is. """ - self._storage[self._scope_key(context, partition)] = data + self._storage[self._partition_key(context, partition)] = data - def _read_scope(self, context: IOContext, partition: Partition | None) -> Any: - """Retrieve one scope's data. + def read_partition(self, context: IOContext, partition: Partition | None) -> Any: + """Retrieve one partition's data. Args: context: IO context whose asset supplies the table, dataset, and @@ -50,18 +50,18 @@ def _read_scope(self, context: IOContext, partition: Partition | None) -> Any: Raises: DataNotFoundError: If no data exists for the resolved key. """ - key = self._scope_key(context, partition) + key = self._partition_key(context, partition) if key not in self._storage: raise DataNotFoundError(f"No data found in memory for: {key}") return self._storage[key] - def _scope_key(self, context: IOContext, partition: Partition | None) -> str: - """Build the storage key for a scope. + def _partition_key(self, context: IOContext, partition: Partition | None) -> str: + """Build the storage key for a partition. Args: context: IO context whose asset supplies the table, dataset, and partitioning. - partition: The scope's partition, or ``None`` for the + partition: The partition, or ``None`` for the unpartitioned whole. Returns: diff --git a/packages/interloper-core/src/interloper/destination/partitioned.py b/packages/interloper-core/src/interloper/destination/partitioned.py deleted file mode 100644 index 2b6c58d7..00000000 --- a/packages/interloper-core/src/interloper/destination/partitioned.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Partition-aware destination template: the three-way scope dispatch, once.""" - -from __future__ import annotations - -from typing import Any - -from interloper.destination.base import Destination -from interloper.destination.context import IOContext -from interloper.partitioning.base import Partition, PartitionWindow - - -class PartitionedDestination(Destination): - """Destination base implementing the partition dispatch once. - - The ``None / Partition / PartitionWindow`` branching used to be - hand-rolled by every destination — the source of the window-duplication - bug where each partition directory received the *full* dataset. - Subclasses implement two scope hooks and are partition-correct by - construction: - - - :meth:`_write_scope` — store data for one scope (``partition=None`` - means the unpartitioned whole). - - :meth:`_read_scope` — load one scope. - - Window writes are split per partition through the data's registered - representation; window reads return one result per partition. Data whose - representation cannot be recognized is passed to the write hook as-is - (it cannot be split). - - Backends whose scoping semantics differ from per-scope storage may - override :meth:`write` or :meth:`read` wholesale instead of implementing - the corresponding hook — :class:`DatabaseDestination` does this for - writes (scoped deletes followed by a single insert). - """ - - def _write_scope(self, context: IOContext, partition: Partition | None, data: Any) -> None: - """Store *data* for a single scope. - - Args: - context: IO context carrying the target asset and the effective schema. - partition: The partition being stored, or ``None`` for the - unpartitioned whole. - data: The scope's slice of the data to store. - - Raises: - NotImplementedError: Subclasses implement this or override ``write``. - """ - raise NotImplementedError(f"{type(self).__name__} must implement _write_scope() or override write().") - - def _read_scope(self, context: IOContext, partition: Partition | None) -> Any: - """Load a single scope. - - Args: - context: IO context carrying the target asset and the effective schema. - partition: The partition to load, or ``None`` for the unpartitioned - whole. - - Raises: - NotImplementedError: Subclasses implement this or override ``read``. - """ - raise NotImplementedError(f"{type(self).__name__} must implement _read_scope() or override read().") - - def write(self, context: IOContext, data: Any) -> None: - """Write data, splitting partition windows per partition. - - Args: - context: IO context carrying the target asset, the partition scope, - and the effective schema. - data: The data to write, in its native representation. - """ - scope = context.partition_or_window - if scope is None: - self._write_scope(context, None, data) - elif isinstance(scope, PartitionWindow): - assert context.asset.partitioning - column = context.asset.partitioning.column - for partition in scope: - self._write_scope(context, partition, partition.slice(data, column)) - else: - assert isinstance(scope, Partition) - self._write_scope(context, scope, data) - - def read(self, context: IOContext) -> Any: - """Read data for the context's scope. - - Args: - context: IO context carrying the target asset, the partition scope, - and the effective schema. - - Returns: - The scope's data; partition windows return one result per - partition, in window order. - """ - scope = context.partition_or_window - if scope is None: - return self._read_scope(context, None) - if isinstance(scope, PartitionWindow): - return [self._read_scope(context, partition) for partition in scope] - assert isinstance(scope, Partition) - return self._read_scope(context, scope) diff --git a/packages/interloper-core/tests/destination/test_base.py b/packages/interloper-core/tests/destination/test_base.py index 370958a9..f71b7e12 100644 --- a/packages/interloper-core/tests/destination/test_base.py +++ b/packages/interloper-core/tests/destination/test_base.py @@ -4,12 +4,16 @@ # component class declares a relation, and the collector needs it as a real # class, not a lazy string. -from typing import Any +import datetime +from typing import Any, ClassVar import pytest import interloper as il +from interloper.destination import IOContext from interloper.destination.base import DestinationDefinition +from interloper.partitioning.base import Partition +from interloper.partitioning.time import TimePartition, TimePartitionWindow class FakeConnection(il.Connection): @@ -113,3 +117,133 @@ def write(self, context: Any, data: Any) -> None: # pragma: no cover with pytest.raises(TypeError, match="not a @fetch_field_provider"): Unmarked.definition() + + +# -- Partition dispatch ---------------------------------------------------------- + +class RecordingPartitions(il.Destination): + """Destination capturing every partition-hook call.""" + + calls: ClassVar[list[tuple[str, Any, Any]]] = [] + + def model_post_init(self, context: Any) -> None: + super().model_post_init(context) + object.__setattr__(self, "calls", []) + + def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None: + self.calls.append(("write", partition.id if partition else None, data)) + + def read_partition(self, context: IOContext, partition: Partition | None) -> Any: + self.calls.append(("read", partition.id if partition else None, None)) + return {"partition": partition.id if partition else None} + + +@il.asset(partitioning=il.TimePartitionConfig(column="date")) +def partitioned_asset(context: il.ExecutionContext) -> list: # noqa: D103 + return [] + + +@il.asset +def plain_asset() -> list: # noqa: D103 + return [] + + +def io_context(asset: il.Asset, partition_or_window=None) -> IOContext: # noqa: D103 + return IOContext(asset=asset, partition_or_window=partition_or_window) + + +class TestWriteDispatch: + """The three-way write dispatch with window splitting.""" + + def test_unpartitioned_write_is_one_call(self): + destination = RecordingPartitions(id="d") + destination.write(io_context(plain_asset()), [{"a": 1}]) + assert destination.calls == [("write", None, [{"a": 1}])] + + def test_partition_write_passes_data_unsplit(self): + destination = RecordingPartitions(id="d") + rows = [{"date": "2024-01-01"}, {"date": "2024-01-02"}] + destination.write(io_context(partitioned_asset(), TimePartition(datetime.date(2024, 1, 1))), rows) + assert destination.calls == [("write", "2024-01-01", rows)] + + def test_window_write_splits_per_partition(self): + destination = RecordingPartitions(id="d") + rows = [ + {"date": "2024-01-01", "v": 1}, + {"date": "2024-01-02", "v": 2}, + {"date": "2024-01-02", "v": 3}, + ] + window = TimePartitionWindow(datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)) + destination.write(io_context(partitioned_asset(), window), rows) + by_partition = {partition: data for kind, partition, data in destination.calls} + assert by_partition["2024-01-01"] == [{"date": "2024-01-01", "v": 1}] + assert by_partition["2024-01-02"] == [{"date": "2024-01-02", "v": 2}, {"date": "2024-01-02", "v": 3}] + + def test_monthly_window_slices_rows_by_period(self): + # Rows carry daily dates; each monthly partition's slice is its whole + # month, which id equality on the period start would miss entirely. + @il.asset(partitioning=il.TimePartitionConfig(column="date", granularity=il.TimeGranularity.MONTH)) + def monthly(context: il.ExecutionContext) -> list: + return [] + + destination = RecordingPartitions(id="d") + rows = [ + {"date": "2024-01-15", "v": 1}, + {"date": "2024-02-10", "v": 2}, + {"date": "2024-02-20", "v": 3}, + ] + window = TimePartitionWindow( + datetime.date(2024, 1, 1), datetime.date(2024, 2, 1), il.TimeGranularity.MONTH + ) + destination.write(io_context(monthly(), window), rows) + by_partition = {partition: data for kind, partition, data in destination.calls} + assert by_partition["2024-01"] == [{"date": "2024-01-15", "v": 1}] + assert by_partition["2024-02"] == [{"date": "2024-02-10", "v": 2}, {"date": "2024-02-20", "v": 3}] + + def test_window_write_splits_dataframes_natively(self): + pd = pytest.importorskip("pandas") + + destination = RecordingPartitions(id="d") + df = pd.DataFrame([{"date": "2024-01-01", "v": 1}, {"date": "2024-01-02", "v": 2}]) + window = TimePartitionWindow(datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)) + destination.write(io_context(partitioned_asset(), window), df) + for _, _, data in destination.calls: + assert isinstance(data, pd.DataFrame) + assert len(data) == 1 + + def test_window_write_passes_unsplittable_data_as_is(self): + destination = RecordingPartitions(id="d") + sentinel = object() + window = TimePartitionWindow(datetime.date(2024, 1, 1), datetime.date(2024, 1, 1)) + destination.write(io_context(partitioned_asset(), window), sentinel) + assert destination.calls == [("write", "2024-01-01", sentinel)] + + +class TestReadDispatch: + """The three-way read dispatch.""" + + def test_unpartitioned_read(self): + assert RecordingPartitions(id="d").read(io_context(plain_asset())) == {"partition": None} + + def test_partition_read(self): + partition = TimePartition(datetime.date(2024, 1, 2)) + result = RecordingPartitions(id="d").read(io_context(partitioned_asset(), partition)) + assert result == {"partition": "2024-01-02"} + + def test_window_read_returns_one_result_per_partition(self): + window = TimePartitionWindow(datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)) + result = RecordingPartitions(id="d").read(io_context(partitioned_asset(), window)) + assert {r["partition"] for r in result} == {"2024-01-01", "2024-01-02"} + + +class TestHookContract: + """The partition hooks are the contract; the templates say which one is missing.""" + + def test_missing_hooks_raise_naming_the_hook(self): + class Bare(il.Destination): + pass + + with pytest.raises(NotImplementedError, match="Bare must implement write_partition"): + Bare(id="b").write(io_context(plain_asset()), []) + with pytest.raises(NotImplementedError, match="Bare must implement read_partition"): + Bare(id="b").read(io_context(plain_asset())) diff --git a/packages/interloper-core/tests/destination/test_context.py b/packages/interloper-core/tests/destination/test_context.py new file mode 100644 index 00000000..d320c89b --- /dev/null +++ b/packages/interloper-core/tests/destination/test_context.py @@ -0,0 +1,59 @@ +"""Tests for ``interloper.destination.context``: the partitions an IO context spells out.""" + +import datetime + +import pytest + +import interloper as il +from interloper.destination import IOContext +from interloper.errors import ConfigError +from interloper.partitioning.time import TimePartition, TimePartitionWindow + + +@il.asset(partitioning=il.TimePartitionConfig(column="date")) +def daily(context: il.ExecutionContext) -> list: # noqa: D103 + return [] + + +@il.asset +def whole() -> list: # noqa: D103 + return [] + + +class TestPartitions: + def test_the_unpartitioned_whole_is_one_partition_none(self): + context = IOContext(asset=whole()) + assert context.partitions == [None] + assert context.window is False + + def test_a_partition_is_itself(self): + partition = TimePartition(datetime.date(2026, 9, 7)) + context = IOContext(asset=daily(), partition_or_window=partition) + assert context.partitions == [partition] + assert context.window is False + + def test_a_window_is_its_partitions_in_window_order(self): + window = TimePartitionWindow(datetime.date(2026, 9, 7), datetime.date(2026, 9, 9)) + context = IOContext(asset=daily(), partition_or_window=window) + assert [p.id if p else None for p in context.partitions] == ["2026-09-09", "2026-09-08", "2026-09-07"] + assert context.window is True + + +class TestSlices: + def test_a_single_partition_receives_the_data_whole(self): + rows = [{"date": "2026-09-07"}, {"date": "2026-09-08"}] + assert IOContext(asset=whole()).slices(rows) == [(None, rows)] + partition = TimePartition(datetime.date(2026, 9, 7)) + assert IOContext(asset=daily(), partition_or_window=partition).slices(rows) == [(partition, rows)] + + def test_a_window_slices_on_the_partition_column(self): + rows = [{"date": "2026-09-07", "v": 1}, {"date": "2026-09-08", "v": 2}, {"date": "2026-09-08", "v": 3}] + window = TimePartitionWindow(datetime.date(2026, 9, 7), datetime.date(2026, 9, 8)) + pairs = IOContext(asset=daily(), partition_or_window=window).slices(rows) + sliced = {p.id if p else None: chunk for p, chunk in pairs} + assert sliced == {"2026-09-07": [rows[0]], "2026-09-08": rows[1:]} + + def test_a_window_over_an_unpartitioned_asset_is_a_config_error(self): + window = TimePartitionWindow(datetime.date(2026, 9, 7), datetime.date(2026, 9, 8)) + with pytest.raises(ConfigError, match="not partitioned"): + IOContext(asset=whole(), partition_or_window=window).slices([]) diff --git a/packages/interloper-core/tests/destination/test_csv.py b/packages/interloper-core/tests/destination/test_csv.py index 4a660666..2f6c0047 100644 --- a/packages/interloper-core/tests/destination/test_csv.py +++ b/packages/interloper-core/tests/destination/test_csv.py @@ -53,7 +53,7 @@ def test_dataframe_write_accepted(self, tmp_path): dest.write(context, pd.DataFrame([{"a": 1}])) assert dest.read(context) == [{"a": "1"}] - def test_missing_scope_raises_data_not_found_error(self, tmp_path): + def test_missing_partition_raises_data_not_found_error(self, tmp_path): dest = CSVDestination(id="csv", base_path=str(tmp_path)) context = IOContext(asset=plain_asset()) with pytest.raises(DataNotFoundError, match="Data file not found"): diff --git a/packages/interloper-core/tests/destination/test_file.py b/packages/interloper-core/tests/destination/test_file.py index 0bb72199..8e299cee 100644 --- a/packages/interloper-core/tests/destination/test_file.py +++ b/packages/interloper-core/tests/destination/test_file.py @@ -79,7 +79,7 @@ def test_a_rewrite_replaces_the_previous_data(self, tmp_path: Path): class TestWindows: - """Window writes are split per partition by ``PartitionedDestination``.""" + """Window writes are split per partition by ``Destination``.""" def test_tabular_rows_are_split_by_the_partition_column(self, tmp_path: Path): dest = FileDestination(id="file", base_path=str(tmp_path)) diff --git a/packages/interloper-core/tests/destination/test_partitioned.py b/packages/interloper-core/tests/destination/test_partitioned.py deleted file mode 100644 index 10dacaf0..00000000 --- a/packages/interloper-core/tests/destination/test_partitioned.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Tests for ``interloper.destination.partitioned``.""" - -import datetime -from typing import Any, ClassVar - -import pytest - -import interloper as il -from interloper.destination import IOContext, PartitionedDestination -from interloper.partitioning.base import Partition -from interloper.partitioning.time import TimePartition, TimePartitionWindow - - -class RecordingScopes(PartitionedDestination): - """Destination capturing every scope-hook call.""" - - calls: ClassVar[list[tuple[str, Any, Any]]] = [] - - def model_post_init(self, context: Any) -> None: - super().model_post_init(context) - object.__setattr__(self, "calls", []) - - def _write_scope(self, context: IOContext, partition: Partition | None, data: Any) -> None: - self.calls.append(("write", partition.id if partition else None, data)) - - def _read_scope(self, context: IOContext, partition: Partition | None) -> Any: - self.calls.append(("read", partition.id if partition else None, None)) - return {"scope": partition.id if partition else None} - - -@il.asset(partitioning=il.TimePartitionConfig(column="date")) -def partitioned_asset(context: il.ExecutionContext) -> list: # noqa: D103 - return [] - - -@il.asset -def plain_asset() -> list: # noqa: D103 - return [] - - -def io_context(asset: il.Asset, scope=None) -> IOContext: # noqa: D103 - return IOContext(asset=asset, partition_or_window=scope) - - -class TestWriteDispatch: - """The three-way write dispatch with window splitting.""" - - def test_unpartitioned_write_is_one_scope(self): - destination = RecordingScopes(id="d") - destination.write(io_context(plain_asset()), [{"a": 1}]) - assert destination.calls == [("write", None, [{"a": 1}])] - - def test_partition_write_passes_data_unsplit(self): - destination = RecordingScopes(id="d") - rows = [{"date": "2024-01-01"}, {"date": "2024-01-02"}] - destination.write(io_context(partitioned_asset(), TimePartition(datetime.date(2024, 1, 1))), rows) - assert destination.calls == [("write", "2024-01-01", rows)] - - def test_window_write_splits_per_partition(self): - destination = RecordingScopes(id="d") - rows = [ - {"date": "2024-01-01", "v": 1}, - {"date": "2024-01-02", "v": 2}, - {"date": "2024-01-02", "v": 3}, - ] - window = TimePartitionWindow(datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)) - destination.write(io_context(partitioned_asset(), window), rows) - by_scope = {scope: data for kind, scope, data in destination.calls} - assert by_scope["2024-01-01"] == [{"date": "2024-01-01", "v": 1}] - assert by_scope["2024-01-02"] == [{"date": "2024-01-02", "v": 2}, {"date": "2024-01-02", "v": 3}] - - def test_monthly_window_slices_rows_by_period(self): - # Rows carry daily dates; each monthly partition's slice is its whole - # month, which id equality on the period start would miss entirely. - @il.asset(partitioning=il.TimePartitionConfig(column="date", granularity=il.TimeGranularity.MONTH)) - def monthly(context: il.ExecutionContext) -> list: - return [] - - destination = RecordingScopes(id="d") - rows = [ - {"date": "2024-01-15", "v": 1}, - {"date": "2024-02-10", "v": 2}, - {"date": "2024-02-20", "v": 3}, - ] - window = TimePartitionWindow( - datetime.date(2024, 1, 1), datetime.date(2024, 2, 1), il.TimeGranularity.MONTH - ) - destination.write(io_context(monthly(), window), rows) - by_scope = {scope: data for kind, scope, data in destination.calls} - assert by_scope["2024-01"] == [{"date": "2024-01-15", "v": 1}] - assert by_scope["2024-02"] == [{"date": "2024-02-10", "v": 2}, {"date": "2024-02-20", "v": 3}] - - def test_window_write_splits_dataframes_natively(self): - pd = pytest.importorskip("pandas") - - destination = RecordingScopes(id="d") - df = pd.DataFrame([{"date": "2024-01-01", "v": 1}, {"date": "2024-01-02", "v": 2}]) - window = TimePartitionWindow(datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)) - destination.write(io_context(partitioned_asset(), window), df) - for _, _, data in destination.calls: - assert isinstance(data, pd.DataFrame) - assert len(data) == 1 - - def test_window_write_passes_unsplittable_data_as_is(self): - destination = RecordingScopes(id="d") - sentinel = object() - window = TimePartitionWindow(datetime.date(2024, 1, 1), datetime.date(2024, 1, 1)) - destination.write(io_context(partitioned_asset(), window), sentinel) - assert destination.calls == [("write", "2024-01-01", sentinel)] - - -class TestReadDispatch: - """The three-way read dispatch.""" - - def test_unpartitioned_read(self): - assert RecordingScopes(id="d").read(io_context(plain_asset())) == {"scope": None} - - def test_partition_read(self): - result = RecordingScopes(id="d").read(io_context(partitioned_asset(), TimePartition(datetime.date(2024, 1, 2)))) - assert result == {"scope": "2024-01-02"} - - def test_window_read_returns_one_result_per_partition(self): - window = TimePartitionWindow(datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)) - result = RecordingScopes(id="d").read(io_context(partitioned_asset(), window)) - assert {r["scope"] for r in result} == {"2024-01-01", "2024-01-02"} - - -class TestHookContract: - """Hooks are required unless the corresponding template is overridden.""" - - def test_missing_hooks_raise_with_guidance(self): - class Bare(PartitionedDestination): - pass - - with pytest.raises(NotImplementedError, match="_write_scope.*or override write"): - Bare(id="b").write(io_context(plain_asset()), []) - with pytest.raises(NotImplementedError, match="_read_scope.*or override read"): - Bare(id="b").read(io_context(plain_asset())) diff --git a/packages/interloper-google-cloud/src/interloper_google_cloud/gcs/destination.py b/packages/interloper-google-cloud/src/interloper_google_cloud/gcs/destination.py index 44d627b1..09ac17cf 100644 --- a/packages/interloper-google-cloud/src/interloper_google_cloud/gcs/destination.py +++ b/packages/interloper-google-cloud/src/interloper_google_cloud/gcs/destination.py @@ -11,7 +11,7 @@ from google.cloud.exceptions import NotFound from google.oauth2 import service_account from interloper.destination import IOContext, destination -from interloper.destination.partitioned import PartitionedDestination +from interloper.destination.base import Destination from interloper.errors import DataNotFoundError from interloper.partitioning import Partition from interloper.representation import Representation @@ -21,7 +21,7 @@ from interloper_google_cloud.connection import GoogleCloudConnection from interloper_google_cloud.gcs.formats import FORMATS, FileFormat -# Custom blob metadata key carrying the scope's row count, so +# Custom blob metadata key carrying the partition's row count, so # partition_row_counts introspects from a single list call without downloads. _ROW_COUNT_METADATA_KEY = "row_count" @@ -32,10 +32,10 @@ icon="icon:gcs", tags=["Cloud"], ) -class GCSDestination(PartitionedDestination): +class GCSDestination(Destination): """Google Cloud Storage destination. - Writes one object per scope in a hive-partitioned layout:: + Writes one object per partition in a hive-partitioned layout:: gs://{bucket}/{prefix}/{dataset}/{table}/data.{ext} gs://{bucket}/{prefix}/{dataset}/{table}/{column}={partition}/data.{ext} @@ -43,7 +43,7 @@ class GCSDestination(PartitionedDestination): Following the hive convention, the partition column lives in the *path only*: it is dropped from partitioned file contents on write (external readers like BigQuery external tables and DuckDB reject a duplicate - partition column) and re-injected from the partition scope on read, so + partition column) and re-injected from the partition on read, so interloper round-trips stay lossless. """ @@ -110,15 +110,15 @@ def _asset_prefix(self, context: IOContext) -> str: return "/".join(part.strip("/") for part in parts if part and part.strip("/")) def _blob_name(self, context: IOContext, partition: Partition | None) -> str: - """Build the object name for a scope. + """Build the object name for a partition. Args: context: The IO context naming the asset. - partition: The partition being addressed, when scoped. + partition: The partition being addressed, or ``None`` for the whole. Returns: ``.../data.{ext}``, inside a ``{column}={id}`` segment for - partition scopes. + partitions. """ parts = [self._asset_prefix(context)] if partition is not None: @@ -128,14 +128,14 @@ def _blob_name(self, context: IOContext, partition: Partition | None) -> str: return "/".join(parts) def _effective_specs(self, context: IOContext, partition: Partition | None) -> list[FieldSpec] | None: - """Return the field specs for a scope's file contents. + """Return the field specs for a partition's file contents. - Partition scopes exclude the partition column (its value lives in the + Partitions exclude the partition column (its value lives in the object path). Args: context: The IO context carrying the schema. - partition: The partition being addressed, when scoped. + partition: The partition being addressed, or ``None`` for the whole. Returns: The specs, or ``None`` when the context carries no schema. @@ -148,17 +148,17 @@ def _effective_specs(self, context: IOContext, partition: Partition | None) -> l specs = [spec for spec in specs if spec.name != context.asset.partitioning.column] return specs - # -- PartitionedDestination hooks -------------------------------------------- + # -- Partition hooks ------------------------------------------------------- - def _write_scope(self, context: IOContext, partition: Partition | None, data: Any) -> None: - """Serialize one scope's data and upload it, overwriting the scope's object. + def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None: + """Serialize one partition's data and upload it, overwriting its object. The row count is stamped as blob metadata so introspection never has to download data. Args: context: The IO context naming the asset. - partition: The partition being written, when scoped. + partition: The partition being written, or ``None`` for the whole. data: The rows to write. """ @@ -173,22 +173,22 @@ def _write_scope(self, context: IOContext, partition: Partition | None, data: An blob.metadata = {_ROW_COUNT_METADATA_KEY: str(len(rows))} blob.upload_from_string(payload, content_type=self._format.content_type) - def _read_scope(self, context: IOContext, partition: Partition | None) -> list[dict[str, Any]]: - """Download and parse one scope's object. + def read_partition(self, context: IOContext, partition: Partition | None) -> list[dict[str, Any]]: + """Download and parse one partition's object. - The partition column is re-injected from the scope, and rows are + The partition column is re-injected from the partition, and rows are reconciled against the context schema when one is set (restoring the declared types — text formats read everything back as strings). Args: context: The IO context carrying the schema. - partition: The partition being read, when scoped. + partition: The partition being read, or ``None`` for the whole. Returns: Rows as a list of dicts. Raises: - DataNotFoundError: If the scope's object does not exist. + DataNotFoundError: If the partition's object does not exist. """ name = self._blob_name(context, partition) try: @@ -228,10 +228,10 @@ def partition_row_counts(self, context: IOContext) -> dict[str, int]: counts: dict[str, int] = {} for blob in self.client.list_blobs(self.bucket, prefix=prefix): - scope = blob.name[len(prefix) :].split("/", 1)[0] - if not scope.startswith(f"{column}="): + segment = blob.name[len(prefix) :].split("/", 1)[0] + if not segment.startswith(f"{column}="): continue - value = scope.split("=", 1)[1] + value = segment.split("=", 1)[1] row_count = (blob.metadata or {}).get(_ROW_COUNT_METADATA_KEY) if row_count is None: row_count = len(self._format.deserialize(blob.download_as_bytes())) diff --git a/plugins/interloper/skills/interloper-destination/SKILL.md b/plugins/interloper/skills/interloper-destination/SKILL.md index ff645a33..b5907b25 100644 --- a/plugins/interloper/skills/interloper-destination/SKILL.md +++ b/plugins/interloper/skills/interloper-destination/SKILL.md @@ -7,16 +7,17 @@ description: Use when writing a custom Interloper destination (files, object sto ## Overview -Two base classes do the partition bookkeeping; you implement scope-level IO. -`il.PartitionedDestination` for files and objects: one write per partition, windows split for -you. `DatabaseDestination` for SQL: delete-then-insert by partition range. Subclassing bare -`il.Destination` means handling `context.partition_or_window` yourself, and the usual result is -a destination that clobbers every partition on each write. +The base class does the partition bookkeeping; you implement partition-level IO, where the +partition is `None` for the whole of an unpartitioned asset. `il.Destination` for files and objects: +implement `write_partition` and `read_partition`, windows are split and gathered for you. +`DatabaseDestination` for SQL: delete-then-insert by partition range, a window in one batch. +Overriding `write`/`read` wholesale is for a backend whose storage is not per partition, +and is how a destination ends up clobbering every partition on each write when done by accident. Reference: https://docs.interloper.dev/guide/destinations/ ## Recipe -1. **Files and objects**: implement `_write_scope` and `_read_scope` for one partition (or +1. **Files and objects**: implement `write_partition` and `read_partition` for one partition (or `None` for unpartitioned assets), and `partition_row_counts`: ```py @@ -27,7 +28,7 @@ Reference: https://docs.interloper.dev/guide/destinations/ from interloper.representation import Representation @il.destination(name="JSONL files") - class JSONLDestination(il.PartitionedDestination): + class JSONLDestination(il.Destination): base_path: str = "" def _path(self, context: il.IOContext, partition: il.Partition | None) -> Path: @@ -36,13 +37,13 @@ Reference: https://docs.interloper.dev/guide/destinations/ return base / "data.jsonl" return base / f"{context.asset.partitioning.column}={partition.id}" / "data.jsonl" - def _write_scope(self, context, partition, data) -> None: + def write_partition(self, context, partition, data) -> None: rows = Representation.of(data).to_records(data) # list[dict] from any representation path = self._path(context, partition) path.parent.mkdir(parents=True, exist_ok=True) path.write_text("".join(json.dumps(row, default=str) + "\n" for row in rows)) - def _read_scope(self, context, partition): + def read_partition(self, context, partition): rows = [json.loads(line) for line in self._path(context, partition).read_text().splitlines() if line] return context.schema.reconcile(rows) if context.schema is not None else rows @@ -67,7 +68,7 @@ Reference: https://docs.interloper.dev/guide/destinations/ delete-then-insert. Ranges are half-open `[start, end)` with `dt.date` bounds; store dates so the comparison works (ISO text sorts lexically). The default `WriteDisposition.REPLACE` deletes the range before inserting, so a - rewrite never duplicates. The base read does not restore types: override `_read_scope` to + rewrite never duplicates. The base read does not restore types: override `read_partition` to call `context.schema.reconcile(rows)` when a schema is present. 3. **Verify** with a two-asset daily source: @@ -106,7 +107,7 @@ Reference: https://docs.interloper.dev/guide/destinations/ - Subclassing `il.Destination` and writing `{table}.json`: a partition rewrite wipes the others and a window lands in one file. -- Returning raw strings from `_read_scope`; the dependent asset does arithmetic on text. +- Returning raw strings from `read_partition`; the dependent asset does arithmetic on text. - Dates stored in a format that does not compare with the `dt.date` bounds of the range hooks. - Forgetting `allow_window=True` on the asset's `TimePartitionConfig` when testing window writes; the error is a `PartitionError`, not a destination problem. diff --git a/plugins/interloper/skills/interloper-upgrade/SKILL.md b/plugins/interloper/skills/interloper-upgrade/SKILL.md index bf75fccd..d4b7e11d 100644 --- a/plugins/interloper/skills/interloper-upgrade/SKILL.md +++ b/plugins/interloper/skills/interloper-upgrade/SKILL.md @@ -83,6 +83,7 @@ Changelog (raw, complete): https://raw.githubusercontent.com/digitl-cloud/interl | spec `resources: {connection: {...}}` | the relation's own name: `connection: {...}` | | spec `upstreams: {orders: [id]}` | `orders: {ref: id}`, a list for a `many` relation | | `DependencyNotFoundError` | `ConfigError` from `validate_relations` | + | `class X(il.PartitionedDestination)` with `_write_scope` / `_read_scope` | `class X(il.Destination)` with `write_partition` / `read_partition`; the partition dispatch is the base class's | A relation left unbound is resolved when it is read: an explicit `default=`, else the target class built from the environment. So a connection that used to resolve through the cascade