No public description

metrics.py

@@ -1100,13 +1100,13 @@ def __neg__(self):
     return CompositeMetric(lambda x, _: -x, '-{}', (self, -1))
 
   def __div__(self, other):
-    return CompositeMetric(lambda x, y: x / y, '{} / {}', (self, other))
+    return CompositeMetric(_safe_divide, '{} / {}', (self, other))
 
   def __truediv__(self, other):
     return self.__div__(other)
 
   def __rdiv__(self, other):
-    return CompositeMetric(lambda x, y: x / y, '{} / {}', (other, self))
+    return CompositeMetric(_safe_divide, '{} / {}', (other, self))
 
   def __rtruediv__(self, other):
     return self.__rdiv__(other)
@@ -1803,6 +1803,35 @@ def get_fingerprint(self, attr_to_exclude=()):
     return super(CompositeMetric, self).get_fingerprint(attr_to_exclude)
 
 
+def _safe_divide(x, y):
+  """Divides x by y, returning NaN instead of +/-inf on division by zero."""
+  if isinstance(x, sql.Column) or isinstance(y, sql.Column):
+    return x / y
+  if isinstance(x, (pd.Series, pd.DataFrame)):
+    x = x.astype(float)
+  if isinstance(y, (pd.Series, pd.DataFrame)):
+    y = y.astype(float)
+    y = y.where(y != 0, np.nan)
+  elif np.ndim(y) == 0:
+    if pd.isna(y) or y == 0:
+      if np.ndim(x) == 0:
+        return np.nan
+      y = np.nan
+  else:
+    y = np.asarray(y, dtype=float)
+    y = np.where(y == 0, np.nan, y)
+  with np.errstate(divide='ignore', invalid='ignore'):
+    try:
+      res = x / y
+    except ZeroDivisionError:
+      return np.nan
+  if isinstance(res, (pd.Series, pd.DataFrame)):
+    return res.replace([np.inf, -np.inf], np.nan)
+  if np.ndim(res) == 0:
+    return np.nan if np.isinf(res) else res
+  return np.where(np.isinf(res), np.nan, res)
+
+
 class Ratio(CompositeMetric):
   """Syntactic sugar for Sum('A') / Sum('B')."""
 
@@ -1812,7 +1841,7 @@ def __init__(self,
                name: Optional[Text] = None,
                where: Optional[Text] = None):
     super(Ratio, self).__init__(
-        lambda x, y: x / y,
+        _safe_divide,
         '{} / {}', (Sum(numerator), Sum(denominator)),
         where=where)
     self.numerator = numerator

metrics_test.py

@@ -18,6 +18,7 @@
 
 import copy
 import inspect
+import warnings
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -532,6 +533,42 @@ def test_ratio(self):
     expected = expected.compute_on(self.df)
     testing.assert_frame_equal(output, expected)
 
+  def test_ratio_and_div_zero_denominator(self):
+    df = pd.DataFrame({
+        'X': [0.0, 2.0, 0.0],
+        'Y': [0.0, 0.0, 0.0],
+        'grp': ['a', 'b', 'b'],
+    })
+    for m in (
+        metrics.Ratio('X', 'Y'),
+        metrics.Sum('X') / metrics.Sum('Y'),
+        1.0 / metrics.Sum('Y'),
+    ):
+      with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter('always', RuntimeWarning)
+        scalar_out = m.compute_on(df, return_dataframe=False)
+        series_out = m.compute_on(df, split_by='grp', return_dataframe=False)
+      self.assertEmpty(w)
+      self.assertTrue(np.isnan(scalar_out))
+      self.assertTrue(series_out.isna().all())
+
+  def test_ratio_zero_denominator_partial_split_by_and_object_dtype(self):
+    df = pd.DataFrame({
+        'X': [10, 20, 0, 5],
+        'Y': [2, 3, 0, 0],
+        'grp': ['a', 'a', 'b', 'c'],
+    }).astype(object)
+    for m in (metrics.Ratio('X', 'Y'), metrics.Sum('X') / metrics.Sum('Y')):
+      with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter('always', RuntimeWarning)
+        res = m.compute_on(df, split_by='grp')
+      self.assertEmpty(w)
+      expected = pd.DataFrame(
+          {m.name: [6.0, np.nan, np.nan]},
+          index=pd.Index(['a', 'b', 'c'], name='grp'),
+      )
+      testing.assert_frame_equal(res, expected)
+
   def test_to_dataframe(self):
     metric = 5 + metrics.Sum('X')
     output = metric.compute_on(self.df)

operations.py

@@ -2235,6 +2235,7 @@ def get_stderrs_or_ci_half_width(self, bucket_estimates):
 
   @staticmethod
   def get_stderrs(bucket_estimates):
+    bucket_estimates = bucket_estimates.replace([np.inf, -np.inf], np.nan)
     dof = bucket_estimates.count(axis=1) - 1
     return bucket_estimates.std(1), dof
 
@@ -2781,6 +2782,17 @@ def can_precompute(self):
     return False
 
 
+def _iter_leaves_with_split_by(metric, split_by):
+  """Yields (leaf_metric, accumulated_split_by) for every leaf path."""
+  curr_split_by = list(split_by) + list(metric.extra_split_by)
+  if not metric.children:
+    yield metric, curr_split_by
+    return
+  for c in metric.children:
+    if utils.is_metric(c):
+      yield from _iter_leaves_with_split_by(c, curr_split_by)
+
+
 class Jackknife(MetricWithCI):
   """Class for Jackknife estimates of standard errors.
 
@@ -2824,7 +2836,7 @@ def compute_slices(self, df, split_by=None):
     Metrics in self can be expressed by Sum or Count so we apply the trick by
     1. replace self with an equivalent tree whose leaf Metrics are all Sum or
     Count.
-    2. computes it on self.unit + split_by.
+    2. computes the leaf Metrics on self.unit + split_by.
     3. loop through the cache and find all Sum and Count results in #2. Use them
     to compute the LOOs and save to cache.
     4. call super().compute_slices() which will just hits the cached results.
@@ -2842,7 +2854,13 @@ def compute_slices(self, df, split_by=None):
       util, df = utils.get_fully_expanded_equivalent_metric_tree(self, df)
       return self.compute_util_metric_on(util, df, split_by)
 
-    self.compute_child(df, (split_by or []) + [self.unit])
+    child_with_pushed_filters = utils.push_filters_to_leaf(
+        self.children[0], is_root=False
+    )
+    for leaf, leaf_split_by in _iter_leaves_with_split_by(
+        child_with_pushed_filters, (split_by or []) + [self.unit]
+    ):
+      self.compute_util_metric_on(leaf, df, leaf_split_by)
     precomputed = self.find_all_in_cache_by_metric_type(metric=metrics.Sum)
     precomputed.update(
         self.find_all_in_cache_by_metric_type(metric=metrics.Count)
@@ -2971,7 +2989,14 @@ def get_samples(self, df, split_by=None, return_cache_key=False):
   @staticmethod
   def get_stderrs(bucket_estimates):
     stderrs, dof = super(Jackknife, Jackknife).get_stderrs(bucket_estimates)
-    return stderrs * dof / np.sqrt(dof + 1), dof
+    valid = dof > 0
+    denom = np.where(valid, np.sqrt(dof + 1), np.nan)
+    se = np.where(valid, stderrs * dof / denom, np.nan)
+    if isinstance(stderrs, pd.Series):
+      se = pd.Series(se, index=stderrs.index, name=stderrs.name)
+    elif isinstance(stderrs, pd.DataFrame):
+      se = pd.DataFrame(se, index=stderrs.index, columns=stderrs.columns)
+    return se, dof
 
   def compute_children_sql(
       self, table, split_by, execute, mode=None, batch_size=None
@@ -4566,7 +4591,8 @@ def __init__(self, child, func, sql_func, name_tmpl, **kwargs):
     self.sql_func = sql_func
 
   def compute_on_children(self, children, split_by):
-    new_df = self.func(children)
+    with np.errstate(divide='ignore', invalid='ignore'):
+      new_df = self.func(children)
     new_df = copy_meterstick_metadata(children, new_df)
     return new_df
 

operations_test.py

@@ -18,6 +18,7 @@
 from __future__ import print_function
 
 import copy
+import warnings
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -987,6 +988,94 @@ def test_jackknife_one_dof(self):
     )
     testing.assert_frame_equal(output, expected)
 
+  def test_jackknife_zero_denominator_in_one_loo_replicate(self):
+    df = pd.DataFrame({
+        'x': [2.0, 0.0, 0.0, 0.0],
+        'y': [4.0, 0.0, 0.0, 0.0],
+        'cookie': [0, 1, 2, 3],
+    })
+    jk = operations.Jackknife('cookie', metrics.Ratio('x', 'y'))
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter('always', RuntimeWarning)
+      output = jk.compute_on(df)
+    self.assertEmpty(w)
+    # Dropping cookie 0 gives 0/0 -> NaN; dropping 1, 2, or 3 gives 2/4 = 0.5.
+    expected = pd.DataFrame(
+        [[0.5, 0.0]],
+        columns=pd.MultiIndex.from_product(
+            [['sum(x) / sum(y)'], ['Value', 'Jackknife SE']],
+            names=['Metric', None],
+        ),
+    )
+    testing.assert_frame_equal(output, expected)
+
+  def test_jackknife_inf_loo_replicate_filtered(self):
+    df = pd.DataFrame({
+        'x': [6.0, 0.0, 0.0, 0.0],
+        'y': [2.0, 1.0, 2.0, 3.0],
+        'cookie': [0, 1, 2, 3],
+    })
+    m = operations.LogTransform(metrics.Ratio('x', 'y'))
+    jk = operations.Jackknife('cookie', m)
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter('always', RuntimeWarning)
+      output = jk.compute_on(df)
+    self.assertEmpty(w)
+    # Point estimate: log(6 / 8).
+    # LOO replicates:
+    # - drop 0: log(0 / 6) = -inf -> filtered to NaN
+    # - drop 1: log(6 / 7)
+    # - drop 2: log(6 / 6)
+    # - drop 3: log(6 / 5)
+    valid_loos = [np.log(6.0 / 7.0), np.log(6.0 / 6.0), np.log(6.0 / 5.0)]
+    expected_se = np.std(valid_loos, ddof=1) * 2.0 / np.sqrt(3.0)
+    expected = pd.DataFrame(
+        [[np.log(6.0 / 8.0), expected_se]],
+        columns=pd.MultiIndex.from_product(
+            [['Ln(sum(x) / sum(y))'], ['Value', 'Jackknife SE']],
+            names=['Metric', None],
+        ),
+    )
+    testing.assert_frame_equal(output, expected)
+
+  def test_jackknife_log_transform_zero_unit_no_runtime_warning(self):
+    df = pd.DataFrame({
+        'x': [0.0, 2.0, 4.0, 6.0],
+        'y': [1.0, 2.0, 2.0, 3.0],
+        'cookie': [0, 1, 2, 3],
+    })
+    m = operations.LogTransform(metrics.Ratio('x', 'y'))
+    opt = operations.Jackknife('cookie', m)
+    no_opt = operations.Jackknife('cookie', m, enable_optimization=False)
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter('always', RuntimeWarning)
+      output = opt.compute_on(df)
+      expected = no_opt.compute_on(df)
+    self.assertEmpty(w)
+    testing.assert_frame_equal(output, expected)
+
+  def test_jackknife_leaf_precompute_extra_split_by_and_where(self):
+    df = pd.DataFrame({
+        'x': [0.0, 2.0, 4.0, 6.0, 0.0, 3.0, 5.0, 7.0],
+        'y': [1.0, 2.0, 2.0, 3.0, 1.0, 2.0, 2.0, 3.0],
+        'cond': ['A', 'A', 'A', 'A', 'B', 'B', 'B', 'B'],
+        'grp': ['g1', 'g1', 'g1', 'g1', 'g1', 'g1', 'g1', 'g1'],
+        'cookie': [0, 1, 2, 3, 0, 1, 2, 3],
+    })
+    r = metrics.Ratio('x', 'y', where='y > 0')
+    for op in (
+        operations.LogTransform(operations.PercentChange('cond', 'A', r) + 100),
+        operations.LogTransform(operations.Distribution('cond', r)),
+    ):
+      opt = operations.Jackknife('cookie', op)
+      no_opt = operations.Jackknife('cookie', op, enable_optimization=False)
+      with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter('always', RuntimeWarning)
+        output = opt.compute_on(df, split_by='grp')
+        expected = no_opt.compute_on(df, split_by='grp')
+      self.assertEmpty(w)
+      testing.assert_frame_equal(output, expected)
+
   @parameterized.named_parameters(*PRECOMPUTABLE_METRICS_JK)
   def test_optimization(self, m):
     opt = operations.Jackknife('cookie', m, 0.9)