partition_gardener 0.3.0

data/lib/partition_gardener/strategy/hash_branches.rb

@@ -0,0 +1,161 @@
+module PartitionGardener
+  module Strategy
+    class HashBranches
+      include Naming
+      include RequiresDefaultPartition
+      include CursorColumns
+
+      DEFAULT_MODULUS = 32
+
+      def initialize(config)
+        @config = config
+      end
+
+      def active_window
+        {start: 0, end: modulus}
+      end
+
+      def build_plan
+        heatmap = collect_heatmap(active_window)
+        hot_buckets = hot_buckets_in_window(heatmap, active_window)
+
+        segments = (0...modulus).map do |remainder|
+          hot = hot_buckets.include?(remainder)
+          Plan::Segment.new(
+            name: hash_partition_name(remainder, hot: hot),
+            range_start: {modulus: modulus, remainder: remainder},
+            range_end: nil,
+            kind: hot ? :hot_bucket : :archive
+          )
+        end
+
+        Plan::Result.new(segments: segments, hot_buckets: hot_buckets)
+      end
+
+      def attached_tail_segments
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          next if partition.default
+          next unless hash_partition_name?(partition.name)
+
+          remainder = hash_remainder_from_partition_name(partition.name)
+          next if remainder.nil?
+
+          hot = hot_bucket_partition_name?(partition.name)
+          Plan::Segment.new(
+            name: partition.name,
+            range_start: {modulus: modulus, remainder: remainder},
+            range_end: nil,
+            kind: hot ? :hot_bucket : :archive
+          )
+        end
+      end
+
+      def collect_heatmap(_window)
+        bucket_counts = HashRouting.collect_bucket_counts(@config)
+
+        {bucket_counts: bucket_counts, default_rows: default_row_count}
+      end
+
+      def hot_buckets_in_window(heatmap, _window)
+        heatmap[:bucket_counts].filter_map do |remainder, row_count|
+          remainder if row_count >= split_row_threshold
+        end.sort
+      end
+
+      def managed_tail_partition_names(window: active_window)
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          next if partition.default
+          next unless hash_partition_name?(partition.name)
+
+          partition.name
+        end
+      end
+
+      def tail_slot_name?(_partition_name)
+        false
+      end
+
+      def current_and_future_where_condition(window: active_window)
+        "FALSE"
+      end
+
+      def default_partition_required?
+        false
+      end
+
+      def rebalance_default_drain_where_condition(window: active_window)
+        "FALSE"
+      end
+
+      def default_partition_drain_where_condition(window: active_window)
+        "TRUE"
+      end
+
+      def bucket_where_condition(_bucket)
+        "TRUE"
+      end
+
+      def archive_bucket?(_bucket)
+        false
+      end
+
+      def future_bucket?(_bucket)
+        false
+      end
+
+      def archive_bucket_from_partition_name(partition_name)
+        hash_remainder_from_partition_name(partition_name)
+      end
+
+      def segment_for_values_clause(segment)
+        "WITH (modulus #{segment.range_start[:modulus]}, remainder #{segment.range_start[:remainder]})"
+      end
+
+      private
+
+      def table_name
+        @config[:table_name]
+      end
+
+      def connection
+        Connection.connection
+      end
+
+      def modulus
+        @config.fetch(:hash_modulus, DEFAULT_MODULUS)
+      end
+
+      def split_row_threshold
+        threshold = @config[:split_row_threshold]
+        threshold.nil? ? FUTURE_MONTH_PARTITION_ROW_THRESHOLD : threshold
+      end
+
+      def hash_partition_name(remainder, hot:)
+        prefix = hot ? "h" : "a"
+        format("#{table_name}_#{prefix}_%02d", remainder)
+      end
+
+      def hash_partition_name?(partition_name)
+        partition_name.match?(/^#{Regexp.escape(table_name)}_[ha]_\d+$/)
+      end
+
+      def hot_bucket_partition_name?(partition_name)
+        partition_name.match?(/^#{Regexp.escape(table_name)}_h_\d+$/)
+      end
+
+      def hash_remainder_from_partition_name(partition_name)
+        match = partition_name.match(/^#{Regexp.escape(table_name)}_[ha]_(\d+)$/)
+        return unless match
+
+        match[1].to_i
+      end
+
+      def default_row_count
+        default_name = default_partition_name(table_name)
+        return 0 unless Connection.partition_exists?(default_name)
+
+        Connection.count_rows_in_partition_table(default_name)
+      end
+    end
+  end
+end

data/lib/partition_gardener/strategy/integer_range.rb

@@ -0,0 +1,261 @@
+module PartitionGardener
+  module Strategy
+    class IntegerRange
+      include Naming
+      include RequiresDefaultPartition
+      include CursorColumns
+
+      DEFAULT_ACTIVE_ID_WIDTH = 10_000_000
+      DEFAULT_CURRENT_BAND_SIZE = 1_000_000
+      DEFAULT_ARCHIVE_BAND_SIZE = 10_000_000
+
+      def initialize(config)
+        @config = config
+      end
+
+      def active_window
+        active_lo = @config.fetch(:active_id_lo) { 0 }
+        active_hi = active_lo + active_id_width
+        {start: active_lo, end: active_hi}
+      end
+
+      def build_plan
+        window = active_window
+        heatmap = collect_heatmap(window)
+        hot_buckets = hot_buckets_in_window(heatmap, window)
+
+        segments = Layout::ZoneSegments.build_filler_and_hot_segments(
+          table_name: table_name,
+          buckets: each_bucket_in_range(window[:start], window[:end]),
+          hot_buckets: hot_buckets,
+          active_start: window[:start],
+          active_end: window[:end],
+          hot_bucket_name: method(:hot_bucket_partition_name),
+          bucket_end: ->(bucket) { [bucket + current_band_size, window[:end]].min }
+        )
+
+        Plan::Result.new(segments: segments, hot_buckets: hot_buckets)
+      end
+
+      def attached_tail_segments
+        window = active_window
+
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          next if partition.default
+          next unless managed_tail_partition?(partition.name, window: window)
+
+          Plan::Segment.new(
+            name: partition.name,
+            range_start: partition.range_start,
+            range_end: partition.range_end,
+            kind: segment_kind_for(partition.name, window: window)
+          )
+        end
+      end
+
+      def collect_heatmap(window)
+        bucket_counts = Hash.new(0)
+        dedicated_partition_counts = {}
+        dedicated_partitions = hot_bucket_partitions_in_window(window).to_h
+
+        heatmap_source_partitions(window).each do |partition_name|
+          counts = counts_by_bucket_in_partition(partition_name, window)
+          counts.each { |bucket, count| bucket_counts[bucket] += count }
+
+          bucket = dedicated_partitions[partition_name]
+          dedicated_partition_counts[bucket] = counts.values.sum if bucket
+        end
+
+        {
+          bucket_counts: bucket_counts,
+          default_rows: default_row_count,
+          dedicated_partition_counts: dedicated_partition_counts
+        }
+      end
+
+      def hot_buckets_in_window(heatmap, window)
+        hot_buckets = Set.new
+
+        heatmap[:bucket_counts].each do |bucket, row_count|
+          next unless bucket_in_window?(bucket, window)
+          next unless row_count >= split_row_threshold
+
+          hot_buckets << bucket
+        end
+
+        dedicated_counts = heatmap.fetch(:dedicated_partition_counts, {})
+        hot_bucket_partitions_in_window(window).each do |_partition_name, bucket|
+          row_count = dedicated_counts[bucket]
+          next if row_count.nil?
+
+          if row_count >= split_row_threshold
+            hot_buckets << bucket
+          else
+            hot_buckets.delete(bucket)
+          end
+        end
+
+        hot_buckets.sort
+      end
+
+      def managed_tail_partition_names(window: active_window)
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          next if partition.default
+          next unless managed_tail_partition?(partition.name, window: window)
+
+          partition.name
+        end
+      end
+
+      def tail_slot_name?(partition_name)
+        partition_name == current_partition_name(table_name) ||
+          partition_name == open_partition_name(table_name) ||
+          partition_name == future_partition_name(table_name) ||
+          partition_name.match?(/^#{Regexp.escape(table_name)}_open_\d+$/)
+      end
+
+      def current_and_future_where_condition(window: active_window)
+        partition_key_column = @config[:partition_key_column]
+        "#{partition_key_column} >= #{window[:start]}"
+      end
+
+      def bucket_where_condition(bucket)
+        partition_key_column = @config[:partition_key_column]
+        end_range = bucket + current_band_size
+        "#{partition_key_column} >= #{bucket} AND #{partition_key_column} < #{end_range}"
+      end
+
+      def archive_bucket?(bucket)
+        bucket < active_window[:start]
+      end
+
+      def future_bucket?(bucket)
+        bucket >= active_window[:end]
+      end
+
+      def archive_bucket_from_partition_name(partition_name)
+        match = partition_name.match(/^#{Regexp.escape(table_name)}_ids_(\d+)_(\d+)$/)
+        return unless match
+
+        match[1].to_i
+      end
+
+      def segment_for_values_clause(segment)
+        if segment.range_end == :max
+          "FROM (#{segment.range_start}) TO (MAXVALUE)"
+        else
+          "FROM (#{segment.range_start}) TO (#{segment.range_end})"
+        end
+      end
+
+      def hot_bucket_partition_name(bucket)
+        end_range = [bucket + current_band_size, active_window[:end]].min
+        "#{table_name}_ids_#{bucket}_#{end_range}"
+      end
+
+      def current_band_size
+        @config.fetch(:current_band_size, DEFAULT_CURRENT_BAND_SIZE)
+      end
+
+      private
+
+      def table_name
+        @config[:table_name]
+      end
+
+      def connection
+        Connection.connection
+      end
+
+      def active_id_width
+        @config.fetch(:active_id_width, DEFAULT_ACTIVE_ID_WIDTH)
+      end
+
+      def split_row_threshold
+        @config.fetch(:split_row_threshold, FUTURE_MONTH_PARTITION_ROW_THRESHOLD)
+      end
+
+      def each_bucket_in_range(range_start, range_end_exclusive)
+        buckets = []
+        cursor = range_start
+        while cursor < range_end_exclusive
+          buckets << cursor
+          cursor += current_band_size
+        end
+        buckets
+      end
+
+      def bucket_in_window?(bucket, window)
+        bucket >= window[:start] && bucket < window[:end]
+      end
+
+      def managed_tail_partition?(partition_name, window:)
+        return true if tail_slot_name?(partition_name)
+
+        bucket = archive_bucket_from_partition_name(partition_name)
+        bucket && bucket_in_window?(bucket, window)
+      end
+
+      def segment_kind_for(partition_name, window:)
+        return :future if partition_name == future_partition_name(table_name)
+        return :filler if tail_slot_name?(partition_name)
+
+        :hot_bucket
+      end
+
+      def heatmap_source_partitions(window)
+        names = tail_slot_names_for_heatmap.select { |name| Connection.partition_attached?(table_name, name) }
+        names.concat(hot_bucket_partitions_in_window(window).map(&:first))
+        names.uniq
+      end
+
+      def tail_slot_names_for_heatmap
+        [
+          current_partition_name(table_name),
+          open_partition_name(table_name),
+          future_partition_name(table_name)
+        ] + open_slot_partition_names
+      end
+
+      def open_slot_partition_names
+        Connection.attached_partitions(table_name).map(&:name).select do |partition_name|
+          partition_name.match?(/^#{Regexp.escape(table_name)}_open_\d+$/)
+        end
+      end
+
+      def hot_bucket_partitions_in_window(window)
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          bucket = archive_bucket_from_partition_name(partition.name)
+          next unless bucket
+          next unless bucket_in_window?(bucket, window)
+
+          [partition.name, bucket]
+        end
+      end
+
+      def counts_by_bucket_in_partition(partition_name, window)
+        partition_key_column = @config[:partition_key_column]
+        band = current_band_size
+
+        sql = <<~SQL
+          SELECT (#{connection.quote_column_name(partition_key_column)} / #{band}) * #{band} AS bucket,
+                 COUNT(*)::int AS row_count
+          FROM #{Connection.quoted_table(partition_name)}
+          GROUP BY 1
+        SQL
+
+        connection.execute(sql).each_with_object({}) do |row, counts|
+          bucket = row["bucket"].to_i
+          counts[bucket] = row["row_count"].to_i if bucket_in_window?(bucket, window)
+        end
+      end
+
+      def default_row_count
+        default_name = default_partition_name(table_name)
+        return 0 unless Connection.partition_exists?(default_name)
+
+        Connection.count_rows_in_partition_table(default_name)
+      end
+    end
+  end
+end

data/lib/partition_gardener/strategy/list_split.rb

@@ -0,0 +1,125 @@
+module PartitionGardener
+  module Strategy
+    class ListSplit
+      include Naming
+      include RequiresDefaultPartition
+      include CursorColumns
+
+      def initialize(config)
+        @config = config
+      end
+
+      def active_window
+        {start: nil, end: nil}
+      end
+
+      def build_plan
+        segments = branches.map do |branch|
+          Plan::Segment.new(
+            name: branch_partition_name(branch),
+            range_start: branch.fetch(:value),
+            range_end: nil,
+            kind: :branch
+          )
+        end
+
+        Plan::Result.new(segments: segments, hot_buckets: [])
+      end
+
+      def attached_tail_segments
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          next if partition.default
+
+          branch = branch_for_partition_name(partition.name)
+          next unless branch
+
+          Plan::Segment.new(
+            name: partition.name,
+            range_start: branch.fetch(:value),
+            range_end: nil,
+            kind: :branch
+          )
+        end
+      end
+
+      def collect_heatmap(_window)
+        {bucket_counts: {}, default_rows: default_row_count}
+      end
+
+      def hot_buckets_in_window(_heatmap, _window)
+        []
+      end
+
+      def managed_tail_partition_names
+        branches.map { |branch| branch_partition_name(branch) }
+      end
+
+      def tail_slot_name?(_partition_name)
+        false
+      end
+
+      def current_and_future_where_condition(window: active_window)
+        "TRUE"
+      end
+
+      def rebalance_default_drain_where_condition(window: active_window)
+        "TRUE"
+      end
+
+      def bucket_where_condition(branch_value)
+        branch = branches.find { |entry| entry.fetch(:value) == branch_value }
+        branch.fetch(:where_condition)
+      end
+
+      def archive_bucket?(_bucket)
+        false
+      end
+
+      def future_bucket?(_bucket)
+        false
+      end
+
+      def archive_bucket_from_partition_name(partition_name)
+        branch = branch_for_partition_name(partition_name)
+        branch&.fetch(:value)
+      end
+
+      def segment_for_values_clause(segment)
+        value = segment.range_start
+        formatted = format_list_value(value)
+        "IN (#{formatted})"
+      end
+
+      private
+
+      def table_name
+        @config[:table_name]
+      end
+
+      def branches
+        @config.fetch(:branches)
+      end
+
+      def branch_partition_name(branch)
+        "#{table_name}_#{branch.fetch(:name)}"
+      end
+
+      def branch_for_partition_name(partition_name)
+        branches.find { |branch| branch_partition_name(branch) == partition_name }
+      end
+
+      def format_list_value(value)
+        return "NULL" if value.nil?
+
+        "'#{value}'"
+      end
+
+      def default_row_count
+        default_name = default_partition_name(table_name)
+        return 0 unless Connection.partition_exists?(default_name)
+
+        Connection.count_rows_in_partition_table(default_name)
+      end
+    end
+  end
+end

data/lib/partition_gardener/strategy/requires_default_partition.rb

@@ -0,0 +1,19 @@
+module PartitionGardener
+  module Strategy
+    module RequiresDefaultPartition
+      def default_partition_required?
+        true
+      end
+
+      # Rows in default matching this condition are staged during tail rebalance.
+      def rebalance_default_drain_where_condition(window: active_window)
+        current_and_future_where_condition(window: window)
+      end
+
+      # Rows moved out of default at the end of maintenance (target: empty default).
+      def default_partition_drain_where_condition(window: active_window)
+        current_and_future_where_condition(window: window)
+      end
+    end
+  end
+end

data/lib/partition_gardener/strategy.rb

@@ -0,0 +1,26 @@
+module PartitionGardener
+  module Strategy
+    class << self
+      def for(config)
+        build(config)
+      end
+
+      def build(config)
+        case config.fetch(:layout, :sliding_window)
+        when :integer_window
+          IntegerRange.new(config)
+        when :hash_branches
+          HashBranches.new(config)
+        when :list_split
+          ListSplit.new(config)
+        when :calendar_year
+          DateRange.new(config.merge(bucket: :year))
+        when :premake_monthly, :rolling_current
+          DateRange.new(config)
+        else
+          DateRange.new(config)
+        end
+      end
+    end
+  end
+end