search-engine-for-typesense 30.1.6.6 → 30.1.6.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 685951054b6b65954ff6517a08176c8bec124c197980dec5fcfc0ebbd19f6691
-  data.tar.gz: 12e290d5e45325cb031b06f87d158b33300fdc8682d3a417ba7191940e35211e
+  metadata.gz: 17890bf666142ee9e0cd293c1ea5fc13dcf4089c50bb1e721708b7c9972f5e15
+  data.tar.gz: 84866d6a6d238f29e3eda054f445a88780c98ab9eff75673f6930c044a336a02
 SHA512:
-  metadata.gz: f7e7e8ab135067d43fd2cffed2dfc0fdeac7f3861c267dc47d5131c2600a7e5adee5b9267d5b326517cfe32c526a0e67d13d614312183c59ff0e496def671990
-  data.tar.gz: dcb0366fa3ddd2dbe03bc50dec6032093761eba588ae2e88116e2b939cbc8f62ea442c0857448b3d69d20cdf7eccc98b0f5d9d1ee2b2d6ce593453539d889b97
+  metadata.gz: 01aa100a2d5a612c56a85cd6f4e02d232476911ceea18cd7f024904a477874ab492d675d0ad9624212a4891e15e0a17371cc2c9314b38bb2bcb332e36212faba
+  data.tar.gz: fb6ba4067d04afdb1cc0bc2236c0040832fc6632d8cd0320080c77b7b9b6fa8d9b8ce6dca99ab725757ca4124107f67570ad2281e03404193c746258becef127

data/lib/search_engine/base/index_maintenance.rb

@@ -270,9 +270,9 @@ module SearchEngine
           if compiled
             parts = Array(compiled.partitions)
             max_p = compiled.max_parallel.to_i
-            return __se_index_partitions_seq!(parts, into) if max_p <= 1 || parts.size <= 1
+            return __se_index_partitions_seq!(parts, into, compiled) if max_p <= 1 || parts.size <= 1
 
-            __se_index_partitions_parallel!(parts, into, max_p)
+            __se_index_partitions_parallel!(parts, into, max_p, compiled)
           else
             summary = SearchEngine::Indexer.rebuild_partition!(self, partition: nil, into: into)
             __se_build_index_result([summary])
@@ -314,10 +314,10 @@ module SearchEngine
 
       class_methods do
         # Sequential processing of partition list with live progress rendering.
-        def __se_index_partitions_seq!(parts, into)
-          estimate = __se_per_partition_estimate(parts.size)
+        def __se_index_partitions_seq!(parts, into, compiled)
+          docs_estimates = __se_per_partition_docs_estimates(parts, compiled)
           renderer = SearchEngine::Logging::LiveRenderer.new(
-            labels: parts.map(&:inspect), partitions: parts, per_partition_estimate: estimate
+            labels: parts.map(&:inspect), partitions: parts, per_partition_docs_estimates: docs_estimates
           )
           renderer.start
 
@@ -351,12 +351,12 @@ module SearchEngine
 
       class_methods do
         # Parallel processing via bounded thread pool with live progress rendering.
-        def __se_index_partitions_parallel!(parts, into, max_p)
+        def __se_index_partitions_parallel!(parts, into, max_p, compiled)
           require 'concurrent-ruby'
 
-          estimate = __se_per_partition_estimate(parts.size)
+          docs_estimates = __se_per_partition_docs_estimates(parts, compiled)
           renderer = SearchEngine::Logging::LiveRenderer.new(
-            labels: parts.map(&:inspect), partitions: parts, per_partition_estimate: estimate
+            labels: parts.map(&:inspect), partitions: parts, per_partition_docs_estimates: docs_estimates
           )
           renderer.start
 
@@ -414,11 +414,27 @@ module SearchEngine
       end
 
       class_methods do
-        # Heuristic per-partition batch estimate for progress bars.
+        # Build an array of per-partition doc estimates.
+        # Tries per-partition counting (auto-detected from AR fetch result),
+        # fills any nils with the equal-split heuristic fallback.
+        #
+        # @param parts [Array] partition keys
+        # @param compiled [SearchEngine::Partitioner::Compiled, nil]
+        # @return [Array<Integer, nil>]
+        def __se_per_partition_docs_estimates(parts, compiled)
+          estimates = parts.map { |part| compiled&.partition_doc_count(part) }
+
+          heuristic = __se_heuristic_docs_estimate(parts.size)
+          parts.each_index.map { |i| estimates[i] || heuristic }
+        rescue StandardError
+          Array.new(parts.size, nil)
+        end
+
+        # Equal-split heuristic: total_docs / partition_count.
         # @param partition_count [Integer]
         # @return [Integer, nil]
-        def __se_per_partition_estimate(partition_count)
-          total = SearchEngine::Indexer::BulkImport.estimate_total_batches(self)
+        def __se_heuristic_docs_estimate(partition_count)
+          total = SearchEngine::Indexer::BulkImport.estimate_total_docs(self)
           return nil unless total
 
           (total.to_f / partition_count).ceil
@@ -426,7 +442,7 @@ module SearchEngine
           nil
         end
 
-        private :__se_per_partition_estimate
+        private :__se_per_partition_docs_estimates, :__se_heuristic_docs_estimate
       end
 
       class_methods do

data/lib/search_engine/indexer/bulk_import.rb

@@ -72,7 +72,7 @@ module SearchEngine
           failed_total = 0
           failed_batches_total = 0
           batches_total = 0
-          # Capture start time before processing any batches to measure total wall-clock duration
+          source_batches_done = 0
           started_at = monotonic_ms
 
           docs_enum.each do |raw_batch|
@@ -95,11 +95,13 @@ module SearchEngine
               batches << stats
               validate_soft_batch_size!(batch_size, stats[:docs_count])
               log_batch(stats, batches_total) if log_batches
-              on_batch&.call(
-                batches_done: batches_total, docs_total: docs_total,
-                success_total: success_total, failed_total: failed_total
-              )
             end
+
+            source_batches_done += 1
+            on_batch&.call(
+              batches_done: source_batches_done, docs_total: docs_total,
+              success_total: success_total, failed_total: failed_total
+            )
           end
 
           # Calculate total duration as wall-clock time from start to finish (not sum of batch durations)
@@ -195,6 +197,7 @@ module SearchEngine
             failed_total: 0,
             failed_batches_total: 0,
             batches_total: 0,
+            source_batches_done: 0,
             idx_counter: -1,
             started_at: monotonic_ms,
             mtx: Mutex.new,
@@ -314,12 +317,11 @@ module SearchEngine
         # @return [void]
         def process_single_batch_parallel(raw_batch:, into:, action:, retry_policy:, batch_size:, log_batches:,
                                           shared_state:)
-          # Each thread gets its own resources
           thread_client = SearchEngine.client
           thread_buffer = +''
           thread_idx = shared_state[:mtx].synchronize { shared_state[:idx_counter] += 1 }
 
-          snapshots = begin
+          snapshot = begin
             stats_list = import_batch_with_handling(
               client: thread_client,
               collection: into,
@@ -332,6 +334,8 @@ module SearchEngine
 
             shared_state[:mtx].synchronize do
               aggregate_stats(stats_list, shared_state, batch_size, log_batches)
+              shared_state[:source_batches_done] += 1
+              progress_snapshot(shared_state)
             end
           rescue StandardError => error
             docs_count = begin
@@ -346,27 +350,26 @@ module SearchEngine
               err_msg = "  batch_index=#{thread_idx} → error=#{error.class}: #{error.message.to_s[0, 200]}"
               warn(SearchEngine::Logging::Color.apply(err_msg, :red))
               aggregate_stats([failure_stat], shared_state, batch_size, log_batches)
+              shared_state[:source_batches_done] += 1
+              progress_snapshot(shared_state)
             end
           end
 
-          on_batch = shared_state[:on_batch]
-          snapshots&.each { |snap| on_batch&.call(**snap) }
+          shared_state[:on_batch]&.call(**snapshot) if snapshot
         end
 
         # Aggregate batch statistics thread-safely into shared state.
         #
         # Must be called within a mutex synchronization block. Updates counters,
         # appends to batches array, validates batch size, and optionally logs.
-        # Returns counter snapshots (one per stats entry) for firing callbacks
-        # outside the lock.
         #
         # @param stats_list [Array<Hash>] array of stats hashes from batch processing
         # @param shared_state [Hash] shared state hash to update (must be mutex-protected)
         # @param batch_size [Integer, nil] soft guard for logging when exceeded
         # @param log_batches [Boolean] whether to log each batch as it completes
-        # @return [Array<Hash>] counter snapshots suitable for on_batch callbacks
+        # @return [void]
         def aggregate_stats(stats_list, shared_state, batch_size, log_batches)
-          stats_list.map do |stats|
+          stats_list.each do |stats|
             shared_state[:docs_total] += stats[:docs_count].to_i
             shared_state[:success_total] += stats[:success_count].to_i
             shared_state[:failed_total] += stats[:failure_count].to_i
@@ -375,13 +378,23 @@ module SearchEngine
             shared_state[:batches] << stats
             validate_soft_batch_size!(batch_size, stats[:docs_count])
             log_batch(stats, shared_state[:batches_total]) if log_batches
-            {
-              batches_done: shared_state[:batches_total], docs_total: shared_state[:docs_total],
-              success_total: shared_state[:success_total], failed_total: shared_state[:failed_total]
-            }
           end
         end
 
+        # Build a progress snapshot from shared state for the on_batch callback.
+        # Must be called within a mutex synchronization block.
+        #
+        # @param shared_state [Hash] shared state hash (must be mutex-protected)
+        # @return [Hash] progress counters keyed by :batches_done, :docs_total, etc.
+        def progress_snapshot(shared_state)
+          {
+            batches_done: shared_state[:source_batches_done],
+            docs_total: shared_state[:docs_total],
+            success_total: shared_state[:success_total],
+            failed_total: shared_state[:failed_total]
+          }
+        end
+
         # Build a Summary object from aggregated shared state.
         #
         # Calculates total duration and constructs a Summary with all aggregated
@@ -426,14 +439,14 @@ module SearchEngine
           enum.is_a?(Enumerator) ? enum : enum.each
         end
 
-        # Estimate total batch count for progress logging.
+        # Estimate total source record count for the given model class.
         #
-        # Attempts to estimate batch count for ActiveRecord sources by counting records
-        # and dividing by batch_size. Returns nil for other source types or when estimation fails.
+        # Shared foundation for batch and doc estimates. Performs a model.count
+        # with a soft timeout to avoid blocking on slow tables.
         #
         # @param klass [Class] a {SearchEngine::Base} subclass
-        # @return [Integer, nil] estimated total batch count or nil if not estimable
-        def estimate_total_batches(klass)
+        # @return [Integer, nil] record count or nil if not estimable
+        def estimate_source_record_count(klass)
           return nil if SearchEngine.config.indexer.estimate_progress == false
           return nil unless klass.is_a?(Class)
 
@@ -447,22 +460,39 @@ module SearchEngine
           model = source_def.dig(:options, :model)
           return nil unless model.respond_to?(:count)
 
-          batch_size = source_def.dig(:options, :batch_size)
-          batch_size ||= SearchEngine.config.sources.active_record.batch_size
-          batch_size = batch_size.to_i
-          return nil unless batch_size.positive?
+          count_with_timeout(model, 10)
+        rescue StandardError
+          nil
+        end
 
-          begin
-            total_records = count_with_timeout(model, 10)
-            return nil unless total_records&.positive?
+        # Estimate total batch count for progress logging.
+        #
+        # @param klass [Class] a {SearchEngine::Base} subclass
+        # @return [Integer, nil] estimated total batch count or nil if not estimable
+        def estimate_total_batches(klass)
+          total_records = estimate_source_record_count(klass)
+          return nil unless total_records&.positive?
 
-            (total_records.to_f / batch_size).ceil
-          rescue StandardError
-            nil
-          end
+          batch_size = batch_size_for_klass(klass)
+          return nil unless batch_size&.positive?
+
+          (total_records.to_f / batch_size).ceil
+        rescue StandardError
+          nil
         end
 
-        public :estimate_total_batches
+        # Estimate total document count for doc-based progress tracking.
+        #
+        # @param klass [Class] a {SearchEngine::Base} subclass
+        # @return [Integer, nil] estimated total docs or nil if not estimable
+        def estimate_total_docs(klass)
+          count = estimate_source_record_count(klass)
+          count&.positive? ? count : nil
+        rescue StandardError
+          nil
+        end
+
+        public :estimate_total_batches, :estimate_total_docs
 
         # Thread-based soft timeout for model.count, avoiding Timeout.timeout
         # which can corrupt ActiveRecord connection state.
@@ -497,6 +527,14 @@ module SearchEngine
           klass.instance_variable_get(:@__mapper_dsl__)
         end
 
+        def batch_size_for_klass(klass)
+          dsl = mapper_dsl_for_klass(klass)
+          source_def = dsl&.dig(:source)
+          batch_size = source_def&.dig(:options, :batch_size)
+          batch_size ||= SearchEngine.config.sources.active_record.batch_size
+          batch_size.to_i
+        end
+
         # Import a single batch with error handling and recursive 413 splitting.
         #
         # Public wrapper that delegates to the internal method with batch_index set to nil,

data/lib/search_engine/logging/live_renderer.rb

@@ -16,7 +16,7 @@ module SearchEngine
     # using {PartitionProgress.line}, preserving CI/pipe compatibility.
     #
     # @example
-    #   renderer = LiveRenderer.new(labels: parts.map(&:inspect), per_partition_estimate: 50)
+    #   renderer = LiveRenderer.new(labels: parts.map(&:inspect), per_partition_docs_estimate: 5000)
     #   renderer.start
     #   parts.each_with_index do |part, i|
     #     renderer[i].start
@@ -34,9 +34,12 @@ module SearchEngine
 
       # @param labels [Array<String>] display label for each slot (partition key)
       # @param partitions [Array, nil] raw partition values for non-TTY output (defaults to labels)
-      # @param per_partition_estimate [Integer, nil] estimated batches per partition (for progress bars)
+      # @param per_partition_docs_estimates [Array<Integer, nil>, nil] per-slot doc estimates (takes priority)
+      # @param per_partition_docs_estimate [Integer, nil] uniform doc estimate for all slots (fallback)
+      # @param per_partition_estimate [Integer, nil] deprecated batch-based estimate (last resort fallback)
       # @param io [IO] output stream (defaults to $stdout)
-      def initialize(labels:, partitions: nil, per_partition_estimate: nil, io: $stdout)
+      def initialize(labels:, partitions: nil, per_partition_docs_estimates: nil,
+                     per_partition_docs_estimate: nil, per_partition_estimate: nil, io: $stdout)
         @io = io
         @tty = Color.enabled?
         @mutex = Mutex.new
@@ -46,8 +49,13 @@ module SearchEngine
         @rendered_once = false
         nontty_cb = @tty ? nil : method(:flush_nontty_slot)
         raw = partitions || labels
+        per_slot = per_partition_docs_estimates || []
+        global_est = per_partition_docs_estimate || per_partition_estimate
         @slots = labels.each_with_index.map do |label, idx|
-          Slot.new(label: label, partition: raw[idx], estimate: per_partition_estimate, on_done: nontty_cb)
+          Slot.new(
+            label: label, partition: raw[idx], docs_estimate: per_slot[idx] || global_est,
+            on_done: nontty_cb
+          )
         end
         @viewport = resolve_viewport
       end
@@ -238,12 +246,12 @@ module SearchEngine
 
         # @param label [String] partition display label (e.g. partition key inspect)
         # @param partition [Object, nil] raw partition value for non-TTY output (defaults to label)
-        # @param estimate [Integer, nil] estimated total batches for progress bar
+        # @param docs_estimate [Integer, nil] estimated total docs for doc-based progress bar
         # @param on_done [Proc, nil] callback invoked after finish/finish_error (non-TTY flush)
-        def initialize(label:, partition: nil, estimate: nil, on_done: nil)
+        def initialize(label:, partition: nil, docs_estimate: nil, on_done: nil)
           @label = label
           @partition = partition.nil? ? label : partition
-          @estimate = estimate
+          @docs_estimate = docs_estimate
           @on_done = on_done
           @state = :pending
           @batches_done = 0
@@ -387,15 +395,16 @@ module SearchEngine
         end
 
         def build_progress_part
-          if @estimate&.positive? && @batches_done.positive?
-            ratio = @batches_done.to_f / @estimate
+          if @docs_estimate&.positive? && @docs_total.positive?
+            ratio = @docs_total.to_f / @docs_estimate
             pct = [100, (ratio * 100).round].min
             filled = [(ratio * BAR_WIDTH).round, BAR_WIDTH].min
             empty = BAR_WIDTH - filled
             bar = "\u2588" * filled + "\u2591" * empty
-            "#{bar}  #{pct}%  #{@batches_done}/#{@estimate} batches  (#{@docs_total} docs)"
+            batch_info = @batches_done.positive? ? "#{@batches_done} batches  " : ''
+            "#{bar}  #{pct}%  #{batch_info}(#{@docs_total}/#{@docs_estimate} docs)"
           elsif @batches_done.positive?
-            "#{@batches_done} batches, #{@docs_total} docs"
+            "#{@batches_done} batches  (#{@docs_total} docs)"
           else
             ''
           end

data/lib/search_engine/partitioner.rb

@@ -76,8 +76,50 @@ module SearchEngine
         end
       end
 
+      # Auto-detect doc count for a partition from the partition_fetch result.
+      # Counts the same scope that partition_fetch will enumerate — single source of truth.
+      #
+      # @param partition [Object] partition key
+      # @return [Integer, nil] doc count or nil when unavailable
+      def partition_doc_count(partition)
+        count = auto_count_from_fetch(partition)
+        count.is_a?(Integer) && count.positive? ? count : nil
+      rescue StandardError
+        nil
+      end
+
       private
 
+      # Call partition_fetch_proc and try to extract a countable relation.
+      # Calling the proc only builds a lazy AR object — no batch queries fire.
+      # The single COUNT query is the only DB cost.
+      def auto_count_from_fetch(partition)
+        return nil unless @partition_fetch_proc
+
+        result = @partition_fetch_proc.call(partition)
+        extract_count(result)
+      rescue StandardError
+        nil
+      end
+
+      def extract_count(result)
+        relation = countable_relation_from(result)
+        relation&.count
+      rescue StandardError
+        nil
+      end
+
+      def countable_relation_from(result)
+        if defined?(ActiveRecord::Batches::BatchEnumerator) &&
+           result.is_a?(ActiveRecord::Batches::BatchEnumerator)
+          return result.instance_variable_get(:@relation)
+        end
+
+        return result if defined?(ActiveRecord::Relation) && result.is_a?(ActiveRecord::Relation)
+
+        nil
+      end
+
       def validate_hook_arity!(proc_obj, name:)
         ar = proc_obj.arity
         return if ar == 1 || ar.negative?

data/lib/search_engine/version.rb

@@ -3,5 +3,5 @@
 module SearchEngine
   # Current gem version.
   # @return [String]
-  VERSION = '30.1.6.6'
+  VERSION = '30.1.6.7'
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: search-engine-for-typesense
 version: !ruby/object:Gem::Version
-  version: 30.1.6.6
+  version: 30.1.6.7
 platform: ruby
 authors:
 - Nikita Shkoda