Allow the preperiod metric and postperiod metric in CUPED/PrePostChange to have same names.

PiperOrigin-RevId: 771202290

Author: tcya

operations.py

@@ -799,6 +799,10 @@ class Adjust(metrics.Metric):
 
       def compute_slices(self, df, split_by: Optional[List[Text]] = None):
         child = df.iloc[:, :len_child]
+        prefix = utils.get_unique_prefix(child)
+        df.columns = list(child.columns) + [
+            prefix + c for c in df.columns[len_child:]
+        ]
         covariate = df.iloc[:, len_child:]
         if len(covariate.columns) > 1:
           return super(Adjust, self).compute_slices(df, split_by)
@@ -826,7 +830,7 @@ def compute_children_sql(self, table, split_by, execute, mode=None):
     child = super(PrePostChange,
                   self).compute_children_sql(table, split_by, execute, mode)
     covariates = child.iloc[:, -self.k_covariates:]
-    child = child.iloc[:, :self.k_covariates]
+    child = child.iloc[:, :-self.k_covariates]
     return self.adjust_value(child, covariates, split_by)
 
   def get_change_raw_sql(
@@ -1045,6 +1049,10 @@ class Adjust(metrics.Metric):
 
       def compute_slices(self, df, split_by: Optional[List[Text]] = None):
         child = df.iloc[:, :len_child]
+        prefix = utils.get_unique_prefix(child)
+        df.columns = list(child.columns) + [
+            prefix + c for c in df.columns[len_child:]
+        ]
         covariate = df.iloc[:, len_child:]
         if len(covariate.columns) > 1:
           return super(Adjust, self).compute_slices(df, split_by)
@@ -1073,7 +1081,7 @@ def compute_children_sql(self, table, split_by, execute, mode=None):
     child = super(CUPED, self).compute_children_sql(table, split_by, execute,
                                                     mode)
     covariates = child.iloc[:, -self.k_covariates:]
-    child = child.iloc[:, :self.k_covariates]
+    child = child.iloc[:, :-self.k_covariates]
     return self.adjust_value(child, covariates, split_by)
 
   def get_change_raw_sql(

operations_test.py

@@ -256,7 +256,7 @@ def test_nxx(self):
     testing.assert_frame_equal(output, expected)
 
 
-class PrePostChangeTests(absltest.TestCase):
+class PrePostChangeTests(parameterized.TestCase):
   n = 40
   df = pd.DataFrame({
       'x': np.random.choice(range(20), n),
@@ -428,7 +428,7 @@ def test_complex(self):
     testing.assert_frame_equal(output, expected)
 
 
-class CUPEDTests(absltest.TestCase):
+class CUPEDTests(parameterized.TestCase):
   n = 40
   df = pd.DataFrame({
       'x': np.random.choice(range(20), n),
@@ -2362,6 +2362,33 @@ def test_different_metrics_have_different_fingerprints(self):
     )
     self.assertLen(fingerprints, len(distinct_ops))
 
+  @parameterized.parameters([operations.CUPED, operations.PrePostChange])
+  def test_cuped_prepost_with_duplicate_names_one_covariate(self, op):
+    s = metrics.Sum('x')
+    cov = metrics.Sum('x', where='x>0.1', name='foo')
+    cov_dup = metrics.Sum('x', where='x>0.1')
+    jk = operations.Jackknife('cookie', confidence=0.9)
+    op_dup = op('grp', 1, s, cov_dup, 'grp4')
+    op = op('grp', 1, s, cov, 'grp4')
+    output = jk(op_dup).compute_on(self.df).display(return_formatted_df=True)
+    expected = jk(op).compute_on(self.df).display(return_formatted_df=True)
+    testing.assert_frame_equal(output, expected)
+
+  @parameterized.parameters([operations.CUPED, operations.PrePostChange])
+  def test_cuped_prepost_with_duplicate_names_multiple_covariates(self, op):
+    s = metrics.Sum('x', name='foo')
+    cov_dup = [
+        metrics.Sum('x', where='x>0.1', name='foo'),
+        metrics.Sum('y', name='foo'),
+    ]
+    cov = [metrics.Sum('x', where='x>0.1'), metrics.Sum('y')]
+    jk = operations.Jackknife('cookie', confidence=0.9)
+    op_dup = op('grp', 1, s, cov_dup, 'grp4')
+    op = op('grp', 1, s, cov, 'grp4')
+    output = jk(op_dup).compute_on(self.df).display(return_formatted_df=True)
+    expected = jk(op).compute_on(self.df).display(return_formatted_df=True)
+    testing.assert_frame_equal(output, expected)
+
 
 if __name__ == '__main__':
   absltest.main()