Improved dataset merging logic.

Started working on improving the dataset and feature set merging logic.
This commit improves dataset merging. It also includes updated tests for
the changed code. Apart from that, it also includes some minor
refactoring and code cleanup.

Author: Sebastijan-Dominis

docs/roadmap.md

@@ -24,7 +24,7 @@ _Last updated: 2026-03-23_
     - Status: done (2026-03-23)
     - Importance: medium
 - [ ] Improve dataset and feature set merging logic
-    - Status: planned
+    - Status: in progress
     - Importance: medium
 - [ ] Create infer.py
     - Status: planned

ml/data/merge/merge_dataset_into_main.py

@@ -3,19 +3,52 @@
 import logging
 from pathlib import Path
 
+import networkx as nx
 import pandas as pd
 
 from ml.data.validation.validate_data import validate_data
-from ml.exceptions import DataError
+from ml.exceptions import ConfigError, DataError
+from ml.feature_freezing.freeze_strategies.tabular.config.models import (
+    DatasetConfig,
+    MergeHow,
+    MergeValidate,
+)
 from ml.utils.loaders import load_json
 
 logger = logging.getLogger(__name__)
 
+def build_dataset_dag(datasets: list[DatasetConfig]) -> list[str]:
+    """Build DAG and return topological merge order."""
+    G = nx.DiGraph()
+    for ds in datasets:
+        G.add_node(ds.name)
+
+    for i, ds1 in enumerate(datasets):
+        keys1 = set(ds1.merge_key if isinstance(ds1.merge_key, list) else [ds1.merge_key])
+        for j, ds2 in enumerate(datasets):
+            if i >= j:
+                continue
+            keys2 = set(ds2.merge_key if isinstance(ds2.merge_key, list) else [ds2.merge_key])
+            if keys1 & keys2:
+                G.add_edge(ds1.name, ds2.name)
+
+    try:
+        merge_order = list(nx.topological_sort(G))
+    except nx.NetworkXUnfeasible:
+        msg = "Cyclic merge dependency detected among datasets."
+        logger.error(msg)
+        raise ConfigError(msg)
+
+    logger.debug(f"DAG-based merge order: {merge_order}")
+    return merge_order
+
 def merge_dataset_into_main(
     data: pd.DataFrame,
     df: pd.DataFrame,
     *,
-    merge_key: str,
+    merge_key: str | list[str],
+    merge_how: MergeHow = "inner",
+    merge_validate: MergeValidate = "m:m",
     dataset_name: str,
     dataset_version: str,
     dataset_snapshot_path: Path,
@@ -26,7 +59,9 @@ def merge_dataset_into_main(
     Args:
         data: Main dataframe accumulated from previous merges.
         df: New dataset dataframe to merge into ``data``.
-        merge_key: Key column used for inner merge.
+        merge_key: Key column(s) used for inner merge.
+        merge_how: Type of merge to perform.
+        merge_validate: Merge validation for row explosion protection.
         dataset_name: Logical dataset name used for logging/errors.
         dataset_version: Dataset version string.
         dataset_snapshot_path: Snapshot directory containing dataset metadata.
@@ -44,31 +79,58 @@ def merge_dataset_into_main(
         dataset-level merge diagnostics/warnings.
     """
 
-    if merge_key not in df.columns:
-        msg = f"Dataset {dataset_name} is missing merge key '{merge_key}' column."
+    missing_keys_in_df = set(merge_key) - set(df.columns)
+    if missing_keys_in_df:
+        msg = f"Dataset {dataset_name} is missing merge key(s): {missing_keys_in_df}"
         logger.error(msg)
         raise DataError(msg)
-    if merge_key not in data.columns and not data.empty:
-        msg = f"Merge key '{merge_key}' column from dataset {dataset_name} not found in the main dataset for merging."
+
+    if not data.empty:
+        missing_keys_in_main = set(merge_key) - set(data.columns)
+        if missing_keys_in_main:
+            msg = f"Merge key(s) {missing_keys_in_main} not found in main dataset when merging {dataset_name}"
+            logger.error(msg)
+            raise DataError(msg)
+
+    original_rows = len(data) if not data.empty else 0
+
+    overlapping_cols = set(data.columns) & set(df.columns) - set(merge_key)
+    if overlapping_cols:
+        logger.warning(f"Dropping overlapping columns {overlapping_cols} from dataset {dataset_name} before merge")
+        df = df.drop(columns=overlapping_cols)
+
+    # Merge with row explosion protection
+    try:
+        if data.empty:
+            data = df
+        else:
+            data = pd.merge(
+                data,
+                df,
+                how=merge_how,
+                on=merge_key,
+                validate=merge_validate,
+                suffixes=("", "_dup"),
+            )
+    except pd.errors.MergeError as e:
+        msg = f"Merge failed for dataset {dataset_name} v{dataset_version}: {e}"
         logger.error(msg)
         raise DataError(msg)
 
+    if len(data) < original_rows:
+        logger.warning(
+            f"Row count decreased after merging {dataset_name}: {original_rows} -> {len(data)}"
+        )
+
+    if merge_key:
+        data = data.sort_values(by=merge_key).reset_index(drop=True)
+
     dataset_metadata = load_json(dataset_snapshot_path / "metadata.json")
     data_hash = validate_data(data_path=dataset_path, metadata=dataset_metadata)
 
-    logger.debug(f"Starting to merge dataset {dataset_name} v{dataset_version} snapshot {dataset_snapshot_path.name} with shape {df.shape} into the main dataset with shape {data.shape}")
-
-    if data.empty:
-        data = df
-    else:
-        overlapping_cols = set(data.columns) & set(df.columns) - {merge_key}
-        if overlapping_cols:
-            logger.warning(f"Overlapping columns found in dataset {dataset_name}: {overlapping_cols}. Dropping these columns from the new dataset before merge.")
-            df = df.drop(columns=overlapping_cols)
-        data = pd.merge(data, df, how="inner", on=merge_key)
-        if data.empty:
-            msg = f"Merged dataset is empty after merging with {dataset_name}. Please check the '{merge_key}' alignment across datasets."
-            logger.error(msg)
-            raise DataError(msg)
+    logger.debug(
+        f"Merged dataset {dataset_name} v{dataset_version} snapshot {dataset_snapshot_path.name} "
+        f"into main dataframe, resulting shape: {data.shape}"
+    )
 
     return data, data_hash

ml/feature_freezing/freeze_strategies/tabular/config/models.py

@@ -11,15 +11,70 @@
 
 logger = logging.getLogger(__name__)
 
+MergeHow = Literal["left", "right", "inner", "outer", "cross"]
+
+MergeValidate = Literal[
+    "one_to_one",
+    "1:1",
+    "one_to_many",
+    "1:m",
+    "many_to_one",
+    "m:1",
+    "many_to_many",
+    "m:m",
+]
+
 class DatasetConfig(BaseModel):
     """Source dataset definition used for feature freezing ingestion."""
 
     ref: str = Field("data/processed", description="Reference path for the dataset, e.g., 'data/processed'")
     name: str = Field(..., description="Name of the dataset, e.g., 'hotel_bookings'")
     version: str = Field(..., description="Version of the dataset, e.g., 'v1'")
     format: Literal["csv", "parquet"]
-    merge_key: str = Field("row_id", description="Key to merge datasets on, default is 'row_id'")
-    path_suffix: str = Field("data.{format}", description="Suffix for the dataset file, supports {format} placeholder")
+    merge_key: str | list[str] = Field(
+        "row_id", description="Key(s) to merge datasets on, default is 'row_id'"
+    )
+    merge_how: MergeHow = Field(
+        "inner", description="Merge type, default 'inner'"
+    )
+    merge_validate: MergeValidate = Field(
+        "m:m", description="Merge validation for row explosion protection"
+    )
+    path_suffix: str = Field(
+        "data.{format}", description="Suffix for the dataset file, supports {format} placeholder"
+    )
+
+    @field_validator("merge_key", mode="before")
+    def ensure_merge_key_list(cls, v):
+        """Always convert merge_key to a list internally."""
+        if isinstance(v, str):
+            return [v]
+        if isinstance(v, list) and all(isinstance(i, str) for i in v):
+            return v
+        raise ConfigError("merge_key must be a string or a list of strings")
+
+    @field_validator("merge_how")
+    def validate_merge_how(cls, v):
+        if v not in {"inner", "left", "right", "outer"}:
+            raise ConfigError(f"Invalid merge_how: {v}")
+        return v
+
+    @field_validator("merge_validate")
+    def normalize_merge_validate(cls, v):
+        """Normalize merge_validate to pandas-compatible strings."""
+        mapping = {
+            "one_to_one": "1:1",
+            "1:1": "1:1",
+            "one_to_many": "1:m",
+            "1:m": "1:m",
+            "many_to_one": "m:1",
+            "m:1": "m:1",
+            "many_to_many": "m:m",
+            "m:m": "m:m",
+        }
+        if v not in mapping:
+            raise ConfigError(f"Invalid merge_validate: {v}")
+        return mapping[v]
 
 class FeatureRolesConfig(BaseModel):
     """Feature role partitioning for validation and downstream usage."""

ml/feature_freezing/utils/data_loader.py

@@ -5,7 +5,7 @@
 
 import pandas as pd
 
-from ml.data.merge.merge_dataset_into_main import merge_dataset_into_main
+from ml.data.merge.merge_dataset_into_main import build_dataset_dag, merge_dataset_into_main
 from ml.exceptions import ConfigError, DataError
 from ml.feature_freezing.freeze_strategies.tabular.config.models import TabularFeaturesConfig
 from ml.snapshot_bindings.extraction.get_snapshot_binding import get_and_validate_snapshot_binding
@@ -16,10 +16,12 @@
 
 logger = logging.getLogger(__name__)
 
+
+
 def load_data_with_lineage(
-        config: TabularFeaturesConfig,
-        snapshot_binding_key: str | None
-    ) -> tuple[pd.DataFrame, list[DataLineageEntry]]:
+    config: TabularFeaturesConfig,
+    snapshot_binding_key: str | None = None
+) -> tuple[pd.DataFrame, list[DataLineageEntry]]:
     """Load configured datasets, merge them, and build lineage entries.
 
     Args:
@@ -30,9 +32,6 @@ def load_data_with_lineage(
         dataset lineage metadata.
     """
 
-    data = pd.DataFrame()
-    data_lineage = []
-
     if not config.data:
         msg = "No datasets specified in the configuration."
         logger.error(msg)
@@ -41,70 +40,71 @@ def load_data_with_lineage(
     datasets_binding = None
     if snapshot_binding_key:
         snapshot_binding_config = get_and_validate_snapshot_binding(
-            snapshot_binding_key,
-            expect_dataset_bindings=True
+            snapshot_binding_key, expect_dataset_bindings=True
         )
         datasets_binding = snapshot_binding_config.datasets
 
-    for dataset in config.data:
-        dataset_versions = datasets_binding.get(dataset.name) if datasets_binding else None
-        dataset_snapshot_binding = dataset_versions.get(dataset.version) if dataset_versions else None
+    dataset_dict = {ds.name: ds for ds in config.data}
+    merge_order = build_dataset_dag(config.data)
+
+    merged_data = pd.DataFrame()
+    data_lineage: list[DataLineageEntry] = []
+
+    for ds_name in merge_order:
+        ds = dataset_dict[ds_name]
+        dataset_versions = datasets_binding.get(ds.name) if datasets_binding else None
+        dataset_snapshot_binding = dataset_versions.get(ds.version) if dataset_versions else None
+
         if snapshot_binding_key:
             if not dataset_snapshot_binding:
-                msg = f"No snapshot binding found for dataset {dataset.name} version {dataset.version} under snapshot binding key {snapshot_binding_key}."
+                msg = f"No snapshot binding found for {ds.name} {ds.version} under {snapshot_binding_key}"
                 logger.error(msg)
                 raise ConfigError(msg)
             dataset_snapshot = dataset_snapshot_binding.snapshot
-            dataset_snapshot_path = Path(dataset.ref) / dataset.name / dataset.version / dataset_snapshot
+            dataset_snapshot_path = Path(ds.ref) / ds.name / ds.version / dataset_snapshot
         else:
-            dataset_path = Path(dataset.ref) / dataset.name / dataset.version
-            dataset_snapshot_path = get_latest_snapshot_path(dataset_path)
+            dataset_path_base = Path(ds.ref) / ds.name / ds.version
+            dataset_snapshot_path = get_latest_snapshot_path(dataset_path_base)
 
-        if dataset.format not in HASH_LOADER_REGISTRY:
-            msg = f"Unsupported data format for loading and hashing: {dataset.format}"
-            logger.error(msg)
-            raise ConfigError(msg)
-
-        dataset_path = dataset_snapshot_path / dataset.path_suffix.format(format=dataset.format)
+        dataset_path = dataset_snapshot_path / ds.path_suffix.format(format=ds.format)
         if not dataset_path.exists():
             msg = f"Dataset file not found at expected path: {dataset_path}"
             logger.error(msg)
             raise DataError(msg)
-        df = read_data(dataset.format, dataset_path)
-        merge_key = dataset.merge_key
+        df = read_data(ds.format, dataset_path)
 
-        data, data_hash = merge_dataset_into_main(
-            data=data,
+        merged_data, data_hash = merge_dataset_into_main(
+            data=merged_data,
             df=df,
-            merge_key=merge_key,
-            dataset_name=dataset.name,
-            dataset_version=dataset.version,
+            merge_key=ds.merge_key,
+            merge_how=ds.merge_how,
+            merge_validate=ds.merge_validate,
+            dataset_name=ds.name,
+            dataset_version=ds.version,
             dataset_snapshot_path=dataset_snapshot_path,
             dataset_path=dataset_path,
         )
 
-        loader_validation_hash = HASH_LOADER_REGISTRY[dataset.format](dataset_path)
+        loader_validation_hash = HASH_LOADER_REGISTRY[ds.format](dataset_path)
 
         entry = DataLineageEntry(
-            ref=dataset.ref,
-            name=dataset.name,
-            version=dataset.version,
-            format=dataset.format,
-            path_suffix=dataset.path_suffix,
-            merge_key=dataset.merge_key,
+            ref=ds.ref,
+            name=ds.name,
+            version=ds.version,
+            format=ds.format,
+            path_suffix=ds.path_suffix,
+            merge_key=ds.merge_key,
+            merge_how=ds.merge_how,
+            merge_validate=ds.merge_validate,
             snapshot_id=dataset_snapshot_path.name,
             path=dataset_path.as_posix(),
             loader_validation_hash=loader_validation_hash,
             data_hash=data_hash,
             row_count=len(df),
             column_count=len(df.columns),
         )
-
         data_lineage.append(entry)
+        logger.debug(f"Loaded dataset {ds.name} {ds.version} snapshot {dataset_snapshot_path.name} loader hash {loader_validation_hash}")
 
-        logger.debug(f"Loaded dataset {dataset.name} {dataset.version} snapshot {dataset_snapshot_path.name} "
-             f"file {dataset_path} loader hash {loader_validation_hash}")
-
-    logger.info(f"Completed loading {len(config.data)} dataframes. Final merged dataframe shape: {data.shape}. Lineage: {data_lineage}")
-
-    return data, data_lineage
+    logger.info(f"Completed loading {len(config.data)} datasets. Final merged dataframe shape: {merged_data.shape}.")
+    return merged_data, data_lineage