<commit_before>require "./abstract_table"
require "../metrics"

module Storage
  struct PhysicalTable < AbstractTable
    MUTEX_SHARDS = 2 << 12
    @@mutexes : Array(Mutex) = (0...MUTEX_SHARDS).map { Mutex.new }

    property read_docs_on_table = PerSecondMetric.new
    property written_docs_on_table = PerSecondMetric.new

    def initialize(@manager : Manager, @table : Manager::Table)
    end

    def primary_key
      @table.info.primary_key
    end

    def get(key)
      @read_docs_on_table.add(1)
      @manager.kv.get_row(@table.info.id, ReQL.encode_key(key)) do |data|
        if data
          hash = ReQL::Datum.unserialize(IO::Memory.new(data)).hash_value
          hash[primary_key] = key
          hash
        end
      end
    end

    def replace(key, durability : ReQL::Durability? = nil)
      key_data = ReQL.encode_key(key)

      @@mutexes[key.hash % MUTEX_SHARDS].synchronize do
        existing_row = @manager.kv.get_row(@table.info.id, key_data) do |existing_row_data|
          if existing_row_data.nil?
            nil
          else
            hash = ReQL::Datum.unserialize(IO::Memory.new(existing_row_data)).hash_value
            hash[primary_key] = key
            hash
          end
        end

        new_row = yield existing_row

        if new_row.nil?
          unless existing_row.nil?
            @written_docs_on_table.add(1)

            if @table.indices.empty?
              @manager.kv.delete_row(@table.info.id, key_data, durability || @table.info.durability)
            else
              batch = @manager.kv.create_batch
              batch.delete_row(@table.info.id, key_data)

              @table.indices.values.each do |index|
                update_index_data(batch, index, key_data, existing_row, new_row)
              end

              @manager.kv.write_batch(batch, durability || @table.info.durability)
            end
          end
        else
          if existing_row != new_row
            @written_docs_on_table.add(1)

            if @table.indices.empty?
              @manager.kv.set_row(@table.info.id, key_data, ReQL::Datum.new(new_row).serialize, durability || @table.info.durability)
            else
              batch = @manager.kv.create_batch

              new_row.delete(primary_key)
              batch.set_row(@table.info.id, key_data, ReQL::Datum.new(new_row).serialize)
              new_row[primary_key] = key

              @table.indices.values.each do |index|
                update_index_data(batch, index, key_data, existing_row, new_row)
              end

              @manager.kv.write_batch(batch, durability || @table.info.durability)
            end
          end
        end
      end
    end

    def scan
      @manager.kv.each_row(@table.info.id) do |key, data|
        @read_docs_on_table.add(1)
        hash = ReQL::Datum.unserialize(IO::Memory.new(data)).hash_value
        hash[primary_key] = ReQL.decode_key(key)
        yield hash
      end
    end

    private def index_values(evaluator, index, row) : Set({ReQL::Datum, Int32})
      computed = row.try { |row| index.info.function.eval(evaluator, {ReQL::Datum.new(row)}).as_datum rescue nil }

      if computed.nil? || computed == nil
        return Set({ReQL::Datum, Int32}).new
      end

      if index.info.multi
        set = Set({ReQL::Datum, Int32}).new
        hash = Hash(ReQL::Datum, Int32).new { 0 }
        computed.array_or_set_value?.try &.each do |val|
          set << {val, hash[val] += 1} unless val == nil
        end

        set
      else
        # https://github.com/crystal-lang/crystal/issues/8868
        [{computed.as(ReQL::Datum), 1}].to_set
      end
    end

    def update_index_data(t, index, key_data, old_row, new_row)
      evaluator = ReQL::Evaluator.new(@manager)
      evaluator.now = Time.unix(0)

      old_values = index_values(evaluator, index, old_row)
      new_values = index_values(evaluator, index, new_row)

      # Remove values that are now missing
      (old_values - new_values).each do |(value, counter)|
        @written_docs_on_table.add(1)
        t.delete_index_entry(@table.info.id, index.info.id, ReQL.encode_key(value), counter, key_data)
      end

      # Add new values
      (new_values - old_values).each do |(value, counter)|
        @written_docs_on_table.add(1)
        t.set_index_entry(@table.info.id, index.info.id, ReQL.encode_key(value), counter, key_data)
      end
    end

    def create_index(name : String, function : ReQL::Func, multi : Bool)
      info = KeyValueStore::IndexInfo.new
      info.name = name
      info.table = @table.info.id
      info.function = function
      info.multi = multi
      index = Manager::Index.new(info)
      @manager.lock.synchronize do
        if @table.indices.has_key?(name)
          raise ReQL::OpFailedError.new("Index `#{name}` already exists on table `#{@table.db_name}.#{@table.info.name}`")
        end
        @manager.kv.save_index(info)
        @table.indices[name] = index
      end

      spawn build_index(index)
    end

    def build_index(index : Manager::Index)
      @manager.kv.build_index(@table.info.id, index.info.id) do |builder|
        @manager.kv.each_row(@table.info.id) do |key, data|
          @read_docs_on_table.add(1)
          hash = ReQL::Datum.unserialize(IO::Memory.new(data)).hash_value
          hash[primary_key] = ReQL.decode_key(key)
          update_index_data(builder, index, key, nil, hash)
        end
      end

      @manager.lock.synchronize do
        # Check if index was dropped while it was building (doesn't exist anymore or was recreated)
        unless @table.indices[index.info.name]? == index<commit_msg>Improve reporting on written keys during index building<commit_after>require "./abstract_table"
require "../metrics"

module Storage
  struct PhysicalTable < AbstractTable
    MUTEX_SHARDS = 2 << 12
    @@mutexes : Array(Mutex) = (0...MUTEX_SHARDS).map { Mutex.new }

    property read_docs_on_table = PerSecondMetric.new
    property written_docs_on_table = PerSecondMetric.new

    def initialize(@manager : Manager, @table : Manager::Table)
    end

    def primary_key
      @table.info.primary_key
    end

    def get(key)
      @read_docs_on_table.add(1)
      @manager.kv.get_row(@table.info.id, ReQL.encode_key(key)) do |data|
        if data
          hash = ReQL::Datum.unserialize(IO::Memory.new(data)).hash_value
          hash[primary_key] = key
          hash
        end
      end
    end

    def replace(key, durability : ReQL::Durability? = nil)
      key_data = ReQL.encode_key(key)

      @@mutexes[key.hash % MUTEX_SHARDS].synchronize do
        existing_row = @manager.kv.get_row(@table.info.id, key_data) do |existing_row_data|
          if existing_row_data.nil?
            nil
          else
            hash = ReQL::Datum.unserialize(IO::Memory.new(existing_row_data)).hash_value
            hash[primary_key] = key
            hash
          end
        end

        new_row = yield existing_row

        if new_row.nil?
          unless existing_row.nil?
            @written_docs_on_table.add(1)

            if @table.indices.empty?
              @manager.kv.delete_row(@table.info.id, key_data, durability || @table.info.durability)
            else
              batch = @manager.kv.create_batch
              batch.delete_row(@table.info.id, key_data)

              @table.indices.values.each do |index|
                update_index_data(batch, index, key_data, existing_row, new_row)
              end

              @manager.kv.write_batch(batch, durability || @table.info.durability)
            end
          end
        else
          if existing_row != new_row
            @written_docs_on_table.add(1)

            if @table.indices.empty?
              @manager.kv.set_row(@table.info.id, key_data, ReQL::Datum.new(new_row).serialize, durability || @table.info.durability)
            else
              batch = @manager.kv.create_batch

              new_row.delete(primary_key)
              batch.set_row(@table.info.id, key_data, ReQL::Datum.new(new_row).serialize)
              new_row[primary_key] = key

              @table.indices.values.each do |index|
                update_index_data(batch, index, key_data, existing_row, new_row)
              end

              @manager.kv.write_batch(batch, durability || @table.info.durability)
            end
          end
        end
      end
    end

    def scan
      @manager.kv.each_row(@table.info.id) do |key, data|
        @read_docs_on_table.add(1)
        hash = ReQL::Datum.unserialize(IO::Memory.new(data)).hash_value
        hash[primary_key] = ReQL.decode_key(key)
        yield hash
      end
    end

    private def index_values(evaluator, index, row) : Set({ReQL::Datum, Int32})
      computed = row.try { |row| index.info.function.eval(evaluator, {ReQL::Datum.new(row)}).as_datum rescue nil }

      if computed.nil? || computed == nil
        return Set({ReQL::Datum, Int32}).new
      end

      if index.info.multi
        set = Set({ReQL::Datum, Int32}).new
        hash = Hash(ReQL::Datum, Int32).new { 0 }
        computed.array_or_set_value?.try &.each do |val|
          set << {val, hash[val] += 1} unless val == nil
        end

        set
      else
        # https://github.com/crystal-lang/crystal/issues/8868
        [{computed.as(ReQL::Datum), 1}].to_set
      end
    end

    def update_index_data(t, index, key_data, old_row, new_row)
      evaluator = ReQL::Evaluator.new(@manager)
      evaluator.now = Time.unix(0)

      old_values = index_values(evaluator, index, old_row)
      new_values = index_values(evaluator, index, new_row)

      # Remove values that are now missing
      (old_values - new_values).each do |(value, counter)|
        @written_docs_on_table.add(1)
        t.delete_index_entry(@table.info.id, index.info.id, ReQL.encode_key(value), counter, key_data)
      end

      # Add new values
      (new_values - old_values).each do |(value, counter)|
        @written_docs_on_table.add(1)
        t.set_index_entry(@table.info.id, index.info.id, ReQL.encode_key(value), counter, key_data)
      end
    end

    def create_index(name : String, function : ReQL::Func, multi : Bool)
      info = KeyValueStore::IndexInfo.new
      info.name = name
      info.table = @table.info.id
      info.function = function
      info.multi = multi
      index = Manager::Index.new(info)
      @manager.lock.synchronize do
        if @table.indices.has_key?(name)
          raise ReQL::OpFailedError.new("Index `#{name}` already exists on table `#{@table.db_name}.#{@table.info.name}`")
        end
        @manager.kv.save_index(info)
        @table.indices[name] = index
      end

      spawn build_index(index)
    end

    def build_index(index : Manager::Index)
      @manager.kv.build_index(@table.info.id, index.info.id) do |builder|
        evaluator = ReQL::Evaluator.new(@manager)
        evaluator.now = Time.unix(0)

        @manager.kv.each_row(@table.info.id) do |key, data|
          @read_docs_on_table.add(1)
          hash = ReQL::Datum.unserialize(IO::Memory.new(data)).hash_value
          hash[primary_key] = ReQL.decode_key(key)

          values = index_values(evaluator, index, hash)
          values.each do |(value, counter)|
            @written_docs_on_table.add(1)
            builder.set_index_entry(@table.info.id, index.info.id, ReQL.encode_key(value), counter, key)
          end
        end
      end

      @manager.lock.synchronize do
        # Check if index was dropped while it was building (doesn't exist anymore or was recreated)
        unless @table.indices[index.info.name]? == index
{'fraction_non_alphanumeric': 0.11411359724612737, 'fraction_numerical': 0.004819277108433735, 'mean_word_length': 2.2939342403628116, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2538384', 'n_tokens_mistral': 3788, 'n_tokens_neox': 3709, 'n_words': 710}