Add ClickHouse limit_by and deduplicate dataset methods

Gemfile.lock

@@ -11,7 +11,7 @@ GIT
 PATH
   remote: .
   specs:
-    umbrellio-utils (1.15.0)
+    umbrellio-utils (1.16.0)
       memery (~> 1)
 
 GEM

README.md

@@ -115,6 +115,37 @@ end
 Utils::Constants.useful_method #=> "Just string"
 ```
 
+### ClickHouse deduplication
+
+Datasets built with `CH.from` understand ClickHouse's `LIMIT n BY`:
+
+```ruby
+CH.from(:events).order(Sequel.desc(:version)).limit_by(:user_id, rows: 3)
+#=> SELECT * FROM "events" ORDER BY "version" DESC LIMIT 3 BY "user_id"
+```
+
+On a `ReplacingMergeTree` table, `#deduplicate` uses that to collapse row versions by hand rather than relying on the `final` setting, which merges the entire table even for a point lookup:
+
+```ruby
+CH.from(:external_operations_distributed)
+  .where(order_id: 42)   # inside the dedup subquery
+  .deduplicate           # boundary
+  .order(:created_at)    # outside
+```
+
+The sorting key, version column and `is_deleted` column are read from `system.tables` and cached per process; `Distributed` tables are resolved through to the local table they wrap. Deduplicated datasets are sent with `final: 0` automatically, since running FINAL inside the subquery would be both slow and redundant. An explicit `final:` passed to `query` / `count` always wins.
+
+`#deduplicate` is a boundary, and which side a filter lands on matters:
+
+- **Before it — immutable selectors only** (primary keys, foreign keys). Filtering a *mutable* column first can match a superseded row version and resurrect a row that FINAL would have dropped.
+- **After it — everything else**, including any filter on a column that changes over a row's lifetime.
+
+`is_deleted` is applied after the boundary for exactly that reason, and is omitted when the engine declares no such column.
+
+Two chain methods are handled rather than passed through, because the dedup subquery has to control them: the subquery always projects `SELECT *` (so the outer query can still filter on `is_deleted` and on columns you did not select) and any projection you set is re-applied outside it; and the version ordering leads the subquery's `ORDER BY`, since it decides which row survives, with any ordering you set kept after it as a tiebreaker.
+
+`#deduplicate` raises on a table it cannot collapse — a non-Replacing engine, a Replacing engine declared without a version argument, or a dataset that is not a single table source (joined, multi-source, or a subquery).
+
 ### Instrumentation
 
 The gem ships a set of opt-in files for collecting GVL and allocation stats

lib/umbrellio_utils/click_house.rb

@@ -10,12 +10,13 @@ module ClickHouse
     extend self
 
     autoload :Backends, "umbrellio_utils/click_house/backends"
+    autoload :TableMetadata, "umbrellio_utils/click_house/table_metadata"
 
     VALID_BACKENDS = %i[legacy native].freeze
 
     DELEGATED = %i[
       execute query query_value query_each count insert
-      from describe_table server_version tables
+      from describe_table server_version tables table_metadata
       create_database drop_database db_name config
       truncate_table! drop_table! optimize_table! on_cluster
       parse_value pg_table_connection populate_temp_table! with_temp_table

lib/umbrellio_utils/click_house/backends/base.rb

@@ -13,13 +13,103 @@ module Backends
       class Base
         include Singleton
 
-        # ClickHouse uses C-style escape sequences in string literals, so
-        # backslashes must be doubled. Sequel's default (Postgres) escaping
-        # only escapes single-quotes.
-        module ClickHouseStringEscaping
+        # ClickHouse-specific dataset behaviour: string escaping plus grammar
+        # Sequel's Postgres dataset doesn't know about.
+        module ClickHouseDatasetMethods
+          # ClickHouse uses C-style escape sequences in string literals, so
+          # backslashes must be doubled. Sequel's default (Postgres) escaping
+          # only escapes single-quotes.
           def literal_string_append(sql, str)
             sql << "'" << str.gsub("\\") { "\\\\" }.gsub("'", "''") << "'"
           end
+
+          # ClickHouse `LIMIT n BY expr, ...` — keeps the first n rows per
+          # distinct combination of the expressions, applied after ORDER BY.
+          def limit_by(*exprs, rows: 1)
+            raise Sequel::Error, "limit_by requires at least one expression" if exprs.empty?
+            clone(limit_by: { exprs:, rows: })
+          end
+
+          # `LIMIT n BY` precedes the regular LIMIT/OFFSET in ClickHouse.
+          def select_limit_sql(sql)
+            if (limit_by = @opts[:limit_by])
+              sql << " LIMIT "
+              literal_append(sql, limit_by[:rows])
+              sql << " BY "
+              expression_list_append(sql, limit_by[:exprs])
+            end
+
+            super
+          end
+
+          # Collapse a ReplacingMergeTree's row versions by hand instead of
+          # relying on the `final` setting, which merges the whole table.
+          #
+          # This is a boundary: everything chained BEFORE it goes inside the
+          # dedup subquery, everything chained AFTER applies to the result.
+          # Only immutable selectors belong before it — filtering a mutable
+          # column first can match a superseded version and resurrect a row
+          # that FINAL would have dropped. `is_deleted` is applied after the
+          # boundary for the same reason.
+          def deduplicate
+            table_name, db_name = deduplication_source
+            meta = ClickHouse.table_metadata(table_name, **db_name)
+
+            unless meta.replacing?
+              raise Sequel::Error,
+                    "#{table_name} is a #{meta.engine}; deduplicate needs a ReplacingMergeTree"
+            end
+
+            unless meta.version
+              raise Sequel::Error,
+                    "#{table_name} declares no version column; deduplicate needs one"
+            end
+
+            wrapped = ClickHouse.from(wrap_source(deduplicated_source(meta)))
+              .clone(ch_dedup: true)
+              .then { |ds| @opts[:select] ? ds.select(*@opts[:select]) : ds }
+
+            meta.is_deleted ? wrapped.where(meta.is_deleted => 0) : wrapped
+          end
+
+          private
+
+          # The subquery must project every column, both so the outer query can
+          # filter on `is_deleted` and so a caller's projection still resolves;
+          # that projection is re-applied outside instead.
+          #
+          # The version ordering leads, since it decides which row survives —
+          # any ordering the caller set is kept after it as a tiebreaker.
+          def deduplicated_source(meta)
+            clone(select: nil)
+              .order(Sequel.desc(meta.version), *@opts[:order])
+              .limit_by(*meta.sorting_key.map { |expr| Sequel.lit(expr) })
+          end
+
+          def wrap_source(inner)
+            source = Array(@opts[:from]).first
+            source.is_a?(Sequel::SQL::AliasedExpression) ? inner.as(source.alias) : inner
+          end
+
+          # => [table_name, {} | { db_name: ... }]
+          def deduplication_source
+            sources = Array(@opts[:from])
+            source = sources.first
+            source = source.expression if source.is_a?(Sequel::SQL::AliasedExpression)
+
+            if sources.size != 1 || @opts[:join]
+              raise Sequel::Error, "deduplicate requires a single table source"
+            end
+
+            case source
+            when Sequel::SQL::QualifiedIdentifier
+              [source.column, { db_name: source.table }]
+            when Sequel::SQL::Identifier
+              [source.value, {}]
+            else
+              raise Sequel::Error, "deduplicate requires a single table source"
+            end
+          end
         end
 
         # Concrete backends implement the low-level ops (execute / query /
@@ -36,11 +126,23 @@ def from(source, db_name: self.db_name)
             else
               DB.from(source)
             end
-          ds.clone(ch: true).with_extend(ClickHouseStringEscaping)
+          ds.clone(ch: true).with_extend(ClickHouseDatasetMethods)
         end
 
-        def count(dataset)
-          query_value(dataset.select(SQL.ch_count))
+        def count(dataset, **)
+          query_value(dataset.select(SQL.ch_count), **)
+        end
+
+        # Sorting key / version / is_deleted of a ReplacingMergeTree table.
+        # Distributed tables carry none of these, so they are resolved through
+        # to the local table they wrap. Memoized per process, like the layout
+        # it describes: a table's engine does not change under a running app.
+        def table_metadata(table_name, db_name: self.db_name)
+          key = [db_name.to_s, table_name.to_s]
+          @table_metadata_cache ||= {}
+          return @table_metadata_cache[key] if @table_metadata_cache.key?(key)
+
+          @table_metadata_cache[key] = build_table_metadata(*key)
         end
 
         def db_name
@@ -144,6 +246,43 @@ def with_temp_table(
 
         protected
 
+        # Every read path needs both halves, and pairing them here keeps a new
+        # one from silently reinstating session-wide FINAL inside a dedup
+        # subquery — the exact cost `deduplicate` exists to remove.
+        def prepare_query(dataset, opts)
+          [sql_for(dataset), settings_for(dataset, opts)]
+        end
+
+        # `final` is usually a session-wide default, which would make a
+        # deduplicated query merge the whole table inside its own subquery —
+        # slow and redundant, since the subquery already collapses versions.
+        # An explicit `final:` from the caller always wins.
+        def settings_for(dataset, opts)
+          return opts if opts.key?(:final)
+          return opts unless dataset.is_a?(Sequel::Dataset) && dataset.opts[:ch_dedup]
+
+          opts.merge(final: 0)
+        end
+
+        def build_table_metadata(db_name, table_name)
+          row = query(
+            from(:tables, db_name: :system)
+              .where(database: db_name, name: table_name)
+              .select(:engine, :engine_full, :sorting_key),
+          ).first
+
+          unless row
+            raise ClickHouse::TableMetadata::UnknownTable, "#{db_name}.#{table_name} not found"
+          end
+
+          if row[:engine] == "Distributed"
+            database, table = ClickHouse::TableMetadata.distributed_target(row[:engine_full])
+            return table_metadata(table, db_name: database)
+          end
+
+          ClickHouse::TableMetadata.parse(**row)
+        end
+
         def log_errors(sql)
           yield
         rescue self.class::SERVER_ERROR => e

lib/umbrellio_utils/click_house/backends/legacy.rb

@@ -16,15 +16,15 @@ def execute(sql, host: nil, **opts)
         end
 
         def query(dataset, host: nil, **opts)
-          sql = sql_for(dataset)
+          sql, settings = prepare_query(dataset, opts)
           log_errors(sql) do
-            select_all(sql, host:, **opts).map { |x| Misc::StrictHash[x.symbolize_keys] }
+            select_all(sql, host:, **settings).map { |x| Misc::StrictHash[x.symbolize_keys] }
           end
         end
 
         def query_value(dataset, host: nil, **opts)
-          sql = sql_for(dataset)
-          log_errors(sql) { select_value(sql, host:, **opts) }
+          sql, settings = prepare_query(dataset, opts)
+          log_errors(sql) { select_value(sql, host:, **settings) }
         end
 
         def query_each(dataset, host: nil, **, &)

lib/umbrellio_utils/click_house/backends/native.rb

@@ -29,18 +29,18 @@ def execute(sql, host: nil, **opts) # rubocop:disable Lint/UnusedMethodArgument
         end
 
         def query(dataset, host: nil, **opts) # rubocop:disable Lint/UnusedMethodArgument
-          sql = sql_for(dataset)
-          log_errors(sql) { pool.query(sql, settings: opts) }
+          sql, settings = prepare_query(dataset, opts)
+          log_errors(sql) { pool.query(sql, settings:) }
         end
 
         def query_value(dataset, host: nil, **opts) # rubocop:disable Lint/UnusedMethodArgument
-          sql = sql_for(dataset)
-          log_errors(sql) { pool.query_value(sql, settings: opts) }
+          sql, settings = prepare_query(dataset, opts)
+          log_errors(sql) { pool.query_value(sql, settings:) }
         end
 
         def query_each(dataset, host: nil, **opts, &) # rubocop:disable Lint/UnusedMethodArgument
-          sql = sql_for(dataset)
-          log_errors(sql) { pool.query_each(sql, settings: opts, &) }
+          sql, settings = prepare_query(dataset, opts)
+          log_errors(sql) { pool.query_each(sql, settings:, &) }
         end
 
         def insert(table_name, db_name: self.db_name, rows: [])

lib/umbrellio_utils/click_house/table_metadata.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+module UmbrellioUtils
+  module ClickHouse
+    # Layout of a MergeTree-family table, read out of `system.tables`.
+    #
+    # ReplacingMergeTree collapses rows sharing the full sorting key, keeping
+    # the one with the highest version and dropping it entirely when the
+    # is_deleted column is set. `Dataset#deduplicate` reproduces that by hand,
+    # so it needs all three parts; `version` and `is_deleted` are only
+    # meaningful on a Replacing engine, and nil elsewhere.
+    TableMetadata = Struct.new(:engine, :sorting_key, :version, :is_deleted)
+
+    # Reopened rather than declared with a block: constants defined inside a
+    # `Struct.new do ... end` block leak to the enclosing lexical scope.
+    class TableMetadata
+      class UnknownTable < StandardError
+      end
+
+      REPLICATED_ARGS_COUNT = 2 # zookeeper path + replica name
+
+      def replacing?
+        engine.include?("Replacing")
+      end
+
+      class << self
+        # Built first so it can answer #replacing? for itself — the engine
+        # test lives in one place only.
+        def parse(engine:, engine_full:, sorting_key:)
+          meta = new(engine, split_args(sorting_key), nil, nil)
+          return meta unless meta.replacing?
+
+          args = engine_args(engine_full)
+          args = args.drop(REPLICATED_ARGS_COUNT) if engine.start_with?("Replicated")
+          meta.version, meta.is_deleted = args[0]&.to_sym, args[1]&.to_sym
+
+          meta
+        end
+
+        # Distributed('cluster', 'database', 'table'[, sharding_key])
+        def distributed_target(engine_full)
+          _cluster, database, table = engine_args(engine_full)
+          [unquote(database), unquote(table)]
+        end
+
+        # Arguments of the leading engine call, or [] when the engine takes none.
+        # Only a parenthesis directly after the engine name counts — later
+        # clauses such as `PARTITION BY toYYYYMM(created_at)` must not be read.
+        def engine_args(engine_full)
+          open_index = engine_full.index("(")
+          return [] unless open_index
+          return [] unless engine_full[0...open_index].match?(/\A\w+\z/)
+
+          split_args(engine_full[(open_index + 1)...close_index(engine_full, open_index)])
+        end
+
+        # Split on top-level commas only: sorting keys hold function calls and
+        # engine arguments hold quoted paths, both of which may contain commas.
+        def split_args(source)
+          args = []
+          current = +""
+
+          scan(source) do |char, depth, in_string|
+            if char == "," && depth.zero? && !in_string
+              args << current.strip
+              current = +""
+            else
+              current << char
+            end
+          end
+
+          args << current.strip
+          args.reject(&:empty?)
+        end
+
+        private
+
+        # Single lexer for both scans, so they can't disagree about what counts
+        # as a parenthesis: a ZooKeeper path may legally contain one inside
+        # quotes, and treating it as structure truncates the argument list.
+        def scan(source)
+          depth = 0
+          in_string = false
+
+          source.to_s.each_char.with_index do |char, index|
+            case char
+            when "'" then in_string = !in_string
+            when "(" then depth += 1 unless in_string
+            when ")" then depth -= 1 unless in_string
+            end
+
+            yield(char, depth, in_string, index)
+          end
+        end
+
+        def close_index(source, open_index)
+          scan(source[open_index..]) do |_char, depth, in_string, index|
+            return open_index + index if depth.zero? && !in_string
+          end
+
+          raise ArgumentError, "unbalanced parentheses in engine: #{source}"
+        end
+
+        def unquote(value)
+          value.to_s.delete_prefix("'").delete_suffix("'")
+        end
+      end
+    end
+  end
+end