CountCmpEvaluator.cpp

Go to the documentation of this file.

/**
 * @file CountCmpEvaluator.cpp
 * @brief Implementation of the Poisson-binomial pre-pass.
 *        See @c CountCmpEvaluator.h for the full docstring.
 */
#include "CountCmpEvaluator.h"
 
#include <algorithm>
#include <climits>
#include <map>
#include <vector>
 
#include "Aggregation.h"          // ComparisonOperator + getAggregationOperator
#include "CmpEvaluatorCommon.h"   // matchAggCmp, computeRefCounts, contributorProb
 
namespace provsql {
 
namespace {
 
/* Partial Poisson-binomial PMF : compute @c dp[j] = Pr(exactly @c j
 * successes among the @c N input Bernoullis) for @c j in @c [0, jmax]
 * only.  @c jmax is clamped to @c N.  Cost : @c O(N x jmax).  Rolling
 * 1-D array, iterate j downward so each read references the
 * not-yet-updated previous row. */
static std::vector<double> partialPMF(const std::vector<double> &p,
                                      std::size_t jmax)
{
  const std::size_t N = p.size();
  jmax = std::min(jmax, N);
  std::vector<double> dp(jmax + 1, 0.0);
  dp[0] = 1.0;
  for (std::size_t i = 0; i < N; ++i) {
    const double pi = p[i];
    const double qi = 1.0 - pi;
    /* Cap inner loop at min(i+1, jmax) : entries beyond i are still
     * zero and entries beyond jmax we never sum. */
    const std::size_t upper = std::min(jmax, i + 1);
    for (std::size_t j = upper; j >= 1; --j) {
      dp[j] = dp[j] * qi + dp[j - 1] * pi;
    }
    dp[0] *= qi;
  }
  return dp;
}
 
/* Probability that the empty world occurs : @c prod_i (1 - p_i).
 * Always needed for SQL HAVING semantics (the empty group never
 * satisfies). */
static double probZero(const std::vector<double> &p)
{
  double q = 1.0;
  for (double pi : p) q *= (1.0 - pi);
  return q;
}
 
/* Probability that at least @c T of the @c N Bernoullis succeed.
 * Dispatches on which side of @c T is closer to the boundary to keep
 * the partial DP at @c O(N x min(T, N - T + 1)).
 *  - If @c T-1 <= N-T (lower tail is smaller) : compute the lower
 *    partial PMF up to @c T-1 and return @c 1 - sum.
 *  - Otherwise (upper tail is smaller) : invert the Bernoullis,
 *    @c Y_i = 1 - X_i, and use @c Pr(B >= T) = Pr(sum Y <= N - T) ;
 *    the partial PMF on @c Y is computed up to @c N - T. */
static double probAtLeast(const std::vector<double> &p, int T)
{
  const int N = static_cast<int>(p.size());
  if (T <= 0) return 1.0;
  if (T >  N) return 0.0;
 
  if (T - 1 <= N - T) {
    auto dp = partialPMF(p, static_cast<std::size_t>(T - 1));
    double sum = 0.0;
    for (int j = 0; j <= T - 1; ++j) sum += dp[j];
    return 1.0 - sum;
  } else {
    std::vector<double> q(N);
    for (int i = 0; i < N; ++i) q[i] = 1.0 - p[i];
    auto dp = partialPMF(q, static_cast<std::size_t>(N - T));
    double sum = 0.0;
    for (int j = 0; j <= N - T; ++j) sum += dp[j];
    return sum;
  }
}
 
/* Probability that at most @c T of the @c N Bernoullis succeed.
 * Same smaller-side dispatch as @c probAtLeast : if @c T is closer
 * to 0 compute the lower partial PMF and sum ; if @c T is closer to
 * @c N invert and compute the upper tail's complement. */
static double probAtMost(const std::vector<double> &p, int T)
{
  const int N = static_cast<int>(p.size());
  if (T <  0) return 0.0;
  if (T >= N) return 1.0;
 
  if (T <= N - 1 - T) {
    auto dp = partialPMF(p, static_cast<std::size_t>(T));
    double sum = 0.0;
    for (int j = 0; j <= T; ++j) sum += dp[j];
    return sum;
  } else {
    std::vector<double> q(N);
    for (int i = 0; i < N; ++i) q[i] = 1.0 - p[i];
    auto dp = partialPMF(q, static_cast<std::size_t>(N - 1 - T));
    double sum = 0.0;
    for (int j = 0; j <= N - 1 - T; ++j) sum += dp[j];
    return 1.0 - sum;
  }
}
 
/* Probability that exactly @c T of the @c N Bernoullis succeed.
 * Same smaller-side dispatch : @c Pr(B = T) = @c Pr(sum Y = N - T)
 * with @c Y_i = 1 - X_i, computed at whichever side has the smaller
 * partial PMF. */
static double probEqual(const std::vector<double> &p, int T)
{
  const int N = static_cast<int>(p.size());
  if (T < 0 || T > N) return 0.0;
 
  if (T <= N - T) {
    auto dp = partialPMF(p, static_cast<std::size_t>(T));
    return dp[T];
  } else {
    std::vector<double> q(N);
    for (int i = 0; i < N; ++i) q[i] = 1.0 - p[i];
    auto dp = partialPMF(q, static_cast<std::size_t>(N - T));
    return dp[N - T];
  }
}
 
/* Does the empty-group count (0) satisfy "0 op C"?  For a scalar
 * aggregation (no GROUP BY) the empty input is a real possible world --
 * one row whose COUNT is 0 -- so a true-on-empty predicate (= 0, < k,
 * <= k, ...) selects it.  The grouped cdfForOperator excludes the empty
 * world unconditionally, so we add probZero back exactly when 0 op C
 * holds. */
static bool zeroSatisfies(ComparisonOperator op, int C)
{
  switch (op) {
    case ComparisonOperator::GE: return 0 >= C;
    case ComparisonOperator::GT: return 0 >  C;
    case ComparisonOperator::LE: return 0 <= C;
    case ComparisonOperator::LT: return 0 <  C;
    case ComparisonOperator::EQ: return 0 == C;
    case ComparisonOperator::NE: return 0 != C;
  }
  return false;
}
 
/* Map operator + threshold to @c Pr(B op C) under SQL HAVING
 * semantics : the empty-group case (@c B = 0) is excluded regardless
 * of operator, matching @c count_enum's @c if (m < 1) m = 1 clamp
 * and its @c x >= 1 enumeration lower bound -- correct for a GROUPED
 * aggregate (the empty group is no row).  For a scalar aggregate
 * (@p is_scalar) the empty world is real, so probZero is added back when
 * @c zeroSatisfies.
 *
 * Each branch picks at most two of probAtLeast / probAtMost /
 * probEqual / probZero, each O(N x min(C, N-C)) ; the whole
 * dispatch is therefore O(N x min(C, N-C)) per cmp. */
static double cdfForOperator(const std::vector<double> &p,
                             ComparisonOperator op,
                             int C, bool is_scalar)
{
  const int N = static_cast<int>(p.size());
  double r = 0.0;
  switch (op) {
    case ComparisonOperator::GE: {
      /* sizes >= max(C, 1) ; the clamp excludes the empty world for
       * GE 0 / GE -K cases.  No further pZero subtraction needed
       * because the [eff_lo, N] range starts at 1 or above. */
      r = probAtLeast(p, std::max(C, 1));
      break;
    }
    case ComparisonOperator::GT: {
      r = probAtLeast(p, std::max(C + 1, 1));
      break;
    }
    case ComparisonOperator::LE: {
      /* sizes [1, min(C, N)] = Pr(B <= min(C, N)) - Pr(B = 0). */
      const int T = std::min(C, N);
      r = (T < 1) ? 0.0 : probAtMost(p, T) - probZero(p);
      break;
    }
    case ComparisonOperator::LT: {
      const int T = std::min(C - 1, N);
      r = (T < 1) ? 0.0 : probAtMost(p, T) - probZero(p);
      break;
    }
    case ComparisonOperator::EQ: {
      r = (C < 1 || C > N) ? 0.0 : probEqual(p, C);
      break;
    }
    case ComparisonOperator::NE: {
      /* sizes [1, N] \ {C} = (1 - Pr(B = 0)) - (Pr(B = C) if 1<=C<=N). */
      const double nonempty = 1.0 - probZero(p);
      const double eq = (C >= 1 && C <= N) ? probEqual(p, C) : 0.0;
      r = nonempty - eq;
      break;
    }
  }
  if (is_scalar && zeroSatisfies(op, C))
    r += probZero(p);
  return r;
}
 
/* The values of a count satisfying "count op C", as the interval
 * [lo, hi] ; false for NE, whose set is not an interval. */
static bool countBounds(ComparisonOperator op, long C, long &lo, long &hi)
{
  lo = LONG_MIN;
  hi = LONG_MAX;
  switch (op) {
    case ComparisonOperator::GE: lo = C;     return true;
    case ComparisonOperator::GT: lo = C + 1; return true;
    case ComparisonOperator::LE: hi = C;     return true;
    case ComparisonOperator::LT: hi = C - 1; return true;
    case ComparisonOperator::EQ: lo = hi = C; return true;
    case ComparisonOperator::NE: return false;
  }
  return false;
}
 
/* Pr(lo <= B <= hi), the empty world (B = 0) counted only for a scalar
 * aggregation, as in cdfForOperator.  O(N x hi). */
static double probInRange(const std::vector<double> &p, long lo, long hi,
                          bool is_scalar)
{
  const long N = static_cast<long>(p.size());
  lo = std::max(lo, is_scalar ? 0L : 1L);
  hi = std::min(hi, N);
  if (lo > hi)
    return 0.0;
  const auto dp = partialPMF(p, static_cast<std::size_t>(hi));
  double r = 0.0;
  for (long j = lo; j <= hi; ++j)
    r += dp[static_cast<std::size_t>(j)];
  return r;
}
 
/* Two comparisons on one count, conjoined: "m < rank <= m + k" of an
 * OFFSET m LIMIT k, or "count > 1 AND count <= 3" on the aggregate of a
 * subquery.  The count is shared, which the single-comparison pass above
 * refuses (it would couple the two comparisons), but the two only meet in
 * one times gate: their conjunction is the event that the count lies in the
 * intersection of their intervals, a single Poisson-binomial range.  The
 * first comparison becomes a Bernoulli of that probability, the second
 * 𝟙.  Conditions: each comparison has that times gate as its only parent,
 * both read the count through the same constant arithmetic, whose topmost
 * gate (or the count itself) has these two parents only, the rest of the
 * chain and the contributors being private as in the single case. */
static unsigned resolveCountRangePairs(GenericCircuit &gc,
                                       const std::vector<unsigned> &ref)
{
  struct Candidate {
    gate_t cmp;
    AggCmpMatch match;
  };
  const auto nb = gc.getNbGates();
  std::map<gate_t, std::vector<gate_t>> parents;
  std::map<gate_t, std::vector<Candidate>> by_top;
  unsigned resolved = 0;
 
  for (std::size_t i = 0; i < nb; ++i) {
    auto g = static_cast<gate_t>(i);
    for (gate_t w : gc.getWires(g))
      if (gc.getGateType(w) == gate_cmp)
        parents[w].push_back(g);
  }
  for (const auto &kv : parents) {
    AggCmpMatch match;
    if (kv.second.size() != 1 || !matchAggCmp(gc, kv.first, match) ||
        match.agg_kind != AggregationOperator::COUNT)
      continue;
    if (!std::all_of(match.ms.begin(), match.ms.end(),
                     [](long m) { return m == 1; }))
      continue;
    const gate_t top = match.via.empty() ? match.agg : match.via.front();
    by_top[top].push_back({kv.first, std::move(match)});
  }
 
  for (const auto &kv : by_top) {
    const auto &pair = kv.second;
    if (pair.size() != 2 || ref[static_cast<std::size_t>(kv.first)] != 2)
      continue;
    const AggCmpMatch &m0 = pair[0].match, &m1 = pair[1].match;
    const gate_t parent = parents[pair[0].cmp].front();
    if (parents[pair[1].cmp].front() != parent ||
        gc.getGateType(parent) != gate_times ||
        m0.agg != m1.agg || m0.via != m1.via)
      continue;
 
    bool sound = true;
    for (std::size_t i = 1; i < m0.via.size(); ++i)
      if (ref[static_cast<std::size_t>(m0.via[i])] != 1)
        sound = false;
    if (!m0.via.empty() && ref[static_cast<std::size_t>(m0.agg)] != 1)
      sound = false;
    std::vector<double> p;
    p.reserve(m0.ks.size());
    for (std::size_t i = 0; sound && i < m0.ks.size(); ++i) {
      if (ref[static_cast<std::size_t>(m0.semimods[i])] != 1) {
        sound = false;
        break;
      }
      double pi = contributorProb(gc, m0.ks[i], ref, sound);
      p.push_back(pi);
    }
    if (!sound)
      continue;
 
    long lo0, hi0, lo1, hi1;
    if (!countBounds(m0.op, m0.C, lo0, hi0) ||
        !countBounds(m1.op, m1.C, lo1, hi1))
      continue;
    const bool is_scalar =
      (gc.getInfos(m0.agg).second & PROVSQL_AGG_SCALAR_FLAG) != 0;
    double pr = probInRange(p, std::max(lo0, lo1), std::min(hi0, hi1),
                            is_scalar);
    if (pr < 0.0) pr = 0.0;
    if (pr > 1.0) pr = 1.0;
 
    gc.resolveCmpToBernoulli(pair[0].cmp, pr);
    gc.resolveCmpToBernoulli(pair[1].cmp, 1.0);
    resolved += 2;
  }
  return resolved;
}
 
}  // namespace
 
unsigned runCountCmpEvaluator(GenericCircuit &gc)
{
  unsigned resolved = 0;
  const auto nb = gc.getNbGates();
 
  /* Snapshot the cmp-gate ids so in-place rewrites don't affect the
   * iteration : same pattern as runAnalyticEvaluator. */
  std::vector<gate_t> cmps;
  for (std::size_t i = 0; i < nb; ++i) {
    auto g = static_cast<gate_t>(i);
    if (gc.getGateType(g) == gate_cmp)
      cmps.push_back(g);
  }
  if (cmps.empty()) return 0;
 
  /* Reference counts are computed once and not updated as we resolve
   * cmps : resolveCmpToBernoulli only clears the cmp's wires (it does
   * not touch any other gate), so children's ref counts are unchanged
   * with respect to the rest of the circuit.  The snapshot reflects
   * the pre-pass state, which is what we need to certify "no outside
   * reachability" for each candidate's input leaves. */
  auto ref = computeRefCounts(gc);
 
  for (gate_t cmp : cmps) {
    if (gc.getGateType(cmp) != gate_cmp) continue;  /* defensive */
 
    AggCmpMatch match;
    if (!matchAggCmp(gc, cmp, match))
      continue;
 
    /* COUNT(*) over unit-weighted contributors only.  matchAggCmp has
     * already remapped SUM-of-1s to COUNT; a genuine COUNT with a
     * non-unit weight, or a SUM / MIN / MAX / AVG aggregate, is out of
     * this pre-pass's scope and is left for its own evaluator or for
     * provsql_having. */
    if (match.agg_kind != AggregationOperator::COUNT) continue;
    {
      bool all_one = true;
      for (long m : match.ms) if (m != 1) { all_one = false; break; }
      if (!all_one) continue;
    }
 
    const gate_t agg = match.agg;
    const auto &semimods = match.semimods;
    const auto &ks = match.ks;
    const ComparisonOperator op = match.op;
    const int C = static_cast<int>(match.C);
 
    /* Independence certification.  The contributors are independent
     * Bernoulli trials -- the precondition for the Poisson-binomial --
     * exactly when each contributor's sub-circuit
     *
     *     K_i -> semimod_i -> gate_agg -> cmp
     *
     * is a private read-once tree, sharing no randomness with another
     * contributor or with the rest of the circuit.  We check:
     *
     *  1. ref_count[gate_agg] == 1 : the aggregate is consumed by this
     *     cmp alone (catches HAVING COUNT(*) >= a AND COUNT(*) <= b
     *     over a shared count, which would couple the cmps).
     *  2. ref_count[semimod_i] == 1 : the wrapper is consumed by
     *     gate_agg alone.
     *  3. Every randomness-bearing gate inside K_i has ref_count == 1
     *     (verified by @c contributorProb as it recurses).  A single
     *     condition that simultaneously gives: leaf sets pairwise
     *     disjoint across contributors (a shared gate would have
     *     ref >= 2), no reuse outside the cmp (an external parent would
     *     push ref >= 2), and read-once-ness within a contributor (a
     *     leaf used twice would have ref >= 2) -- so the contributor's
     *     marginal is its read-once probability and the contributors
     *     are mutually independent.  Generalises the previous
     *     "K_i is a single gate_input" rule to arbitrary products /
     *     sums of private leaves (e.g. the bid * expertise row of a
     *     join), and bails (leaving the cmp for provsql_having) on any
     *     unsupported gate type.
     *
     * Constants on the path (semimod's M = gate_value(1), the
     * const_side gate_one + gate_value(C), and any gate_one / gate_zero
     * inside a contributor) carry no randomness, so their ref counts
     * are not constrained. */
    if (!aggPrivateToCmp(match, ref)) continue;
    bool sound = true;
    std::vector<double> p;
    p.reserve(ks.size());
    for (std::size_t i = 0; i < ks.size(); ++i) {
      if (ref[static_cast<std::size_t>(semimods[i])] != 1) { sound = false; break; }
      double pi = contributorProb(gc, ks[i], ref, sound);
      if (!sound) break;
      p.push_back(pi);
    }
    if (!sound) continue;
 
    /* Scalar aggregation (no GROUP BY): the empty input is a real world,
     * flagged in the agg gate's info2 high bit by provenance_aggregate. */
    const bool is_scalar =
      (gc.getInfos(agg).second & PROVSQL_AGG_SCALAR_FLAG) != 0;
 
    /* Run the smaller-side dispatch over the contributor marginals. */
    double pr = cdfForOperator(p, op, C, is_scalar);
 
    /* Defensive clamp against floating-point roundoff. */
    if (pr < 0.0) pr = 0.0;
    if (pr > 1.0) pr = 1.0;
 
    gc.resolveCmpToBernoulli(cmp, pr);
    ++resolved;
  }
 
  resolved += resolveCountRangePairs(gc, ref);
  return resolved;
}
 
}  // namespace provsql