Add UNION optimization for keyset pagination

- Add UNION optimization for keyset pagination
- Add an EachBatch like iterator class

app/models/concerns/from_set_operator.rb

@@ -10,8 +10,8 @@ module FromSetOperator
 
     raise "Trying to redefine method '#{method(method_name)}'" if methods.include?(method_name)
 
-    define_method(method_name) do |members, remove_duplicates: true, alias_as: table_name|
-      operator_sql = operator.new(members, remove_duplicates: remove_duplicates).to_sql
+    define_method(method_name) do |members, remove_duplicates: true, remove_order: true, alias_as: table_name|
+      operator_sql = operator.new(members, remove_duplicates: remove_duplicates, remove_order: remove_order).to_sql
 
       from(Arel.sql("(#{operator_sql}) #{alias_as}"))
     end

app/models/issue.rb

@@ -108,6 +108,7 @@ class Issue < ApplicationRecord
   scope :order_due_date_desc, -> { reorder(::Gitlab::Database.nulls_last_order('due_date', 'DESC')) }
   scope :order_closest_future_date, -> { reorder(Arel.sql('CASE WHEN issues.due_date >= CURRENT_DATE THEN 0 ELSE 1 END ASC, ABS(CURRENT_DATE - issues.due_date) ASC')) }
   scope :order_relative_position_asc, -> { reorder(::Gitlab::Database.nulls_last_order('relative_position', 'ASC')) }
+  scope :order_relative_position_desc, -> { reorder(::Gitlab::Database.nulls_first_order('relative_position', 'DESC')) }
   scope :order_closed_date_desc, -> { reorder(closed_at: :desc) }
   scope :order_created_at_desc, -> { reorder(created_at: :desc) }
   scope :order_severity_asc, -> { includes(:issuable_severity).order('issuable_severities.severity ASC NULLS FIRST') }

lib/gitlab/database.rb

@@ -123,6 +123,16 @@ module Gitlab
       # ignore - happens when Rake tasks yet have to create a database, e.g. for testing
     end
 
+    def self.nulls_order(field, direction = :asc, nulls_order = :nulls_last)
+      raise ArgumentError unless [:nulls_last, :nulls_first].include?(nulls_order)
+      raise ArgumentError unless [:asc, :desc].include?(direction)
+
+      case nulls_order
+      when :nulls_last then nulls_last_order(field, direction)
+      when :nulls_first then nulls_first_order(field, direction)
+      end
+    end
+
     def self.nulls_last_order(field, direction = 'ASC')
       Arel.sql("#{field} #{direction} NULLS LAST")
     end

lib/gitlab/pagination/keyset/iterator.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    module Keyset
+      class Iterator
+        def initialize(scope:, use_union_optimization: false)
+          @scope = scope
+          @order = Gitlab::Pagination::Keyset::Order.extract_keyset_order_object(scope)
+          @use_union_optimization = use_union_optimization
+        end
+
+        # rubocop: disable CodeReuse/ActiveRecord
+        def each_batch(of: 1000)
+          cursor_attributes = {}
+
+          loop do
+            current_scope = scope.dup.limit(of)
+            relation = order
+              .apply_cursor_conditions(current_scope, cursor_attributes, { use_union_optimization: @use_union_optimization })
+              .reorder(order)
+              .limit(of)
+
+            yield relation
+
+            last_record = relation.last
+            break unless last_record
+
+            cursor_attributes = order.cursor_attributes_for_node(last_record)
+          end
+        end
+        # rubocop: enable CodeReuse/ActiveRecord
+
+        private
+
+        attr_reader :scope, :order
+      end
+    end
+  end
+end

lib/gitlab/pagination/keyset/order.rb

@@ -135,7 +135,7 @@ module Gitlab
         #
         # (id < 3 AND created_at IS NULL) OR (created_at IS NOT NULL)
         def build_where_values(values)
-          return if values.blank?
+          return [] if values.blank?
 
           verify_incoming_values!(values)
 
@@ -156,13 +156,26 @@ module Gitlab
             end
           end
 
-          build_or_query(where_values)
+          where_values
+        end
+
+        def where_values_with_or_query(values)
+          build_or_query(build_where_values(values.with_indifferent_access))
         end
 
         # rubocop: disable CodeReuse/ActiveRecord
-        def apply_cursor_conditions(scope, values = {})
+        def apply_cursor_conditions(scope, values = {}, options = { use_union_optimization: false })
+          values ||= {}
+          transformed_values = values.with_indifferent_access
           scope = apply_custom_projections(scope)
-          scope.where(build_where_values(values.with_indifferent_access))
+
+          where_values = build_where_values(transformed_values)
+
+          if options[:use_union_optimization] && where_values.size > 1
+            build_union_query(scope, where_values).reorder(self)
+          else
+            scope.where(build_or_query(where_values)) # rubocop: disable CodeReuse/ActiveRecord
+          end
         end
         # rubocop: enable CodeReuse/ActiveRecord
 
@@ -212,11 +225,19 @@ module Gitlab
         end
 
         def build_or_query(expressions)
-          or_expression = expressions.reduce { |or_expression, expression| Arel::Nodes::Or.new(or_expression, expression) }
+          return [] if expressions.blank?
 
+          or_expression = expressions.reduce { |or_expression, expression| Arel::Nodes::Or.new(or_expression, expression) }
           Arel::Nodes::Grouping.new(or_expression)
         end
 
+        def build_union_query(scope, where_values)
+          scopes = where_values.map do |where_value|
+            scope.dup.where(where_value).reorder(self) # rubocop: disable CodeReuse/ActiveRecord
+          end
+          scope.model.from_union(scopes, remove_duplicates: false, remove_order: false)
+        end
+
         def to_sql_literal(column_definitions)
           column_definitions.map do |column_definition|
             if column_definition.order_expression.respond_to?(:to_sql)

spec/factories/issues.rb

@@ -13,6 +13,14 @@ FactoryBot.define do
       confidential { true }
     end
 
+    trait :with_asc_relative_position do
+      sequence(:relative_position) { |n| n * 1000 }
+    end
+
+    trait :with_desc_relative_position do
+      sequence(:relative_position) { |n| -n * 1000 }
+    end
+
     trait :opened do
       state_id { Issue.available_states[:opened] }
     end

spec/lib/gitlab/pagination/keyset/iterator_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Pagination::Keyset::Iterator do
+  let_it_be(:project) { create(:project) }
+  let_it_be(:issue_list_with_same_pos) { create_list(:issue, 3, project: project, relative_position: 100, updated_at: 1.day.ago) }
+  let_it_be(:issue_list_with_null_pos) { create_list(:issue, 3, project: project, relative_position: nil, updated_at: 1.day.ago) }
+  let_it_be(:issue_list_with_asc_pos) { create_list(:issue, 3, :with_asc_relative_position, project: project, updated_at: 1.day.ago) }
+
+  let(:klass) { Issue }
+  let(:column) { 'relative_position' }
+  let(:direction) { :asc }
+  let(:reverse_direction) { ::Gitlab::Pagination::Keyset::ColumnOrderDefinition::REVERSED_ORDER_DIRECTIONS[direction] }
+  let(:nulls_position) { :nulls_last }
+  let(:reverse_nulls_position) { ::Gitlab::Pagination::Keyset::ColumnOrderDefinition::REVERSED_NULL_POSITIONS[nulls_position] }
+  let(:custom_reorder) do
+    Gitlab::Pagination::Keyset::Order.build([
+      Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+        attribute_name: column,
+         column_expression: klass.arel_table[column],
+         order_expression: ::Gitlab::Database.nulls_order(column, direction, nulls_position),
+         reversed_order_expression: ::Gitlab::Database.nulls_order(column, reverse_direction, reverse_nulls_position),
+         order_direction: direction,
+         nullable: nulls_position,
+         distinct: false
+      ),
+      Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+        attribute_name: 'id',
+         order_expression: klass.arel_table[:id].send(direction),
+         add_to_projections: true
+      )
+    ])
+  end
+
+  let(:scope) { project.issues.reorder(custom_reorder) }
+
+  subject { described_class.new(scope: scope) }
+
+  describe '.each_batch' do
+    it 'yields an ActiveRecord::Relation when a block is given' do
+      subject.each_batch(of: 1) do |relation|
+        expect(relation).to be_a_kind_of(ActiveRecord::Relation)
+      end
+    end
+
+    it 'accepts a custom batch size' do
+      count = 0
+
+      subject.each_batch(of: 2) { |relation| count += relation.count(:all) }
+
+      expect(count).to eq(9)
+    end
+
+    it 'allows updating of the yielded relations' do
+      time = Time.current
+
+      subject.each_batch(of: 2) do |relation|
+        relation.update_all(updated_at: time)
+      end
+
+      expect(Issue.where(updated_at: time).count).to eq(9)
+    end
+
+    context 'with ordering direction' do
+      context 'when ordering asc' do
+        it 'orders ascending by default, including secondary order column' do
+          positions = []
+
+          subject.each_batch(of: 2) { |rel| positions.concat(rel.pluck(:relative_position, :id)) }
+
+          expect(positions).to eq(project.issues.order_relative_position_asc.order(id: :asc).pluck(:relative_position, :id))
+        end
+      end
+
+      context 'when reversing asc order' do
+        let(:scope) { project.issues.order(custom_reorder.reversed_order) }
+
+        it 'orders in reverse of ascending' do
+          positions = []
+
+          subject.each_batch(of: 2) { |rel| positions.concat(rel.pluck(:relative_position, :id)) }
+
+          expect(positions).to eq(project.issues.order_relative_position_desc.order(id: :desc).pluck(:relative_position, :id))
+        end
+      end
+
+      context 'when asc order, with nulls first' do
+        let(:nulls_position) { :nulls_first }
+
+        it 'orders ascending with nulls first' do
+          positions = []
+
+          subject.each_batch(of: 2) { |rel| positions.concat(rel.pluck(:relative_position, :id)) }
+
+          expect(positions).to eq(project.issues.reorder(::Gitlab::Database.nulls_first_order('relative_position', 'ASC')).order(id: :asc).pluck(:relative_position, :id))
+        end
+      end
+
+      context 'when ordering desc' do
+        let(:direction) { :desc }
+        let(:nulls_position) { :nulls_last }
+
+        it 'orders descending' do
+          positions = []
+
+          subject.each_batch(of: 2) { |rel| positions.concat(rel.pluck(:relative_position, :id)) }
+
+          expect(positions).to eq(project.issues.reorder(::Gitlab::Database.nulls_last_order('relative_position', 'DESC')).order(id: :desc).pluck(:relative_position, :id))
+        end
+      end
+
+      context 'when ordering by columns are repeated twice' do
+        let(:direction) { :desc }
+        let(:column) { :id }
+
+        it 'orders descending' do
+          positions = []
+
+          subject.each_batch(of: 2) { |rel| positions.concat(rel.pluck(:id)) }
+
+          expect(positions).to eq(project.issues.reorder(id: :desc).pluck(:id))
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/pagination/keyset/order_spec.rb

@@ -3,76 +3,77 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Pagination::Keyset::Order do
-  let(:table) { Arel::Table.new(:my_table) }
-  let(:order) { nil }
+  describe 'paginate over items correctly' do
+    let(:table) { Arel::Table.new(:my_table) }
+    let(:order) { nil }
 
-  def run_query(query)
-    ActiveRecord::Base.connection.execute(query).to_a
-  end
+    def run_query(query)
+      ActiveRecord::Base.connection.execute(query).to_a
+    end
 
-  def build_query(order:, where_conditions: nil, limit: nil)
-    <<-SQL
+    def build_query(order:, where_conditions: nil, limit: nil)
+      <<-SQL
     SELECT id, year, month
       FROM (#{table_data}) my_table (id, year, month) 
       WHERE #{where_conditions || '1=1'} 
       ORDER BY #{order} 
       LIMIT #{limit || 999};
-    SQL
-  end
-
-  def iterate_and_collect(order:, page_size:, where_conditions: nil)
-    all_items = []
-
-    loop do
-      paginated_items = run_query(build_query(order: order, where_conditions: where_conditions, limit: page_size))
-      break if paginated_items.empty?
-
-      all_items.concat(paginated_items)
-      last_item = paginated_items.last
-      cursor_attributes = order.cursor_attributes_for_node(last_item)
-      where_conditions = order.build_where_values(cursor_attributes).to_sql
+      SQL
     end
 
-    all_items
-  end
+    def iterate_and_collect(order:, page_size:, where_conditions: nil)
+      all_items = []
 
-  subject do
-    run_query(build_query(order: order))
-  end
+      loop do
+        paginated_items = run_query(build_query(order: order, where_conditions: where_conditions, limit: page_size))
+        break if paginated_items.empty?
+
+        all_items.concat(paginated_items)
+        last_item = paginated_items.last
+        cursor_attributes = order.cursor_attributes_for_node(last_item)
+        where_conditions = order.where_values_with_or_query(cursor_attributes).to_sql
+      end
 
-  shared_examples 'order examples' do
-    it { expect(subject).to eq(expected) }
+      all_items
+    end
 
-    context 'when paginating forwards' do
-      subject { iterate_and_collect(order: order, page_size: 2) }
+    subject do
+      run_query(build_query(order: order))
+    end
 
+    shared_examples 'order examples' do
       it { expect(subject).to eq(expected) }
 
-      context 'with different page size' do
-        subject { iterate_and_collect(order: order, page_size: 5) }
+      context 'when paginating forwards' do
+        subject { iterate_and_collect(order: order, page_size: 2) }
 
         it { expect(subject).to eq(expected) }
-      end
-    end
 
-    context 'when paginating backwards' do
-      subject do
-        last_item = expected.last
-        cursor_attributes = order.cursor_attributes_for_node(last_item)
-        where_conditions = order.reversed_order.build_where_values(cursor_attributes)
+        context 'with different page size' do
+          subject { iterate_and_collect(order: order, page_size: 5) }
 
-        iterate_and_collect(order: order.reversed_order, page_size: 2, where_conditions: where_conditions.to_sql)
+          it { expect(subject).to eq(expected) }
+        end
       end
 
-      it do
-        expect(subject).to eq(expected.reverse[1..-1]) # removing one item because we used it to calculate cursor data for the "last" page in subject
+      context 'when paginating backwards' do
+        subject do
+          last_item = expected.last
+          cursor_attributes = order.cursor_attributes_for_node(last_item)
+          where_conditions = order.reversed_order.where_values_with_or_query(cursor_attributes)
+
+          iterate_and_collect(order: order.reversed_order, page_size: 2, where_conditions: where_conditions.to_sql)
+        end
+
+        it do
+          expect(subject).to eq(expected.reverse[1..-1]) # removing one item because we used it to calculate cursor data for the "last" page in subject
+        end
       end
     end
-  end
 
-  context 'when ordering by a distinct column' do
-    let(:table_data) do
-      <<-SQL
+    context 'when ordering by a distinct column' do
+      let(:table_data) do
+        <<-SQL
       VALUES (1,  0, 0),
              (2,  0, 0),
              (3,  0, 0),
@@ -82,41 +83,41 @@ RSpec.describe Gitlab::Pagination::Keyset::Order do
              (7,  0, 0),
              (8,  0, 0),
              (9,  0, 0)
-      SQL
-    end
+        SQL
+      end
 
-    let(:order) do
-      Gitlab::Pagination::Keyset::Order.build([
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'id',
-          column_expression: table['id'],
-          order_expression: table['id'].desc,
-          nullable: :not_nullable,
-          distinct: true
-        )
-      ])
-    end
+      let(:order) do
+        Gitlab::Pagination::Keyset::Order.build([
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'id',
+            column_expression: table['id'],
+            order_expression: table['id'].desc,
+            nullable: :not_nullable,
+            distinct: true
+          )
+        ])
+      end
 
-    let(:expected) do
-      [
-        { "id" => 9, "year" => 0, "month" => 0 },
-        { "id" => 8, "year" => 0, "month" => 0 },
-        { "id" => 7, "year" => 0, "month" => 0 },
-        { "id" => 6, "year" => 0, "month" => 0 },
-        { "id" => 5, "year" => 0, "month" => 0 },
-        { "id" => 4, "year" => 0, "month" => 0 },
-        { "id" => 3, "year" => 0, "month" => 0 },
-        { "id" => 2, "year" => 0, "month" => 0 },
-        { "id" => 1, "year" => 0, "month" => 0 }
-      ]
-    end
+      let(:expected) do
+        [
+          { "id" => 9, "year" => 0, "month" => 0 },
+          { "id" => 8, "year" => 0, "month" => 0 },
+          { "id" => 7, "year" => 0, "month" => 0 },
+          { "id" => 6, "year" => 0, "month" => 0 },
+          { "id" => 5, "year" => 0, "month" => 0 },
+          { "id" => 4, "year" => 0, "month" => 0 },
+          { "id" => 3, "year" => 0, "month" => 0 },
+          { "id" => 2, "year" => 0, "month" => 0 },
+          { "id" => 1, "year" => 0, "month" => 0 }
+        ]
+      end
 
-    it_behaves_like 'order examples'
-  end
+      it_behaves_like 'order examples'
+    end
 
-  context 'when ordering by two non-nullable columns and a distinct column' do
-    let(:table_data) do
-      <<-SQL
+    context 'when ordering by two non-nullable columns and a distinct column' do
+      let(:table_data) do
+        <<-SQL
       VALUES (1,  2010, 2),
              (2,  2011, 1),
              (3,  2009, 2),
@@ -126,55 +127,55 @@ RSpec.describe Gitlab::Pagination::Keyset::Order do
              (7,  2010, 3),
              (8,  2012, 4),
              (9,  2013, 5)
-      SQL
-    end
+        SQL
+      end
 
-    let(:order) do
-      Gitlab::Pagination::Keyset::Order.build([
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'year',
-          column_expression: table['year'],
-          order_expression: table['year'].asc,
-          nullable: :not_nullable,
-          distinct: false
-        ),
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'month',
-          column_expression: table['month'],
-          order_expression: table['month'].asc,
-          nullable: :not_nullable,
-          distinct: false
-        ),
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'id',
-          column_expression: table['id'],
-          order_expression: table['id'].asc,
-          nullable: :not_nullable,
-          distinct: true
-        )
-      ])
-    end
+      let(:order) do
+        Gitlab::Pagination::Keyset::Order.build([
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'year',
+            column_expression: table['year'],
+            order_expression: table['year'].asc,
+            nullable: :not_nullable,
+            distinct: false
+          ),
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'month',
+            column_expression: table['month'],
+            order_expression: table['month'].asc,
+            nullable: :not_nullable,
+            distinct: false
+          ),
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'id',
+            column_expression: table['id'],
+            order_expression: table['id'].asc,
+            nullable: :not_nullable,
+            distinct: true
+          )
+        ])
+      end
 
-    let(:expected) do
-      [
-        { 'year' => 2009, 'month' => 2, 'id' => 3 },
-        { 'year' => 2009, 'month' => 2, 'id' => 6 },
-        { 'year' => 2010, 'month' => 2, 'id' => 1 },
-        { 'year' => 2010, 'month' => 3, 'id' => 7 },
-        { 'year' => 2011, 'month' => 1, 'id' => 2 },
-        { 'year' => 2011, 'month' => 1, 'id' => 4 },
-        { 'year' => 2011, 'month' => 1, 'id' => 5 },
-        { 'year' => 2012, 'month' => 4, 'id' => 8 },
-        { 'year' => 2013, 'month' => 5, 'id' => 9 }
-      ]
-    end
+      let(:expected) do
+        [
+          { 'year' => 2009, 'month' => 2, 'id' => 3 },
+          { 'year' => 2009, 'month' => 2, 'id' => 6 },
+          { 'year' => 2010, 'month' => 2, 'id' => 1 },
+          { 'year' => 2010, 'month' => 3, 'id' => 7 },
+          { 'year' => 2011, 'month' => 1, 'id' => 2 },
+          { 'year' => 2011, 'month' => 1, 'id' => 4 },
+          { 'year' => 2011, 'month' => 1, 'id' => 5 },
+          { 'year' => 2012, 'month' => 4, 'id' => 8 },
+          { 'year' => 2013, 'month' => 5, 'id' => 9 }
+        ]
+      end
 
-    it_behaves_like 'order examples'
-  end
+      it_behaves_like 'order examples'
+    end
 
-  context 'when ordering by nullable columns and a distinct column' do
-    let(:table_data) do
-      <<-SQL
+    context 'when ordering by nullable columns and a distinct column' do
+      let(:table_data) do
+        <<-SQL
       VALUES (1,  2010, null),
              (2,  2011, 2),
              (3,  null, null),
@@ -186,61 +187,61 @@ RSpec.describe Gitlab::Pagination::Keyset::Order do
              (9,  null, 2),
              (10, null, null),
              (11, 2010, 2)
-      SQL
-    end
+        SQL
+      end
 
-    let(:order) do
-      Gitlab::Pagination::Keyset::Order.build([
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'year',
-          column_expression: table['year'],
-          order_expression: Gitlab::Database.nulls_last_order('year', :asc),
-          reversed_order_expression: Gitlab::Database.nulls_first_order('year', :desc),
-          order_direction: :asc,
-          nullable: :nulls_last,
-          distinct: false
-        ),
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'month',
-          column_expression: table['month'],
-          order_expression: Gitlab::Database.nulls_last_order('month', :asc),
-          reversed_order_expression: Gitlab::Database.nulls_first_order('month', :desc),
-          order_direction: :asc,
-          nullable: :nulls_last,
-          distinct: false
-        ),
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'id',
-          column_expression: table['id'],
-          order_expression: table['id'].asc,
-          nullable: :not_nullable,
-          distinct: true
-        )
-      ])
-    end
+      let(:order) do
+        Gitlab::Pagination::Keyset::Order.build([
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'year',
+            column_expression: table['year'],
+            order_expression: Gitlab::Database.nulls_last_order('year', :asc),
+            reversed_order_expression: Gitlab::Database.nulls_first_order('year', :desc),
+            order_direction: :asc,
+            nullable: :nulls_last,
+            distinct: false
+          ),
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'month',
+            column_expression: table['month'],
+            order_expression: Gitlab::Database.nulls_last_order('month', :asc),
+            reversed_order_expression: Gitlab::Database.nulls_first_order('month', :desc),
+            order_direction: :asc,
+            nullable: :nulls_last,
+            distinct: false
+          ),
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'id',
+            column_expression: table['id'],
+            order_expression: table['id'].asc,
+            nullable: :not_nullable,
+            distinct: true
+          )
+        ])
+      end
 
-    let(:expected) do
-      [
-        { "id" => 7, "year" => 2010, "month" => 2 },
-        { "id" => 11, "year" => 2010, "month" => 2 },
-        { "id" => 1, "year" => 2010, "month" => nil },
-        { "id" => 5, "year" => 2010, "month" => nil },
-        { "id" => 2, "year" => 2011, "month" => 2 },
-        { "id" => 6, "year" => 2011, "month" => 2 },
-        { "id" => 8, "year" => 2012, "month" => 2 },
-        { "id" => 9, "year" => nil, "month" => 2 },
-        { "id" => 4, "year" => nil, "month" => 5 },
-        { "id" => 3, "year" => nil, "month" => nil },
-        { "id" => 10, "year" => nil, "month" => nil }
-      ]
-    end
+      let(:expected) do
+        [
+          { "id" => 7, "year" => 2010, "month" => 2 },
+          { "id" => 11, "year" => 2010, "month" => 2 },
+          { "id" => 1, "year" => 2010, "month" => nil },
+          { "id" => 5, "year" => 2010, "month" => nil },
+          { "id" => 2, "year" => 2011, "month" => 2 },
+          { "id" => 6, "year" => 2011, "month" => 2 },
+          { "id" => 8, "year" => 2012, "month" => 2 },
+          { "id" => 9, "year" => nil, "month" => 2 },
+          { "id" => 4, "year" => nil, "month" => 5 },
+          { "id" => 3, "year" => nil, "month" => nil },
+          { "id" => 10, "year" => nil, "month" => nil }
+        ]
+      end
 
-    it_behaves_like 'order examples'
-  end
+      it_behaves_like 'order examples'
+    end
 
-  context 'when ordering by nullable columns with nulls first ordering and a distinct column' do
-    let(:table_data) do
-      <<-SQL
+    context 'when ordering by nullable columns with nulls first ordering and a distinct column' do
+      let(:table_data) do
+        <<-SQL
       VALUES (1,  2010, null),
              (2,  2011, 2),
              (3,  null, null),
@@ -252,61 +253,61 @@ RSpec.describe Gitlab::Pagination::Keyset::Order do
              (9,  null, 2),
              (10, null, null),
              (11, 2010, 2)
-      SQL
-    end
+        SQL
+      end
 
-    let(:order) do
-      Gitlab::Pagination::Keyset::Order.build([
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'year',
-          column_expression: table['year'],
-          order_expression: Gitlab::Database.nulls_first_order('year', :asc),
-          reversed_order_expression: Gitlab::Database.nulls_last_order('year', :desc),
-          order_direction: :asc,
-          nullable: :nulls_first,
-          distinct: false
-        ),
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'month',
-          column_expression: table['month'],
-          order_expression: Gitlab::Database.nulls_first_order('month', :asc),
-          order_direction: :asc,
-          reversed_order_expression: Gitlab::Database.nulls_last_order('month', :desc),
-          nullable: :nulls_first,
-          distinct: false
-        ),
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'id',
-          column_expression: table['id'],
-          order_expression: table['id'].asc,
-          nullable: :not_nullable,
-          distinct: true
-        )
-      ])
-    end
+      let(:order) do
+        Gitlab::Pagination::Keyset::Order.build([
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'year',
+            column_expression: table['year'],
+            order_expression: Gitlab::Database.nulls_first_order('year', :asc),
+            reversed_order_expression: Gitlab::Database.nulls_last_order('year', :desc),
+            order_direction: :asc,
+            nullable: :nulls_first,
+            distinct: false
+          ),
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'month',
+            column_expression: table['month'],
+            order_expression: Gitlab::Database.nulls_first_order('month', :asc),
+            order_direction: :asc,
+            reversed_order_expression: Gitlab::Database.nulls_last_order('month', :desc),
+            nullable: :nulls_first,
+            distinct: false
+          ),
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'id',
+            column_expression: table['id'],
+            order_expression: table['id'].asc,
+            nullable: :not_nullable,
+            distinct: true
+          )
+        ])
+      end
 
-    let(:expected) do
-      [
-        { "id" => 3, "year" => nil, "month" => nil },
-        { "id" => 10, "year" => nil, "month" => nil },
-        { "id" => 9, "year" => nil, "month" => 2 },
-        { "id" => 4, "year" => nil, "month" => 5 },
-        { "id" => 1, "year" => 2010, "month" => nil },
-        { "id" => 5, "year" => 2010, "month" => nil },
-        { "id" => 7, "year" => 2010, "month" => 2 },
-        { "id" => 11, "year" => 2010, "month" => 2 },
-        { "id" => 2, "year" => 2011, "month" => 2 },
-        { "id" => 6, "year" => 2011, "month" => 2 },
-        { "id" => 8, "year" => 2012, "month" => 2 }
-      ]
-    end
+      let(:expected) do
+        [
+          { "id" => 3, "year" => nil, "month" => nil },
+          { "id" => 10, "year" => nil, "month" => nil },
+          { "id" => 9, "year" => nil, "month" => 2 },
+          { "id" => 4, "year" => nil, "month" => 5 },
+          { "id" => 1, "year" => 2010, "month" => nil },
+          { "id" => 5, "year" => 2010, "month" => nil },
+          { "id" => 7, "year" => 2010, "month" => 2 },
+          { "id" => 11, "year" => 2010, "month" => 2 },
+          { "id" => 2, "year" => 2011, "month" => 2 },
+          { "id" => 6, "year" => 2011, "month" => 2 },
+          { "id" => 8, "year" => 2012, "month" => 2 }
+        ]
+      end
 
-    it_behaves_like 'order examples'
-  end
+      it_behaves_like 'order examples'
+    end
 
-  context 'when ordering by non-nullable columns with mixed directions and a distinct column' do
-    let(:table_data) do
-      <<-SQL
+    context 'when ordering by non-nullable columns with mixed directions and a distinct column' do
+      let(:table_data) do
+        <<-SQL
       VALUES (1,  2010, 0),
              (2,  2011, 0),
              (3,  2010, 0),
@@ -318,158 +319,216 @@ RSpec.describe Gitlab::Pagination::Keyset::Order do
              (9,  2013, 0),
              (10, 2014, 0),
              (11, 2013, 0)
-      SQL
-    end
+        SQL
+      end
 
-    let(:order) do
-      Gitlab::Pagination::Keyset::Order.build([
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'year',
-          column_expression: table['year'],
-          order_expression: table['year'].asc,
-          nullable: :not_nullable,
-          distinct: false
-        ),
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'id',
-          column_expression: table['id'],
-          order_expression: table['id'].desc,
-          nullable: :not_nullable,
-          distinct: true
-        )
-      ])
-    end
+      let(:order) do
+        Gitlab::Pagination::Keyset::Order.build([
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'year',
+            column_expression: table['year'],
+            order_expression: table['year'].asc,
+            nullable: :not_nullable,
+            distinct: false
+          ),
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'id',
+            column_expression: table['id'],
+            order_expression: table['id'].desc,
+            nullable: :not_nullable,
+            distinct: true
+          )
+        ])
+      end
 
-    let(:expected) do
-      [
-        { "id" => 7, "year" => 2010, "month" => 0 },
-        { "id" => 4, "year" => 2010, "month" => 0 },
-        { "id" => 3, "year" => 2010, "month" => 0 },
-        { "id" => 1, "year" => 2010, "month" => 0 },
-        { "id" => 8, "year" => 2011, "month" => 0 },
-        { "id" => 2, "year" => 2011, "month" => 0 },
-        { "id" => 6, "year" => 2012, "month" => 0 },
-        { "id" => 5, "year" => 2012, "month" => 0 },
-        { "id" => 11, "year" => 2013, "month" => 0 },
-        { "id" => 9, "year" => 2013, "month" => 0 },
-        { "id" => 10, "year" => 2014, "month" => 0 }
-      ]
-    end
+      let(:expected) do
+        [
+          { "id" => 7, "year" => 2010, "month" => 0 },
+          { "id" => 4, "year" => 2010, "month" => 0 },
+          { "id" => 3, "year" => 2010, "month" => 0 },
+          { "id" => 1, "year" => 2010, "month" => 0 },
+          { "id" => 8, "year" => 2011, "month" => 0 },
+          { "id" => 2, "year" => 2011, "month" => 0 },
+          { "id" => 6, "year" => 2012, "month" => 0 },
+          { "id" => 5, "year" => 2012, "month" => 0 },
+          { "id" => 11, "year" => 2013, "month" => 0 },
+          { "id" => 9, "year" => 2013, "month" => 0 },
+          { "id" => 10, "year" => 2014, "month" => 0 }
+        ]
+      end
 
-    it 'takes out a slice between two cursors' do
-      after_cursor = { "id" => 8, "year" => 2011 }
-      before_cursor = { "id" => 5, "year" => 2012 }
+      it 'takes out a slice between two cursors' do
+        after_cursor = { "id" => 8, "year" => 2011 }
+        before_cursor = { "id" => 5, "year" => 2012 }
 
-      after_conditions = order.build_where_values(after_cursor)
-      reversed = order.reversed_order
-      before_conditions = reversed.build_where_values(before_cursor)
+        after_conditions = order.where_values_with_or_query(after_cursor)
+        reversed = order.reversed_order
+        before_conditions = reversed.where_values_with_or_query(before_cursor)
 
-      query = build_query(order: order, where_conditions: "(#{after_conditions.to_sql}) AND (#{before_conditions.to_sql})", limit: 100)
+        query = build_query(order: order, where_conditions: "(#{after_conditions.to_sql}) AND (#{before_conditions.to_sql})", limit: 100)
 
-      expect(run_query(query)).to eq([
-        { "id" => 2, "year" => 2011, "month" => 0 },
-        { "id" => 6, "year" => 2012, "month" => 0 }
-      ])
+        expect(run_query(query)).to eq([
+          { "id" => 2, "year" => 2011, "month" => 0 },
+          { "id" => 6, "year" => 2012, "month" => 0 }
+        ])
+      end
     end
-  end
 
-  context 'when the passed cursor values do not match with the order definition' do
-    let(:order) do
-      Gitlab::Pagination::Keyset::Order.build([
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'year',
-          column_expression: table['year'],
-          order_expression: table['year'].asc,
-          nullable: :not_nullable,
-          distinct: false
-        ),
-        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-          attribute_name: 'id',
-          column_expression: table['id'],
-          order_expression: table['id'].desc,
-          nullable: :not_nullable,
-          distinct: true
-        )
-      ])
-    end
+    context 'when the passed cursor values do not match with the order definition' do
+      let(:order) do
+        Gitlab::Pagination::Keyset::Order.build([
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'year',
+            column_expression: table['year'],
+            order_expression: table['year'].asc,
+            nullable: :not_nullable,
+            distinct: false
+          ),
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: 'id',
+            column_expression: table['id'],
+            order_expression: table['id'].desc,
+            nullable: :not_nullable,
+            distinct: true
+          )
+        ])
+      end
 
-    context 'when values are missing' do
-      it 'raises error'  do
-        expect { order.build_where_values(id: 1) }.to raise_error(/Missing items: year/)
+      context 'when values are missing' do
+        it 'raises error'  do
+          expect { order.build_where_values(id: 1) }.to raise_error(/Missing items: year/)
+        end
       end
-    end
 
-    context 'when extra values are present' do
-      it 'raises error' do
-        expect { order.build_where_values(id: 1, year: 2, foo: 3) }.to raise_error(/Extra items: foo/)
+      context 'when extra values are present' do
+        it 'raises error' do
+          expect { order.build_where_values(id: 1, year: 2, foo: 3) }.to raise_error(/Extra items: foo/)
+        end
       end
-    end
 
-    context 'when values are missing and extra values are present' do
-      it 'raises error'  do
-        expect { order.build_where_values(year: 2, foo: 3) }.to raise_error(/Extra items: foo\. Missing items: id/)
+      context 'when values are missing and extra values are present' do
+        it 'raises error'  do
+          expect { order.build_where_values(year: 2, foo: 3) }.to raise_error(/Extra items: foo\. Missing items: id/)
+        end
       end
-    end
 
-    context 'when no values are passed' do
-      it 'returns nil' do
-        expect(order.build_where_values({})).to eq(nil)
+      context 'when no values are passed' do
+        it 'returns empty array' do
+          expect(order.build_where_values({})).to eq([])
+        end
       end
     end
-  end
 
-  context 'extract and apply cursor attributes' do
-    let(:model) { Project.new(id: 100) }
-    let(:scope) { Project.all }
+    context 'extract and apply cursor attributes' do
+      let(:model) { Project.new(id: 100) }
+      let(:scope) { Project.all }
 
-    shared_examples 'cursor attribute examples' do
-      describe '#cursor_attributes_for_node' do
-        it { expect(order.cursor_attributes_for_node(model)).to eq({ id: '100' }.with_indifferent_access) }
-      end
+      shared_examples 'cursor attribute examples' do
+        describe '#cursor_attributes_for_node' do
+          it { expect(order.cursor_attributes_for_node(model)).to eq({ id: '100' }.with_indifferent_access) }
+        end
+
+        describe '#apply_cursor_conditions' do
+          context 'when params with string keys are passed' do
+            subject(:sql) { order.apply_cursor_conditions(scope, { 'id' => '100' }).to_sql }
 
-      describe '#apply_cursor_conditions' do
-        context 'when params with string keys are passed' do
-          subject(:sql) { order.apply_cursor_conditions(scope, { 'id' => '100' }).to_sql }
+            it { is_expected.to include('"projects"."id" < 100)') }
+          end
 
-          it { is_expected.to include('"projects"."id" < 100)') }
+          context 'when params with symbol keys are passed' do
+            subject(:sql) { order.apply_cursor_conditions(scope, { id: '100' }).to_sql }
+
+            it { is_expected.to include('"projects"."id" < 100)') }
+          end
         end
+      end
 
-        context 'when params with symbol keys are passed' do
-          subject(:sql) { order.apply_cursor_conditions(scope, { id: '100' }).to_sql }
+      context 'when string attribute name is given' do
+        let(:order) do
+          Gitlab::Pagination::Keyset::Order.build([
+            Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+              attribute_name: 'id',
+              order_expression: Project.arel_table['id'].desc,
+              nullable: :not_nullable,
+              distinct: true
+            )
+          ])
+        end
 
-          it { is_expected.to include('"projects"."id" < 100)') }
+        it_behaves_like 'cursor attribute examples'
+      end
+
+      context 'when symbol attribute name is given' do
+        let(:order) do
+          Gitlab::Pagination::Keyset::Order.build([
+            Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+              attribute_name: :id,
+              order_expression: Project.arel_table['id'].desc,
+              nullable: :not_nullable,
+              distinct: true
+            )
+          ])
         end
+
+        it_behaves_like 'cursor attribute examples'
       end
     end
+  end
 
-    context 'when string attribute name is given' do
-      let(:order) do
-        Gitlab::Pagination::Keyset::Order.build([
-          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-            attribute_name: 'id',
-            order_expression: Project.arel_table['id'].desc,
-            nullable: :not_nullable,
-            distinct: true
-          )
-        ])
+  describe 'UNION optimization' do
+    let_it_be(:five_months_ago) { 5.months.ago }
+
+    let_it_be(:user_1) { create(:user, created_at: five_months_ago) }
+    let_it_be(:user_2) { create(:user, created_at: five_months_ago) }
+    let_it_be(:user_3) { create(:user, created_at: 1.month.ago) }
+    let_it_be(:user_4) { create(:user, created_at: 2.months.ago) }
+
+    let(:expected_results) { [user_3, user_4, user_2, user_1] }
+    let(:scope) { User.order(created_at: :desc, id: :desc) }
+    let(:keyset_aware_scope) { Gitlab::Pagination::Keyset::SimpleOrderBuilder.build(scope).first }
+    let(:iterator_options) { { scope: keyset_aware_scope } }
+
+    subject(:items) do
+      [].tap do |collector|
+        Gitlab::Pagination::Keyset::Iterator.new(**iterator_options).each_batch(of: 2) do |models|
+          collector.concat(models)
+        end
       end
+    end
 
-      it_behaves_like 'cursor attribute examples'
+    context 'when UNION optimization is off' do
+      it 'returns items in the correct order' do
+        iterator_options[:use_union_optimization] = false
+
+        expect(items).to eq(expected_results)
+      end
     end
 
-    context 'when symbol attribute name is given' do
-      let(:order) do
-        Gitlab::Pagination::Keyset::Order.build([
-          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
-            attribute_name: :id,
-            order_expression: Project.arel_table['id'].desc,
-            nullable: :not_nullable,
-            distinct: true
-          )
-        ])
+    context 'when UNION optimization is on' do
+      before do
+        iterator_options[:use_union_optimization] = true
       end
 
-      it_behaves_like 'cursor attribute examples'
+      it 'returns items in the correct order' do
+        expect(items).to eq(expected_results)
+      end
+
+      it 'calls Gitlab::SQL::Union' do
+        expect_next_instances_of(Gitlab::SQL::Union, 2) do |instance|
+          expect(instance.send(:remove_order)).to eq(false) # Do not remove order from the queries
+          expect(instance.send(:remove_duplicates)).to eq(false) # Do not deduplicate the results
+        end
+
+        items
+      end
+
+      it 'builds UNION query' do
+        cursor_attributes = { created_at: five_months_ago, id: user_2.id }
+        order = Gitlab::Pagination::Keyset::Order.extract_keyset_order_object(keyset_aware_scope)
+
+        query = order.apply_cursor_conditions(scope, cursor_attributes, use_union_optimization: true).to_sql
+        expect(query).to include('UNION ALL')
+      end
     end
   end
 end