import collections.abc
import itertools
from abc import ABC, abstractmethod, abstractproperty
from collections import defaultdict
from datetime import datetime
from typing import (
Collection,
Dict,
List,
Mapping,
NamedTuple,
Optional,
Sequence,
Tuple,
Type,
Union,
cast,
)
import dagster._check as check
from dagster._annotations import PublicAttr, experimental, public
from dagster._core.definitions.multi_dimensional_partitions import (
MultiPartitionKey,
MultiPartitionsDefinition,
)
from dagster._core.definitions.partition import (
PartitionsDefinition,
PartitionsSubset,
StaticPartitionsDefinition,
)
from dagster._core.definitions.time_window_partitions import TimeWindowPartitionsDefinition
from dagster._core.instance import DynamicPartitionsStore
from dagster._serdes import whitelist_for_serdes
from dagster._utils.cached_method import cached_method
from dagster._utils.warnings import disable_dagster_warnings
class UpstreamPartitionsResult(NamedTuple):
"""Represents the result of mapping a PartitionsSubset to the corresponding
partitions in another PartitionsDefinition.
partitions_subset (PartitionsSubset): The resulting partitions subset that was
mapped to. Only contains partitions for existent partitions, filtering out nonexistent partitions.
required_but_nonexistent_partition_keys (Sequence[str]): A list containing invalid partition keys in to_partitions_def
that partitions in from_partitions_subset were mapped to.
"""
partitions_subset: PartitionsSubset
required_but_nonexistent_partition_keys: Sequence[str]
[docs]class PartitionMapping(ABC):
"""Defines a correspondence between the partitions in an asset and the partitions in an asset
that it depends on.
Overriding PartitionMapping outside of Dagster is not supported. The abstract methods of this
class may change at any time.
"""
[docs] @public
@abstractmethod
def get_downstream_partitions_for_partitions(
self,
upstream_partitions_subset: PartitionsSubset,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> PartitionsSubset:
"""Returns the subset of partition keys in the downstream asset that use the data in the given
partition key subset of the upstream asset.
Args:
upstream_partitions_subset (Union[PartitionKeyRange, PartitionsSubset]): The
subset of partition keys in the upstream asset.
downstream_partitions_def (PartitionsDefinition): The partitions definition for the
downstream asset.
"""
[docs] @public
@abstractmethod
def get_upstream_mapped_partitions_result_for_partitions(
self,
downstream_partitions_subset: Optional[PartitionsSubset],
downstream_partitions_def: Optional[PartitionsDefinition],
upstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> UpstreamPartitionsResult:
"""Returns a UpstreamPartitionsResult object containing the partition keys the downstream
partitions subset was mapped to in the upstream partitions definition.
Valid upstream partitions will be included in UpstreamPartitionsResult.partitions_subset.
Invalid upstream partitions will be included in UpstreamPartitionsResult.required_but_nonexistent_partition_keys.
For example, if an upstream asset is time-partitioned and starts in June 2023, and the
downstream asset is time-partitioned and starts in May 2023, this function would return a
UpstreamPartitionsResult(PartitionsSubset("2023-06-01"), required_but_nonexistent_partition_keys=["2023-05-01"])
when downstream_partitions_subset contains 2023-05-01 and 2023-06-01.
"""
@abstractproperty
def description(self) -> str:
"""A human-readable description of the partition mapping, displayed in the Dagster UI."""
raise NotImplementedError()
[docs]@whitelist_for_serdes
class IdentityPartitionMapping(PartitionMapping, NamedTuple("_IdentityPartitionMapping", [])):
"""Expects that the upstream and downstream assets are partitioned in the same way, and maps
partitions in the downstream asset to the same partition in the upstream asset.
"""
def get_upstream_mapped_partitions_result_for_partitions(
self,
downstream_partitions_subset: Optional[PartitionsSubset],
downstream_partitions_def: Optional[PartitionsDefinition],
upstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> UpstreamPartitionsResult:
if downstream_partitions_subset is None:
check.failed("downstream asset is not partitioned")
if downstream_partitions_def == upstream_partitions_def:
return UpstreamPartitionsResult(downstream_partitions_subset, [])
upstream_partition_keys = set(
upstream_partitions_def.get_partition_keys(
dynamic_partitions_store=dynamic_partitions_store
)
)
downstream_partition_keys = set(downstream_partitions_subset.get_partition_keys())
return UpstreamPartitionsResult(
upstream_partitions_def.subset_with_partition_keys(
list(upstream_partition_keys & downstream_partition_keys)
),
list(downstream_partition_keys - upstream_partition_keys),
)
def get_downstream_partitions_for_partitions(
self,
upstream_partitions_subset: PartitionsSubset,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> PartitionsSubset:
if upstream_partitions_subset is None:
check.failed("upstream asset is not partitioned")
if upstream_partitions_def == downstream_partitions_def:
return upstream_partitions_subset
upstream_partition_keys = set(upstream_partitions_subset.get_partition_keys())
downstream_partition_keys = set(
downstream_partitions_def.get_partition_keys(
dynamic_partitions_store=dynamic_partitions_store
)
)
return downstream_partitions_def.empty_subset().with_partition_keys(
list(downstream_partition_keys & upstream_partition_keys)
)
@property
def description(self) -> str:
return (
"Assumes upstream and downstream assets share the same partitions definition. "
"Maps each partition in the downstream asset to the same partition in the upstream asset."
)
[docs]@whitelist_for_serdes
class AllPartitionMapping(PartitionMapping, NamedTuple("_AllPartitionMapping", [])):
"""Maps every partition in the downstream asset to every partition in the upstream asset.
Commonly used in the case when the downstream asset is not partitioned, in which the entire
downstream asset depends on all partitions of the usptream asset.
"""
def get_upstream_mapped_partitions_result_for_partitions(
self,
downstream_partitions_subset: Optional[PartitionsSubset],
downstream_partitions_def: Optional[PartitionsDefinition],
upstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> UpstreamPartitionsResult:
upstream_subset = upstream_partitions_def.subset_with_all_partitions(
current_time=current_time, dynamic_partitions_store=dynamic_partitions_store
)
return UpstreamPartitionsResult(upstream_subset, [])
def get_downstream_partitions_for_partitions(
self,
upstream_partitions_subset: PartitionsSubset,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> PartitionsSubset:
if upstream_partitions_subset is None:
check.failed("upstream asset is not partitioned")
if len(upstream_partitions_subset) == 0:
return downstream_partitions_def.empty_subset()
return downstream_partitions_def.subset_with_all_partitions(
current_time=current_time, dynamic_partitions_store=dynamic_partitions_store
)
@property
def description(self) -> str:
return "Each downstream partition depends on all partitions of the upstream asset."
[docs]@whitelist_for_serdes
class LastPartitionMapping(PartitionMapping, NamedTuple("_LastPartitionMapping", [])):
"""Maps all dependencies to the last partition in the upstream asset.
Commonly used in the case when the downstream asset is not partitioned, in which the entire
downstream asset depends on the last partition of the upstream asset.
"""
def get_upstream_mapped_partitions_result_for_partitions(
self,
downstream_partitions_subset: Optional[PartitionsSubset],
downstream_partitions_def: Optional[PartitionsDefinition],
upstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> UpstreamPartitionsResult:
last = upstream_partitions_def.get_last_partition_key(
current_time=current_time, dynamic_partitions_store=dynamic_partitions_store
)
upstream_subset = upstream_partitions_def.empty_subset()
if last is not None:
upstream_subset = upstream_subset.with_partition_keys([last])
return UpstreamPartitionsResult(upstream_subset, [])
def get_downstream_partitions_for_partitions(
self,
upstream_partitions_subset: PartitionsSubset,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> PartitionsSubset:
last_upstream_partition = upstream_partitions_def.get_last_partition_key(
current_time=current_time, dynamic_partitions_store=dynamic_partitions_store
)
if last_upstream_partition and last_upstream_partition in upstream_partitions_subset:
return downstream_partitions_def.subset_with_all_partitions(
current_time=current_time, dynamic_partitions_store=dynamic_partitions_store
)
else:
return downstream_partitions_def.empty_subset()
@property
def description(self) -> str:
return "Each downstream partition depends on the last partition of the upstream asset."
[docs]@whitelist_for_serdes
class SpecificPartitionsPartitionMapping(
PartitionMapping,
NamedTuple(
"_SpecificPartitionsPartitionMapping", [("partition_keys", PublicAttr[Sequence[str]])]
),
):
"""Maps to a specific subset of partitions in the upstream asset.
Example:
.. code-block:: python
from dagster import SpecificPartitionsPartitionMapping, StaticPartitionsDefinition, asset
@asset(partitions_def=StaticPartitionsDefinition(["a", "b", "c"]))
def upstream():
...
@asset(
ins={
"upstream": AssetIn(partition_mapping=SpecificPartitionsPartitionMapping(["a"]))
}
)
def a_downstream(upstream):
...
"""
def get_upstream_mapped_partitions_result_for_partitions(
self,
downstream_partitions_subset: Optional[PartitionsSubset],
downstream_partitions_def: Optional[PartitionsDefinition],
upstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> UpstreamPartitionsResult:
return UpstreamPartitionsResult(
upstream_partitions_def.subset_with_partition_keys(self.partition_keys), []
)
def get_downstream_partitions_for_partitions(
self,
upstream_partitions_subset: PartitionsSubset,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> PartitionsSubset:
# if any of the partition keys in this partition mapping are contained within the upstream
# partitions subset, then all partitions of the downstream asset are dependencies
if any(key in upstream_partitions_subset for key in self.partition_keys):
return downstream_partitions_def.subset_with_all_partitions(
dynamic_partitions_store=dynamic_partitions_store
)
return downstream_partitions_def.empty_subset()
@property
def description(self) -> str:
return f"Each downstream partition depends on the following upstream partitions: {self.partition_keys}"
class DimensionDependency(NamedTuple):
partition_mapping: PartitionMapping
upstream_dimension_name: Optional[str] = None
downstream_dimension_name: Optional[str] = None
class BaseMultiPartitionMapping(ABC):
@abstractmethod
def get_dimension_dependencies(
self,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
) -> Sequence[DimensionDependency]: ...
def get_partitions_def(
self, partitions_def: PartitionsDefinition, dimension_name: Optional[str]
) -> PartitionsDefinition:
if isinstance(partitions_def, MultiPartitionsDefinition):
if not isinstance(dimension_name, str):
check.failed("Expected dimension_name to be a string")
return partitions_def.get_partitions_def_for_dimension(dimension_name)
return partitions_def
def _get_dependency_partitions_subset(
self,
a_partitions_def: PartitionsDefinition,
a_partitions_subset: PartitionsSubset,
b_partitions_def: PartitionsDefinition,
a_upstream_of_b: bool,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
current_time: Optional[datetime] = None,
) -> Union[UpstreamPartitionsResult, PartitionsSubset]:
"""Given two partitions definitions a_partitions_def and b_partitions_def that have a dependency
relationship (a_upstream_of_b is True if a_partitions_def is upstream of b_partitions_def),
and a_partition_keys, a list of partition keys in a_partitions_def, returns a list of
partition keys in the partitions definition b_partitions_def that are
dependencies of the partition keys in a_partition_keys.
"""
a_partition_keys_by_dimension = defaultdict(set)
if isinstance(a_partitions_def, MultiPartitionsDefinition):
for partition_key in a_partitions_subset.get_partition_keys():
key = a_partitions_def.get_partition_key_from_str(partition_key)
for dimension_name, key in key.keys_by_dimension.items():
a_partition_keys_by_dimension[dimension_name].add(key)
else:
for partition_key in a_partitions_subset.get_partition_keys():
a_partition_keys_by_dimension[None].add(partition_key)
# Maps the dimension name and key of a partition in a_partitions_def to the list of
# partition keys in b_partitions_def that are dependencies of that partition
dep_b_keys_by_a_dim_and_key: Dict[Optional[str], Dict[Optional[str], List[str]]] = (
defaultdict(lambda: defaultdict(list))
)
required_but_nonexistent_upstream_partitions = set()
b_dimension_partitions_def_by_name: Dict[Optional[str], PartitionsDefinition] = (
{
dimension.name: dimension.partitions_def
for dimension in b_partitions_def.partitions_defs
}
if isinstance(b_partitions_def, MultiPartitionsDefinition)
else {None: b_partitions_def}
)
if a_upstream_of_b:
# a_partitions_def is upstream of b_partitions_def, so we need to map the
# dimension names of a_partitions_def to the corresponding dependent dimensions of
# b_partitions_def
a_dim_to_dependency_b_dim = {
dimension_mapping.upstream_dimension_name: (
dimension_mapping.downstream_dimension_name,
dimension_mapping.partition_mapping,
)
for dimension_mapping in self.get_dimension_dependencies(
a_partitions_def, b_partitions_def
)
}
for a_dim_name, keys in a_partition_keys_by_dimension.items():
if a_dim_name in a_dim_to_dependency_b_dim:
(
b_dim_name,
dimension_mapping,
) = a_dim_to_dependency_b_dim[a_dim_name]
a_dimension_partitions_def = self.get_partitions_def(
a_partitions_def, a_dim_name
)
b_dimension_partitions_def = self.get_partitions_def(
b_partitions_def, b_dim_name
)
for key in keys:
# if downstream dimension mapping exists, for a given key, get the list of
# downstream partition keys that are dependencies of that key
dep_b_keys_by_a_dim_and_key[a_dim_name][key] = list(
dimension_mapping.get_downstream_partitions_for_partitions(
a_dimension_partitions_def.empty_subset().with_partition_keys(
[key]
),
a_dimension_partitions_def,
b_dimension_partitions_def,
current_time=current_time,
dynamic_partitions_store=dynamic_partitions_store,
).get_partition_keys()
)
else:
# a_partitions_def is downstream of b_partitions_def, so we need to map the
# dimension names of a_partitions_def to the corresponding dependency dimensions of
# b_partitions_def
a_dim_to_dependency_b_dim = {
dimension_mapping.downstream_dimension_name: (
dimension_mapping.upstream_dimension_name,
dimension_mapping.partition_mapping,
)
for dimension_mapping in self.get_dimension_dependencies(
b_partitions_def, a_partitions_def
)
}
for a_dim_name, keys in a_partition_keys_by_dimension.items():
if a_dim_name in a_dim_to_dependency_b_dim:
(
b_dim_name,
partition_mapping,
) = a_dim_to_dependency_b_dim[a_dim_name]
a_dimension_partitions_def = self.get_partitions_def(
a_partitions_def, a_dim_name
)
b_dimension_partitions_def = self.get_partitions_def(
b_partitions_def, b_dim_name
)
for key in keys:
mapped_partitions_result = (
partition_mapping.get_upstream_mapped_partitions_result_for_partitions(
a_dimension_partitions_def.empty_subset().with_partition_keys(
[key]
),
a_dimension_partitions_def,
b_dimension_partitions_def,
current_time=current_time,
dynamic_partitions_store=dynamic_partitions_store,
)
)
dep_b_keys_by_a_dim_and_key[a_dim_name][key] = list(
mapped_partitions_result.partitions_subset.get_partition_keys()
)
required_but_nonexistent_upstream_partitions.update(
set(mapped_partitions_result.required_but_nonexistent_partition_keys)
)
b_partition_keys = set()
mapped_a_dim_names = a_dim_to_dependency_b_dim.keys()
mapped_b_dim_names = [mapping[0] for mapping in a_dim_to_dependency_b_dim.values()]
unmapped_b_dim_names = list(
set(b_dimension_partitions_def_by_name.keys()) - set(mapped_b_dim_names)
)
for key in a_partitions_subset.get_partition_keys():
for b_key_values in itertools.product(
*(
[
dep_b_keys_by_a_dim_and_key[dim_name][
(
cast(MultiPartitionsDefinition, a_partitions_def)
.get_partition_key_from_str(key)
.keys_by_dimension[dim_name]
if dim_name
else key
)
]
for dim_name in mapped_a_dim_names
]
),
*[
b_dimension_partitions_def_by_name[dim_name].get_partition_keys(
dynamic_partitions_store=dynamic_partitions_store, current_time=current_time
)
for dim_name in unmapped_b_dim_names
],
):
b_partition_keys.add(
MultiPartitionKey(
{
cast(str, (mapped_b_dim_names + unmapped_b_dim_names)[i]): key
for i, key in enumerate(b_key_values)
}
)
if len(b_key_values) > 1
else b_key_values[0]
)
mapped_subset = b_partitions_def.empty_subset().with_partition_keys(b_partition_keys)
if a_upstream_of_b:
return mapped_subset
else:
return UpstreamPartitionsResult(
mapped_subset,
required_but_nonexistent_partition_keys=list(
required_but_nonexistent_upstream_partitions
),
)
def get_upstream_mapped_partitions_result_for_partitions(
self,
downstream_partitions_subset: Optional[PartitionsSubset],
downstream_partitions_def: Optional[PartitionsDefinition],
upstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> UpstreamPartitionsResult:
if downstream_partitions_subset is None:
check.failed("downstream asset is not partitioned")
result = self._get_dependency_partitions_subset(
check.not_none(downstream_partitions_def),
downstream_partitions_subset,
cast(MultiPartitionsDefinition, upstream_partitions_def),
a_upstream_of_b=False,
dynamic_partitions_store=dynamic_partitions_store,
current_time=current_time,
)
if not isinstance(result, UpstreamPartitionsResult):
check.failed("Expected UpstreamPartitionsResult")
return result
def get_downstream_partitions_for_partitions(
self,
upstream_partitions_subset: PartitionsSubset,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> PartitionsSubset:
if upstream_partitions_subset is None:
check.failed("upstream asset is not partitioned")
result = self._get_dependency_partitions_subset(
upstream_partitions_def,
upstream_partitions_subset,
cast(MultiPartitionsDefinition, downstream_partitions_def),
a_upstream_of_b=True,
dynamic_partitions_store=dynamic_partitions_store,
)
if isinstance(result, UpstreamPartitionsResult):
check.failed("Expected PartitionsSubset")
return result
[docs]@experimental
@whitelist_for_serdes
class MultiToSingleDimensionPartitionMapping(
BaseMultiPartitionMapping,
PartitionMapping,
NamedTuple(
"_MultiToSingleDimensionPartitionMapping", [("partition_dimension_name", Optional[str])]
),
):
"""Defines a correspondence between an single-dimensional partitions definition
and a MultiPartitionsDefinition. The single-dimensional partitions definition must be
a dimension of the MultiPartitionsDefinition.
This class handles the case where the upstream asset is multipartitioned and the
downstream asset is single dimensional, and vice versa.
For a partition key X, this partition mapping assumes that any multi-partition key with
X in the selected dimension is a dependency.
Args:
partition_dimension_name (Optional[str]): The name of the partition dimension in the
MultiPartitionsDefinition that matches the single-dimension partitions definition.
"""
def __new__(cls, partition_dimension_name: Optional[str] = None):
return super(MultiToSingleDimensionPartitionMapping, cls).__new__(
cls,
partition_dimension_name=check.opt_str_param(
partition_dimension_name, "partition_dimension_name"
),
)
@property
def description(self) -> str:
return (
"Assumes that the single-dimension partitions definition is a dimension of the "
"multi-partitions definition. For a single-dimension partition key X, any "
"multi-partition key with X in the matching dimension is a dependency."
)
def get_dimension_dependencies(
self,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
) -> Sequence[DimensionDependency]:
infer_mapping_result = _get_infer_single_to_multi_dimension_deps_result(
upstream_partitions_def, downstream_partitions_def
)
if not infer_mapping_result.can_infer:
check.invariant(isinstance(infer_mapping_result.inference_failure_reason, str))
check.failed(cast(str, infer_mapping_result.inference_failure_reason))
return [cast(DimensionDependency, infer_mapping_result.dimension_dependency)]
@whitelist_for_serdes
class DimensionPartitionMapping(
NamedTuple(
"_DimensionPartitionMapping",
[
("dimension_name", str),
("partition_mapping", PartitionMapping),
],
)
):
"""A helper class for MultiPartitionMapping that defines a partition mapping used to calculate
the dependent partition keys in the selected downstream MultiPartitions definition dimension.
Args:
dimension_name (str): The name of the dimension in the downstream MultiPartitionsDefinition.
partition_mapping (PartitionMapping): The partition mapping object used to calculate
the downstream dimension partitions from the upstream dimension partitions and vice versa.
"""
def __new__(
cls,
dimension_name: str,
partition_mapping: PartitionMapping,
):
return super(DimensionPartitionMapping, cls).__new__(
cls,
dimension_name=check.str_param(dimension_name, "dimension_name"),
partition_mapping=check.inst_param(
partition_mapping, "partition_mapping", PartitionMapping
),
)
[docs]@experimental
@whitelist_for_serdes
class MultiPartitionMapping(
BaseMultiPartitionMapping,
PartitionMapping,
NamedTuple(
"_MultiPartitionMapping",
[("downstream_mappings_by_upstream_dimension", Mapping[str, DimensionPartitionMapping])],
),
):
"""Defines a correspondence between two MultiPartitionsDefinitions.
Accepts a mapping of upstream dimension name to downstream DimensionPartitionMapping, representing
the explicit correspondence between the upstream and downstream MultiPartitions dimensions
and the partition mapping used to calculate the downstream partitions.
Examples:
.. code-block:: python
weekly_abc = MultiPartitionsDefinition(
{
"abc": StaticPartitionsDefinition(["a", "b", "c"]),
"weekly": WeeklyPartitionsDefinition("2023-01-01"),
}
)
daily_123 = MultiPartitionsDefinition(
{
"123": StaticPartitionsDefinition(["1", "2", "3"]),
"daily": DailyPartitionsDefinition("2023-01-01"),
}
)
MultiPartitionMapping(
{
"abc": DimensionPartitionMapping(
dimension_name="123",
partition_mapping=StaticPartitionMapping({"a": "1", "b": "2", "c": "3"}),
),
"weekly": DimensionPartitionMapping(
dimension_name="daily",
partition_mapping=TimeWindowPartitionMapping(),
)
}
)
For upstream or downstream dimensions not explicitly defined in the mapping, Dagster will
assume an `AllPartitionsMapping`, meaning that all upstream partitions in those dimensions
will be mapped to all downstream partitions in those dimensions.
Examples:
.. code-block:: python
weekly_abc = MultiPartitionsDefinition(
{
"abc": StaticPartitionsDefinition(["a", "b", "c"]),
"daily": DailyPartitionsDefinition("2023-01-01"),
}
)
daily_123 = MultiPartitionsDefinition(
{
"123": StaticPartitionsDefinition(["1", "2", "3"]),
"daily": DailyPartitionsDefinition("2023-01-01"),
}
)
MultiPartitionMapping(
{
"daily": DimensionPartitionMapping(
dimension_name="daily",
partition_mapping=IdentityPartitionMapping(),
)
}
)
# Will map `daily_123` partition key {"123": "1", "daily": "2023-01-01"} to the upstream:
# {"abc": "a", "daily": "2023-01-01"}
# {"abc": "b", "daily": "2023-01-01"}
# {"abc": "c", "daily": "2023-01-01"}
Args:
downstream_mappings_by_upstream_dimension (Mapping[str, DimensionPartitionMapping]): A
mapping that defines an explicit correspondence between one dimension of the upstream
MultiPartitionsDefinition and one dimension of the downstream MultiPartitionsDefinition.
Maps a string representing upstream dimension name to downstream DimensionPartitionMapping,
containing the downstream dimension name and partition mapping.
"""
def __new__(
cls, downstream_mappings_by_upstream_dimension: Mapping[str, DimensionPartitionMapping]
):
return super(MultiPartitionMapping, cls).__new__(
cls,
downstream_mappings_by_upstream_dimension=check.mapping_param(
downstream_mappings_by_upstream_dimension,
"downstream_mappings_by_upstream_dimension",
key_type=str,
value_type=DimensionPartitionMapping,
),
)
@property
def description(self) -> str:
return "\n ".join(
[
(
f"Upstream dimension '{upstream_dim}' mapped to downstream dimension "
f"'{downstream_mapping.dimension_name}' using {type(downstream_mapping.partition_mapping).__name__}."
)
for upstream_dim, downstream_mapping in self.downstream_mappings_by_upstream_dimension.items()
]
)
def get_dimension_dependencies(
self,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
) -> Sequence[DimensionDependency]:
self._check_all_dimensions_accounted_for(
upstream_partitions_def,
downstream_partitions_def,
)
return [
DimensionDependency(
mapping.partition_mapping,
upstream_dimension_name=upstream_dimension,
downstream_dimension_name=mapping.dimension_name,
)
for upstream_dimension, mapping in self.downstream_mappings_by_upstream_dimension.items()
]
def _check_all_dimensions_accounted_for(
self,
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
) -> None:
if any(
not isinstance(partitions_def, MultiPartitionsDefinition)
for partitions_def in (upstream_partitions_def, downstream_partitions_def)
):
check.failed(
"Both partitions defs provided to a MultiPartitionMapping must be multi-partitioned"
)
upstream_dimension_names = {
dim.name
for dim in cast(MultiPartitionsDefinition, upstream_partitions_def).partitions_defs
}
dimension_names = {
dim.name
for dim in cast(MultiPartitionsDefinition, downstream_partitions_def).partitions_defs
}
for (
upstream_dimension_name,
dimension_mapping,
) in self.downstream_mappings_by_upstream_dimension.items():
if upstream_dimension_name not in upstream_dimension_names:
check.failed(
"Dimension mapping has an upstream dimension name that is not in the upstream "
"partitions def"
)
if dimension_mapping.dimension_name not in dimension_names:
check.failed(
"Dimension mapping has a downstream dimension name that is not in the"
" downstream partitions def"
)
upstream_dimension_names.remove(upstream_dimension_name)
dimension_names.remove(dimension_mapping.dimension_name)
[docs]@whitelist_for_serdes
class StaticPartitionMapping(
PartitionMapping,
NamedTuple(
"_StaticPartitionMapping",
[
(
"downstream_partition_keys_by_upstream_partition_key",
PublicAttr[Mapping[str, Union[str, Collection[str]]]],
)
],
),
):
"""Define an explicit correspondence between two StaticPartitionsDefinitions.
Args:
downstream_partition_keys_by_upstream_partition_key (Dict[str, str | Collection[str]]):
The single or multi-valued correspondence from upstream keys to downstream keys.
"""
def __init__(
self,
downstream_partition_keys_by_upstream_partition_key: Mapping[
str, Union[str, Collection[str]]
],
):
check.mapping_param(
downstream_partition_keys_by_upstream_partition_key,
"downstream_partition_keys_by_upstream_partition_key",
key_type=str,
value_type=(str, collections.abc.Collection),
)
# cache forward and reverse mappings
self._mapping = defaultdict(set)
for (
upstream_key,
downstream_keys,
) in downstream_partition_keys_by_upstream_partition_key.items():
self._mapping[upstream_key] = (
{downstream_keys} if isinstance(downstream_keys, str) else set(downstream_keys)
)
self._inverse_mapping = defaultdict(set)
for upstream_key, downstream_keys in self._mapping.items():
for downstream_key in downstream_keys:
self._inverse_mapping[downstream_key].add(upstream_key)
@cached_method
def _check_upstream(self, *, upstream_partitions_def: StaticPartitionsDefinition):
"""Validate that the mapping from upstream to downstream is only defined on upstream keys."""
check.inst_param(
upstream_partitions_def,
"upstream_partitions_def",
StaticPartitionsDefinition,
"StaticPartitionMapping can only be defined between two StaticPartitionsDefinitions",
)
upstream_keys = upstream_partitions_def.get_partition_keys()
extra_keys = set(self._mapping.keys()).difference(upstream_keys)
if extra_keys:
raise ValueError(
f"mapping source partitions not in the upstream partitions definition: {extra_keys}"
)
@cached_method
def _check_downstream(self, *, downstream_partitions_def: StaticPartitionsDefinition):
"""Validate that the mapping from upstream to downstream only maps to downstream keys."""
check.inst_param(
downstream_partitions_def,
"downstream_partitions_def",
StaticPartitionsDefinition,
"StaticPartitionMapping can only be defined between two StaticPartitionsDefinitions",
)
downstream_keys = downstream_partitions_def.get_partition_keys()
extra_keys = set(self._inverse_mapping.keys()).difference(downstream_keys)
if extra_keys:
raise ValueError(
"mapping target partitions not in the downstream partitions definition:"
f" {extra_keys}"
)
def get_downstream_partitions_for_partitions(
self,
upstream_partitions_subset: PartitionsSubset,
upstream_partitions_def: StaticPartitionsDefinition,
downstream_partitions_def: StaticPartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> PartitionsSubset:
self._check_downstream(downstream_partitions_def=downstream_partitions_def)
downstream_subset = downstream_partitions_def.empty_subset()
downstream_keys = set()
for key in upstream_partitions_subset.get_partition_keys():
downstream_keys.update(self._mapping[key])
return downstream_subset.with_partition_keys(downstream_keys)
def get_upstream_mapped_partitions_result_for_partitions(
self,
downstream_partitions_subset: Optional[PartitionsSubset],
downstream_partitions_def: Optional[PartitionsDefinition],
upstream_partitions_def: StaticPartitionsDefinition,
current_time: Optional[datetime] = None,
dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
) -> UpstreamPartitionsResult:
self._check_upstream(upstream_partitions_def=upstream_partitions_def)
upstream_subset = upstream_partitions_def.empty_subset()
if downstream_partitions_subset is None:
return UpstreamPartitionsResult(upstream_subset, [])
upstream_keys = set()
for key in downstream_partitions_subset.get_partition_keys():
upstream_keys.update(self._inverse_mapping[key])
return UpstreamPartitionsResult(upstream_subset.with_partition_keys(upstream_keys), [])
@property
def description(self) -> str:
return (
f"Maps upstream partitions to their downstream dependencies according to the "
f"following mapping: \n{self.downstream_partition_keys_by_upstream_partition_key}"
)
class InferSingleToMultiDimensionDepsResult(
NamedTuple(
"_InferSingleToMultiDimensionDepsResult",
[
("can_infer", bool),
("inference_failure_reason", Optional[str]),
("dimension_dependency", Optional[DimensionDependency]),
],
)
):
def __new__(
cls,
can_infer: bool,
inference_failure_reason: Optional[str] = None,
dimension_dependency: Optional[DimensionDependency] = None,
):
if can_infer and dimension_dependency is None:
check.failed("dimension_dependency must be provided if can_infer is True")
if not can_infer and inference_failure_reason is None:
check.failed("inference_failure_reason must be provided if can_infer is False")
return super(InferSingleToMultiDimensionDepsResult, cls).__new__(
cls,
can_infer,
inference_failure_reason,
dimension_dependency,
)
def _get_infer_single_to_multi_dimension_deps_result(
upstream_partitions_def: PartitionsDefinition,
downstream_partitions_def: PartitionsDefinition,
partition_dimension_name: Optional[str] = None,
) -> InferSingleToMultiDimensionDepsResult:
from dagster._core.definitions.time_window_partition_mapping import TimeWindowPartitionMapping
upstream_is_multipartitioned = isinstance(upstream_partitions_def, MultiPartitionsDefinition)
multipartitions_defs = [
partitions_def
for partitions_def in [upstream_partitions_def, downstream_partitions_def]
if isinstance(partitions_def, MultiPartitionsDefinition)
]
if len(multipartitions_defs) != 1:
return InferSingleToMultiDimensionDepsResult(
False,
"Can only use MultiToSingleDimensionPartitionMapping when upstream asset is"
" multipartitioned and the downstream asset is single dimensional, or vice versa."
f" Instead received {len(multipartitions_defs)} multi-partitioned assets.",
)
multipartitions_def = cast(MultiPartitionsDefinition, next(iter(multipartitions_defs)))
single_dimension_partitions_def = next(
iter(
{
upstream_partitions_def,
downstream_partitions_def,
}
- set(multipartitions_defs)
)
)
filtered_multipartition_dims = (
multipartitions_def.partitions_defs
if partition_dimension_name is None
else [
dim
for dim in multipartitions_def.partitions_defs
if dim.name == partition_dimension_name
]
)
if partition_dimension_name:
if len(filtered_multipartition_dims) != 1:
return InferSingleToMultiDimensionDepsResult(
False,
f"Provided partition dimension name {partition_dimension_name} not found in"
f" multipartitions definition {multipartitions_def}.",
)
matching_dimension_defs = [
dimension_def
for dimension_def in filtered_multipartition_dims
if dimension_def.partitions_def == single_dimension_partitions_def
]
if len(matching_dimension_defs) == 1:
return InferSingleToMultiDimensionDepsResult(
True,
dimension_dependency=DimensionDependency(
IdentityPartitionMapping(),
upstream_dimension_name=(
matching_dimension_defs[0].name if upstream_is_multipartitioned else None
),
downstream_dimension_name=(
matching_dimension_defs[0].name if not upstream_is_multipartitioned else None
),
),
)
elif len(matching_dimension_defs) > 1:
return InferSingleToMultiDimensionDepsResult(
False,
"partition dimension name must be specified when multiple dimensions of the"
" MultiPartitionsDefinition match the single dimension partitions def",
)
time_dimensions = [
dimension_def
for dimension_def in filtered_multipartition_dims
if isinstance(dimension_def.partitions_def, TimeWindowPartitionsDefinition)
]
if len(time_dimensions) == 1 and isinstance(
single_dimension_partitions_def, TimeWindowPartitionsDefinition
):
return InferSingleToMultiDimensionDepsResult(
True,
dimension_dependency=DimensionDependency(
TimeWindowPartitionMapping(),
upstream_dimension_name=(
time_dimensions[0].name if upstream_is_multipartitioned else None
),
downstream_dimension_name=(
time_dimensions[0].name if not upstream_is_multipartitioned else None
),
),
)
return InferSingleToMultiDimensionDepsResult(
False,
"MultiToSingleDimensionPartitionMapping can only be used when: \n(a) The single dimensional"
" partitions definition is a dimension of the MultiPartitionsDefinition.\n(b) The single"
" dimensional partitions definition is a TimeWindowPartitionsDefinition and the"
" MultiPartitionsDefinition has a single time dimension.",
)
def infer_partition_mapping(
partition_mapping: Optional[PartitionMapping],
downstream_partitions_def: Optional[PartitionsDefinition],
upstream_partitions_def: Optional[PartitionsDefinition],
) -> PartitionMapping:
from .time_window_partition_mapping import TimeWindowPartitionMapping
if partition_mapping is not None:
return partition_mapping
elif upstream_partitions_def and downstream_partitions_def:
if _get_infer_single_to_multi_dimension_deps_result(
upstream_partitions_def, downstream_partitions_def
).can_infer:
with disable_dagster_warnings():
return MultiToSingleDimensionPartitionMapping()
elif isinstance(upstream_partitions_def, TimeWindowPartitionsDefinition) and isinstance(
downstream_partitions_def, TimeWindowPartitionsDefinition
):
return TimeWindowPartitionMapping()
else:
return IdentityPartitionMapping()
else:
return AllPartitionMapping()
def get_builtin_partition_mapping_types() -> Tuple[Type[PartitionMapping], ...]:
from dagster._core.definitions.time_window_partition_mapping import TimeWindowPartitionMapping
return (
AllPartitionMapping,
IdentityPartitionMapping,
LastPartitionMapping,
SpecificPartitionsPartitionMapping,
StaticPartitionMapping,
TimeWindowPartitionMapping,
MultiToSingleDimensionPartitionMapping,
MultiPartitionMapping,
)