diff --git a/docs/docs/in_depth/join_data.md b/docs/docs/in_depth/join_data.md index 95281035..f84b7c5e 100644 --- a/docs/docs/in_depth/join_data.md +++ b/docs/docs/in_depth/join_data.md @@ -51,7 +51,7 @@ Join Types specify how two datasets are merged based on their keys. The framewor - Inner Join, - Left Join, - Outer Join, -- Right Join (use sparingly; prefer left joins when possible). +- Right Join (use sparingly; prefer left joins when possible). The declared left feature group's data is the merge engine's left argument, whichever compute framework the join executes in. - ASOF Join (point-in-time / as-of: equi match on the by-keys, nearest time match on the time columns). ```python diff --git a/mloda/core/prepare/execution_plan.py b/mloda/core/prepare/execution_plan.py index 4985898a..19cbb716 100644 --- a/mloda/core/prepare/execution_plan.py +++ b/mloda/core/prepare/execution_plan.py @@ -17,7 +17,7 @@ from mloda.core.prepare.joinstep_collection import JoinStepCollection from mloda.core.prepare.graph.graph import Graph from mloda.core.prepare.resolve_graph import PlannedQueue -from mloda.core.prepare.resolve_links import LinkFrameworkTrekker, LinkTrekker +from mloda.core.prepare.resolve_links import LinkFrameworkTrekker, LinkTrekker, inheritance_distance from mloda.core.core.step.feature_group_step import FeatureGroupStep from mloda.core.core.step.join_step import JoinStep from mloda.core.core.step.transform_frame_work_step import TransformFrameworkStep @@ -42,6 +42,13 @@ def _filter_options_sort_key(single_filter: SingleFilter) -> tuple[str, str]: ) +def _nearest_frameworks(frameworks_by_distance: dict[int, set[type[ComputeFramework]]]) -> set[type[ComputeFramework]]: + """A declared side is held by its closest subclasses only; farther ones answer for a different side.""" + if not frameworks_by_distance: + return set() + return frameworks_by_distance[min(frameworks_by_distance)] + + class ExecutionPlan: def __init__( self, @@ -439,13 +446,31 @@ def run_link( destination_framework_uuids: set[UUID] = set() source_framework_uuids: set[UUID] = set() + left_frameworks_by_distance: dict[int, set[type[ComputeFramework]]] = defaultdict(set) + right_frameworks_by_distance: dict[int, set[type[ComputeFramework]]] = defaultdict(set) + for uuid in required_uuids: - if graph.get_nodes()[uuid].feature.get_compute_framework() == destination_framework: + node = graph.get_nodes()[uuid] + node_framework = node.feature.get_compute_framework() + + if node_framework == destination_framework: destination_framework_uuids.add(uuid) - if graph.get_nodes()[uuid].feature.get_compute_framework() == source_framework: + if node_framework == source_framework: source_framework_uuids.add(uuid) + # Links match polymorphically, so a subclass of a declared side counts as that side, ranked by distance. + if issubclass(node.feature_group_class, link.left_feature_group): + left_distance = inheritance_distance(node.feature_group_class, link.left_feature_group) + left_frameworks_by_distance[left_distance].add(node_framework) + + if issubclass(node.feature_group_class, link.right_feature_group): + right_distance = inheritance_distance(node.feature_group_class, link.right_feature_group) + right_frameworks_by_distance[right_distance].add(node_framework) + + declared_left_frameworks = _nearest_frameworks(left_frameworks_by_distance) + declared_right_frameworks = _nearest_frameworks(right_frameworks_by_distance) + # The order shows which items should be added first. # Thus, we need to make sure that higher ordered links are calculated first. for k, v in link_trekker.order.items(): @@ -483,13 +508,33 @@ def run_link( required_uuids, destination_framework_uuids, source_framework_uuids, - swap_merge_sides, + self.swap_merge_sides_by_declared_side( + destination_framework, declared_left_frameworks, declared_right_frameworks, swap_merge_sides + ), ) # This makes sure that we do not write on the same datasets due to overlapping joins at once. self.joinstep_collection.add(js) return js + @staticmethod + def swap_merge_sides_by_declared_side( + destination_framework: type[ComputeFramework], + declared_left_frameworks: set[type[ComputeFramework]], + declared_right_frameworks: set[type[ComputeFramework]], + fallback: bool, + ) -> bool: + """The declared left group's data must stay the merge engine's left argument, wherever the join runs.""" + holds_left = destination_framework in declared_left_frameworks + holds_right = destination_framework in declared_right_frameworks + + if holds_left and not holds_right: + return False + if holds_right and not holds_left: + return True + # Self links and sides sharing one framework are not decidable from the declared sides. + return fallback + def find_fg_per_uuid( self, pre_execution_plan: list[LinkFrameworkTrekker | FeatureGroupStep], uuid: UUID ) -> type[FeatureGroup]: diff --git a/mloda/core/prepare/resolve_links.py b/mloda/core/prepare/resolve_links.py index d0a25247..ca306a4e 100644 --- a/mloda/core/prepare/resolve_links.py +++ b/mloda/core/prepare/resolve_links.py @@ -11,6 +11,12 @@ LinkFrameworkTrekker = tuple[Link, type[ComputeFramework], type[ComputeFramework]] +def inheritance_distance(child: type, parent: type) -> int: + """Steps from child to parent in the MRO, or 9999 if parent is not in child's hierarchy.""" + mro: tuple[type, ...] = getattr(child, "__mro__", ()) + return mro.index(parent) if parent in mro else 9999 + + class LinkTrekker: """This class is used to keep track of Links and which children depend on this link.""" @@ -353,16 +359,7 @@ def _find_matching_links( return self._select_most_specific_links(polymorphic_matches, left_fg, right_fg) def _inheritance_distance(self, child: type, parent: type) -> int: - """Calculate the inheritance distance from child to parent in the MRO. - - Returns the number of steps in the Method Resolution Order from child to parent. - Returns a large number if parent is not in child's MRO. - """ - try: - mro = child.__mro__ - return mro.index(parent) - except (ValueError, AttributeError): - return 9999 # Not in hierarchy + return inheritance_distance(child, parent) def _select_most_specific_links(self, links: list[Link], left_fg: type, right_fg: type) -> list[Link]: """Select links that are most specific (closest in inheritance hierarchy). diff --git a/tests/test_core/test_integration/test_core/test_link_planner_orientation_characterization.py b/tests/test_core/test_integration/test_core/test_link_planner_orientation_characterization.py index 0aa60c42..290c3adf 100644 --- a/tests/test_core/test_integration/test_core/test_link_planner_orientation_characterization.py +++ b/tests/test_core/test_integration/test_core/test_link_planner_orientation_characterization.py @@ -293,11 +293,6 @@ def _run_pair( RIGHT_JOIN_ROWS = ["k3|L3|k3|R3", "k4|L4|k4|R4", f"{MISSING}|{MISSING}|k5|R5"] CHAIN_ROWS = ["k4|L4|R4|T4", "k3|L3|R3|T3"] -RIGHT_SIDE_BINDING_REASON = ( - "plain RIGHT joins bind the merge arguments to the resolved frameworks instead of the " - "declared sides, so the declared left index is looked up in the right group's data" -) - @MODES_WITH_MULTIPROCESSING def test_inner_join_declared_orientation_keeps_left_group_first( @@ -348,17 +343,7 @@ def test_right_join_keeps_every_right_row_for_a_child_declaring_the_left_framewo @MODES_SYNC_THREADING -def test_right_join_raises_for_a_child_on_the_right_framework( - modes: set[ParallelizationMode], flight_server: Any -) -> None: - # The exception type is incidental: the column-semantics guard reaches the key column before the merge does. - with pytest.raises((KeyError, ValueError), match="oc_b_left_key"): - _run_pair(PAIR_B, "right", OrientCharInvertedChild, modes, flight_server) - - -@pytest.mark.xfail(strict=True, reason=RIGHT_SIDE_BINDING_REASON) -@MODES_SYNC_THREADING -def test_right_join_should_keep_every_right_row_for_a_child_on_the_right_framework( +def test_right_join_keeps_every_right_row_for_a_child_on_the_right_framework( modes: set[ParallelizationMode], flight_server: Any ) -> None: rows = _run_pair(PAIR_B, "right", OrientCharInvertedChild, modes, flight_server) diff --git a/tests/test_core/test_integration/test_core/test_right_join_side_binding.py b/tests/test_core/test_integration/test_core/test_right_join_side_binding.py new file mode 100644 index 00000000..4f681ba4 --- /dev/null +++ b/tests/test_core/test_integration/test_core/test_right_join_side_binding.py @@ -0,0 +1,313 @@ +"""A RIGHT join binds the declared left group as the left merge argument; surplus left keys expose a swap.""" + +from typing import Any, Optional + +import pytest + +from mloda.provider import BaseInputData +from mloda.provider import ComputeFramework +from mloda.provider import DataCreator +from mloda.provider import FeatureGroup +from mloda.provider import FeatureSet +from mloda.user import Feature +from mloda.user import FeatureName +from mloda.user import Index +from mloda.user import JoinSpec, Link +from mloda.user import Options +from mloda.user import ParallelizationMode +from mloda.user import PluginCollector +from mloda.user import mloda +from mloda_plugins.compute_framework.base_implementations.pandas.dataframe import PandasDataFrame +from mloda_plugins.compute_framework.base_implementations.pyarrow.table import PyArrowTable + + +LEFT_KEYS = ["k3", "k2", "k1"] +LEFT_PAYLOADS = ["l3", "l2", "l1"] +RIGHT_KEYS = ["k1", "k2", "k9"] +RIGHT_PAYLOADS = ["r1", "r2", "r9"] + +MISSING = "-" + +# k3 has no right partner and drops out; k9 has no left partner and survives null-padded. +RIGHT_JOIN_ROWS = ["k1|l1|k1|r1", "k2|l2|k2|r2", f"{MISSING}|{MISSING}|k9|r9"] + +SIBLING_KEYS = RIGHT_KEYS +SIBLING_PAYLOADS = ["s1", "s2", "s9"] + +# The sibling covers every right key, so its inner join only widens the rows the RIGHT join already produced. +RIGHT_JOIN_ROWS_WITH_SIBLING = ["k1|l1|k1|r1|s1", "k2|l2|k2|r2|s2", f"{MISSING}|{MISSING}|k9|r9|s9"] + +MODES = pytest.mark.parametrize("modes", [{ParallelizationMode.SYNC}, {ParallelizationMode.THREADING}]) + + +def _cell(value: Any) -> str: + """Render one joined cell so an unmatched value reads the same in every framework.""" + if value is None: + return MISSING + text = str(value) + return MISSING if text in ("nan", "None", "") else text + + +def _columns(data: Any) -> dict[str, list[Any]]: + if hasattr(data, "column_names"): + return {name: data.column(name).to_pylist() for name in data.column_names} + return {name: list(data[name]) for name in data.columns} + + +def _joined_rows(data: Any, prefix: str) -> list[str]: + columns = _columns(data) + return [ + "|".join(_cell(value) for value in row) + for row in zip( + columns[f"{prefix}_left_key"], + columns[f"{prefix}_left_payload"], + columns[f"{prefix}_right_key"], + columns[f"{prefix}_right_payload"], + ) + ] + + +def _pair_features(prefix: str) -> set[Feature]: + return { + Feature(name=f"{prefix}_left_key"), + Feature(name=f"{prefix}_left_payload"), + Feature(name=f"{prefix}_right_key"), + Feature(name=f"{prefix}_right_payload"), + } + + +def _packed_rows(results: Any, column: str) -> list[str]: + matching = [frame for frame in results if column in _columns(frame)] + assert len(matching) == 1, f"Expected exactly one result frame carrying {column}, got {len(matching)}." + return [str(value) for value in _columns(matching[0])[column]] + + +class RightBindLeftInArrow(FeatureGroup): + @classmethod + def input_data(cls) -> Optional[BaseInputData]: + return DataCreator(supports_features={"rjsb_left_key", "rjsb_left_payload"}) + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {"rjsb_left_key": LEFT_KEYS, "rjsb_left_payload": LEFT_PAYLOADS} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PyArrowTable} + + +class RightBindRightInPandas(FeatureGroup): + @classmethod + def input_data(cls) -> Optional[BaseInputData]: + return DataCreator(supports_features={"rjsb_right_key", "rjsb_right_payload"}) + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {"rjsb_right_key": RIGHT_KEYS, "rjsb_right_payload": RIGHT_PAYLOADS} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PandasDataFrame} + + +class RightBindChild(FeatureGroup): + """Runs in the declared right group's framework, which is where the RIGHT join executes.""" + + def input_features(self, options: Options, feature_name: FeatureName) -> Optional[set[Feature]]: + return _pair_features("rjsb") + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {cls.get_class_name(): _joined_rows(data, "rjsb")} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PandasDataFrame} + + +class RightBindPolyBase(FeatureGroup): + """Declared left side of a link whose declared right side derives from it.""" + + @classmethod + def input_data(cls) -> Optional[BaseInputData]: + return DataCreator(supports_features={"rjsbp_left_key", "rjsbp_left_payload"}) + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {"rjsbp_left_key": LEFT_KEYS, "rjsbp_left_payload": LEFT_PAYLOADS} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PyArrowTable} + + +class RightBindPolyDerived(RightBindPolyBase): + """Declared right side, and a subclass of the declared left side, so it answers to both sides.""" + + @classmethod + def input_data(cls) -> Optional[BaseInputData]: + return DataCreator(supports_features={"rjsbp_right_key", "rjsbp_right_payload"}) + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {"rjsbp_right_key": RIGHT_KEYS, "rjsbp_right_payload": RIGHT_PAYLOADS} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PandasDataFrame} + + +class RightBindPolyChild(FeatureGroup): + def input_features(self, options: Options, feature_name: FeatureName) -> Optional[set[Feature]]: + return _pair_features("rjsbp") + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {cls.get_class_name(): _joined_rows(data, "rjsbp")} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PandasDataFrame} + + +class RightBindAncestorBase(FeatureGroup): + """Declared left side of the RIGHT join, and the base class of an unrelated ancestor of the same child.""" + + @classmethod + def input_data(cls) -> Optional[BaseInputData]: + return DataCreator(supports_features={"rjsba_left_key", "rjsba_left_payload"}) + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {"rjsba_left_key": LEFT_KEYS, "rjsba_left_payload": LEFT_PAYLOADS} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PyArrowTable} + + +class RightBindAncestorRight(FeatureGroup): + @classmethod + def input_data(cls) -> Optional[BaseInputData]: + return DataCreator(supports_features={"rjsba_right_key", "rjsba_right_payload"}) + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {"rjsba_right_key": RIGHT_KEYS, "rjsba_right_payload": RIGHT_PAYLOADS} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PandasDataFrame} + + +class RightBindAncestorSibling(RightBindAncestorBase): + """Subclasses the declared left side but runs in the destination framework and joins on its own link.""" + + @classmethod + def input_data(cls) -> Optional[BaseInputData]: + return DataCreator(supports_features={"rjsba_sibling_key", "rjsba_sibling_payload"}) + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + return {"rjsba_sibling_key": SIBLING_KEYS, "rjsba_sibling_payload": SIBLING_PAYLOADS} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PandasDataFrame} + + +class RightBindAncestorChild(FeatureGroup): + def input_features(self, options: Options, feature_name: FeatureName) -> Optional[set[Feature]]: + return _pair_features("rjsba") | {Feature(name="rjsba_sibling_payload")} + + @classmethod + def calculate_feature(cls, data: Any, features: FeatureSet) -> Any: + siblings = _columns(data)["rjsba_sibling_payload"] + rows = [f"{row}|{_cell(value)}" for row, value in zip(_joined_rows(data, "rjsba"), siblings)] + return {cls.get_class_name(): rows} + + @classmethod + def compute_framework_rule(cls) -> set[type[ComputeFramework]]: + return {PandasDataFrame} + + +# MULTIPROCESSING is left out: cross-framework joins fail in the transform hop for unrelated reasons. +@MODES +def test_right_join_keeps_every_right_row_and_drops_unmatched_left_rows( + modes: set[ParallelizationMode], flight_server: Any +) -> None: + link = Link.right( + JoinSpec(RightBindLeftInArrow, Index(("rjsb_left_key",))), + JoinSpec(RightBindRightInPandas, Index(("rjsb_right_key",))), + ) + + results = mloda.run_all( + [Feature(name=RightBindChild.get_class_name())], + links={link}, + compute_frameworks={PandasDataFrame, PyArrowTable}, + plugin_collector=PluginCollector.enabled_feature_groups( + {RightBindLeftInArrow, RightBindRightInPandas, RightBindChild} + ), + flight_server=flight_server if ParallelizationMode.MULTIPROCESSING in modes else None, + parallelization_modes=modes, + ) + rows = _packed_rows(results, RightBindChild.get_class_name()) + + assert len(rows) == len(RIGHT_KEYS) + assert sorted(rows) == sorted(RIGHT_JOIN_ROWS) + + +@MODES +def test_right_join_binds_the_declared_left_side_when_the_declared_right_side_subclasses_it( + modes: set[ParallelizationMode], flight_server: Any +) -> None: + """The right node answers both issubclass tests, which must not cost it its right-argument position.""" + link = Link.right( + JoinSpec(RightBindPolyBase, Index(("rjsbp_left_key",))), + JoinSpec(RightBindPolyDerived, Index(("rjsbp_right_key",))), + ) + + results = mloda.run_all( + [Feature(name=RightBindPolyChild.get_class_name())], + links={link}, + compute_frameworks={PandasDataFrame, PyArrowTable}, + plugin_collector=PluginCollector.enabled_feature_groups( + {RightBindPolyBase, RightBindPolyDerived, RightBindPolyChild} + ), + flight_server=flight_server if ParallelizationMode.MULTIPROCESSING in modes else None, + parallelization_modes=modes, + ) + rows = _packed_rows(results, RightBindPolyChild.get_class_name()) + + assert len(rows) == len(RIGHT_KEYS) + assert sorted(rows) == sorted(RIGHT_JOIN_ROWS) + + +@MODES +def test_right_join_binds_the_declared_left_side_when_a_sibling_subclass_is_a_second_ancestor( + modes: set[ParallelizationMode], flight_server: Any +) -> None: + """A subclass of the declared left side running in the destination framework must not rebind the sides.""" + link = Link.right( + JoinSpec(RightBindAncestorBase, Index(("rjsba_left_key",))), + JoinSpec(RightBindAncestorRight, Index(("rjsba_right_key",))), + ) + sibling_link = Link.inner( + JoinSpec(RightBindAncestorSibling, Index(("rjsba_sibling_key",))), + JoinSpec(RightBindAncestorRight, Index(("rjsba_right_key",))), + ) + + results = mloda.run_all( + [Feature(name=RightBindAncestorChild.get_class_name())], + links={link, sibling_link}, + compute_frameworks={PandasDataFrame, PyArrowTable}, + plugin_collector=PluginCollector.enabled_feature_groups( + {RightBindAncestorBase, RightBindAncestorRight, RightBindAncestorSibling, RightBindAncestorChild} + ), + flight_server=flight_server if ParallelizationMode.MULTIPROCESSING in modes else None, + parallelization_modes=modes, + ) + rows = _packed_rows(results, RightBindAncestorChild.get_class_name()) + + assert len(rows) == len(RIGHT_KEYS) + assert sorted(rows) == sorted(RIGHT_JOIN_ROWS_WITH_SIBLING) diff --git a/tests/test_core/test_prepare/test_link_planner_plan_characterization.py b/tests/test_core/test_prepare/test_link_planner_plan_characterization.py index 6f715b95..eb729603 100644 --- a/tests/test_core/test_prepare/test_link_planner_plan_characterization.py +++ b/tests/test_core/test_prepare/test_link_planner_plan_characterization.py @@ -585,7 +585,7 @@ def test_a_hand_built_inconsistent_trekker_key_yields_an_inconsistent_joinstep() def test_a_right_join_plans_the_joinstep_where_the_declared_right_side_runs() -> None: - """RIGHT swaps the frameworks before the trekker lookup, so only the merge order still follows the trekker.""" + """RIGHT swaps the frameworks, so the destination holds the declared right side and the merge sides swap back.""" declared = _pair_scenario(link_factory=Link.right, child_cfw=PandasDataFrame) _trek(declared, PyArrowTable, PandasDataFrame) inverted = _pair_scenario(link_factory=Link.right, child_cfw=PandasDataFrame) @@ -599,7 +599,7 @@ def test_a_right_join_plans_the_joinstep_where_the_declared_right_side_runs() -> assert declared_step.source_framework is PyArrowTable assert declared_step.destination_framework_uuids == {declared.right_uuid} assert declared_step.source_framework_uuids == {declared.left_uuid} - assert declared_step.swap_merge_sides is False + assert declared_step.swap_merge_sides is True assert isinstance(inverted_step, JoinStep) assert inverted_step.destination_framework is PandasDataFrame @@ -609,6 +609,35 @@ def test_a_right_join_plans_the_joinstep_where_the_declared_right_side_runs() -> assert inverted_step.swap_merge_sides is True +def test_a_right_join_reached_through_a_reversed_key_keeps_the_declared_merge_sides() -> None: + """Hand built: no planner path reverses this key, so keep the shape; a reversed key must not swap the sides.""" + planned = _pair_scenario(link_factory=Link.right) + _trek(planned, PandasDataFrame, PyArrowTable) + + join_step = _run(planned, PandasDataFrame, PyArrowTable) + + assert isinstance(join_step, JoinStep) + assert join_step.destination_framework is PyArrowTable + assert join_step.source_framework is PandasDataFrame + assert join_step.destination_framework_uuids == {planned.left_uuid} + assert join_step.source_framework_uuids == {planned.right_uuid} + assert join_step.swap_merge_sides is False + + +def test_a_left_join_inverted_after_queueing_swaps_the_merge_sides() -> None: + planned = _pair_scenario(link_factory=Link.left, child_cfw=PandasDataFrame) + _trek(planned, PandasDataFrame, PyArrowTable) + + join_step = _run(planned, PyArrowTable, PandasDataFrame) + + assert isinstance(join_step, JoinStep) + assert join_step.destination_framework is PandasDataFrame + assert join_step.source_framework is PyArrowTable + assert join_step.destination_framework_uuids == {planned.right_uuid} + assert join_step.source_framework_uuids == {planned.left_uuid} + assert join_step.swap_merge_sides is True + + # Fresh interpreters are slow to start, so this one needs more than the suite-wide per-test budget. @pytest.mark.timeout(60) def test_fresh_interpreters_plan_the_same_link_orientation() -> None: