work on Polars join

Author: JohnMount

data_algebra/polars_model.py

@@ -884,12 +884,18 @@ def _natural_join_step(self, op: data_algebra.data_ops_types.OperatorPlatform, *
         how = op.jointype.lower()
         if how == "full":
             how = "outer"
-        coalesce_columns = (
-            set(op.sources[0].columns_produced()).intersection(op.sources[1].columns_produced()) 
-            - set(op.on_a))
         if how != "right":
+            coalesce_columns = (
+                set(op.sources[0].columns_produced()).intersection(op.sources[1].columns_produced()) 
+                - set(op.on_a))
+            orphan_keys = [c for c in op.on_b if c not in set(op.on_a)]
+            input_right = inputs[1]
+            if len(orphan_keys) > 0:
+                input_right = input_right.with_columns([
+                    pl.col(c).alias(f"{c}_da_join_tmp_key") for c in orphan_keys
+                ])
             res = inputs[0].join(
-                inputs[1],
+                input_right,
                 left_on=op.on_a,
                 right_on=op.on_b,
                 how=how,
@@ -903,10 +909,21 @@ def _natural_join_step(self, op: data_algebra.data_ops_types.OperatorPlatform, *
                         .alias(c)
                     for c in coalesce_columns
                 ])
+            if len(orphan_keys) > 0:
+                res = res.rename({f"{c}_da_join_tmp_key": c for c in orphan_keys})
         else:
             # simulate right join with left join
+            coalesce_columns = (
+                set(op.sources[0].columns_produced()).intersection(op.sources[1].columns_produced()) 
+                - set(op.on_b))
+            orphan_keys = [c for c in op.on_a if c not in set(op.on_b)]
+            input_right = inputs[0]
+            if len(orphan_keys) > 0:
+                input_right = input_right.with_columns([
+                    pl.col(c).alias(f"{c}_da_join_tmp_key") for c in orphan_keys
+                ])
             res = inputs[1].join(
-                inputs[0],
+                input_right,
                 left_on=op.on_b,
                 right_on=op.on_a,
                 how="left",
@@ -920,6 +937,8 @@ def _natural_join_step(self, op: data_algebra.data_ops_types.OperatorPlatform, *
                         .alias(c)
                     for c in coalesce_columns
                 ])
+            if len(orphan_keys) > 0:
+                res = res.rename({f"{c}_da_join_tmp_key": c for c in orphan_keys})
         res = res.select(op.columns_produced())
         return res
 

tests/test_expression_expectations.py

@@ -48,7 +48,7 @@ def test_expression_expectations_1():
                         valid_for_empty=False,
                         empty_produces_empty=False,
                         models_to_skip=to_skip,
-                        try_on_Polars=False,  # TODO: turn this on
+                        try_on_Polars=False,  # TODO: complete coverage, and turn this on
                     )
     for op, op_class, exp, ops, expect in u_results:
         # re-run, but don't check value

tests/test_idioms.py

@@ -247,7 +247,7 @@ def test_idiom_cross_join():
     )
     data_algebra.test_util.check_transform(
         ops=ops, data={table_name_d: d, table_name_e: e}, expect=expect,
-        try_on_Polars=False,  # TODO: turn this on
+        try_on_Polars=False,  # TODO: get empty case to match and turn this on
     )
 
 

tests/test_join_conditions.py

@@ -60,5 +60,4 @@ def test_join_conditions_on_join():
     assert data_algebra.test_util.equivalent_frames(res, expect)
     data_algebra.test_util.check_transform(
         ops=ops, data={"d1": d1, "d2": d2}, expect=expect,
-        try_on_Polars=False,  # TODO: turn this on, to get it to work we have to copy columns with name changes in join conditions
     )