text: """Bloom filtering on fragments for equality queries."""
  import base64
  import itertools
  import struct
  from datetime import datetime
  from hashlib import sha256
  from math import ceil, exp, log
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
  
  from psycopg2.sql import SQL, Identifier
  
  from splitgraph.config import SPLITGRAPH_META_SCHEMA
  from splitgraph.core.output import pretty_size
  from splitgraph.core.types import Changeset
  from splitgraph.engine.postgres.engine import SG_UD_FLAG
  
  if TYPE_CHECKING:
      from splitgraph.engine.postgres.engine import PsycopgEngine
  
  
  def _hash_value(value: Union[datetime, int, str, None]) -> Tuple[bytes, bytes]:
      if value is None:
          # See explanation in generate_bloom_index for why this isn't
          # a completely terrible idea.
          value = "NULL"
      return (
          sha256(str(value).encode("utf-8")).digest(),
          sha256((str(value) + "salt").encode("utf-8")).digest(),
      )
  
  
  def generate_bloom_index(
      engine: "PsycopgEngine",
      object_id: str,
      changeset: Optional[Changeset],
      column: str,
      probability: Optional[float] = None,
      size: Optional[int] = None,
  ) -> Tuple[int, str]:
      """
      Generates a bloom filter signature for a given column and a given fragment. Bloom filters
      can answer queries asking whether an item is definitely not in a given set or possibly can be.
  
      The tradeoff is between the probability of a false positive (item said to be in the set when
      it actually isn't) and the size of the filter.
  
      Bloom filters also have an extra parameter, k, or the number of bits in the signature that
      a certain item flips. This parameter has an optimal value for a given number of distinct items
      or a probability and so isn't explicitly passed by the user.
  
      :param engine: Object engine the fragment is cached in.
      :param object_id: Fragment ID
      :param changeset: Optional, if specified, the old column values are included in the index.
      :param column: Column name to generate the index on.
      :param probability: Probability of a false positive. Either this or the size of the filter must
          be specified, but not both.
      :param size: Size of the filter, in bytes.
      :return: Dictionary to be inserted into the index.
      """
  
      if not (probability is None) ^ (size is None):
          raise ValueError("One of probability or size must be specified, but not both!")
  
      # We need k hash functions to generate a signature for every item, which we can construct
      # by taking a linear combination of two hash functions. The first hash function is a simple sha of the
      # column, the second one is a sha of the column + a deterministic salt.
      # First, we outsource the actual hashing to Postgres.
  
      # NULLs are interesting since a sha of a NULL is a NULL again. Here, we
      # turn them into strings with "NULL". This does mean they will collide if
      # we are indexing a string column with an actual "NULL" in it, but in this case
      # it will only mean chunks with NULLs will be fetched for a query with "NULL"
      # and vice versa, which doesn't break anything (this is just a preflight optimisation).
  
      digest_query = SQL(
          "SELECT digest(coalesce({0}::text, 'NULL'), 'sha256'), "
          "digest(coalesce({0}::text, 'NULL') || 'salt', 'sha256') "
          "FROM {1}.{2} o WHERE o.{3} = true"
      ).format(
          Identifier(column),
          Identifier(SPLITGRAPH_META_SCHEMA),
          Identifier(object_id),
          Identifier(SG_UD_FLAG),
      )
  
      digests = engine.run_sql(digest_query)
  
      # Add digests of the old values in the changeset for this column.
      if changeset:
          for _, old_row, _ in changeset.values():
              if column in old_row:
                  # Note that this will add this value to the digest
                  # even if it didn't get changed (say the row got overwritten
                  # but another column got changed). This means that the value
                  # in this case will be counted twice (both the new value
                  # that we hashed before and the old value which we're
                  # hashing now which is the same). This is not an issue
                  # since we deduplicate our digests and the same digest
                  # will set the same bits in the filter to 1, but something
                  # to keep in mind.
                  digests.append(_hash_value(old_row[column]))
  
      # Count the number of distinct items and determine the size (if needed) and optimal number
      # of hash functions.
      distinct_items = list(set(digests))
  
      if probability:
          # The formula gives the number of bits in the array, but we divide it by
          # 8 since we'll be using a byte array for operations + to store the signature.
          size = int(ceil(-len(distinct_items) * log(probability) / log(2) ** 2 / 8))
  
      # For mypy: we've asserted previously that either probability or size are specified
      # and we calculate the size from the probability, so size is no longer Optional[int].
      size = cast(int, size)
  
      size_bits = size * 8
      no_funcs = int(ceil(log(2) * size_bits / len(distinct_items)))
  
      # Generate the filter
      result = bytearray(size)
      for hash_1, hash_2 in distinct_items:
          hash_1 = int.from_bytes(hash_1, byteorder="big")
          hash_2 = int.from_bytes(hash_2, byteorder="big")
          for i in range(no_funcs):
              hash_i = (hash_1 + i * hash_2) % size_bits
              result[hash_i // 8] |= 1 << hash_i % 8
  
      return no_funcs, base64.b64encode(result).decode("ascii")
  
  
  def describe(index_tuple: Tuple[int, str]) -> str:
      """
      Returns a pretty-printed summary of the bloom filter
  
      :param index_tuple: Tuple of (k, base64-encoded fingerprint)
          returned by generate_bloom_index
      :return: String
      """
  
      k, bloom_filter = index_tuple
      footprint = len(bloom_filter)
      bloom_filter = base64.b64decode(bloom_filter)
  
      # Get filter size in bytes (this is not the in-db footprint of the base64 value)
      filter_size = len(bloom_filter)
  
      # Calculate the number of set bits (used to approximate number of items)
      set_bits = 0
      for n in struct.unpack("B" * (len(bloom_filter)), bloom_filter):
          set_bits += bin(n).count("1")
  
      approx_items = -(filter_size * 8) / k * log(1 - set_bits / filter_size / 8)
  
      # Calculate the rough false positive probability
      probability = (1 - exp(-k * approx_items / filter_size / 8)) ** k
  
      return "k=%d, size %s, approx. %d item(s), false positive probability %.1f%%" % (
          k,
          pretty_size(footprint),
          int(approx_items),
          probability * 100,
      )
  
  
  def _prepare_bloom_quals(quals: Any) -> List[List[Tuple[str, int, int]]]:
      """
      Convert list of qualifiers in CNF (ANDed OR-clauses where each clause is "column, operator, value")
      to prepare it for querying the bloom filter:
  
      * Clauses where operator isn't equality are set to True (we can't make a judgement on anything
        but exact matches).
      * Clauses with equality are converted to (column, hash_1, hash_2) so that the sought value
        isn't re-hashed for every fragment
      * OR-clauses where one operator is True are set to True completely (e.g. if we query
        a = 5 OR b > 6, the bloom filter can't say with certainty that there are no rows
        with b > 6 in the fragment and so we have to inspect it)
      * The toplevel AND-clause has all True values removed.
  
      :param quals: Quals in CNF.
      :return: Transformed list of quals
      """
  
      def _process_qual(qual):
          column, operator, value = qual
          if operator != "=":
              return True
  
          hash_1, hash_2 = _hash_value(value)
          hash_1 = int.from_bytes(hash_1, byteorder="big")
          hash_2 = int.from_bytes(hash_2, byteorder="big")
  
          return column, hash_1, hash_2
  
      def _process_or(quals):
          result = []
          for qual in quals:
              qual = _process_qual(qual)
              if qual is True:
                  # anything OR True is True
                  return True
              result.append(qual)
          return result
  
      result = []
      for or_quals in quals:
          or_quals = _process_or(or_quals)
          if or_quals is not True:
              result.append(or_quals)
  
      return result
  
  
  def _match(qual: Tuple[str, int, int], bloom_index: Dict[str, Tuple[int, bytes]]) -> bool:
      """
      Checks whether a processed qual (column, hash_1, hash_2) can match a fragment with
      a given index.
  
      :param qual:
      :param bloom_index:
      """
  
      column, hash_1, hash_2 = qual
      if column not in bloom_index:
          # No index info for this column -- might match
          return True
  
      no_funcs, bloom_filter = bloom_index[column]
      size_bits = len(bloom_filter) * 8
      for i in range(no_funcs):
          hash_i = (hash_1 + i * hash_2) % size_bits
          if not bloom_filter[hash_i // 8] & (1 << hash_i % 8):
              # If at least one position in the filter isn't filled,
              # this qualifier can't be met by anything in the fragment.
              return False
  
      return True
  
  
  def filter_bloom_index(engine: "PsycopgEngine", object_ids: List[str], quals: Any) -> List[str]:
      """
      Runs a bloom filter on given qualifiers using the given objects' previously-generated
      fingerprints.
  
      :param engine: Object engine
      :param object_ids: Object IDs
      :param quals: List of qualifiers
      :return: List of object IDs that might match the qualifiers in `quals` (including
          IDs that don't have a bloom index).
      """
      if not object_ids:
          return object_ids
  
      quals = _prepare_bloom_quals(quals)
      # If we don't have any equalities in quals or quals collapse to something
      # that the bloom filter can't make a judgement about, do nothing.
      if not quals:
          return object_ids
  
      # Load the index: my SQLfu isn't strong enough to create a query that takes
      # care of varying values of K and varying signature sizes.
      bloom_index = engine.run_sql(
          SQL(  # nosec
              "SELECT object_id, index -> 'bloom' FROM {}.{} WHERE object_id IN ("
              + ",".join(itertools.repeat("%s", len(object_ids)))
              + ")"
          ).format(Identifier(SPLITGRAPH_META_SCHEMA), Identifier("objects")),
          object_ids,
      )
  
      bloom_index = {
          o: {col: (i[0], base64.b64decode(i[1])) for col, i in index.items()}
          for o, index in bloom_index
          if index
      }
  
      dropped = []
  
      for object_id in object_ids:
          if object_id not in bloom_index:
              continue
  
          and_result = True
          for or_quals in quals:
              or_result = False
              for or_qual in or_quals:
                  if _match(or_qual, bloom_index[object_id]):
                      or_result = True
                      break
              if not or_result:
                  # One of the subclauses discarded this fragment.
                  and_result = False
                  break
  
          if not and_result:
              dropped.append(object_id)
  
      return [o for o in object_ids if o not in dropped]
  