Let merge_datasets_as_delayed merge >2 datasets, filter by predicates and subset column

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
     ----------
     """

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,84 +22,39 @@ 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
     # 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,
+    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