Query container

CMakeLists.txt

@@ -105,6 +105,7 @@ target_link_libraries(pisa PUBLIC # TODO(michal): are there any of these we can
     spdlog
     fmt::fmt
     range-v3
+    nlohmann_json::nlohmann_json
 )
 target_include_directories(pisa PUBLIC external)
 

include/pisa/query.hpp

@@ -0,0 +1,149 @@
+#pragma once
+
+#include <functional>
+#include <istream>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include <gsl/span>
+
+namespace pisa {
+
+struct QueryContainerInner;
+
+struct ResolvedTerm {
+    std::uint32_t id;
+    std::string term;
+};
+
+using TermProcessorFn = std::function<std::optional<std::string>(std::string)>;
+using ParseFn = std::function<std::vector<ResolvedTerm>(std::string const&)>;
+
+class QueryContainer;
+
+/// QueryRequest is a special container that maintains important invariants, such as sorted term
+/// IDs, and also has some additional data, like term weights, etc.
+class QueryRequest {
+  public:
+    explicit QueryRequest(QueryContainer const& data, std::size_t k);
+
+    [[nodiscard]] auto term_ids() const -> gsl::span<std::uint32_t const>;
+    [[nodiscard]] auto threshold() const -> std::optional<float>;
+    [[nodiscard]] auto k() const -> std::optional<float>;
+
+  private:
+    std::size_t m_k;
+    std::optional<float> m_threshold{};
+    std::vector<std::uint32_t> m_term_ids{};
+};
+
+class QueryContainer {
+  public:
+    QueryContainer(QueryContainer const&);
+    QueryContainer(QueryContainer&&) noexcept;
+    QueryContainer& operator=(QueryContainer const&);
+    QueryContainer& operator=(QueryContainer&&) noexcept;
+    ~QueryContainer();
+
+    [[nodiscard]] auto operator==(QueryContainer const& other) const noexcept -> bool;
+
+    /// Constructs a query from a raw string.
+    [[nodiscard]] static auto raw(std::string query_string) -> QueryContainer;
+
+    /// Constructs a query from a list of terms.
+    ///
+    /// \param terms            List of terms
+    /// \param term_processor   Function executed for each term before stroring them,
+    ///                         e.g., stemming or filtering. This function returns
+    ///                         `std::optional<std::string>`, and all `std::nullopt` values
+    ///                         will be filtered out.
+    [[nodiscard]] static auto
+    from_terms(std::vector<std::string> terms, std::optional<TermProcessorFn> term_processor)
+        -> QueryContainer;
+
+    /// Constructs a query from a list of term IDs.
+    [[nodiscard]] static auto from_term_ids(std::vector<std::uint32_t> term_ids) -> QueryContainer;
+
+    /// Constructs a query from a JSON object.
+    [[nodiscard]] static auto from_json(std::string_view json_string) -> QueryContainer;
+
+    [[nodiscard]] auto to_json() const -> std::string;
+
+    /// Constructs a query from a colon-separated format:
+    ///
+    /// ```
+    /// id:raw query string
+    /// ```
+    /// or
+    /// ```
+    /// raw query string
+    /// ```
+    [[nodiscard]] static auto from_colon_format(std::string_view line) -> QueryContainer;
+
+    // Accessors
+
+    [[nodiscard]] auto id() const noexcept -> std::optional<std::string> const&;
+    [[nodiscard]] auto string() const noexcept -> std::optional<std::string> const&;
+    [[nodiscard]] auto terms() const noexcept -> std::optional<std::vector<std::string>> const&;
+    [[nodiscard]] auto term_ids() const noexcept -> std::optional<std::vector<std::uint32_t>> const&;
+    [[nodiscard]] auto threshold(std::size_t k) const noexcept -> std::optional<float>;
+    [[nodiscard]] auto thresholds() const noexcept
+        -> std::vector<std::pair<std::size_t, float>> const&;
+
+    /// Sets the raw string.
+    [[nodiscard]] auto string(std::string) -> QueryContainer&;
+
+    /// Parses the raw query with the given parser.
+    ///
+    /// \throws std::domain_error   when raw string is not set
+    auto parse(ParseFn parse_fn) -> QueryContainer&;
+
+    /// Sets the query score threshold for `k`.
+    ///
+    /// If another threshold for the same `k` exists, it will be replaced,
+    /// and `true` will be returned. Otherwise, `false` will be returned.
+    auto add_threshold(std::size_t k, float score) -> bool;
+
+    /// Returns a query ready to be used for retrieval.
+    [[nodiscard]] auto query(std::size_t k) const -> QueryRequest;
+
+  private:
+    QueryContainer();
+    std::unique_ptr<QueryContainerInner> m_data;
+};
+
+enum class Format { Json, Colon };
+
+class QueryReader {
+  public:
+    /// Open reader from file.
+    static auto from_file(std::string const& file) -> QueryReader;
+    /// Open reader from stdin.
+    static auto from_stdin() -> QueryReader;
+
+    /// Read next query or return `nullopt` if stream has ended.
+    [[nodiscard]] auto next() -> std::optional<QueryContainer>;
+
+    /// Execute `fn(q)` for each query `q`.
+    template <typename Fn>
+    void for_each(Fn&& fn)
+    {
+        auto query = next();
+        while (query) {
+            fn(std::move(*query));
+            query = next();
+        }
+    }
+
+  private:
+    explicit QueryReader(std::unique_ptr<std::istream> stream, std::istream& stream_ref);
+
+    std::unique_ptr<std::istream> m_stream;
+    std::istream& m_stream_ref;
+    std::string m_line_buf{};
+    std::optional<Format> m_format{};
+};
+
+}  // namespace pisa

include/pisa/query/query_parser.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <string>
+
+#include "query.hpp"
+#include "term_resolver.hpp"
+
+namespace pisa {
+
+/// Parses a query string to processed terms.
+class QueryParser {
+  public:
+    explicit QueryParser(TermResolver term_processor);
+    /// Given a query string, it returns a list of (possibly processed) terms.
+    ///
+    /// Possible transformations of terms include lower-casing and stemming.
+    /// Some terms could be also removed, e.g., because they are on a list of
+    /// stop words. The exact implementation depends on the term processor
+    /// passed to the constructor.
+    auto operator()(std::string const&) -> std::vector<ResolvedTerm>;
+
+  private:
+    TermResolver m_term_resolver;
+};
+
+}  // namespace pisa

include/pisa/query/term_resolver.hpp

@@ -0,0 +1,53 @@
+#pragma once
+
+#include <functional>
+#include <optional>
+#include <string>
+
+#include "query.hpp"
+
+namespace pisa {
+
+/// Thrown if expected resolver but none found.
+struct MissingResolverError {
+};
+
+using TermResolver = std::function<std::optional<ResolvedTerm>(std::string)>;
+
+struct StandardTermResolverParams;
+
+/// Provides a standard implementation of `TermResolver`.
+class StandardTermResolver {
+  public:
+    StandardTermResolver(
+        std::string const& term_lexicon_path,
+        std::optional<std::string> const& stopwords_filename,
+        std::optional<std::string> const& stemmer_type);
+    StandardTermResolver(StandardTermResolver const&);
+    StandardTermResolver(StandardTermResolver&&) noexcept;
+    StandardTermResolver& operator=(StandardTermResolver const&);
+    StandardTermResolver& operator=(StandardTermResolver&&) noexcept;
+    ~StandardTermResolver();
+
+    [[nodiscard]] auto operator()(std::string token) const -> std::optional<ResolvedTerm>;
+
+  private:
+    [[nodiscard]] auto is_stopword(std::uint32_t const term) const -> bool;
+
+    std::unique_ptr<StandardTermResolverParams> m_self;
+};
+
+/// Reads queries from `query_file`, resolves them with `term_resolver`, filters by
+/// query length (number of resolved terms in the query), and prints the selected
+/// queries to `out`.
+///
+/// \throws MissingResolverError  When no resolver passed but queries don't have IDs resolved.
+//
+void filter_queries(
+    std::optional<std::string> const& query_file,
+    std::optional<TermResolver> term_resolver,
+    std::size_t min_query_len,
+    std::size_t max_query_len,
+    std::ostream& out);
+
+}  // namespace pisa