query_execution/ProbabilityStore.hpp - incubator-retired-quickstep - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, 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.
  **/

 #ifndef QUICKSTEP_QUERY_EXECUTION_PROBABILITY_STORE_HPP_
 #define QUICKSTEP_QUERY_EXECUTION_PROBABILITY_STORE_HPP_

 #include <algorithm>
 #include <cstddef>
 #include <random>
 #include <unordered_map>
 #include <vector>

 #include "utility/Macros.hpp"

 #include "glog/logging.h"

 namespace quickstep {

 /**
  * @brief A class that stores the probabilities of objects. We use an integer field
  *        called "key" to identify each object.
  *        A probability is expressed as a fraction. All the objects share a common denominator.
  **/
 class ProbabilityStore {
  public:
   /**
    * @brief Constructor.
    **/
   ProbabilityStore()
       : common_denominator_(1.0), dist_(0.0, 1.0), engine_(std::random_device()()) {}

   /**
    * @brief Get the number of objects in the store.
    **/
   inline std::size_t getNumObjects() const {
     DCHECK_EQ(individual_probabilities_.size(), cumulative_probabilities_.size());
     return individual_probabilities_.size();
   }

   /**
    * @brief Get the common denominator.
    */
   inline std::size_t getDenominator() const {
     return common_denominator_;
   }

   /**
    * @brief Check if an object with a given key is present.
    * @param key The key of the given object.
    * @return True if the object is present, false otherwise.
    */
   inline bool hasObject(const std::size_t key) const {
     return (individual_probabilities_.find(key) != individual_probabilities_.end());
   }

   /**
    * @brief Add individual (not cumulative) probability for a given object with
    *        updated denominator.
    *
    * @note This function leaves the cumulative probabilities in a consistent
    *       state. An alternative lazy implementation should be written if cost
    *       of calculating cumulative probabilities is high.
    *
    * @param key The key of the given object.
    * @param numerator The numerator for the given object.
    * @param new_denominator The updated denominator for the store.
    **/
   void addOrUpdateObjectNewDenominator(const std::size_t key,
                                        const float numerator,
                                        const float new_denominator) {
     CHECK_GT(new_denominator, 0u);
     DCHECK_LE(numerator, new_denominator);
     common_denominator_ = new_denominator;
     addOrUpdateObjectHelper(key, numerator);
   }

   /**
    * @brief Add or update multiple objects with a new denominator.
    * @param keys The keys to be added/updated.
    * @param numerators The numerators to be added/updated.
    * @param new_denominator The new denominator.
    */
   void addOrUpdateObjectsNewDenominator(
       const std::vector<std::size_t> &keys,
       const std::vector<float> &numerators,
       const float new_denominator) {
     CHECK_GT(new_denominator, 0u);
     common_denominator_ = new_denominator;
     addOrUpdateObjectsHelper(keys, numerators);
   }

   /**
    * @brief Remove an object from the store.
    *
    * @note This function decrements the denominator with the value of the numerator being removed.
    *
    * @param key The key of the object to be removed.
    **/
   void removeObject(const std::size_t key) {
     auto individual_it = individual_probabilities_.find(key);
     DCHECK(individual_it != individual_probabilities_.end());
     const float new_denominator = common_denominator_ - individual_it->second;
     individual_probabilities_.erase(individual_it);
     common_denominator_ = new_denominator;
     updateCumulativeProbabilities();
   }

   /**
    * @brief Get the individual probability (not cumulative) for an object.
    *
    * @param key The key of the object.
    **/
   const float getIndividualProbability(const std::size_t key) const {
     const auto it = individual_probabilities_.find(key);
     DCHECK(it != individual_probabilities_.end());
     DCHECK_NE(0.0, common_denominator_);
     return it->second / common_denominator_;
   }

   /**
    * @brief Update the cumulative probabilities.
    *
    * @note This function should be called upon if there are any updates,
    *       additions or deletions to the individual probabilities.
    * @note An efficient implementation should be written if there are large
    *       number of objects.
    **/
   void updateCumulativeProbabilities() {
     cumulative_probabilities_.clear();
     float cumulative_probability = 0;
     for (const auto p : individual_probabilities_) {
       cumulative_probability += p.second / common_denominator_;
       cumulative_probabilities_.emplace_back(p.first,
                                              cumulative_probability);
     }
     if (!cumulative_probabilities_.empty()) {
       // Adjust the last cumulative probability manually to 1, so that
       // floating addition related rounding issues are ignored.
       cumulative_probabilities_.back().updateProbability(1);
     }
   }

   /**
    * @brief Return a uniformly chosen random key.
    **/
   inline const std::size_t pickRandomKey() const {
     return getKeyForProbability(dist_(engine_));
   }

  private:
   struct ProbabilityInfo {
    public:
     ProbabilityInfo(const std::size_t key, const float probability)
         : key_(key), probability_(probability) {
       DCHECK(probability_ <= 1.0);
     }

     void updateProbability(const float new_probability) {
       probability_ = new_probability;
     }

     const std::size_t key_;
     float probability_;
   };

   /**
    * @brief Get a key for a given cumulative probability.
    *
    * @param key_cumulative_probability The input cumulative probability.
    *
    * @return The object that has a cumulative probability that is greater than
    *         or equal to the input cumulative probability.
    **/
   inline std::size_t getKeyForProbability(
       const float key_cumulative_probability) const {
     DCHECK(!cumulative_probabilities_.empty());
     // It doesn't matter in which order the objects are arranged in the
     // cumulative_probabilities_ vector.
     ProbabilityInfo search_key(0, key_cumulative_probability);
     const auto it = std::upper_bound(
         cumulative_probabilities_.begin(),
         cumulative_probabilities_.end(),
         search_key,
         [](const ProbabilityInfo &a, const ProbabilityInfo &b) {
           return a.probability_ < b.probability_;
         });
     DCHECK(it != std::end(cumulative_probabilities_));
     return it->key_;
   }

   /**
    * @brief Add individual (not cumulative) probability for a given object.
    *
    * @note This function leaves the cumulative probabilities in a consistent
    *       state. An alternative lazy implementation should be written if cost
    *       of calculating cumulative probabilities is high.
    * @note This function may override previously written probability values.
    *
    * @param key The key of the given object.
    * @param numerator The numerator for the given object.
    **/
   void addOrUpdateObjectHelper(const std::size_t key,
                                const float numerator) {
     DCHECK_LE(numerator, common_denominator_);
     individual_probabilities_[key] = numerator;
     updateCumulativeProbabilities();
   }

   /**
    * @brief Add individual (not cumulative) probabilities for given objects.
    *
    * @note This function leaves the cumulative probabilities in a consistent
    *       state. An alternative lazy implementation should be written if cost
    *       of calculating cumulative probabilities is high.
    * @note This function may override previously written probability values.
    *
    * @param keys A vector of keys to be added.
    * @param numerators The numerators of the given objects.
    **/
   void addOrUpdateObjectsHelper(const std::vector<std::size_t> &keys,
                                 const std::vector<float> &numerators) {
     DCHECK_EQ(keys.size(), numerators.size());
     for (std::size_t i = 0; i < keys.size(); ++i) {
       DCHECK_LE(numerators[i], common_denominator_);
       individual_probabilities_[keys[i]] = numerators[i];
     }
     updateCumulativeProbabilities();
   }

   // Key = key of the object.
   // Value = Numerator of the object.
   std::unordered_map<std::size_t, float> individual_probabilities_;
   std::vector<ProbabilityInfo> cumulative_probabilities_;

   float common_denominator_;

   mutable std::uniform_real_distribution<float> dist_;
   mutable std::default_random_engine engine_;

   DISALLOW_COPY_AND_ASSIGN(ProbabilityStore);
 };

 /** @} */

 }  // namespace quickstep

 #endif  // QUICKSTEP_QUERY_EXECUTION_PROBABILITY_STORE_HPP_
	/**
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, 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.
	**/

	#ifndef QUICKSTEP_QUERY_EXECUTION_PROBABILITY_STORE_HPP_
	#define QUICKSTEP_QUERY_EXECUTION_PROBABILITY_STORE_HPP_

	#include <algorithm>
	#include <cstddef>
	#include <random>
	#include <unordered_map>
	#include <vector>

	#include "utility/Macros.hpp"

	#include "glog/logging.h"

	namespace quickstep {

	/**
	* @brief A class that stores the probabilities of objects. We use an integer field
	* called "key" to identify each object.
	* A probability is expressed as a fraction. All the objects share a common denominator.
	**/
	class ProbabilityStore {
	public:
	/**
	* @brief Constructor.
	**/
	ProbabilityStore()
	: common_denominator_(1.0), dist_(0.0, 1.0), engine_(std::random_device()()) {}

	/**
	* @brief Get the number of objects in the store.
	**/
	inline std::size_t getNumObjects() const {
	DCHECK_EQ(individual_probabilities_.size(), cumulative_probabilities_.size());
	return individual_probabilities_.size();
	}

	/**
	* @brief Get the common denominator.
	*/
	inline std::size_t getDenominator() const {
	return common_denominator_;
	}

	/**
	* @brief Check if an object with a given key is present.
	* @param key The key of the given object.
	* @return True if the object is present, false otherwise.
	*/
	inline bool hasObject(const std::size_t key) const {
	return (individual_probabilities_.find(key) != individual_probabilities_.end());
	}

	/**
	* @brief Add individual (not cumulative) probability for a given object with
	* updated denominator.
	*
	* @note This function leaves the cumulative probabilities in a consistent
	* state. An alternative lazy implementation should be written if cost
	* of calculating cumulative probabilities is high.
	*
	* @param key The key of the given object.
	* @param numerator The numerator for the given object.
	* @param new_denominator The updated denominator for the store.
	**/
	void addOrUpdateObjectNewDenominator(const std::size_t key,
	const float numerator,
	const float new_denominator) {
	CHECK_GT(new_denominator, 0u);
	DCHECK_LE(numerator, new_denominator);
	common_denominator_ = new_denominator;
	addOrUpdateObjectHelper(key, numerator);
	}

	/**
	* @brief Add or update multiple objects with a new denominator.
	* @param keys The keys to be added/updated.
	* @param numerators The numerators to be added/updated.
	* @param new_denominator The new denominator.
	*/
	void addOrUpdateObjectsNewDenominator(
	const std::vector<std::size_t> &keys,
	const std::vector<float> &numerators,
	const float new_denominator) {
	CHECK_GT(new_denominator, 0u);
	common_denominator_ = new_denominator;
	addOrUpdateObjectsHelper(keys, numerators);
	}

	/**
	* @brief Remove an object from the store.
	*
	* @note This function decrements the denominator with the value of the numerator being removed.
	*
	* @param key The key of the object to be removed.
	**/
	void removeObject(const std::size_t key) {
	auto individual_it = individual_probabilities_.find(key);
	DCHECK(individual_it != individual_probabilities_.end());
	const float new_denominator = common_denominator_ - individual_it->second;
	individual_probabilities_.erase(individual_it);
	common_denominator_ = new_denominator;
	updateCumulativeProbabilities();
	}

	/**
	* @brief Get the individual probability (not cumulative) for an object.
	*
	* @param key The key of the object.
	**/
	const float getIndividualProbability(const std::size_t key) const {
	const auto it = individual_probabilities_.find(key);
	DCHECK(it != individual_probabilities_.end());
	DCHECK_NE(0.0, common_denominator_);
	return it->second / common_denominator_;
	}

	/**
	* @brief Update the cumulative probabilities.
	*
	* @note This function should be called upon if there are any updates,
	* additions or deletions to the individual probabilities.
	* @note An efficient implementation should be written if there are large
	* number of objects.
	**/
	void updateCumulativeProbabilities() {
	cumulative_probabilities_.clear();
	float cumulative_probability = 0;
	for (const auto p : individual_probabilities_) {
	cumulative_probability += p.second / common_denominator_;
	cumulative_probabilities_.emplace_back(p.first,
	cumulative_probability);
	}
	if (!cumulative_probabilities_.empty()) {
	// Adjust the last cumulative probability manually to 1, so that
	// floating addition related rounding issues are ignored.
	cumulative_probabilities_.back().updateProbability(1);
	}
	}

	/**
	* @brief Return a uniformly chosen random key.
	**/
	inline const std::size_t pickRandomKey() const {
	return getKeyForProbability(dist_(engine_));
	}

	private:
	struct ProbabilityInfo {
	public:
	ProbabilityInfo(const std::size_t key, const float probability)
	: key_(key), probability_(probability) {
	DCHECK(probability_ <= 1.0);
	}

	void updateProbability(const float new_probability) {
	probability_ = new_probability;
	}

	const std::size_t key_;
	float probability_;
	};

	/**
	* @brief Get a key for a given cumulative probability.
	*
	* @param key_cumulative_probability The input cumulative probability.
	*
	* @return The object that has a cumulative probability that is greater than
	* or equal to the input cumulative probability.
	**/
	inline std::size_t getKeyForProbability(
	const float key_cumulative_probability) const {
	DCHECK(!cumulative_probabilities_.empty());
	// It doesn't matter in which order the objects are arranged in the
	// cumulative_probabilities_ vector.
	ProbabilityInfo search_key(0, key_cumulative_probability);
	const auto it = std::upper_bound(
	cumulative_probabilities_.begin(),
	cumulative_probabilities_.end(),
	search_key,
	[](const ProbabilityInfo &a, const ProbabilityInfo &b) {
	return a.probability_ < b.probability_;
	});
	DCHECK(it != std::end(cumulative_probabilities_));
	return it->key_;
	}

	/**
	* @brief Add individual (not cumulative) probability for a given object.
	*
	* @note This function leaves the cumulative probabilities in a consistent
	* state. An alternative lazy implementation should be written if cost
	* of calculating cumulative probabilities is high.
	* @note This function may override previously written probability values.
	*
	* @param key The key of the given object.
	* @param numerator The numerator for the given object.
	**/
	void addOrUpdateObjectHelper(const std::size_t key,
	const float numerator) {
	DCHECK_LE(numerator, common_denominator_);
	individual_probabilities_[key] = numerator;
	updateCumulativeProbabilities();
	}

	/**
	* @brief Add individual (not cumulative) probabilities for given objects.
	*
	* @note This function leaves the cumulative probabilities in a consistent
	* state. An alternative lazy implementation should be written if cost
	* of calculating cumulative probabilities is high.
	* @note This function may override previously written probability values.
	*
	* @param keys A vector of keys to be added.
	* @param numerators The numerators of the given objects.
	**/
	void addOrUpdateObjectsHelper(const std::vector<std::size_t> &keys,
	const std::vector<float> &numerators) {
	DCHECK_EQ(keys.size(), numerators.size());
	for (std::size_t i = 0; i < keys.size(); ++i) {
	DCHECK_LE(numerators[i], common_denominator_);
	individual_probabilities_[keys[i]] = numerators[i];
	}
	updateCumulativeProbabilities();
	}

	// Key = key of the object.
	// Value = Numerator of the object.
	std::unordered_map<std::size_t, float> individual_probabilities_;
	std::vector<ProbabilityInfo> cumulative_probabilities_;

	float common_denominator_;

	mutable std::uniform_real_distribution<float> dist_;
	mutable std::default_random_engine engine_;

	DISALLOW_COPY_AND_ASSIGN(ProbabilityStore);
	};

	/** @} */

	} // namespace quickstep

	#endif // QUICKSTEP_QUERY_EXECUTION_PROBABILITY_STORE_HPP_