Fix: Cube index validation (JDASoftwareGroup#413)

The issue occured when updating a cube with multiple datasets, which also have different dimension columns.

This fixes issue JDASoftwareGroup#413.

CHANGES.rst

@@ -6,6 +6,7 @@ Version 3.19.1 (2021-02-XX)
 ===========================
 
 * Allow ``pyarrow==3`` as a dependency.
+* Fix an issue with the cube index validation introduced in v3.19.0 (#413).
 
 Version 3.19.0 (2021-02-12)
 ===========================

kartothek/io/dask/common_cube.py

@@ -66,17 +66,16 @@ def ensure_valid_cube_indices(
     `index_columns` and `suppress_index_on` fields adjusted to reflect the
     existing datasets.
     """
-    required_indices = set(cube.index_columns)
-    suppress_index_on = set(cube.suppress_index_on)
+    dataset_indices = []
     for ds in existing_datasets.values():
         for internal_table in ds.table_meta:
             dataset_columns = set(ds.table_meta[internal_table].names)
-            table_indices = required_indices & dataset_columns
+            table_indices = cube.index_columns & dataset_columns
             compatible_indices = _ensure_compatible_indices(ds, table_indices)
             if compatible_indices:
-                dataset_indices = set(compatible_indices)
-                suppress_index_on -= dataset_indices
-                required_indices |= dataset_indices
+                dataset_indices.append(set(compatible_indices))
+    required_indices = cube.index_columns.union(*dataset_indices)
+    suppress_index_on = cube.suppress_index_on.difference(*dataset_indices)
     # Need to remove dimension columns since they *are* technically indices but
     # the cube interface class declares them as not indexed just to add them
     # later on, assuming it is not blacklisted

tests/io/dask/test_common_cube.py

@@ -0,0 +1,51 @@
+from kartothek.core.cube.cube import Cube
+from kartothek.core.dataset import DatasetMetadata
+from kartothek.io.dask.common_cube import ensure_valid_cube_indices
+
+
+class FakeSeedTableMetadata:
+    names = ["d1", "d2", "p", "i"]
+
+
+class FakeExtraTableMetadata:
+    names = ["d1", "p", "i"]
+
+
+def test_cube_indices_are_validated():
+    source_metadata = DatasetMetadata.from_dict(
+        {
+            "dataset_uuid": "source",
+            "dataset_metadata_version": 4,
+            "table_meta": {"table": FakeSeedTableMetadata()},
+            "partition_keys": ["p"],
+            "indices": {
+                "d1": {"1": ["part_1"]},
+                "d2": {"1": ["part_1"]},
+                "i": {"1": ["part_1"]},
+            },
+        }
+    )
+    extra_metadata = DatasetMetadata.from_dict(
+        {
+            "dataset_uuid": "extra",
+            "dataset_metadata_version": 4,
+            "table_meta": {"table": FakeExtraTableMetadata()},
+            "partition_keys": ["p"],
+            "indices": {
+                "i": {"1": ["part_1"]},
+            },
+        }
+    )
+    cube = Cube(
+        dimension_columns=["d1", "d2"],
+        partition_columns=["p"],
+        uuid_prefix="cube",
+        seed_dataset="source",
+        index_columns=["i"],
+    )
+
+    validated_cube = ensure_valid_cube_indices(
+        {"source": source_metadata, "extra": extra_metadata}, cube
+    )
+
+    assert validated_cube == cube