tantiny-in-memory 1.0.0

data/lib/tantiny/index.rb

@@ -0,0 +1,165 @@
+# frozen_string_literal: true
+
+module Tantiny
+  class Index
+    LOCKFILE = ".tantiny.lock"
+    DEFAULT_WRITER_MEMORY = 5_000_000 # 5MB
+    DEFAULT_LIMIT = 10
+
+    def self.new(path, **options, &block)
+      default_tokenizer = options[:tokenizer] || Tokenizer.default
+      schema = Schema.new(default_tokenizer, &block)
+
+      object = __new(
+        path.to_s,
+        schema.default_tokenizer,
+        schema.field_tokenizers.transform_keys(&:to_s),
+        schema.text_fields.map(&:to_s),
+        schema.string_fields.map(&:to_s),
+        schema.integer_fields.map(&:to_s),
+        schema.double_fields.map(&:to_s),
+        schema.date_fields.map(&:to_s),
+        schema.facet_fields.map(&:to_s)
+      )
+
+      object.send(:initialize, path, schema, **options)
+
+      object
+    end
+
+    def initialize(path, schema, **options)
+      @path = path
+      @schema = schema
+
+      @indexer_memory = options[:writer_memory] || DEFAULT_WRITER_MEMORY
+      @exclusive_writer = options[:exclusive_writer] || false
+
+      @active_transaction = Concurrent::ThreadLocalVar.new(false)
+      @transaction_semaphore = Mutex.new
+
+      acquire_index_writer if exclusive_writer?
+    end
+
+    attr_reader :schema
+
+    def transaction
+      if inside_transaction?
+        yield
+      else
+        synchronize do
+          open_transaction!
+
+          yield
+
+          close_transaction!
+        end
+      end
+
+      nil
+    end
+
+    def reload
+      __reload
+    end
+
+    def <<(document)
+      transaction do
+        __add_document(
+          resolve(document, schema.id_field).to_s,
+          slice_document(document, schema.text_fields) { |v| v.to_s },
+          slice_document(document, schema.string_fields) { |v| v.to_s },
+          slice_document(document, schema.integer_fields) { |v| v.to_i },
+          slice_document(document, schema.double_fields) { |v| v.to_f },
+          slice_document(document, schema.date_fields) { |v| Helpers.timestamp(v) },
+          slice_document(document, schema.facet_fields) { |v| v.to_s }
+        )
+      end
+    end
+
+    def delete(id)
+      transaction do
+        __delete_document(id.to_s)
+      end
+    end
+
+    def search(query, limit: DEFAULT_LIMIT, **smart_query_options)
+      unless query.is_a?(Query)
+        fields = schema.text_fields
+        query = Query.smart_query(self, fields, query.to_s, **smart_query_options)
+      end
+
+      __search(query, limit)
+    end
+
+    # Shortcuts for creating queries:
+    Query::TYPES.each do |query_type|
+      method_name = "#{query_type}_query"
+      define_method(method_name) do |*args, **kwargs|
+        Query.send(method_name, self, *args, **kwargs)
+      end
+    end
+
+    private
+
+    def slice_document(document, fields, &block)
+      fields.inject({}) do |hash, field|
+        hash.tap { |h| h[field.to_s] = resolve(document, field) }
+      end.compact.transform_values(&block)
+    end
+
+    def resolve(document, field)
+      document.is_a?(Hash) ? document[field] : document.send(field)
+    end
+
+    def acquire_index_writer
+      __acquire_index_writer(@indexer_memory)
+    rescue TantivyError => e
+      case e.message
+      when /Failed to acquire Lockfile/
+        raise IndexWriterBusyError.new
+      else
+        raise
+      end
+    end
+
+    def release_index_writer
+      __release_index_writer
+    end
+
+    def commit
+      __commit
+    end
+
+    def open_transaction!
+      acquire_index_writer unless exclusive_writer?
+
+      @active_transaction.value = true
+    end
+
+    def close_transaction!
+      commit
+
+      release_index_writer unless exclusive_writer?
+
+      @active_transaction.value = false
+    end
+
+    def inside_transaction?
+      @active_transaction.value
+    end
+
+    def exclusive_writer?
+      @exclusive_writer
+    end
+
+    def synchronize(&block)
+      @transaction_semaphore.synchronize do
+        Helpers.with_lock(lockfile_path, &block)
+      end
+    end
+
+    def lockfile_path
+      @lockfile_path ||= File.join(@path, LOCKFILE)
+    end
+  end
+end

data/lib/tantiny/query.rb

@@ -0,0 +1,165 @@
+# frozen_string_literal: true
+
+require "date"
+
+module Tantiny
+  class Query
+    TYPES = %i[
+      all empty term fuzzy_term
+      phrase regex range facet
+      smart prefix
+    ].freeze
+
+    DEFAULT_BOOST = 1.0
+    DEFAULT_FUZZY_DISTANCE = 1
+
+    class << self
+      def conjunction(*queries)
+        # @type var queries: Array[untyped]
+        queries.one? ? queries.first : __conjunction(queries)
+      end
+
+      def disjunction(*queries)
+        # @type var queries: Array[untyped]
+        queries.one? ? queries.first : __disjunction(queries)
+      end
+
+      def all_query(_index = nil)
+        __new_all_query
+      end
+
+      def empty_query(_index = nil)
+        __new_empty_query
+      end
+
+      def term_query(index, fields, term, **options)
+        allowed_fields = text_and_strings(index)
+        construct_query(index, :term, allowed_fields, fields, [term.to_s], **options)
+      end
+
+      def fuzzy_term_query(index, fields, term, distance = DEFAULT_FUZZY_DISTANCE, **options)
+        params = [term.to_s, distance.to_i]
+        allowed_fields = text_and_strings(index)
+        construct_query(index, :fuzzy_term, allowed_fields, fields, params, **options)
+      end
+
+      def phrase_query(index, fields, phrase, **options)
+        queries = [*fields].map do |f|
+          terms = index.schema.tokenizer_for(f).terms(phrase)
+          allowed_fields = index.schema.text_fields
+          construct_query(index, :phrase, allowed_fields, f, [terms], **options)
+        end
+
+        queries.empty? ? empty_query : disjunction(*queries)
+      end
+
+      def regex_query(index, fields, regex, **options)
+        allowed_fields = text_and_strings(index)
+        construct_query(index, :regex, allowed_fields, fields, [regex.to_s], **options)
+      end
+
+      def prefix_query(index, fields, prefix, **options)
+        regex_query(index, fields, Regexp.escape(prefix) + ".*", **options)
+      end
+
+      def range_query(index, fields, range, **options)
+        schema = index.schema
+
+        case range.first
+        when Integer
+          allowed_fields = schema.integer_fields
+          from, to = [range.min, range.max]
+        when Float
+          allowed_fields = schema.double_fields
+          from, to = [range.first, range.last]
+        when Date, DateTime
+          # @type var range: Range[Date | DateTime]
+          allowed_fields = schema.date_fields
+          from, to = [Helpers.timestamp(range.first), Helpers.timestamp(range.last)]
+        else
+          raise UnsupportedRange.new(range.first.class)
+        end
+
+        # @type var allowed_fields: Array[Symbol]
+        construct_query(index, :range, allowed_fields, fields, [from, to], **options)
+      end
+
+      def facet_query(index, field, path, **options)
+        allowed_fields = index.schema.facet_fields
+        construct_query(index, :facet, allowed_fields, field, [path], **options)
+      end
+
+      def smart_query(index, fields, query_string, **options)
+        fuzzy_distance = options[:fuzzy_distance]
+        boost_factor = options.fetch(:boost, DEFAULT_BOOST)
+
+        field_queries = [*fields].map do |field|
+          terms = index.schema.tokenizer_for(field).terms(query_string)
+
+          # See: https://github.com/soutaro/steep/issues/272
+          # @type block: nil | Query
+          next if terms.empty?
+
+          term_queries = terms.map do |term|
+            if fuzzy_distance.nil?
+              term_query(index, field, term)
+            else
+              fuzzy_term_query(index, field, term, fuzzy_distance)
+            end
+          end
+
+          # @type var terms: untyped
+          # @type var term_queries: untyped
+          last_term_query = prefix_query(index, field, terms.last) | term_queries.last
+
+          conjunction(last_term_query, *term_queries[0...-1])
+        end.compact
+
+        disjunction(*field_queries).boost(boost_factor)
+      end
+
+      private
+
+      # Can't use variadic argument `params` here due to:
+      # https://github.com/soutaro/steep/issues/480
+      def construct_query(index, query_type, allowed_fields, fields, params, **options)
+        queries = [*fields].map do |field|
+          supported = allowed_fields.include?(field)
+          raise UnsupportedField.new(field) unless supported
+
+          send("__new_#{query_type}_query", index, field.to_s, *params)
+        end
+
+        return empty_query if fields.empty?
+
+        disjunction(*queries).boost(options.fetch(:boost, DEFAULT_BOOST))
+      end
+
+      def text_and_strings(index)
+        index.schema.text_fields | index.schema.string_fields
+      end
+    end
+
+    def |(other)
+      raise ArgumentError.new("Not a #{self.class}.") unless other.is_a?(self.class)
+
+      self.class.disjunction(self, other)
+    end
+
+    def &(other)
+      raise ArgumentError.new("Not a #{self.class}.") unless other.is_a?(self.class)
+
+      self.class.conjunction(self, other)
+    end
+
+    def !
+      __negation
+    end
+
+    def boost(boost_factor)
+      return self if boost_factor == DEFAULT_BOOST
+
+      __boost(boost_factor.to_f)
+    end
+  end
+end

data/lib/tantiny/schema.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Tantiny
+  class Schema
+    attr_reader :default_tokenizer,
+      :id_field,
+      :text_fields,
+      :string_fields,
+      :integer_fields,
+      :double_fields,
+      :date_fields,
+      :facet_fields,
+      :field_tokenizers
+
+    def initialize(tokenizer, &block)
+      @default_tokenizer = tokenizer
+      @id_field = :id
+      @text_fields = []
+      @string_fields = []
+      @integer_fields = []
+      @double_fields = []
+      @date_fields = []
+      @facet_fields = []
+      @field_tokenizers = {}
+
+      instance_exec(&block)
+    end
+
+    def tokenizer_for(field)
+      field_tokenizers[field] || default_tokenizer
+    end
+
+    private
+
+    def id(key) = @id_field = key
+
+    def string(key) = @string_fields << key
+
+    def integer(key) = @integer_fields << key
+
+    def double(key) = @double_fields << key
+
+    def date(key) = @date_fields << key
+
+    def facet(key) = @facet_fields << key
+
+    def text(key, tokenizer: nil)
+      @field_tokenizers[key] = tokenizer if tokenizer
+
+      @text_fields << key
+    end
+  end
+end

data/lib/tantiny/tokenizer.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Tantiny
+  class Tokenizer
+    def self.default
+      new(:simple)
+    end
+
+    def self.new(kind, **options)
+      case kind
+      when :simple
+        __new_simple_tokenizer
+      when :stemmer
+        language = options[:language] || :en
+        __new_stemmer_tokenizer(language.to_s)
+      when :ngram
+        prefix_only = options.fetch(:prefix_only, false)
+        __new_ngram_tokenizer(options[:min], options[:max], prefix_only)
+      else
+        raise UnknownTokenizer.new(kind)
+      end
+    end
+
+    def terms(string)
+      __extract_terms(string)
+    end
+  end
+end

data/lib/tantiny/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Tantiny
+  VERSION = "1.0.0" # {x-release-please-version}
+end

data/lib/tantiny.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require "ruby-next/language/setup"
+RubyNext::Language.setup_gem_load_path
+
+require "rutie"
+require "thermite/fiddle"
+require "concurrent"
+require "fileutils"
+
+require "tantiny/version"
+require "tantiny/errors"
+require "tantiny/helpers"
+require "tantiny/schema"
+require "tantiny/tokenizer"
+require "tantiny/query"
+require "tantiny/index"
+
+module Tantiny
+  project_dir = File.expand_path("../..", __FILE__)
+
+  Thermite::Fiddle.load_module(
+    "Init_tantiny",
+    cargo_project_path: project_dir,
+    ruby_project_path: project_dir
+  )
+end

data/sig/tantiny/errors.rbs

@@ -0,0 +1,20 @@
+module Tantiny
+  class TantivyError < StandardError
+  end
+
+  class UnknownField < StandardError
+    def initialize: () -> void
+  end
+
+  class UnknownTokenizer < StandardError
+    def initialize: (Symbol tokenizer_type) -> void
+  end
+
+  class UnsupportedRange < StandardError
+    def initialize: (Class range_type) -> void
+  end
+
+  class UnsupportedField < StandardError
+    def initialize: (Symbol field) -> void
+  end
+end

data/sig/tantiny/helpers.rbs

@@ -0,0 +1,8 @@
+
+module Tantiny
+  module Helpers
+    def self.timestamp: ((Date | DateTime) date) -> String
+
+    def self.with_lock: (String lockfile) { (*untyped) -> void } -> void
+  end
+end

data/sig/tantiny/index.rbs

@@ -0,0 +1,103 @@
+module Tantiny
+  class Index
+    LOCKFILE: String
+    DEFAULT_WRITER_MEMORY: Integer
+    DEFAULT_LIMIT: Integer
+
+    def self.new: (
+      String path,
+      **untyped options
+    ) { (*untyped) -> void } -> Index
+
+    def self.__new: (
+      String path,
+      Tokenizer default_tokenizer,
+      Hash[String, Tokenizer] field_tokenizers,
+      Array[String] text_fields,
+      Array[String] string_fields,
+      Array[String] integer_fields,
+      Array[String] double_fields,
+      Array[String] date_fields,
+      Array[String] facet_fields
+    ) -> Index
+
+    def initialize: (
+      String path,
+      Schema schema,
+      **untyped options
+    ) -> void
+
+    attr_reader schema: Schema
+
+    def transaction: () { (*untyped) -> void } -> void
+
+    def reload: () -> void
+    def <<: (untyped document) -> void
+    def delete: (String id) -> void
+
+    def search: (
+      (Query | String) query,
+      ?limit: Integer,
+      **untyped smart_query_options
+    ) -> Array[String]
+
+    def all_query: () -> Query
+    def empty_query: () -> Query
+    def term_query: (fields fields, String term, **untyped options) -> Query
+    def fuzzy_term_query: (fields fields, String term, ?Integer distance, **untyped options) -> Query
+    def phrase_query: (fields fields, String phrase, **untyped options) -> Query
+    def regex_query: (fields fields, String regex, **untyped options) -> Query
+    def prefix_query: (fields fields, String prefix, **untyped options) -> Query
+    def facet_query: (Symbol field, String path, **untyped options) -> Query
+    def range_query: (fields fields, Range[numeric | date] range, **untyped options) -> Query
+    def smart_query: (fields fields, String query_string, **untyped options) -> Query
+
+    def __commit: () -> void
+    def __reload: () -> void
+
+    def __add_document: (
+      String id,
+      Hash[String, String] text_fields,
+      Hash[String, String] string_fields,
+      Hash[String, Integer] integer_fields,
+      Hash[String, Float] double_fields,
+      Hash[String, String] date_fields,
+      Hash[String, String] facet_fields
+    ) -> void
+
+    def __delete_document: (String id) -> void
+
+    def __search: (Query query, Integer limit) -> Array[String]
+
+    def __acquire_index_writer: (Integer overall_memory) -> void
+    def __release_index_writer: () -> void
+
+    private
+
+    def commit: () -> void
+
+    def slice_document: (
+      untyped document,
+      Array[Symbol] fields
+    ) { (untyped v) -> untyped } -> Hash[String, untyped]
+
+    def default_search: (
+      String query_string,
+      Integer limit,
+      ?fuzzy_distance: Integer
+    ) -> Array[String]
+
+    def resolve: (untyped document, Symbol field) -> untyped
+
+    def synchronize: () { (*untyped) -> void } -> void
+    def lockfile_path: () -> String
+
+    def exclusive_writer?: () -> bool
+    def acquire_index_writer: () -> void
+    def release_index_writer: () -> void
+
+    def open_transaction!: () -> void
+    def close_transaction!: () -> void
+    def inside_transaction?: () -> bool
+  end
+end