def align_datasets(left_dataset_uuid, right_dataset_uuid, store, match_how="exact"):
"""
Determine dataset partition alignment
Parameters
----------
left_dataset_uuid : basestring
right_dataset_uuid : basestring
store : KeyValuestore or callable
match_how : basestring or callable, {exact, prefix, all, callable}
Yields
------
list
"""
store = _instantiate_store(store)
left_dataset = DatasetMetadata.load_from_store(uuid=left_dataset_uuid, store=store)
right_dataset = DatasetMetadata.load_from_store(
uuid=right_dataset_uuid, store=store
)
metadata_version = left_dataset.metadata_version
# Loop over the dataset with fewer partitions, treating its keys as
# partition label prefixes
if (
callable(match_how)
or match_how == "left"
or (
match_how == "prefix"
and len(list(left_dataset.partitions.keys())[0])
< len(list(right_dataset.partitions.keys())[0])
)
):
first_dataset = left_dataset
second_dataset = right_dataset
else:
first_dataset = right_dataset
second_dataset = left_dataset
# The del statements are here to reduce confusion below
del left_dataset
del right_dataset
# For every partition in the 'small' dataset, at least one partition match
# needs to be found in the larger dataset.
available_partitions = list(second_dataset.partitions.items())
partition_stack = available_partitions[:]
# TODO: write a test which protects against the following scenario!!
# Sort the partition labels by length of the labels, starting with the
# labels which are the longest. This way we prevent label matching for
# similar partitions, e.g. cluster_100 and cluster_1. This, of course,
# works only as long as the internal loop removes elements which were
# matched already (here improperly called stack)
for l_1 in sorted(first_dataset.partitions, key=len, reverse=True):
p_1 = first_dataset.partitions[l_1]
res = [
MetaPartition.from_partition(
partition=p_1, metadata_version=metadata_version
)
]
for parts in available_partitions:
l_2, p_2 = parts
if callable(match_how) and not match_how(l_1, l_2):
continue
if match_how == "exact" and l_1 != l_2:
continue
elif match_how == "prefix" and not l_2.startswith(l_1):
LOGGER.debug("rejecting (%s, %s)", l_1, l_2)
continue
LOGGER.debug(
"Found alignment between partitions " "(%s, %s) and" "(%s, %s)",
first_dataset.uuid,
p_1.label,
second_dataset.uuid,
p_2.label,
)
res.append(
MetaPartition.from_partition(
partition=p_2, metadata_version=metadata_version
)
)
# In exact or prefix matching schemes, it is expected to only
# find one partition alignment. in this case reduce the size of
# the inner loop
if match_how in ["exact", "prefix"]:
partition_stack.remove((l_2, p_2))
# Need to copy, otherwise remove will alter the loop iterator
available_partitions = partition_stack[:]
if len(res) == 1:
raise RuntimeError(
"No matching partition for {} in dataset {} "
"found".format(p_1, first_dataset)
)
yield res
def dispatch_metapartitions_from_factory(
dataset_factory,
label_filter=None,
concat_partitions_on_primary_index=False,
predicates=None,
store=None,
dispatch_by=None,
) -> Union[Iterator[MetaPartition], Iterator[List[MetaPartition]]]:
if not callable(dataset_factory) and not isinstance(
dataset_factory, DatasetFactory):
raise TypeError("Need to supply a dataset factory!")
if dispatch_by and concat_partitions_on_primary_index:
raise ValueError(
"Both `dispatch_by` and `concat_partitions_on_primary_index` are provided, "
"`concat_partitions_on_primary_index` is deprecated and will be removed in the next major release. "
"Please only provide the `dispatch_by` argument. ")
if concat_partitions_on_primary_index:
warnings.warn(
"The keyword `concat_partitions_on_primary_index` is deprecated and will be removed in the next major release. Use `dispatch_by=dataset_factory.partition_keys` to achieve the same behavior instead.",
DeprecationWarning,
)
dispatch_by = dataset_factory.partition_keys
if dispatch_by and not set(dispatch_by).issubset(
set(dataset_factory.index_columns)):
raise RuntimeError(
f"Dispatch columns must be indexed.\nRequested index: {dispatch_by} but available index columns: {sorted(dataset_factory.index_columns)}"
)
check_predicates(predicates)
# Determine which indices need to be loaded.
index_cols: Set[str] = set()
if dispatch_by:
index_cols |= set(dispatch_by)
if predicates:
predicate_cols = set(columns_in_predicates(predicates))
predicate_index_cols = predicate_cols & set(
dataset_factory.index_columns)
index_cols |= predicate_index_cols
for col in index_cols:
dataset_factory.load_index(col)
base_df = dataset_factory.get_indices_as_dataframe(list(index_cols),
predicates=predicates)
if label_filter:
base_df = base_df[base_df.index.map(label_filter)]
indices_to_dispatch = {
name: ix.unload()
for name, ix in dataset_factory.indices.items()
if isinstance(ix, ExplicitSecondaryIndex)
}
if dispatch_by:
base_df = cast(pd.DataFrame, base_df)
# Group the resulting MetaParitions by partition keys or a subset of those keys
merged_partitions = base_df.groupby(by=list(dispatch_by),
sort=False,
as_index=False)
for group_name, group in merged_partitions:
if not isinstance(group_name, tuple):
group_name = (group_name, )
mps = []
logical_conjunction = list(
zip(dispatch_by, ["=="] * len(dispatch_by), group_name))
for label in group.index.unique():
mps.append(
MetaPartition.from_partition(
partition=dataset_factory.partitions[label],
dataset_metadata=dataset_factory.metadata,
indices=indices_to_dispatch,
metadata_version=dataset_factory.metadata_version,
table_meta=dataset_factory.table_meta,
partition_keys=dataset_factory.partition_keys,
logical_conjunction=logical_conjunction,
))
yield mps
else:
for part_label in base_df.index.unique():
part = dataset_factory.partitions[part_label]
yield MetaPartition.from_partition(
partition=part,
dataset_metadata=dataset_factory.metadata,
indices=indices_to_dispatch,
metadata_version=dataset_factory.metadata_version,
table_meta=dataset_factory.table_meta,
partition_keys=dataset_factory.partition_keys,
)
def dispatch_metapartitions_from_factory(
dataset_factory,
label_filter=None,
concat_partitions_on_primary_index=False,
predicates=None,
store=None,
dispatch_by=None,
) -> Union[Iterator[MetaPartition], Iterator[List[MetaPartition]]]:
if not callable(dataset_factory) and not isinstance(
dataset_factory, DatasetFactory):
raise TypeError("Need to supply a dataset factory!")
if dispatch_by and concat_partitions_on_primary_index:
raise ValueError(
"Both `dispatch_by` and `concat_partitions_on_primary_index` are provided, "
"`concat_partitions_on_primary_index` is deprecated and will be removed in the next major release. "
"Please only provide the `dispatch_by` argument. ")
if concat_partitions_on_primary_index:
warnings.warn(
"The keyword `concat_partitions_on_primary_index` is deprecated and will be removed in the next major release. Use `dispatch_by=dataset_factory.partition_keys` to achieve the same behavior instead.",
DeprecationWarning,
)
dispatch_by = dataset_factory.partition_keys
if dispatch_by and not set(dispatch_by).issubset(
set(dataset_factory.index_columns)):
raise RuntimeError(
f"Dispatch columns must be indexed.\nRequested index: {dispatch_by} but available index columns: {sorted(dataset_factory.index_columns)}"
)
if predicates is not None:
dataset_factory, allowed_labels = _allowed_labels_by_predicates(
predicates, dataset_factory, dispatch_by)
else:
allowed_labels = None
indices_to_dispatch = {
name: ix.copy(index_dct={})
for name, ix in dataset_factory.indices.items()
if isinstance(ix, ExplicitSecondaryIndex)
}
if dispatch_by:
# Build up a DataFrame that contains per row a Partition and its primary index columns.
base_df = None
for part_key in dispatch_by:
dataset_factory.load_index(part_key)
idx = dataset_factory.indices[part_key].index_dct
df = _index_to_dataframe(part_key, idx, allowed_labels)
if base_df is None:
base_df = df
else:
base_df = base_df.merge(df, on=["__partition__"])
assert base_df is not None
base_df = cast(pd.DataFrame, base_df)
# Group the resulting MetaParitions by partition keys or a subset of those keys
merged_partitions = base_df.groupby(by=list(dispatch_by),
sort=False,
as_index=False)
for group_name, group in merged_partitions:
if not isinstance(group_name, tuple):
group_name = (group_name, )
mps = []
logical_conjunction = list(
zip(dispatch_by, ["=="] * len(dispatch_by), group_name))
for label in group.__partition__:
mps.append(
MetaPartition.from_partition(
partition=dataset_factory.partitions[label],
dataset_metadata=dataset_factory.metadata,
indices=indices_to_dispatch,
metadata_version=dataset_factory.metadata_version,
table_meta=dataset_factory.table_meta,
partition_keys=dataset_factory.partition_keys,
logical_conjunction=logical_conjunction,
))
yield mps
else:
if allowed_labels is not None:
partition_labels = allowed_labels
else:
partition_labels = dataset_factory.partitions.keys()
for part_label in partition_labels:
if label_filter is not None:
if not label_filter(part_label):
continue
part = dataset_factory.partitions[part_label]
yield MetaPartition.from_partition(
partition=part,
dataset_metadata=dataset_factory.metadata,
indices=indices_to_dispatch,
metadata_version=dataset_factory.metadata_version,
table_meta=dataset_factory.table_meta,
partition_keys=dataset_factory.partition_keys,
)
def dispatch_metapartitions_from_factory(
dataset_factory: DatasetFactory,
label_filter: Optional[Callable] = None,
concat_partitions_on_primary_index: bool = False,
predicates: PredicatesType = None,
store: Optional[StoreInput] = None,
dispatch_by: Optional[List[str]] = None,
dispatch_metadata: bool = False,
) -> Union[Iterator[MetaPartition], Iterator[List[MetaPartition]]]:
"""
:meta private:
"""
if dispatch_by is not None and concat_partitions_on_primary_index:
raise ValueError(
"Both `dispatch_by` and `concat_partitions_on_primary_index` are provided, "
"`concat_partitions_on_primary_index` is deprecated and will be removed in the next major release. "
"Please only provide the `dispatch_by` argument. "
)
if concat_partitions_on_primary_index:
dispatch_by = dataset_factory.partition_keys
if dispatch_by is not None and not set(dispatch_by).issubset(
set(dataset_factory.index_columns)
):
raise RuntimeError(
f"Dispatch columns must be indexed.\nRequested index: {dispatch_by} but available index columns: {sorted(dataset_factory.index_columns)}"
)
check_predicates(predicates)
# Determine which indices need to be loaded.
index_cols: Set[str] = set()
if dispatch_by:
index_cols |= set(dispatch_by)
if predicates:
predicate_cols = set(columns_in_predicates(predicates))
predicate_index_cols = predicate_cols & set(dataset_factory.index_columns)
index_cols |= predicate_index_cols
for col in index_cols:
dataset_factory.load_index(col)
base_df = dataset_factory.get_indices_as_dataframe(
list(index_cols), predicates=predicates
)
if label_filter:
base_df = base_df[base_df.index.map(label_filter)]
indices_to_dispatch = {
name: ix.unload()
for name, ix in dataset_factory.indices.items()
if isinstance(ix, ExplicitSecondaryIndex)
}
if dispatch_by is not None:
base_df = cast(pd.DataFrame, base_df)
if len(dispatch_by) == 0:
merged_partitions = [((""), base_df)]
else:
# Group the resulting MetaParitions by partition keys or a subset of those keys
merged_partitions = base_df.groupby(
by=list(dispatch_by), sort=True, as_index=False
)
for group_name, group in merged_partitions:
if not isinstance(group_name, tuple):
group_name = (group_name,) # type: ignore
mps = []
logical_conjunction = list(
zip(dispatch_by, ["=="] * len(dispatch_by), group_name)
)
for label in group.index.unique():
mps.append(
MetaPartition.from_partition(
partition=dataset_factory.partitions[label],
dataset_metadata=dataset_factory.metadata
if dispatch_metadata
else None,
indices=indices_to_dispatch if dispatch_metadata else None,
metadata_version=dataset_factory.metadata_version,
table_meta=dataset_factory.table_meta,
partition_keys=dataset_factory.partition_keys,
logical_conjunction=logical_conjunction,
)
)
yield mps
else:
for part_label in base_df.index.unique():
part = dataset_factory.partitions[part_label]
yield MetaPartition.from_partition(
partition=part,
dataset_metadata=dataset_factory.metadata
if dispatch_metadata
else None,
indices=indices_to_dispatch if dispatch_metadata else None,
metadata_version=dataset_factory.metadata_version,
table_meta=dataset_factory.table_meta,
partition_keys=dataset_factory.partition_keys,
)
def dispatch_metapartitions_from_factory(
dataset_factory,
label_filter=None,
concat_partitions_on_primary_index=False,
predicates=None,
store=None,
):
if not callable(dataset_factory) and not isinstance(
dataset_factory, DatasetFactory):
raise TypeError("Need to supply a dataset factory!")
if predicates is not None:
dataset_factory, allowed_labels = _allowed_labels_by_predicates(
predicates, dataset_factory)
else:
allowed_labels = None
indices_to_dispatch = {
name: ix.copy(index_dct={})
for name, ix in dataset_factory.indices.items()
if isinstance(ix, ExplicitSecondaryIndex)
}
if concat_partitions_on_primary_index:
if dataset_factory.explicit_partitions:
dataset_factory = dataset_factory.load_partition_indices()
# Build up a DataFrame that contains per row a Partition and its
# primary index columns.
base_df = None
for part_key in dataset_factory.partition_keys:
idx = dataset_factory.indices[part_key].index_dct
df = _index_to_dataframe(part_key, idx, allowed_labels)
if base_df is None:
base_df = df
else:
base_df = base_df.merge(df, on=["__partition__"])
# Group the resulting MetaParitions by partition keys
merged_partitions = base_df.groupby(dataset_factory.partition_keys)
merged_partitions = merged_partitions["__partition__"].unique()
for row, labels in merged_partitions.iteritems():
mps = []
for label in labels:
mps.append(
MetaPartition.from_partition(
partition=dataset_factory.partitions[label],
dataset_metadata=dataset_factory.metadata,
indices=indices_to_dispatch,
metadata_version=dataset_factory.metadata_version,
table_meta=dataset_factory.table_meta,
partition_keys=dataset_factory.partition_keys,
))
yield mps
else:
if allowed_labels is not None:
partition_labels = allowed_labels
else:
partition_labels = dataset_factory.partitions.keys()
for part_label in partition_labels:
if label_filter is not None:
if not label_filter(part_label):
continue
part = dataset_factory.partitions[part_label]
yield MetaPartition.from_partition(
partition=part,
dataset_metadata=dataset_factory.metadata,
indices=indices_to_dispatch,
metadata_version=dataset_factory.metadata_version,
table_meta=dataset_factory.table_meta,
partition_keys=dataset_factory.partition_keys,
)
def dispatch_metapartitions_from_factory(
dataset_factory: DatasetFactory,
predicates: PredicatesType = None,
dispatch_by: Optional[List[str]] = None,
) -> Union[Iterator[MetaPartition], Iterator[List[MetaPartition]]]:
"""
:meta private:
"""
if dispatch_by is not None and not set(dispatch_by).issubset(
set(dataset_factory.index_columns)
):
raise RuntimeError(
f"Dispatch columns must be indexed.\nRequested index: {dispatch_by} but available index columns: {sorted(dataset_factory.index_columns)}"
)
check_predicates(predicates)
# Determine which indices need to be loaded.
index_cols: Set[str] = set()
if dispatch_by:
index_cols |= set(dispatch_by)
if predicates:
predicate_cols = set(columns_in_predicates(predicates))
predicate_index_cols = predicate_cols & set(dataset_factory.index_columns)
index_cols |= predicate_index_cols
for col in index_cols:
dataset_factory.load_index(col)
base_df = dataset_factory.get_indices_as_dataframe(
list(index_cols), predicates=predicates
)
if dispatch_by is not None:
base_df = cast(pd.DataFrame, base_df)
if len(dispatch_by) == 0:
merged_partitions = [((""), base_df)]
else:
# Group the resulting MetaParitions by partition keys or a subset of those keys
merged_partitions = base_df.groupby(
by=list(dispatch_by), sort=True, as_index=False
)
for group_name, group in merged_partitions:
if not isinstance(group_name, tuple):
group_name = (group_name,) # type: ignore
mps = []
logical_conjunction = list(
zip(dispatch_by, ["=="] * len(dispatch_by), group_name)
)
for label in group.index.unique():
mps.append(
MetaPartition.from_partition(
partition=dataset_factory.partitions[label],
metadata_version=dataset_factory.metadata_version,
schema=dataset_factory.schema,
partition_keys=dataset_factory.partition_keys,
logical_conjunction=logical_conjunction,
table_name=dataset_factory.table_name,
)
)
yield mps
else:
for part_label in base_df.index.unique():
part = dataset_factory.partitions[part_label]
yield MetaPartition.from_partition(
partition=part,
metadata_version=dataset_factory.metadata_version,
schema=dataset_factory.schema,
partition_keys=dataset_factory.partition_keys,
table_name=dataset_factory.table_name,
)