Merge pull request #463 from infehm/add_predicates_to_conditions_converter

contribution / Add predicates to conditions converter

Author: stephan-hesselmann-by

CHANGES.rst

@@ -3,10 +3,11 @@ Changelog
 =========
 
 
-Version 5.0.1 (2021-06-xx)
+Version 5.1.0 (2021-06-xx)
 ==========================
 * Add :meth:`~kartothek.io.eager.copy_dataset` to copy and optionally rename datasets within one store or between stores (eager only)
 * Add renaming option to :meth:`~kartothek.io.eager_cube.copy_cube`
+* Add predicates to cube condition converter to :meth:`~kartothek.utils.predicate_converter`
 
 
 Version 5.0.0 (2021-06-23)

kartothek/utils/predicate_converter.py

@@ -0,0 +1,88 @@
+"""
+Helper module to convert kartothek dataset load predicates into cube conditions.
+"""
+
+from typing import Any, List, Sequence, Tuple
+
+import pandas as pd
+import pyarrow as pa
+
+from kartothek.core.cube.conditions import Condition
+
+
+def write_predicate_as_cube_condition(predicate: Tuple[str, str, Any]) -> Condition:
+    """
+    Rewrites a single io.dask.dataset 'read_dataset_as_ddf' predicate condition as cube condition.
+
+    Remark: This function is restricted by "Condition.from_string" which does not allow for IsInCondition
+    and InIntervalCondition and will throw an error if conditions of those types are passed.
+
+    Parameters
+    ----------
+    predicate: list
+        list containing single predicate definition
+
+    Returns
+    -------
+    condition: Condition
+        cube condition containing the predicate definition
+    """
+    condition_string = None
+    parameter_format_dict = {}
+
+    if len(predicate) != 3:
+        raise ValueError("Please use predicates consisting of exactly 3 entries")
+
+    if type(predicate[2]) == int:
+        condition_string = f"{predicate[0]} {predicate[1]} {str(predicate[2])}"
+        parameter_format_dict[predicate[0]] = pa.int64()
+    if type(predicate[2]) == str:
+        condition_string = f"{predicate[0]} {predicate[1]} {predicate[2]}"
+        parameter_format_dict[predicate[0]] = pa.string()
+    if type(predicate[2]) == pd._libs.tslibs.timestamps.Timestamp:
+        condition_string = (
+            f"{predicate[0]} {predicate[1]} {predicate[2].strftime('%Y-%m-%d')}"
+        )
+        parameter_format_dict[predicate[0]] = pa.timestamp("s")
+    if type(predicate[2]) == bool:
+        condition_string = f"{predicate[0]} {predicate[1]} {str(predicate[2])}"
+        parameter_format_dict[predicate[0]] = pa.bool_()
+    if type(predicate[2]) == float:
+        condition_string = f"{predicate[0]} {predicate[1]} {str(predicate[2])}"
+        parameter_format_dict[predicate[0]] = pa.float64()
+
+    if condition_string is not None:
+        condition = Condition.from_string(condition_string, parameter_format_dict)
+    else:
+        raise TypeError(
+            "Please only enter predicates for parameter values of the following type:"
+            " str, int, float, bool or pandas timestamp, "
+        )
+    return condition
+
+
+def convert_predicates_to_cube_conditions(
+    predicates: List[List[Tuple[str, str, Any]]],
+) -> Sequence[Condition]:
+    """
+    Converts a io.dask.dataset 'read_dataset_as_ddf' predicate to a cube condition
+
+    Parameters
+    ----------
+    predicates: list
+        list containing a list of single predicates
+
+    Returns
+    -------
+    condition: Condition
+        cube condition containing the combined predicate definitions
+    """
+    condition: Any = ()
+    if len(predicates) > 1:
+        raise ValueError(
+            "Cube conditions cannot handle 'or' operators, therefore, "
+            "please pass a predicate list with one element."
+        )
+    for predicate in predicates[0]:
+        condition = condition + (write_predicate_as_cube_condition(predicate),)
+    return condition

tests/utils/test_predicate_converter.py

@@ -0,0 +1,76 @@
+import pandas as pd
+import pytest
+
+from kartothek.core.cube.conditions import C
+from kartothek.utils.predicate_converter import (
+    convert_predicates_to_cube_conditions,
+    write_predicate_as_cube_condition,
+)
+
+
+@pytest.mark.parametrize(
+    "test_input, expected",
+    [
+        (["column", "==", 1], (C("column") == 1)),
+        (["column", "==", 1.1], C("column") == 1.1),
+        (["column", "==", "1"], C("column") == "1"),
+        (
+            ["date", "==", pd.to_datetime("2020-01-01")],
+            (C("date") == pd.Timestamp("2020-01-01 00:00:00")),
+        ),
+        (["column", "==", True], C("column") == True),  # noqa: E712)
+        (["column", "<=", 1], (C("column") <= 1)),
+        (["column", ">=", 1], (C("column") >= 1)),
+        (["column", "!=", 1], (C("column") != 1)),
+    ],
+)
+def test_write_predicate_as_cube_condition(test_input, expected):
+    assert write_predicate_as_cube_condition(test_input) == expected
+
+
+def test_raises_type_error_write_predicate_as_cube_condition():
+    with pytest.raises(
+        TypeError,
+        match="Please only enter predicates for parameter values of the "
+        "following type: str, int, float, bool or pandas timestamp, ",
+    ):
+        write_predicate_as_cube_condition(
+            ("date", "==", pd.to_datetime("2020-01-01").date())
+        )
+
+
+def test_raises_value_error_write_predicate_as_cube_condition():
+    with pytest.raises(
+        ValueError, match="Please use predicates consisting of exactly 3 entries"
+    ):
+        write_predicate_as_cube_condition(("date", "=="))
+
+    with pytest.raises(
+        ValueError, match="Please use predicates consisting of exactly 3 entries"
+    ):
+        write_predicate_as_cube_condition(("date", "==", "date", "=="))
+
+
+@pytest.mark.parametrize(
+    "test_input, expected",
+    [
+        ([[("column", "==", 1)]], (C("column") == 1,)),
+        (
+            [[("column", "==", 1), ("column2", "==", "1")]],
+            (C("column") == 1, C("column2") == "1"),
+        ),
+    ],
+)
+def test_convert_predicates_to_cube_conditions(test_input, expected):
+    assert convert_predicates_to_cube_conditions(test_input) == expected
+
+
+def test_raises_value_error_convert_predicates_to_cube_conditions():
+    with pytest.raises(
+        ValueError,
+        match="Cube conditions cannot handle 'or' operators, therefore, "
+        "please pass a predicate list with one element.",
+    ):
+        convert_predicates_to_cube_conditions(
+            [[("column", "==", 1)], [("column2", "==", 2)]]
+        )