diff --git a/kartothek/io/dask/delayed.py b/kartothek/io/dask/delayed.py index b74f3e21..3ccd9fb1 100644 --- a/kartothek/io/dask/delayed.py +++ b/kartothek/io/dask/delayed.py @@ -11,6 +11,7 @@ from kartothek.core.docs import default_docs from kartothek.core.factory import _ensure_factory from kartothek.core.naming import DEFAULT_METADATA_VERSION +from kartothek.core.urlencode import decode_key from kartothek.core.utils import _check_callable from kartothek.core.uuid import gen_uuid from kartothek.io_components.delete import ( @@ -38,6 +39,7 @@ raise_if_dataset_exists, store_dataset_from_partitions, ) +from kartothek.serialization._generic import check_predicates from ._update import _update_dask_partitions_one_to_one from ._utils import ( @@ -129,17 +131,24 @@ def garbage_collect_dataset__delayed( ) -def _load_and_merge_mps(mp_list, store, label_merger, metadata_merger, merge_tasks): - mp_list = [mp.load_dataframes(store=store) for mp in mp_list] +def _load_and_merge_mps( + mp_list, store, label_merger, metadata_merger, merge_tasks, predicates, columns +): + for i in range(len(mp_list)): + mp = mp_list[i] + uuid, _, _, _ = decode_key(mp.label) + assert uuid in columns.keys() + mp = mp.load_dataframes(store=store, columns=columns[uuid]) + mp_list[i] = mp mp = MetaPartition.merge_metapartitions( mp_list, label_merger=label_merger, metadata_merger=metadata_merger ) mp = mp.concat_dataframes() for task in merge_tasks: - mp = mp.merge_dataframes(**task) + mp = mp.merge_dataframes(**task, predicates=predicates) - return mp + return mp # Why return metapartitions here? @default_docs @@ -208,15 +217,90 @@ def merge_datasets_as_delayed( ... "output_label": 'merged_core_data' ... }, ... ] + """ + _check_callable(store) + + dataset_uuids = ( + [left_dataset_uuid, right_dataset_uuid] + if match_how == "right" + else [left_dataset_uuid, right_dataset_uuid] + ) + match_how = "first" if match_how in ["left", "right"] else match_how + + return merge_many_datasets_as_delayed( + dataset_uuids=dataset_uuids, + store=store, + merge_tasks=merge_tasks, + match_how=match_how, + label_merger=label_merger, + metadata_merger=metadata_merger, + ) + +@default_docs +def merge_many_datasets_as_delayed( + dataset_uuids, + store, + merge_tasks, + match_how="exact", + label_merger=None, + metadata_merger=None, + predicates=None, + columns=None, +): + """ + A dask.delayed graph to perform the merge of two full kartothek datasets. + Parameters + ---------- + dataset_uuids : List[str] + List of UUIDs of datasets (order does not matter in all merge schemas) + match_how : Union[str, Callable] + Define the partition label matching scheme. + Available implementations are: + * first: The first partitions are considered to be + the base partitions and **all** partitions of the + other dataset are joined to the first + partition. This should only be used if one of the + datasets contains very few partitions. + * prefix : The labels of the partitions of the dataset with fewest + partitions are considered to be the prefixes to the + other datasets. + * exact : All partition labels of the first dataset need to have + an exact match in the other datasets + * callable: A callable with signature func(left, right) which + returns a boolean to determine if the partitions match + If True, an exact match of partition labels between the to-be-merged + datasets is required in order to merge. + If False (Default), the partition labels of the dataset with fewer + partitions are interpreted as prefixes. + merge_tasks : List[Dict] + A list of merge tasks. Each item in this list is a dictionary giving + explicit instructions for a specific merge. + Each dict should contain key/values: + * 'output_label' : The table for the merged dataframe + * `merge_func`: A callable with signature + `merge_func(df_list, merge_kwargs)` to + handle the data preprocessing and merging. + Default pandas.concat. + * 'merge_kwargs' : The kwargs to be passed to the `merge_func` + Example: + .. code:: + >>> merge_tasks = [ + ... { + ... "output_label": "table", + "merge_func": pd.concat, + ... "merge_kwargs": {"axis": 1, "copy": False}, + ... }, + ... ] """ _check_callable(store) + check_predicates(predicates) mps = align_datasets( - left_dataset_uuid=left_dataset_uuid, - right_dataset_uuid=right_dataset_uuid, + dataset_uuids=dataset_uuids, store=store, match_how=match_how, + predicates=predicates, ) mps = map_delayed( _load_and_merge_mps, @@ -225,6 +309,8 @@ def merge_datasets_as_delayed( label_merger=label_merger, metadata_merger=metadata_merger, merge_tasks=merge_tasks, + predicates=predicates, + columns=columns, ) return list(mps) @@ -268,7 +354,6 @@ def read_dataset_as_delayed_metapartitions( Parameters ---------- - """ ds_factory = _ensure_factory( dataset_uuid=dataset_uuid, @@ -346,7 +431,6 @@ def read_dataset_as_delayed( """ A collection of dask.delayed objects to retrieve a dataset from store where each partition is loaded as a :class:`~pandas.DataFrame`. - Parameters ---------- """ @@ -386,7 +470,6 @@ def read_table_as_delayed( """ A collection of dask.delayed objects to retrieve a single table from a dataset as partition-individual :class:`~pandas.DataFrame` instances. - You can transform the collection of ``dask.delayed`` objects into a ``dask.dataframe`` using the following code snippet. As older kartothek specifications don't store schema information, this must be provided by @@ -444,7 +527,6 @@ def update_dataset_from_delayed( :class:`~karothek.io.metapartition.MetaPartition`. If you want to use this pipeline step for just deleting partitions without adding new ones you have to give an empty meta partition as input (``[Metapartition(None)]``). - Parameters ---------- """ diff --git a/kartothek/io_components/merge.py b/kartothek/io_components/merge.py index c1752fd2..60ffbffd 100644 --- a/kartothek/io_components/merge.py +++ b/kartothek/io_components/merge.py @@ -1,20 +1,19 @@ import logging +import pathlib -from kartothek.core.dataset import DatasetMetadata -from kartothek.io_components.metapartition import MetaPartition +from kartothek.io_components.read import dispatch_metapartitions from kartothek.io_components.utils import _instantiate_store LOGGER = logging.getLogger(__name__) -def align_datasets(left_dataset_uuid, right_dataset_uuid, store, match_how="exact"): +def align_datasets(dataset_uuids, store, predicates=None, match_how="exact"): """ Determine dataset partition alignment Parameters ---------- - left_dataset_uuid : basestring - right_dataset_uuid : basestring + dataset_uuids : List[basestring] store : KeyValuestore or callable match_how : basestring or callable, {exact, prefix, all, callable} @@ -23,37 +22,6 @@ def align_datasets(left_dataset_uuid, right_dataset_uuid, store, match_how="exac 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 @@ -61,46 +29,32 @@ def align_datasets(left_dataset_uuid, right_dataset_uuid, store, match_how="exac # 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, + mp_0 = dispatch_metapartitions(dataset_uuids[0], store, predicates=predicates) + + for mp_0_i in mp_0: + l_1 = mp_0_i.label + res = [mp_0_i] + for j in range(1, len(dataset_uuids)): + mp_j = dispatch_metapartitions( + dataset_uuids[j], store, predicates=predicates ) - res.append( - MetaPartition.from_partition( - partition=p_2, metadata_version=metadata_version + for mp_j_i in mp_j: + l_j = mp_j_i.label + if callable(match_how) and not match_how(l_1, l_j): + continue + if match_how == "prefix": + raise NotImplementedError + elif match_how == "exact" and not ( + pathlib.Path(l_1).parent == pathlib.Path(l_j).parent + ): + continue + + LOGGER.debug( + "Found alignment between partitions " "(%s, %s) and" "(%s, %s)", + dataset_uuids[0], + l_1, + dataset_uuids[j], + l_j, ) - ) - - # 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) - ) + res.append(mp_j_i) yield res diff --git a/kartothek/io_components/metapartition.py b/kartothek/io_components/metapartition.py index c1c2646c..3dd8e3b0 100644 --- a/kartothek/io_components/metapartition.py +++ b/kartothek/io_components/metapartition.py @@ -137,7 +137,7 @@ def _impl(self, *method_args, **method_kwargs): result = result.add_metapartition(mp, schema_validation=False) if not isinstance(result, MetaPartition): raise ValueError( - "Result for method {} is not a `MetaPartition` but".format( + "Result for method {} is not a `MetaPartition` but {}".format( method.__name__, type(method_return) ) ) @@ -701,9 +701,10 @@ def load_dataframes( # the conditition. # # We separate these predicates into their index and their Parquet part. - split_predicates, has_index_condition = self._split_predicates_in_index_and_content( - predicates - ) + ( + split_predicates, + has_index_condition, + ) = self._split_predicates_in_index_and_content(predicates) filtered_predicates = [] if has_index_condition: @@ -856,7 +857,12 @@ def _reconstruct_index_columns( @_apply_to_list def merge_dataframes( - self, left, right, output_label, merge_func=pd.merge, merge_kwargs=None + self, + output_label, + labels=None, + merge_func=pd.merge, + merge_kwargs=None, + predicates=None, ): """ Merge internal dataframes. @@ -873,10 +879,8 @@ def merge_dataframes( Parameters ---------- - left : basestring - Category of the left dataframe. - right : basestring - Category of the right dataframe. + labels: basestring + List of categories of dataframes. output_label : basestring Category for the newly created dataframe merge_func : callable, optional @@ -896,29 +900,32 @@ def merge_dataframes( if merge_kwargs is None: merge_kwargs = {} - left_df = new_data.pop(left) - right_df = new_data.pop(right) + if labels is None: + labels = list(new_data.keys()) + + dfs = [new_data.pop(label) for label in labels] LOGGER.debug("Merging internal dataframes of %s", self.label) try: - df_merged = merge_func(left_df, right_df, **merge_kwargs) + df_merged = merge_func(dfs, **merge_kwargs) except TypeError: LOGGER.error( - "Tried to merge using %s with\n left:%s\nright:%s\n " "kwargs:%s", + "Tried to merge using %s with\n %s\n " "kwargs:%s", merge_func.__name__, - left_df.head(), - right_df.head(), + [df.head() for df in dfs], merge_kwargs, ) raise + df_merged = filter_df_from_predicates(df_merged, predicates) + new_data[output_label] = df_merged new_table_meta = copy(self.table_meta) # The tables are no longer part of the MetaPartition, thus also drop # their schema. - del new_table_meta[left] - del new_table_meta[right] + for label in labels: + del new_table_meta[label] new_table_meta[output_label] = make_meta( df_merged, origin="{}/{}".format(output_label, self.label), @@ -1474,6 +1481,7 @@ def merge_metapartitions(metapartitions, label_merger=None, metadata_merger=None label=new_label, data=new_data, dataset_metadata=new_ds_meta, + partition_keys=metapartitions[0].partition_keys, metadata_version=new_metadata_version, logical_conjunction=logical_conjunction, ) diff --git a/kartothek/io_components/read.py b/kartothek/io_components/read.py index 7eee0482..915a9f8d 100644 --- a/kartothek/io_components/read.py +++ b/kartothek/io_components/read.py @@ -88,7 +88,7 @@ def dispatch_metapartitions_from_factory( mps.append( MetaPartition.from_partition( partition=dataset_factory.partitions[label], - dataset_metadata=dataset_factory.metadata, + dataset_metadata=dataset_factory.metadata, # Should this not be dataset_factory.dataset_metadata? indices=indices_to_dispatch, metadata_version=dataset_factory.metadata_version, table_meta=dataset_factory.table_meta,