Fix naming issue for MetricList. name_tmpl wasn't respected in compute_on_sql when MetricList only has one child because its result is squeezed to a number in an intermediate step so the column name is lost.

PiperOrigin-RevId: 771168986

Author: tcya

metrics.py

@@ -1115,6 +1115,7 @@ def __init__(self,
     tmpl = name_tmpl or 'MetricList({})'
     if len(children) == 1:
       name = children[0].name
+      name = name_tmpl.format(name) if name_tmpl else name
     else:
       name = tmpl.format(', '.join(m.name for m in children))
     super(MetricList, self).__init__(name, children, where, name_tmpl)
@@ -1154,8 +1155,11 @@ def compute_slices(self, df, split_by=None):
   def compute_on_children(self, children, split_by):
     if isinstance(children, list):
       children = self.to_dataframe(children)
-    if self.name_tmpl:
-      children.columns = [self.name_tmpl.format(c) for c in children.columns]
+    if isinstance(children, pd.DataFrame):
+      if self.name_tmpl:
+        children.columns = [self.name_tmpl.format(c) for c in children.columns]
+    elif not isinstance(children, pd.Series):
+      children = pd.DataFrame({self.name: [children]})
     if self.columns:
       if len(children.columns) != len(self.columns):
         raise ValueError(
@@ -1262,8 +1266,7 @@ def compute_children_sql(self, table, split_by, execute, mode=None):
                 return_dataframe=self.children_return_dataframe,
             )
         )
-    # When there is only one child, returns the result of the child.
-    return children[0] if len(self.children) == 1 else children
+    return children
 
   def get_sql_and_with_clause(self, table, split_by, global_filter, indexes,
                               local_filter, with_data):

metrics_test.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Metrics."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -645,7 +644,7 @@ def test_duplicate_column_names(self):
     testing.assert_frame_equal(output, expected)
 
 
-class TestMetricList(absltest.TestCase):
+class TestMetricList(parameterized.TestCase):
 
   def test_return_list(self):
     df = pd.DataFrame({'X': [0, 1, 2, 3]})
@@ -675,17 +674,26 @@ def test_return_df(self):
         }, columns=['sum(X)', 'mean(X)'])
     testing.assert_frame_equal(output, expected)
 
-  def test_with_name_tmpl(self):
+  @parameterized.parameters(
+      ([metrics.Sum('X')],), ([metrics.Sum('X'), metrics.Mean('X')],)
+  )
+  def test_name_tmpl(self, children):
     df = pd.DataFrame({'X': [0, 1, 2, 3]})
-    ms = [metrics.Sum('X'), metrics.Mean('X')]
-    m = metrics.MetricList(ms, name_tmpl='a{}b')
+    m = metrics.MetricList(children, name_tmpl='a{}b')
     output = m.compute_on(df)
-    expected = pd.DataFrame(
-        data={
-            'asum(X)b': [6],
-            'amean(X)b': [1.5]
-        },
-        columns=['asum(X)b', 'amean(X)b'])
+    expected = metrics.MetricList(children).compute_on(df)
+    expected.columns = [f'a{c}b' for c in expected.columns]
+    testing.assert_frame_equal(output, expected)
+
+  @parameterized.parameters(
+      ([metrics.Sum('X')],), ([metrics.Sum('X'), metrics.Mean('X')],)
+  )
+  def test_nested_name_tmpl(self, children):
+    df = pd.DataFrame({'X': [0, 1, 2, 3]})
+    m = metrics.MetricList(children, name_tmpl='a {}')
+    m = metrics.MetricList(m, name_tmpl='b {}')
+    output = m.compute_on(df)
+    expected = metrics.MetricList(children, name_tmpl='b a {}').compute_on(df)
     testing.assert_frame_equal(output, expected)
 
   def test_operations(self):