<commit_before>(ns cmr.search.services.query-execution.granule-counts-results-feature
  "This enables the :include-granule-counts feature for collection search results. When it is enabled
  collection search results will include the number of granules in each collection that would match
  the same query. It is currently limited to spatial and temporal conditions within collection queries.
  Other types of conditions will not be included in limiting the granule counts."
  (:require [cmr.common.services.errors :as errors]
            [cmr.search.services.query-walkers.condition-extractor :as condition-extractor]
            [cmr.common-app.services.search.query-model :as q]
            [cmr.common-app.services.search.group-query-conditions :as gc]
            [cmr.common-app.services.search.elastic-search-index :as idx]
            [cmr.common-app.services.search.complex-to-simple :as c2s]
            [cmr.common-app.services.search.query-execution :as query-execution]
            [cmr.common-app.services.search.query-to-elastic :as q2e]
            [cmr.search.services.acl-service :as acl-service])
  (:import [cmr.common_app.services.search.query_model
            Query
            ConditionGroup]
           [cmr.search.models.query
            SpatialCondition
            TemporalCondition]))

(defmulti query-results->concept-ids
  "Returns the concept ids from the query results. It is expected that results handlers will
  implement this multi-method"
  (fn [results]
    (:result-format results)))

(defn- valid-parent-condition?
  "The granule count query extractor can extract spatial and temporal conditions from a collection
  query. It only supports extracting spatial and temporal conditions that aren't OR'd or negated
  as well. Returns true if the parent condition is allowed"
  [condition]
  (let [cond-type (type condition)]
    (or (= cond-type Query)
        (and (= cond-type ConditionGroup)
             (= :and (:operation condition))))))

(defn- validate-path-to-condition
  "Validates that the path to the condition we are extracting from the query doesn't contain
  any of the conditions that change it's meaning. This is mainly that the condition is AND'd from
  the top level and not negated. Throws an exception if one is encountered."
  [query condition-path]
  (loop [conditions condition-path]
    (when-let [condition (first conditions)]
      (when-not (valid-parent-condition? condition)
        (errors/internal-error!
          (str "Granule query extractor found spatial or temporal condition within an unexpected "
               "condition. Query:" (pr-str query))))
      (recur (rest conditions)))))

(defn- extract-spatial-and-temporal-conditions
  "Returns a list of the spatial and temporal conditions in the query. Throws an exception if a
  parent of the spatial or temporal condition isn't expected."
  [query]
  (condition-extractor/extract-conditions
    query
    ;; Matches spatial or temporal conditions
    (fn [condition-path condition]
      (let [cond-type (type condition)]
        (when (or (= cond-type SpatialCondition)
                  (= cond-type TemporalCondition))
          ;; Validate that the parent of the spatial or temporal condition is acceptable
          (validate-path-to-condition query condition-path)
          true)))))

(defn- convert-spatial-or-temporal-condition
  "Applies any necessary changes to spatial or temporal condition to make them work with granules"
  [condition]
  (if (= TemporalCondition (type condition))
    ;; Turn off limit to granules so that the correct fields will be searched. limit to granules is a
    ;; collection applicable parameter.
    (assoc condition :limit-to-granules false)
    condition))

(defn extract-granule-count-query
  "Extracts a query to find the number of granules per collection in the results from a collection query
  coll-query - The collection query
  results - the results of the collection query"
  [coll-query results]
  (let [collection-ids (query-results->concept-ids results)
        spatial-temp-conds (->> coll-query
                                extract-spatial-and-temporal-conditions
                                (map convert-spatial-or-temporal-condition))
        condition (if (seq collection-ids)
                    (gc/and-conds (cons (q/string-conditions :collection-concept-id collection-ids true)
                                        spatial-temp-conds))
                    ;; The results were empty so the granule count query doesn't need to find anything.
                    q/match-none)]
    (q/query {:concept-type :granule
              :condition condition
              ;; We don't need any results
              :page-size 0
              :result-format :query-specified
              :aggregations {:granule-counts-by-collection-id
                             {:terms {:field (q2e/query-field->elastic-field
                                              :collection-concept-id :granule)
                                      :size (count collection-ids)}}}})))

(defn- search-results->granule-counts
  "Extracts the granule counts per collection id in a map from the search results from ElasticSearch"
  [elastic-results]
  (let [aggs (get-in elastic-results [:aggregations :granule-counts-by-collection-id :buckets])]
    (into {} (for [{collection-id :key num-granules :doc_count} aggs]
               [collection-id num-granules]))))

;; This find granule counts per collection.
(defmethod query-execution/post-process-query-result-feature :granule-counts
  [context query elastic-results query-results feature]
  (->> query-results
       (extract-granule-count-query query)
       (query-execution/execute-query context)
       search-results->granule-counts
       (assoc query-results :granule-counts-map)))
<commit_msg>CMR-3228: Address review comments.<commit_after>(ns cmr.search.services.query-execution.granule-counts-results-feature
  "This enables the :include-granule-counts feature for collection search results. When it is enabled
  collection search results will include the number of granules in each collection that would match
  the same query. It is currently limited to spatial and temporal conditions within collection queries.
  Other types of conditions will not be included in limiting the granule counts."
  (:require [cmr.common.services.errors :as errors]
            [cmr.search.services.query-walkers.condition-extractor :as condition-extractor]
            [cmr.common-app.services.search.query-model :as q]
            [cmr.common-app.services.search.group-query-conditions :as gc]
            [cmr.common-app.services.search.elastic-search-index :as idx]
            [cmr.common-app.services.search.complex-to-simple :as c2s]
            [cmr.common-app.services.search.query-execution :as query-execution]
            [cmr.common-app.services.search.query-to-elastic :as q2e]
            [cmr.search.services.acl-service :as acl-service])
  (:import [cmr.common_app.services.search.query_model
            Query
            ConditionGroup]
           [cmr.search.models.query
            SpatialCondition
            TemporalCondition]))

(defmulti query-results->concept-ids
  "Returns the concept ids from the query results. It is expected that results handlers will
  implement this multi-method"
  (fn [results]
    (:result-format results)))

(defn- valid-parent-condition?
  "The granule count query extractor can extract spatial and temporal conditions from a collection
  query. It only supports extracting spatial and temporal conditions that aren't OR'd or negated
  as well. Returns true if the parent condition is allowed"
  [condition]
  (let [cond-type (type condition)]
    (or (= cond-type Query)
        (and (= cond-type ConditionGroup)
             (= :and (:operation condition))))))

(defn- validate-path-to-condition
  "Validates that the path to the condition we are extracting from the query doesn't contain
  any of the conditions that change it's meaning. This is mainly that the condition is AND'd from
  the top level and not negated. Throws an exception if one is encountered."
  [query condition-path]
  (loop [conditions condition-path]
    (when-let [condition (first conditions)]
      (when-not (valid-parent-condition? condition)
        (errors/internal-error!
          (str "Granule query extractor found spatial or temporal condition within an unexpected "
               "condition. Query:" (pr-str query))))
      (recur (rest conditions)))))

(defn- extract-spatial-and-temporal-conditions
  "Returns a list of the spatial and temporal conditions in the query. Throws an exception if a
  parent of the spatial or temporal condition isn't expected."
  [query]
  (condition-extractor/extract-conditions
    query
    ;; Matches spatial or temporal conditions
    (fn [condition-path condition]
      (when (or (instance? SpatialCondition condition)
                (instance? TemporalCondition condition))
        ;; Validate that the parent of the spatial or temporal condition is acceptable
        (validate-path-to-condition query condition-path)
        true))))


(defn- convert-spatial-or-temporal-condition
  "Applies any necessary changes to spatial or temporal condition to make them work with granules"
  [condition]
  (if (instance? TemporalCondition condition)
    ;; Turn off limit to granules so that the correct fields will be searched. limit to granules is a
    ;; collection applicable parameter.
    (assoc condition :limit-to-granules false)
    condition))

(defn extract-granule-count-query
  "Extracts a query to find the number of granules per collection in the results from a collection query
  coll-query - The collection query
  results - the results of the collection query"
  [coll-query results]
  (let [collection-ids (query-results->concept-ids results)
        spatial-temp-conds (->> coll-query
                                extract-spatial-and-temporal-conditions
                                (map convert-spatial-or-temporal-condition))
        condition (if (seq collection-ids)
                    (gc/and-conds (cons (q/string-conditions :collection-concept-id collection-ids true)
                                        spatial-temp-conds))
                    ;; The results were empty so the granule count query doesn't need to find anything.
                    q/match-none)]
    (q/query {:concept-type :granule
              :condition condition
              ;; We don't need any results
              :page-size 0
              :result-format :query-specified
              :aggregations {:granule-counts-by-collection-id
                             {:terms {:field (q2e/query-field->elastic-field
                                              :collection-concept-id :granule)
                                      :size (count collection-ids)}}}})))

(defn- search-results->granule-counts
  "Extracts the granule counts per collection id in a map from the search results from ElasticSearch"
  [elastic-results]
  (let [aggs (get-in elastic-results [:aggregations :granule-counts-by-collection-id :buckets])]
    (into {} (for [{collection-id :key num-granules :doc_count} aggs]
               [collection-id num-granules]))))

;; This find granule counts per collection.
(defmethod query-execution/post-process-query-result-feature :granule-counts
  [context query elastic-results query-results feature]
  (->> query-results
       (extract-granule-count-query query)
       (query-execution/execute-query context)
       search-results->granule-counts
       (assoc query-results :granule-counts-map)))

{'fraction_non_alphanumeric': 0.08958046767537826, 'fraction_numerical': 0.0010316368638239339, 'mean_word_length': 2.7130545802744974, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1569841', 'n_tokens_mistral': 3166, 'n_tokens_neox': 2988, 'n_words': 1027}