vectra-client 0.2.2 → 0.3.1

data/lib/vectra/batch.rb

@@ -0,0 +1,148 @@
+# frozen_string_literal: true
+
+require "concurrent"
+
+module Vectra
+  # Batch operations with concurrent processing
+  #
+  # Provides async batch upsert capabilities with configurable concurrency
+  # and automatic chunking of large vector sets.
+  #
+  # @example Async batch upsert
+  #   batch = Vectra::Batch.new(client, concurrency: 4)
+  #   result = batch.upsert_async(
+  #     index: 'my-index',
+  #     vectors: large_vector_array,
+  #     chunk_size: 100
+  #   )
+  #   puts "Upserted: #{result[:upserted_count]}"
+  #
+  class Batch
+    DEFAULT_CONCURRENCY = 4
+    DEFAULT_CHUNK_SIZE = 100
+
+    attr_reader :client, :concurrency
+
+    # Initialize a new Batch processor
+    #
+    # @param client [Client] the Vectra client
+    # @param concurrency [Integer] max concurrent requests (default: 4)
+    def initialize(client, concurrency: DEFAULT_CONCURRENCY)
+      @client = client
+      @concurrency = [concurrency, 1].max
+    end
+
+    # Perform async batch upsert with concurrent requests
+    #
+    # @param index [String] the index name
+    # @param vectors [Array<Hash>] vectors to upsert
+    # @param namespace [String, nil] optional namespace
+    # @param chunk_size [Integer] vectors per chunk (default: 100)
+    # @return [Hash] aggregated result with :upserted_count, :chunks, :errors
+    def upsert_async(index:, vectors:, namespace: nil, chunk_size: DEFAULT_CHUNK_SIZE)
+      chunks = vectors.each_slice(chunk_size).to_a
+      return { upserted_count: 0, chunks: 0, errors: [] } if chunks.empty?
+
+      results = process_chunks_concurrently(chunks) do |chunk|
+        client.upsert(index: index, vectors: chunk, namespace: namespace)
+      end
+
+      aggregate_results(results, vectors.size)
+    end
+
+    # Perform async batch delete with concurrent requests
+    #
+    # @param index [String] the index name
+    # @param ids [Array<String>] IDs to delete
+    # @param namespace [String, nil] optional namespace
+    # @param chunk_size [Integer] IDs per chunk (default: 100)
+    # @return [Hash] aggregated result
+    def delete_async(index:, ids:, namespace: nil, chunk_size: DEFAULT_CHUNK_SIZE)
+      chunks = ids.each_slice(chunk_size).to_a
+      return { deleted_count: 0, chunks: 0, errors: [] } if chunks.empty?
+
+      results = process_chunks_concurrently(chunks) do |chunk|
+        client.delete(index: index, ids: chunk, namespace: namespace)
+      end
+
+      aggregate_delete_results(results, ids.size)
+    end
+
+    # Perform async batch fetch with concurrent requests
+    #
+    # @param index [String] the index name
+    # @param ids [Array<String>] IDs to fetch
+    # @param namespace [String, nil] optional namespace
+    # @param chunk_size [Integer] IDs per chunk (default: 100)
+    # @return [Hash<String, Vector>] merged results
+    def fetch_async(index:, ids:, namespace: nil, chunk_size: DEFAULT_CHUNK_SIZE)
+      chunks = ids.each_slice(chunk_size).to_a
+      return {} if chunks.empty?
+
+      results = process_chunks_concurrently(chunks) do |chunk|
+        client.fetch(index: index, ids: chunk, namespace: namespace)
+      end
+
+      merge_fetch_results(results)
+    end
+
+    private
+
+    def process_chunks_concurrently(chunks)
+      pool = Concurrent::FixedThreadPool.new(concurrency)
+      futures = []
+
+      chunks.each_with_index do |chunk, index|
+        futures << Concurrent::Future.execute(executor: pool) do
+          { index: index, result: yield(chunk), error: nil }
+        rescue StandardError => e
+          { index: index, result: nil, error: e }
+        end
+      end
+
+      # Wait for all futures and collect results
+      results = futures.map(&:value)
+      pool.shutdown
+      pool.wait_for_termination(30)
+
+      results.sort_by { |r| r[:index] }
+    end
+
+    def aggregate_results(results, total_vectors)
+      errors = results.select { |r| r[:error] }.map { |r| r[:error] }
+      successful = results.reject { |r| r[:error] }
+      upserted = successful.sum { |r| r.dig(:result, :upserted_count) || 0 }
+
+      {
+        upserted_count: upserted,
+        total_vectors: total_vectors,
+        chunks: results.size,
+        successful_chunks: successful.size,
+        errors: errors
+      }
+    end
+
+    def aggregate_delete_results(results, total_ids)
+      errors = results.select { |r| r[:error] }.map { |r| r[:error] }
+      successful = results.reject { |r| r[:error] }
+
+      {
+        deleted_count: total_ids - (errors.size * (total_ids / results.size.to_f).ceil),
+        total_ids: total_ids,
+        chunks: results.size,
+        successful_chunks: successful.size,
+        errors: errors
+      }
+    end
+
+    def merge_fetch_results(results)
+      merged = {}
+      results.each do |r|
+        next if r[:error] || r[:result].nil?
+
+        merged.merge!(r[:result])
+      end
+      merged
+    end
+  end
+end

data/lib/vectra/cache.rb

@@ -0,0 +1,261 @@
+# frozen_string_literal: true
+
+require "digest"
+
+module Vectra
+  # Optional caching layer for frequently queried vectors
+  #
+  # Provides in-memory caching with TTL support for query results
+  # and fetched vectors to reduce database load.
+  #
+  # @example Basic usage
+  #   cache = Vectra::Cache.new(ttl: 300, max_size: 1000)
+  #   cached_client = Vectra::CachedClient.new(client, cache: cache)
+  #
+  #   # First call hits the database
+  #   result1 = cached_client.query(index: 'idx', vector: vec, top_k: 10)
+  #
+  #   # Second call returns cached result
+  #   result2 = cached_client.query(index: 'idx', vector: vec, top_k: 10)
+  #
+  class Cache
+    DEFAULT_TTL = 300 # 5 minutes
+    DEFAULT_MAX_SIZE = 1000
+
+    attr_reader :ttl, :max_size
+
+    # Initialize cache
+    #
+    # @param ttl [Integer] time-to-live in seconds (default: 300)
+    # @param max_size [Integer] maximum cache entries (default: 1000)
+    def initialize(ttl: DEFAULT_TTL, max_size: DEFAULT_MAX_SIZE)
+      @ttl = ttl
+      @max_size = max_size
+      @store = {}
+      @timestamps = {}
+      @mutex = Mutex.new
+    end
+
+    # Get value from cache
+    #
+    # @param key [String] cache key
+    # @return [Object, nil] cached value or nil if not found/expired
+    def get(key)
+      @mutex.synchronize do
+        return nil unless @store.key?(key)
+
+        if expired?(key)
+          delete_entry(key)
+          return nil
+        end
+
+        @store[key]
+      end
+    end
+
+    # Set value in cache
+    #
+    # @param key [String] cache key
+    # @param value [Object] value to cache
+    # @return [Object] the cached value
+    def set(key, value)
+      @mutex.synchronize do
+        evict_if_needed
+        @store[key] = value
+        @timestamps[key] = Time.now
+        value
+      end
+    end
+
+    # Get or set value with block
+    #
+    # @param key [String] cache key
+    # @yield block to compute value if not cached
+    # @return [Object] cached or computed value
+    def fetch(key)
+      cached = get(key)
+      return cached unless cached.nil?
+
+      value = yield
+      set(key, value)
+      value
+    end
+
+    # Delete entry from cache
+    #
+    # @param key [String] cache key
+    # @return [Object, nil] deleted value
+    def delete(key)
+      @mutex.synchronize { delete_entry(key) }
+    end
+
+    # Clear all cache entries
+    #
+    # @return [void]
+    def clear
+      @mutex.synchronize do
+        @store.clear
+        @timestamps.clear
+      end
+    end
+
+    # Get cache statistics
+    #
+    # @return [Hash] cache stats
+    def stats
+      @mutex.synchronize do
+        {
+          size: @store.size,
+          max_size: max_size,
+          ttl: ttl,
+          keys: @store.keys
+        }
+      end
+    end
+
+    # Check if key exists and is not expired
+    #
+    # @param key [String] cache key
+    # @return [Boolean]
+    def exist?(key)
+      @mutex.synchronize do
+        return false unless @store.key?(key)
+        return false if expired?(key)
+
+        true
+      end
+    end
+
+    private
+
+    def expired?(key)
+      return true unless @timestamps.key?(key)
+
+      Time.now - @timestamps[key] > ttl
+    end
+
+    def delete_entry(key)
+      @timestamps.delete(key)
+      @store.delete(key)
+    end
+
+    def evict_if_needed
+      # Evict when at or above max_size to make room for new entry
+      return if @store.size < max_size
+
+      # Remove oldest entries (at least 10% of max_size to avoid frequent evictions)
+      entries_to_remove = [(max_size * 0.2).ceil, 1].max
+      oldest_keys = @timestamps.sort_by { |_, v| v }.first(entries_to_remove).map(&:first)
+      oldest_keys.each { |key| delete_entry(key) }
+    end
+  end
+
+  # Client wrapper with caching support
+  #
+  # Wraps a Vectra::Client to add transparent caching for query and fetch operations.
+  #
+  class CachedClient
+    attr_reader :client, :cache
+
+    # Initialize cached client
+    #
+    # @param client [Client] the underlying Vectra client
+    # @param cache [Cache] cache instance (creates default if nil)
+    # @param cache_queries [Boolean] whether to cache query results (default: true)
+    # @param cache_fetches [Boolean] whether to cache fetch results (default: true)
+    def initialize(client, cache: nil, cache_queries: true, cache_fetches: true)
+      @client = client
+      @cache = cache || Cache.new
+      @cache_queries = cache_queries
+      @cache_fetches = cache_fetches
+    end
+
+    # Query with caching
+    #
+    # @see Client#query
+    def query(index:, vector:, top_k: 10, namespace: nil, filter: nil, **options)
+      unless @cache_queries
+        return client.query(index: index, vector: vector, top_k: top_k,
+                            namespace: namespace, filter: filter, **options)
+      end
+
+      key = query_cache_key(index, vector, top_k, namespace, filter)
+      cache.fetch(key) do
+        client.query(index: index, vector: vector, top_k: top_k,
+                     namespace: namespace, filter: filter, **options)
+      end
+    end
+
+    # Fetch with caching
+    #
+    # @see Client#fetch
+    def fetch(index:, ids:, namespace: nil)
+      return client.fetch(index: index, ids: ids, namespace: namespace) unless @cache_fetches
+
+      # Check cache for each ID
+      results = {}
+      uncached_ids = []
+
+      ids.each do |id|
+        key = fetch_cache_key(index, id, namespace)
+        cached = cache.get(key)
+        if cached
+          results[id] = cached
+        else
+          uncached_ids << id
+        end
+      end
+
+      # Fetch uncached IDs
+      if uncached_ids.any?
+        fetched = client.fetch(index: index, ids: uncached_ids, namespace: namespace)
+        fetched.each do |id, vector|
+          key = fetch_cache_key(index, id, namespace)
+          cache.set(key, vector)
+          results[id] = vector
+        end
+      end
+
+      results
+    end
+
+    # Pass through other methods to underlying client
+    def method_missing(method, *, **, &)
+      if client.respond_to?(method)
+        client.public_send(method, *, **, &)
+      else
+        super
+      end
+    end
+
+    def respond_to_missing?(method, include_private = false)
+      client.respond_to?(method, include_private) || super
+    end
+
+    # Invalidate cache entries for an index
+    #
+    # @param index [String] index name
+    def invalidate_index(index)
+      cache.stats[:keys].each do |key|
+        cache.delete(key) if key.start_with?("#{index}:")
+      end
+    end
+
+    # Clear entire cache
+    def clear_cache
+      cache.clear
+    end
+
+    private
+
+    def query_cache_key(index, vector, top_k, namespace, filter)
+      vector_hash = Digest::MD5.hexdigest(vector.to_s)[0, 16]
+      filter_hash = filter ? Digest::MD5.hexdigest(filter.to_s)[0, 8] : "nofilter"
+      "#{index}:q:#{vector_hash}:#{top_k}:#{namespace || 'default'}:#{filter_hash}"
+    end
+
+    def fetch_cache_key(index, id, namespace)
+      "#{index}:f:#{id}:#{namespace || 'default'}"
+    end
+  end
+end