feat(orm): extend store with rows_by_id and range_bounds

lib/prism/orm/store.rb

@@ -453,6 +453,217 @@ module Prism
 
       private
 
+      # `id` is always "indexed" for free since `table_definition(table)[:rows]`
+      # is already a Hash keyed by id - shared by #lookup_by_index's single-id
+      # equality fast-path and #lookup_by_index_in's id IN-list fast-path, which
+      # differed only in that one wrapped a single id and the other a list.
+      def rows_by_id(table, ids)
+        within_transaction(read_only: true) do
+          rows = table_definition(table)[:rows]
+          ids.map { |id| rows[id.to_i]&.dup }.compact
+        end
+      end
+
+      # Finds [lo, hi) into `sorted` (ascending, unique values) covering
+      # every value allowed by `operators`. A tighter bound on the same side
+      # (both :gte and :gt given, say) - the caller (ConditionNode) never
+      # produces that from a single #where, but even if it did, only one
+      # side is honored here; the final Relation#rows `#select` against the
+      # real condition still filters out anything that slipped through.
+      def range_bounds(sorted, operators)
+        lo = 0
+        hi = sorted.length
+        if operators.key?(:gte)
+          lo = sorted.bsearch_index { |v| v >= operators[:gte] } || sorted.length
+        elsif operators.key?(:gt)
+          lo = sorted.bsearch_index { |v| v > operators[:gt] } || sorted.length
+        end
+        if operators.key?(:lte)
+          hi = sorted.bsearch_index { |v| v > operators[:lte] } || sorted.length
+        elsif operators.key?(:lt)
+          hi = sorted.bsearch_index { |v| v >= operators[:lt] } || sorted.length
+        end
+        [lo, hi]
+      end
+
+      def resolve_index_name(definition, name_or_columns)
+        return name_or_columns.to_s if definition[:indexes].key?(name_or_columns.to_s)
+        columns = Array(name_or_columns).map(&:to_s)
+        match = definition[:indexes].values.find { |index| index[:columns] == columns }
+        match && match[:name]
+      end
+
+      # Builds the map/range metadata for an index from the table's current
+      # rows - the same computation #add_index needs for a brand-new index
+      # and #rebuild_indexes_covering needs to bring an existing index back
+      # in sync after a backfill changed every row's value for one of its
+      # columns.
+      def build_index_data(definition, columns, unique:, name:)
+        map = {}
+        definition[:rows].each do |id, row|
+          key = columns.map { |column| row[column] }
+          existing = (map[key] ||= [])
+          if unique && !existing.empty?
+            raise UniqueConstraintViolation,
+                  "cannot create unique index #{name}: duplicate value for #{columns.join(', ')}"
+          end
+          existing << id
+        end
+
+        # Range queries (gt/gte/lt/lte, see #lookup_by_index_range) only
+        # make sense for a single-column index - a composite key has no
+        # single total order Predicates agrees with. range_capable stays
+        # false if the column's values can't all be compared with each
+        # other (e.g. mixed types slipped in), same defensive posture as
+        # everywhere else in this store: an index that can't help safely
+        # just means callers fall back to a full scan, never a wrong one.
+        range_capable = false
+        sorted_values = []
+        if columns.size == 1
+          begin
+            sorted_values = map.keys.map(&:first).compact.sort
+            range_capable = true
+          rescue ArgumentError, TypeError
+            range_capable = false
+            sorted_values = []
+          end
+        end
+
+        # The composite counterpart of range_capable above: a multi-column
+        # index can still answer a range query on its LAST column as long
+        # as every value observed for that last column, within a given
+        # prefix (every column before it), can be totally ordered -
+        # sorted_by_prefix holds one ascending, de-duplicated array of
+        # last-column values per prefix key, exactly what
+        # #lookup_by_index_equality_range binary-searches into. A
+        # single-column index has no prefix to partition by, so this is
+        # only ever attempted for columns.size > 1 (a single-column index
+        # already gets the simpler range_capable path above).
+        composite_range_capable = false
+        sorted_by_prefix = {}
+        if columns.size > 1
+          begin
+            grouped = Hash.new { |hash, key| hash[key] = [] }
+            map.each_key { |key| grouped[key[0..-2]] << key.last }
+            grouped.each { |prefix, values| sorted_by_prefix[prefix] = values.compact.uniq.sort }
+            composite_range_capable = true
+          rescue ArgumentError, TypeError
+            composite_range_capable = false
+            sorted_by_prefix = {}
+          end
+        end
+
+        { name: name, columns: columns, unique: !!unique, map: map,
+          range_capable: range_capable, sorted_values: sorted_values,
+          composite_range_capable: composite_range_capable, sorted_by_prefix: sorted_by_prefix }
+      end
+
+      # Recomputes every index that covers `column` from scratch. Called
+      # after #add_column backfills a value onto every existing row for a
+      # column an index already existed on - see the comment on #add_column
+      # for why a stale map here is a silent-wrong-results bug, not just a
+      # missed optimization.
+      def rebuild_indexes_covering(definition, column)
+        (definition[:indexes] || {}).each do |name, index|
+          next unless index[:columns].include?(column)
+          definition[:indexes][name] = build_index_data(definition, index[:columns], unique: index[:unique], name: name)
+        end
+      end
+
+      def validate_unique_indexes(definition, id, row)
+        (definition[:indexes] || {}).each_value do |index|
+          next unless index[:unique]
+          key = index[:columns].map { |column| row[column] }
+          existing = index[:map][key]
+          next if existing.nil? || existing.empty? || existing == [id]
+          raise UniqueConstraintViolation,
+                "duplicate value for unique index #{index[:name]} on #{index[:columns].join(', ')}"
+        end
+      end
+
+      def index_add(definition, id, row)
+        (definition[:indexes] || {}).each_value do |index|
+          key = index[:columns].map { |column| row[column] }
+          is_new_key = !index[:map].key?(key) || index[:map][key].empty?
+          list = (index[:map][key] ||= [])
+          list << id unless list.include?(id)
+          insert_sorted_value(index, key.first) if is_new_key && index[:range_capable] && !key.first.nil?
+          insert_sorted_prefix_value(index, key) if is_new_key && index[:composite_range_capable] && !key.last.nil?
+        end
+      end
+
+      def index_remove(definition, id, row)
+        (definition[:indexes] || {}).each_value do |index|
+          key = index[:columns].map { |column| row[column] }
+          list = index[:map][key]
+          next unless list
+          list.delete(id)
+          next unless list.empty?
+          index[:map].delete(key)
+          remove_sorted_value(index, key.first) if index[:range_capable] && !key.first.nil?
+          remove_sorted_prefix_value(index, key) if index[:composite_range_capable] && !key.last.nil?
+        end
+      end
+
+      # Keeps sorted_values ascending and unique as rows come and go, so
+      # #lookup_by_index_range never has to re-sort. A comparison failure
+      # (mixed, mutually-incomparable types landing in the same column)
+      # disables range_capable rather than raising - the index still works
+      # for #lookup_by_index equality, it just stops being usable for range
+      # queries, the same safe-degrade posture as #add_index above.
+      def insert_sorted_value(index, value)
+        sorted = index[:sorted_values]
+        idx = sorted.bsearch_index { |v| v >= value } || sorted.length
+        sorted.insert(idx, value) unless sorted[idx] == value
+      rescue ArgumentError
+        index[:range_capable] = false
+        index[:sorted_values] = []
+      end
+
+      def remove_sorted_value(index, value)
+        sorted = index[:sorted_values]
+        idx = sorted.bsearch_index { |v| v >= value }
+        sorted.delete_at(idx) if idx && sorted[idx] == value
+      rescue ArgumentError
+        index[:range_capable] = false
+        index[:sorted_values] = []
+      end
+
+      # Same incremental bookkeeping as #insert_sorted_value/
+      # #remove_sorted_value, but scoped to one prefix key's own sorted
+      # array (there is one such array per distinct prefix, not one for the
+      # whole index) - #lookup_by_index_equality_range only ever binary-
+      # searches the single array for the prefix it was asked about, so
+      # every OTHER prefix's array must stay untouched by a row that
+      # doesn't belong to it.
+      def insert_sorted_prefix_value(index, key)
+        prefix, value = key[0..-2], key.last
+        sorted = (index[:sorted_by_prefix][prefix] ||= [])
+        idx = sorted.bsearch_index { |v| v >= value } || sorted.length
+        sorted.insert(idx, value) unless sorted[idx] == value
+      rescue ArgumentError
+        index[:composite_range_capable] = false
+        index[:sorted_by_prefix] = {}
+      end
+
+      def remove_sorted_prefix_value(index, key)
+        prefix, value = key[0..-2], key.last
+        sorted = index[:sorted_by_prefix][prefix]
+        return unless sorted
+        idx = sorted.bsearch_index { |v| v >= value }
+        sorted.delete_at(idx) if idx && sorted[idx] == value
+        index[:sorted_by_prefix].delete(prefix) if sorted.empty?
+      rescue ArgumentError
+        index[:composite_range_capable] = false
+        index[:sorted_by_prefix] = {}
+      end
+
+      def instrument(operation, table, extra = {})
+        Instrumentation.instrument('sql.prism', { operation: operation, table: table.to_s }.merge(extra)) do
+          yield
+        end
+      end
+
       def within_transaction(read_only: false)
         return yield if @in_transaction