xref: /aosp_15_r20/external/icing/icing/scoring/advanced_scoring/score-expression.h (revision 8b6cd535a057e39b3b86660c4aa06c99747c2136)

// Copyright (C) 2022 Google LLC
//
// Licensed 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 ICING_SCORING_ADVANCED_SCORING_SCORE_EXPRESSION_H_
#define ICING_SCORING_ADVANCED_SCORING_SCORE_EXPRESSION_H_

#include <cstdint>
#include <memory>
#include <string>
#include <string_view>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include "icing/text_classifier/lib3/utils/base/statusor.h"
#include "icing/absl_ports/canonical_errors.h"
#include "icing/index/embed/embedding-query-results.h"
#include "icing/index/hit/doc-hit-info.h"
#include "icing/index/iterator/doc-hit-info-iterator.h"
#include "icing/join/join-children-fetcher.h"
#include "icing/schema/schema-store.h"
#include "icing/scoring/bm25f-calculator.h"
#include "icing/scoring/section-weights.h"
#include "icing/store/document-filter-data.h"
#include "icing/store/document-id.h"
#include "icing/store/document-store.h"
#include "icing/util/status-macros.h"

namespace icing {
namespace lib {

enum class ScoreExpressionType {
  kDouble,
  kDoubleList,
  kDocument,  // Only "this" is considered as document type.
  // TODO(b/326656531): Instead of creating a vector index type, consider
  // changing it to vector type so that the data is the vector directly.
  kVectorIndex,
  kString,
};

// A map from alias schema type to a set of Icing schema types.
using SchemaTypeAliasMap =
    std::unordered_map<std::string, std::unordered_set<std::string>>;

class ScoreExpression {
 public:
  virtual ~ScoreExpression() = default;

  // Evaluate the score expression to double with the current document.
  //
  // RETURNS:
  //   - The evaluated result as a double on success.
  //   - INVALID_ARGUMENT if a non-finite value is reached while evaluating the
  //                      expression.
  //   - INTERNAL if there are inconsistencies.
  virtual libtextclassifier3::StatusOr<double> EvaluateDouble(
      const DocHitInfo& hit_info, const DocHitInfoIterator* query_it) const {
    if (type() == ScoreExpressionType::kDouble) {
      return absl_ports::UnimplementedError(
          "All ScoreExpressions of type double must provide their own "
          "implementation of EvaluateDouble!");
    }
    return absl_ports::InternalError(
        "Runtime type error: the expression should never be evaluated to a "
        "double. There must be inconsistencies in the static type checking.");
  }

  virtual libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
      const DocHitInfo& hit_info, const DocHitInfoIterator* query_it) const {
    if (type() == ScoreExpressionType::kDoubleList) {
      return absl_ports::UnimplementedError(
          "All ScoreExpressions of type double List must provide their own "
          "implementation of EvaluateList!");
    }
    return absl_ports::InternalError(
        "Runtime type error: the expression should never be evaluated to a "
        "double list. There must be inconsistencies in the static type "
        "checking.");
  }

  virtual libtextclassifier3::StatusOr<std::string_view> EvaluateString()
      const {
    if (type() == ScoreExpressionType::kString) {
      return absl_ports::UnimplementedError(
          "All ScoreExpressions of type string must provide their own "
          "implementation of EvaluateString!");
    }
    return absl_ports::InternalError(
        "Runtime type error: the expression should never be evaluated to a "
        "string. There must be inconsistencies in the static type checking.");
  }

  // Indicate the type to which the current expression will be evaluated.
  virtual ScoreExpressionType type() const = 0;

  // Indicate whether the current expression is a constant.
  // Returns true if and only if the object is of ConstantScoreExpression or
  // StringExpression type.
  virtual bool is_constant() const { return false; }
};

class ThisExpression : public ScoreExpression {
 public:
  static std::unique_ptr<ThisExpression> Create() {
    return std::unique_ptr<ThisExpression>(new ThisExpression());
  }

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDocument;
  }

 private:
  ThisExpression() = default;
};

class ConstantScoreExpression : public ScoreExpression {
 public:
  static std::unique_ptr<ConstantScoreExpression> Create(
      double c, ScoreExpressionType type = ScoreExpressionType::kDouble) {
    return std::unique_ptr<ConstantScoreExpression>(
        new ConstantScoreExpression(c, type));
  }

  libtextclassifier3::StatusOr<double> EvaluateDouble(
      const DocHitInfo&, const DocHitInfoIterator*) const override {
    return c_;
  }

  ScoreExpressionType type() const override { return type_; }

  bool is_constant() const override { return true; }

 private:
  explicit ConstantScoreExpression(double c, ScoreExpressionType type)
      : c_(c), type_(type) {}

  double c_;
  ScoreExpressionType type_;
};

class StringExpression : public ScoreExpression {
 public:
  static std::unique_ptr<StringExpression> Create(std::string str) {
    return std::unique_ptr<StringExpression>(
        new StringExpression(std::move(str)));
  }

  libtextclassifier3::StatusOr<std::string_view> EvaluateString()
      const override {
    return str_;
  }

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kString;
  }

  bool is_constant() const override { return true; }

 private:
  explicit StringExpression(std::string str) : str_(std::move(str)) {}
  std::string str_;
};

class OperatorScoreExpression : public ScoreExpression {
 public:
  enum class OperatorType { kPlus, kMinus, kNegative, kTimes, kDiv };

  // RETURNS:
  //   - An OperatorScoreExpression instance on success if not simplifiable.
  //   - A ConstantScoreExpression instance on success if simplifiable.
  //   - FAILED_PRECONDITION on any null pointer in children.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<std::unique_ptr<ScoreExpression>> Create(
      OperatorType op, std::vector<std::unique_ptr<ScoreExpression>> children);

  libtextclassifier3::StatusOr<double> EvaluateDouble(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDouble;
  }

 private:
  explicit OperatorScoreExpression(
      OperatorType op, std::vector<std::unique_ptr<ScoreExpression>> children)
      : op_(op), children_(std::move(children)) {}

  OperatorType op_;
  std::vector<std::unique_ptr<ScoreExpression>> children_;
};

class MathFunctionScoreExpression : public ScoreExpression {
 public:
  enum class FunctionType {
    kLog,
    kPow,
    kMax,
    kMin,
    kLen,
    kSum,
    kAvg,
    kSqrt,
    kAbs,
    kSin,
    kCos,
    kTan,
    kMaxOrDefault,
    kMinOrDefault,
  };

  static const std::unordered_map<std::string, FunctionType> kFunctionNames;

  static const std::unordered_set<FunctionType> kVariableArgumentsFunctions;

  static const std::unordered_set<FunctionType> kListArgumentFunctions;

  // RETURNS:
  //   - A MathFunctionScoreExpression instance on success if not simplifiable.
  //   - A ConstantScoreExpression instance on success if simplifiable.
  //   - FAILED_PRECONDITION on any null pointer in args.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<std::unique_ptr<ScoreExpression>> Create(
      FunctionType function_type,
      std::vector<std::unique_ptr<ScoreExpression>> args);

  libtextclassifier3::StatusOr<double> EvaluateDouble(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDouble;
  }

 private:
  explicit MathFunctionScoreExpression(
      FunctionType function_type,
      std::vector<std::unique_ptr<ScoreExpression>> args)
      : function_type_(function_type), args_(std::move(args)) {}

  FunctionType function_type_;
  std::vector<std::unique_ptr<ScoreExpression>> args_;
};

class ListOperationFunctionScoreExpression : public ScoreExpression {
 public:
  enum class FunctionType { kFilterByRange };

  static const std::unordered_map<std::string, FunctionType> kFunctionNames;

  // RETURNS:
  //   - A ListOperationFunctionScoreExpression instance on success.
  //   - FAILED_PRECONDITION on any null pointer in args.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<std::unique_ptr<ScoreExpression>> Create(
      FunctionType function_type,
      std::vector<std::unique_ptr<ScoreExpression>> args);

  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDoubleList;
  }

 private:
  explicit ListOperationFunctionScoreExpression(
      FunctionType function_type,
      std::vector<std::unique_ptr<ScoreExpression>> args)
      : function_type_(function_type), args_(std::move(args)) {}

  FunctionType function_type_;
  std::vector<std::unique_ptr<ScoreExpression>> args_;
};

class DocumentFunctionScoreExpression : public ScoreExpression {
 public:
  enum class FunctionType {
    kDocumentScore,
    kCreationTimestamp,
    kUsageCount,
    kUsageLastUsedTimestamp,
  };

  static const std::unordered_map<std::string, FunctionType> kFunctionNames;

  // RETURNS:
  //   - A DocumentFunctionScoreExpression instance on success.
  //   - FAILED_PRECONDITION on any null pointer in args.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<
      std::unique_ptr<DocumentFunctionScoreExpression>>
  Create(FunctionType function_type,
         std::vector<std::unique_ptr<ScoreExpression>> args,
         const DocumentStore* document_store, double default_score,
         int64_t current_time_ms);

  libtextclassifier3::StatusOr<double> EvaluateDouble(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDouble;
  }

 private:
  explicit DocumentFunctionScoreExpression(
      FunctionType function_type,
      std::vector<std::unique_ptr<ScoreExpression>> args,
      const DocumentStore* document_store, double default_score,
      int64_t current_time_ms)
      : args_(std::move(args)),
        document_store_(*document_store),
        default_score_(default_score),
        function_type_(function_type),
        current_time_ms_(current_time_ms) {}

  std::vector<std::unique_ptr<ScoreExpression>> args_;
  const DocumentStore& document_store_;
  double default_score_;
  FunctionType function_type_;
  int64_t current_time_ms_;
};

class RelevanceScoreFunctionScoreExpression : public ScoreExpression {
 public:
  static constexpr std::string_view kFunctionName = "relevanceScore";

  // RETURNS:
  //   - A RelevanceScoreFunctionScoreExpression instance on success.
  //   - FAILED_PRECONDITION on any null pointer in args.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<
      std::unique_ptr<RelevanceScoreFunctionScoreExpression>>
  Create(std::vector<std::unique_ptr<ScoreExpression>> args,
         Bm25fCalculator* bm25f_calculator, double default_score);

  libtextclassifier3::StatusOr<double> EvaluateDouble(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDouble;
  }

 private:
  explicit RelevanceScoreFunctionScoreExpression(
      Bm25fCalculator* bm25f_calculator, double default_score)
      : bm25f_calculator_(*bm25f_calculator), default_score_(default_score) {}

  Bm25fCalculator& bm25f_calculator_;
  double default_score_;
};

class ChildrenRankingSignalsFunctionScoreExpression : public ScoreExpression {
 public:
  static constexpr std::string_view kFunctionName = "childrenRankingSignals";

  // RETURNS:
  //   - A ChildrenRankingSignalsFunctionScoreExpression instance on success.
  //   - FAILED_PRECONDITION on any null pointer in children.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<
      std::unique_ptr<ChildrenRankingSignalsFunctionScoreExpression>>
  Create(std::vector<std::unique_ptr<ScoreExpression>> args,
         const DocumentStore& document_store,
         const JoinChildrenFetcher* join_children_fetcher,
         int64_t current_time_ms);

  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDoubleList;
  }

 private:
  explicit ChildrenRankingSignalsFunctionScoreExpression(
      const DocumentStore& document_store,
      const JoinChildrenFetcher& join_children_fetcher, int64_t current_time_ms)
      : document_store_(document_store),
        join_children_fetcher_(join_children_fetcher),
        current_time_ms_(current_time_ms) {}

  const DocumentStore& document_store_;               // Does not own.
  const JoinChildrenFetcher& join_children_fetcher_;  // Does not own.
  int64_t current_time_ms_;
};

class PropertyWeightsFunctionScoreExpression : public ScoreExpression {
 public:
  static constexpr std::string_view kFunctionName = "propertyWeights";

  // RETURNS:
  //   - A PropertyWeightsFunctionScoreExpression instance on success.
  //   - FAILED_PRECONDITION on any null pointer in children.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<
      std::unique_ptr<PropertyWeightsFunctionScoreExpression>>
  Create(std::vector<std::unique_ptr<ScoreExpression>> args,
         const DocumentStore* document_store,
         const SectionWeights* section_weights, int64_t current_time_ms);

  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
      const DocHitInfo& hit_info, const DocHitInfoIterator*) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDoubleList;
  }

 private:
  explicit PropertyWeightsFunctionScoreExpression(
      const DocumentStore* document_store,
      const SectionWeights* section_weights, int64_t current_time_ms)
      : document_store_(*document_store),
        section_weights_(*section_weights),
        current_time_ms_(current_time_ms) {}
  const DocumentStore& document_store_;
  const SectionWeights& section_weights_;
  int64_t current_time_ms_;
};

class GetEmbeddingParameterFunctionScoreExpression : public ScoreExpression {
 public:
  static constexpr std::string_view kFunctionName = "getEmbeddingParameter";

  // RETURNS:
  //   - A GetEmbeddingParameterFunctionScoreExpression instance on success if
  //     not simplifiable.
  //   - A ConstantScoreExpression instance on success if simplifiable.
  //   - FAILED_PRECONDITION on any null pointer in children.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<std::unique_ptr<ScoreExpression>> Create(
      std::vector<std::unique_ptr<ScoreExpression>> args);

  libtextclassifier3::StatusOr<double> EvaluateDouble(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kVectorIndex;
  }

 private:
  explicit GetEmbeddingParameterFunctionScoreExpression(
      std::unique_ptr<ScoreExpression> arg)
      : arg_(std::move(arg)) {}
  std::unique_ptr<ScoreExpression> arg_;
};

class MatchedSemanticScoresFunctionScoreExpression : public ScoreExpression {
 public:
  static constexpr std::string_view kFunctionName = "matchedSemanticScores";

  // RETURNS:
  //   - A MatchedSemanticScoresFunctionScoreExpression instance on success.
  //   - FAILED_PRECONDITION on any null pointer in children.
  //   - INVALID_ARGUMENT on type errors.
  static libtextclassifier3::StatusOr<
      std::unique_ptr<MatchedSemanticScoresFunctionScoreExpression>>
  Create(std::vector<std::unique_ptr<ScoreExpression>> args,
         SearchSpecProto::EmbeddingQueryMetricType::Code default_metric_type,
         const EmbeddingQueryResults* embedding_query_results);

  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDoubleList;
  }

 private:
  explicit MatchedSemanticScoresFunctionScoreExpression(
      std::vector<std::unique_ptr<ScoreExpression>> args,
      SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
      const EmbeddingQueryResults& embedding_query_results)
      : args_(std::move(args)),
        metric_type_(metric_type),
        embedding_query_results_(embedding_query_results) {}

  std::vector<std::unique_ptr<ScoreExpression>> args_;
  const SearchSpecProto::EmbeddingQueryMetricType::Code metric_type_;
  const EmbeddingQueryResults& embedding_query_results_;
};

class GetScorablePropertyFunctionScoreExpression : public ScoreExpression {
 public:
  static constexpr std::string_view kFunctionName = "getScorableProperty";

  // Returns:
  //   - FAILED_PRECONDITION on any null pointer in children.
  //   - INVALID_ARGUMENT on
  //     - |args| type errors.
  //     - alias_schema_type in the scoring expression does not match to any
  //       schema type in the |schema_type_alias_map|.
  //     - any matched schema type not having the specified property_path as a
  //       scorable property.
  static libtextclassifier3::StatusOr<
      std::unique_ptr<GetScorablePropertyFunctionScoreExpression>>
  Create(std::vector<std::unique_ptr<ScoreExpression>> args,
         const DocumentStore* document_store, const SchemaStore* schema_store,
         const SchemaTypeAliasMap& schema_type_alias_map,
         int64_t current_time_ms);

  ScoreExpressionType type() const override {
    return ScoreExpressionType::kDoubleList;
  }

  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
      const DocHitInfo& hit_info,
      const DocHitInfoIterator* query_it) const override;

 private:
  explicit GetScorablePropertyFunctionScoreExpression(
      const DocumentStore* document_store, const SchemaStore* schema_store,
      int64_t current_time_ms,
      std::unordered_set<SchemaTypeId>&& schema_type_ids,
      std::string_view property_path);

  // Returns a set of schema type ids that are matched to the
  // |alias_schema_type| in the scoring expression, based on the
  // |schema_type_alias_map|.
  //
  // For each of the schema type in the returned set, this function also
  // validates that:
  //   - The schema type is valid in the schema store.
  //   - The |property_path| is defined as scorable under the schema type.
  static libtextclassifier3::StatusOr<std::unordered_set<SchemaTypeId>>
  GetAndValidateSchemaTypeIds(std::string_view alias_schema_type,
                              std::string_view property_path,
                              const SchemaTypeAliasMap& schema_type_alias_map,
                              const SchemaStore& schema_store);

  const DocumentStore& document_store_;
  const SchemaStore& schema_store_;
  int64_t current_time_ms_;
  // A doc hit is evaluated by this function only if its schema type id is in
  // this set.
  std::unordered_set<SchemaTypeId> schema_type_ids_;
  std::string property_path_;
};

}  // namespace lib
}  // namespace icing

#endif  // ICING_SCORING_ADVANCED_SCORING_SCORE_EXPRESSION_H_