<commit_before>(ns bzzz.index
  (use bzzz.analyzer)
  (use bzzz.util)
  (use bzzz.const)
  (use bzzz.query)
  (:require [clojure.tools.logging :as log])
  (:import (java.io StringReader File)
           (org.apache.lucene.analysis Analyzer TokenStream)
           (org.apache.lucene.document Document Field Field$Index Field$Store)
           (org.apache.lucene.search.highlight Highlighter QueryScorer
                                               SimpleHTMLFormatter TextFragment)
           (org.apache.lucene.index IndexWriter IndexReader Term
                                    IndexWriterConfig DirectoryReader FieldInfo)
           (org.apache.lucene.search Query ScoreDoc SearcherManager IndexSearcher
                                     Explanation Collector TopScoreDocCollector TopDocsCollector)
           (org.apache.lucene.store NIOFSDirectory Directory)))
(set! *warn-on-reflection* true)
(def root* (atom default-root))
(def identifier* (atom default-identifier))
(def mapping* (atom {}))

(defn acceptable-index-name [name]
  (clojure.string/replace name #"[^a-zA-Z_0-9-]" ""))

(defn root-identifier-path []
  (File. ^File (File. (as-str @root*)) (as-str @identifier*)))

(defn new-index-directory ^Directory [name]
  (let [path ^File (root-identifier-path)]
    (.mkdir path)
    (NIOFSDirectory. (File. path (as-str (acceptable-index-name name))))))

(defn new-index-writer ^IndexWriter [name]
  (IndexWriter. (new-index-directory name)
                (IndexWriterConfig. *version* @analyzer*)))

(defn new-index-reader ^IndexReader [name]
  (with-open [writer (new-index-writer name)]
    (DirectoryReader/open ^IndexWriter writer false)))

(defn analyzed? [name]
  (if (or (substring? "_not_analyzed" name)
          (= name id-field))
    false
    true))

(defn norms? [name]
  (if (or (substring? "_no_norms" name)
          (= name id-field))
    false
    true))

(defn stored? [name]
  (if (and (substring? "_no_store" name)
           (not (= name id-field)))
    false
    true))

(defn indexed? [name]
  (if (and (substring? "_no_index" name)
           (not (= name id-field)))
    false
    true))

(defn- add-field [document key value]
  (let [ str-key (as-str key) ]
    (.add ^Document document
          (Field. str-key (as-str value)
                  (if (stored? str-key)
                    Field$Store/YES
                    Field$Store/NO)
                  (if (indexed? str-key)
                    (case [(analyzed? str-key) (norms? str-key)]
                      [true true] Field$Index/ANALYZED
                      [true false] Field$Index/ANALYZED_NO_NORMS
                      [false true] Field$Index/NOT_ANALYZED
                      [false false] Field$Index/NOT_ANALYZED_NO_NORMS)
                    Field$Index/NO)))))

(defn map->document
  [hmap]
  (let [document (Document.)]
    (doseq [[key value] hmap]
      (add-field document (as-str key) value))
    document))

(defn document->map
  [^Document doc score highlighter ^Explanation explanation]
  (let [m (into {:_score score }
                (for [^Field f (.getFields doc)]
                  [(keyword (.name f)) (.stringValue f)]))
        fragments (highlighter m)]
    (conj
     m
     (when explanation (assoc m :_explain (.toString explanation)))
     (when fragments   (assoc m :_highlight fragments)))))

(defn get-search-manager ^SearcherManager [index]
  (locking mapping*
    (when (nil? (@mapping* index))
      (swap! mapping* assoc index (SearcherManager. (new-index-directory index)
                                                    nil)))
    (@mapping* index)))

(defn refresh-search-managers []
  (locking mapping*
    (doseq [[index ^SearcherManager manager] @mapping*]
      (log/debug "refreshing: " index " " manager)
      (.maybeRefresh manager))))

(defn reset-search-managers []
  (doseq [[name ^SearcherManager manager] @mapping*]
    (log/info "\tclosing: " name " " manager)
    (.close manager))
  (reset! mapping* {}))

(defn bootstrap-indexes []
  (doseq [f (.listFiles (File. (as-str @root*)))]
    (if (.isDirectory ^File f)
      (get-search-manager (.getName ^File f)))))

(defn use-searcher [index callback]
  (let [manager (get-search-manager index)
        searcher (.acquire manager)]
    (try
      (callback searcher)
      (finally (.release manager searcher)))))

(defn use-writer [index callback]
  (let [writer (new-index-writer index)]
    (try
      (callback ^IndexWriter writer)
      (finally
        (.commit writer)
        (.forceMerge writer 1)
        (.close writer)))))

(defn store
  [index maps analyzer]
  (if analyzer
    (reset! analyzer* (parse-analyzer analyzer)))
  (use-writer index (fn [^IndexWriter writer]
                      (doseq [m maps]
                        (if (:id m)
                          (.updateDocument writer ^Term (Term. ^String id-field
                                                               (as-str (:id m))) (map->document m))
                          (.addDocument writer (map->document m))))
                      { index true })))

(defn delete-from-query
  [index input]
  (use-writer index (fn [^IndexWriter writer]
                      ;; method is deleteDocuments(Query...)
                      (let [query (parse-query input (extract-analyzer nil))]
                        (.deleteDocuments writer ^"[Lorg.apache.lucene.search.Query;" (into-array Query [query]))
                        { index (.toString query) }))))

(defn delete-all [index]
  (use-writer index (fn [^IndexWriter writer] (.deleteAll writer))))


(defn fragment->map [^TextFragment fragment]
  {:text (.toString fragment)
   :score (.getScore fragment)
   :frag-num (wall-hack-field TextFragment :fragNum fragment)
   :text-start-pos (wall-hack-field TextFragment :textStartPos fragment)
   :text-end-pos (wall-hack-field TextFragment :textEndPos fragment)})
<commit_msg>just use the object's class<commit_after>(ns bzzz.index
  (use bzzz.analyzer)
  (use bzzz.util)
  (use bzzz.const)
  (use bzzz.query)
  (:require [clojure.tools.logging :as log])
  (:import (java.io StringReader File)
           (org.apache.lucene.analysis Analyzer TokenStream)
           (org.apache.lucene.document Document Field Field$Index Field$Store)
           (org.apache.lucene.search.highlight Highlighter QueryScorer
                                               SimpleHTMLFormatter TextFragment)
           (org.apache.lucene.index IndexWriter IndexReader Term
                                    IndexWriterConfig DirectoryReader FieldInfo)
           (org.apache.lucene.search Query ScoreDoc SearcherManager IndexSearcher
                                     Explanation Collector TopScoreDocCollector TopDocsCollector)
           (org.apache.lucene.store NIOFSDirectory Directory)))
(set! *warn-on-reflection* true)
(def root* (atom default-root))
(def identifier* (atom default-identifier))
(def mapping* (atom {}))

(defn acceptable-index-name [name]
  (clojure.string/replace name #"[^a-zA-Z_0-9-]" ""))

(defn root-identifier-path []
  (File. ^File (File. (as-str @root*)) (as-str @identifier*)))

(defn new-index-directory ^Directory [name]
  (let [path ^File (root-identifier-path)]
    (.mkdir path)
    (NIOFSDirectory. (File. path (as-str (acceptable-index-name name))))))

(defn new-index-writer ^IndexWriter [name]
  (IndexWriter. (new-index-directory name)
                (IndexWriterConfig. *version* @analyzer*)))

(defn new-index-reader ^IndexReader [name]
  (with-open [writer (new-index-writer name)]
    (DirectoryReader/open ^IndexWriter writer false)))

(defn analyzed? [name]
  (if (or (substring? "_not_analyzed" name)
          (= name id-field))
    false
    true))

(defn norms? [name]
  (if (or (substring? "_no_norms" name)
          (= name id-field))
    false
    true))

(defn stored? [name]
  (if (and (substring? "_no_store" name)
           (not (= name id-field)))
    false
    true))

(defn indexed? [name]
  (if (and (substring? "_no_index" name)
           (not (= name id-field)))
    false
    true))

(defn- add-field [document key value]
  (let [ str-key (as-str key) ]
    (.add ^Document document
          (Field. str-key (as-str value)
                  (if (stored? str-key)
                    Field$Store/YES
                    Field$Store/NO)
                  (if (indexed? str-key)
                    (case [(analyzed? str-key) (norms? str-key)]
                      [true true] Field$Index/ANALYZED
                      [true false] Field$Index/ANALYZED_NO_NORMS
                      [false true] Field$Index/NOT_ANALYZED
                      [false false] Field$Index/NOT_ANALYZED_NO_NORMS)
                    Field$Index/NO)))))

(defn map->document
  [hmap]
  (let [document (Document.)]
    (doseq [[key value] hmap]
      (add-field document (as-str key) value))
    document))

(defn document->map
  [^Document doc score highlighter ^Explanation explanation]
  (let [m (into {:_score score }
                (for [^Field f (.getFields doc)]
                  [(keyword (.name f)) (.stringValue f)]))
        fragments (highlighter m)]
    (conj
     m
     (when explanation (assoc m :_explain (.toString explanation)))
     (when fragments   (assoc m :_highlight fragments)))))

(defn get-search-manager ^SearcherManager [index]
  (locking mapping*
    (when (nil? (@mapping* index))
      (swap! mapping* assoc index (SearcherManager. (new-index-directory index)
                                                    nil)))
    (@mapping* index)))

(defn refresh-search-managers []
  (locking mapping*
    (doseq [[index ^SearcherManager manager] @mapping*]
      (log/debug "refreshing: " index " " manager)
      (.maybeRefresh manager))))

(defn reset-search-managers []
  (doseq [[name ^SearcherManager manager] @mapping*]
    (log/info "\tclosing: " name " " manager)
    (.close manager))
  (reset! mapping* {}))

(defn bootstrap-indexes []
  (doseq [f (.listFiles (File. (as-str @root*)))]
    (if (.isDirectory ^File f)
      (get-search-manager (.getName ^File f)))))

(defn use-searcher [index callback]
  (let [manager (get-search-manager index)
        searcher (.acquire manager)]
    (try
      (callback searcher)
      (finally (.release manager searcher)))))

(defn use-writer [index callback]
  (let [writer (new-index-writer index)]
    (try
      (callback ^IndexWriter writer)
      (finally
        (.commit writer)
        (.forceMerge writer 1)
        (.close writer)))))

(defn store
  [index maps analyzer]
  (if analyzer
    (reset! analyzer* (parse-analyzer analyzer)))
  (use-writer index (fn [^IndexWriter writer]
                      (doseq [m maps]
                        (if (:id m)
                          (.updateDocument writer ^Term (Term. ^String id-field
                                                               (as-str (:id m))) (map->document m))
                          (.addDocument writer (map->document m))))
                      { index true })))

(defn delete-from-query
  [index input]
  (use-writer index (fn [^IndexWriter writer]
                      ;; method is deleteDocuments(Query...)
                      (let [query (parse-query input (extract-analyzer nil))]
                        (.deleteDocuments writer ^"[Lorg.apache.lucene.search.Query;" (into-array Query [query]))
                        { index (.toString query) }))))

(defn delete-all [index]
  (use-writer index (fn [^IndexWriter writer] (.deleteAll writer))))


(defn fragment->map [^TextFragment fragment]
  {:text (.toString fragment)
   :score (.getScore fragment)
   :frag-num (wall-hack-field (class fragment) :fragNum fragment)
   :text-start-pos (wall-hack-field (class fragment) :textStartPos fragment)
   :text-end-pos (wall-hack-field (class fragment) :textEndPos fragment)})

{'fraction_non_alphanumeric': 0.13354247922672546, 'fraction_numerical': 0.0005087332541970493, 'mean_word_length': 2.272752497225305, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3580034', 'n_tokens_mistral': 3707, 'n_tokens_neox': 3505, 'n_words': 949}