check project a bit earlier

Author: JohnMount

build/lib/data_algebra/data_ops.py

@@ -331,16 +331,16 @@ def extend(
     ):
         if (ops is None) or (len(ops) < 1):
             return self
-        parsed_ops = data_algebra.expr_rep.parse_assignments_in_context(
-            ops, self, parse_env=parse_env
-        )
-        new_cols_used_in_calc = set(data_algebra.expr_rep.get_columns_used(parsed_ops))
         if partition_by is None:
             partition_by = []
         if order_by is None:
             order_by = []
         if reverse is None:
             reverse = []
+        parsed_ops = data_algebra.expr_rep.parse_assignments_in_context(
+            ops, self, parse_env=parse_env
+        )
+        new_cols_used_in_calc = set(data_algebra.expr_rep.get_columns_used(parsed_ops))
         new_cols_produced_in_calc = set([k for k in parsed_ops.keys()])
         if (partition_by != 1) and (len(partition_by) > 0):
             if len(new_cols_produced_in_calc.intersection(partition_by)) > 0:
@@ -357,6 +357,7 @@ def extend(
                 reverse=reverse,
                 parse_env=parse_env,
             )
+        # see if we can combine nodes
         if isinstance(self, ExtendNode):
             compatible_partition = (partition_by == self.partition_by) or (
                 ((partition_by == 1) or (len(partition_by) <= 0))
@@ -398,6 +399,7 @@ def extend(
                     order_by=order_by,
                     reverse=reverse,
                 )
+        # new node
         return ExtendNode(
             source=self,
             parsed_ops=parsed_ops,
@@ -407,15 +409,19 @@ def extend(
         )
 
     def project(self, ops=None, *, group_by=None, parse_env=None):
-        if ((ops is None) or (len(ops) < 1)) and (
-            (group_by is None) or (len(group_by) < 1)
-        ):
+        if group_by is None:
+            group_by = []
+        if ((ops is None) or (len(ops) < 1)) and (len(group_by) < 1):
             raise ValueError("must have ops or group_by")
-        if self.is_trivial_when_intermediate():
-            return self.sources[0].project(ops, group_by=group_by, parse_env=parse_env)
         parsed_ops = data_algebra.expr_rep.parse_assignments_in_context(
             ops, self, parse_env=parse_env
         )
+        new_cols_used_in_calc = set(data_algebra.expr_rep.get_columns_used(parsed_ops))
+        new_cols_produced_in_calc = set([k for k in parsed_ops.keys()])
+        if len(new_cols_used_in_calc.intersection(group_by)):
+            raise ValueError("can not alter grouping columns")
+        if self.is_trivial_when_intermediate():
+            return self.sources[0].project(ops, group_by=group_by, parse_env=parse_env)
         return ProjectNode(source=self, parsed_ops=parsed_ops, group_by=group_by)
 
     def natural_join(self, b, *, by=None, jointype="INNER"):

coverage.txt

@@ -58,7 +58,7 @@ data_algebra/cdata.py                    232     75    68%
 data_algebra/cdata_impl.py                10      1    90%
 data_algebra/connected_components.py      49      1    98%
 data_algebra/data_model.py                29     13    55%
-data_algebra/data_ops.py                1244    240    81%
+data_algebra/data_ops.py                1250    242    81%
 data_algebra/data_ops_types.py            42     16    62%
 data_algebra/db_model.py                 389     72    81%
 data_algebra/diagram.py                   56     43    23%
@@ -71,7 +71,7 @@ data_algebra/test_util.py                119     17    86%
 data_algebra/util.py                      44     10    77%
 data_algebra/yaml.py                     101     13    87%
 ----------------------------------------------------------
-TOTAL                                   3493    800    77%
+TOTAL                                   3499    802    77%
 
 
-============================== 74 passed in 6.94s ==============================
+============================== 74 passed in 7.19s ==============================

data_algebra/data_ops.py

@@ -331,16 +331,16 @@ def extend(
     ):
         if (ops is None) or (len(ops) < 1):
             return self
-        parsed_ops = data_algebra.expr_rep.parse_assignments_in_context(
-            ops, self, parse_env=parse_env
-        )
-        new_cols_used_in_calc = set(data_algebra.expr_rep.get_columns_used(parsed_ops))
         if partition_by is None:
             partition_by = []
         if order_by is None:
             order_by = []
         if reverse is None:
             reverse = []
+        parsed_ops = data_algebra.expr_rep.parse_assignments_in_context(
+            ops, self, parse_env=parse_env
+        )
+        new_cols_used_in_calc = set(data_algebra.expr_rep.get_columns_used(parsed_ops))
         new_cols_produced_in_calc = set([k for k in parsed_ops.keys()])
         if (partition_by != 1) and (len(partition_by) > 0):
             if len(new_cols_produced_in_calc.intersection(partition_by)) > 0:
@@ -357,6 +357,7 @@ def extend(
                 reverse=reverse,
                 parse_env=parse_env,
             )
+        # see if we can combine nodes
         if isinstance(self, ExtendNode):
             compatible_partition = (partition_by == self.partition_by) or (
                 ((partition_by == 1) or (len(partition_by) <= 0))
@@ -398,6 +399,7 @@ def extend(
                     order_by=order_by,
                     reverse=reverse,
                 )
+        # new node
         return ExtendNode(
             source=self,
             parsed_ops=parsed_ops,
@@ -407,15 +409,19 @@ def extend(
         )
 
     def project(self, ops=None, *, group_by=None, parse_env=None):
-        if ((ops is None) or (len(ops) < 1)) and (
-            (group_by is None) or (len(group_by) < 1)
-        ):
+        if group_by is None:
+            group_by = []
+        if ((ops is None) or (len(ops) < 1)) and (len(group_by) < 1):
             raise ValueError("must have ops or group_by")
-        if self.is_trivial_when_intermediate():
-            return self.sources[0].project(ops, group_by=group_by, parse_env=parse_env)
         parsed_ops = data_algebra.expr_rep.parse_assignments_in_context(
             ops, self, parse_env=parse_env
         )
+        new_cols_used_in_calc = set(data_algebra.expr_rep.get_columns_used(parsed_ops))
+        new_cols_produced_in_calc = set([k for k in parsed_ops.keys()])
+        if len(new_cols_used_in_calc.intersection(group_by)):
+            raise ValueError("can not alter grouping columns")
+        if self.is_trivial_when_intermediate():
+            return self.sources[0].project(ops, group_by=group_by, parse_env=parse_env)
         return ProjectNode(source=self, parsed_ops=parsed_ops, group_by=group_by)
 
     def natural_join(self, b, *, by=None, jointype="INNER"):