Add Dataset#as_set and Dataset#select_set, as Set will be a core class in Ruby 4

Author: jeremyevans

CHANGELOG

@@ -1,5 +1,7 @@
 === master
 
+* Add Dataset#as_set and Dataset#select_set, as Set will be a core class in Ruby 4 (jeremyevans)
+
 * Add support for Database#create_index :only option on PostgreSQL 11+ (KirIgor) (#2343)
 
 * Add support for Database#indexes :invalid option on PostgresSQL (KirIgor, jeremyevans) (#2343)

lib/sequel/core.rb

@@ -1,6 +1,6 @@
 # frozen-string-literal: true
 
-%w'bigdecimal date thread time uri'.each{|f| require f}
+%w'bigdecimal date set thread time uri'.each{|f| require f}
 
 # Top level module for Sequel
 #

lib/sequel/dataset/actions.rb

@@ -11,11 +11,12 @@ class Dataset
 
     # Action methods defined by Sequel that execute code on the database.
     ACTION_METHODS = (<<-METHS).split.map(&:to_sym).freeze
-      << [] all as_hash avg count columns columns! delete each
+      << [] all as_hash as_set avg count columns columns! delete each
       empty? fetch_rows first first! get import insert last
-      map max min multi_insert paged_each select_hash select_hash_groups select_map select_order_map
-      single_record single_record! single_value single_value! sum to_hash to_hash_groups truncate update
-      where_all where_each where_single_value
+      map max min multi_insert paged_each select_hash select_hash_groups
+      select_map select_order_map select_set single_record single_record!
+      single_value single_value! sum to_hash to_hash_groups
+      truncate update where_all where_each where_single_value
     METHS
 
     # The clone options to use when retrieving columns for a dataset.
@@ -51,6 +52,26 @@ def all(&block)
       _all(block){|a| each{|r| a << r}}
     end
 
+    # Returns sets for column values for each record in the dataset.
+    #
+    #   DB[:table].as_set(:id) # SELECT * FROM table
+    #   # => Set[1, 2, 3, ...]
+    #
+    # You can also provide an array of column names, in which case the elements
+    # of the returned set are arrays (not sets):
+    #
+    #   DB[:table].as_set([:id, :name]) # SELECT * FROM table
+    #   # => Set[[1, 'A'], [2, 'B'], [3, 'C'], ...]
+    def as_set(column)
+      return naked.as_set(column) if row_proc
+
+      if column.is_a?(Array)
+        to_set{|r| r.values_at(*column)}
+      else
+        to_set{|r| r[column]}
+      end
+    end
+
     # Returns the average value for the given column/expression.
     # Uses a virtual row block if no argument is given.
     #
@@ -765,6 +786,35 @@ def select_order_map(column=nil, &block)
       _select_map(column, true, &block)
     end
 
+    # Selects the column given (either as an argument or as a block), and
+    # returns a set of all values of that column in the dataset.
+    #
+    #   DB[:table].select_set(:id) # SELECT id FROM table
+    #   # => Set[3, 5, 8, 1, ...]
+    #
+    #   DB[:table].select_set{id * 2} # SELECT (id * 2) FROM table
+    #   # => Set[6, 10, 16, 2, ...]
+    #
+    # You can also provide an array of column names, which returns a set
+    # with array elements (not set elements):
+    #
+    #   DB[:table].select_map([:id, :name]) # SELECT id, name FROM table
+    #   # => Set[[1, 'A'], [2, 'B'], [3, 'C'], ...]
+    #
+    # If you provide an array of expressions, you must be sure that each entry
+    # in the array has an alias that Sequel can determine.
+    def select_set(column=nil, &block)
+      ds = ungraphed.naked
+      columns = Array(column)
+      virtual_row_columns(columns, block)
+      select_cols = order ? columns.map{|c| c.is_a?(SQL::OrderedExpression) ? c.expression : c} : columns
+      if column.is_a?(Array) || (columns.length > 1)
+        ds.select(*select_cols)._select_set_multiple(hash_key_symbols(select_cols))
+      else
+        ds.select(auto_alias_expression(select_cols.first))._select_set_single
+      end
+    end
+    
     # Limits the dataset to one record, and returns the first record in the dataset,
     # or nil if the dataset has no records. Users should probably use +first+ instead of
     # this method. Example:
@@ -1089,6 +1139,17 @@ def _select_map_single
       map{|r| r[k||=r.keys.first]}
     end
 
+    # Return a set of arrays of values given by the symbols in ret_cols.
+    def _select_set_multiple(ret_cols)
+      to_set{|r| r.values_at(*ret_cols)}
+    end
+  
+    # Returns a set of the first value in each row.
+    def _select_set_single
+      k = nil
+      to_set{|r| r[k||=r.keys.first]}
+    end
+  
     # A dataset for returning single values from the current dataset.
     def single_value_ds
       clone(:limit=>1).ungraphed.naked

lib/sequel/model/base.rb

@@ -10,7 +10,7 @@ class Model
     # Class methods for Sequel::Model that implement basic model functionality.
     #
     # * All of the following methods have class methods created that send the method
-    #   to the model's dataset: all, any?, as_hash, avg, count, cross_join, distinct, each,
+    #   to the model's dataset: all, any?, as_hash, as_set, avg, count, cross_join, distinct, each,
     #   each_server, empty?, except, exclude, exclude_having, fetch_rows,
     #   filter, first, first!, for_update, from, from_self, full_join, full_outer_join,
     #   get, graph, grep, group, group_and_count, group_append, group_by, having, import,
@@ -19,7 +19,7 @@ class Model
     #   natural_join, natural_left_join, natural_right_join, offset, order, order_append, order_by,
     #   order_more, order_prepend, paged_each, qualify, reverse, reverse_order, right_join,
     #   right_outer_join, select, select_all, select_append, select_group, select_hash,
-    #   select_hash_groups, select_map, select_more, select_order_map, select_prepend, server,
+    #   select_hash_groups, select_map, select_more, select_order_map, select_prepend, select_set, server,
     #   single_record, single_record!, single_value, single_value!, sum, to_hash, to_hash_groups,
     #   truncate, unfiltered, ungraphed, ungrouped, union, unlimited, unordered, where, where_all,
     #   where_each, where_single_value, with, with_recursive, with_sql

spec/core/dataset_spec.rb

@@ -2314,6 +2314,26 @@ def ==(o) @h == o.h; end
   end
 end
 
+describe "Dataset#as_set" do
+  before do
+    @d = Sequel.mock(:fetch=>[{:a => 1, :b => 2}, {:a => 3, :b => 4}, {:a => 5, :b => 6}])[:items]
+  end
+  
+  it "should return set of column elements if column symbol is given" do
+    @d.as_set(:a).must_equal Set[1, 3, 5]
+  end
+  
+  it "should support set with array elements if an array of column symbols is given" do
+    @d.as_set([:a, :b]).must_equal Set[[1, 2], [3, 4], [5, 6]]
+  end
+  
+  it "should not call the row_proc" do
+    @d = @d.with_row_proc(proc{|r| h = {}; r.keys.each{|k| h[k] = r[k] * 2}; h})
+    @d.as_set(:a).must_equal Set[1, 3, 5]
+    @d.as_set([:a, :b]).must_equal Set[[1, 2], [3, 4], [5, 6]]
+  end
+end
+
 describe "Dataset#distinct" do
   before do
     @db = Sequel.mock
@@ -5420,6 +5440,100 @@ def default_timestamp_format
   end
 end
 
+describe "Sequel::Dataset#select_set" do
+  before do
+    @ds = Sequel.mock(:fetch=>[{:c=>1}, {:c=>2}])[:t]
+  end
+
+  it "should do select and map in one step" do
+    @ds.select_set(:a).must_equal Set[1, 2]
+    @ds.db.sqls.must_equal ['SELECT a FROM t']
+  end
+
+  it "should handle qualified identifiers in arguments" do
+    @ds.select_set(Sequel[:a][:b]).must_equal Set[1, 2]
+    @ds.db.sqls.must_equal ['SELECT a.b FROM t']
+  end
+
+  with_symbol_splitting "should handle implicit qualifiers in arguments" do
+    @ds.select_set(:a__b).must_equal Set[1, 2]
+    @ds.db.sqls.must_equal ['SELECT a.b FROM t']
+  end
+
+  it "should raise if multiple arguments and can't determine alias" do
+    proc{@ds.select_set([Sequel.function(:a), :b])}.must_raise(Sequel::Error)
+    proc{@ds.select_set(Sequel.function(:a)){b}}.must_raise(Sequel::Error)
+    proc{@ds.select_set{[a.function, b]}}.must_raise(Sequel::Error)
+  end
+
+  with_symbol_splitting "should handle implicit aliases in arguments" do
+    @ds.select_set(:a___b).must_equal Set[1, 2]
+    @ds.db.sqls.must_equal ['SELECT a AS b FROM t']
+  end
+
+  it "should handle aliased expressions in arguments" do
+    @ds.select_set(Sequel[:a].as(:b)).must_equal Set[1, 2]
+    @ds.db.sqls.must_equal ['SELECT a AS b FROM t']
+  end
+
+  it "should handle other objects" do
+    @ds.select_set(Sequel.lit("a").as(:b)).must_equal Set[1, 2]
+    @ds.db.sqls.must_equal ['SELECT a AS b FROM t']
+  end
+  
+  it "should handle identifiers with strings" do
+    @ds.select_set([Sequel::SQL::Identifier.new('c'), :c]).must_equal Set[[1, 1], [2, 2]]
+    @ds.db.sqls.must_equal ['SELECT c, c FROM t']
+  end
+  
+  it "should raise an error for plain strings" do
+    proc{@ds.select_set(['c', :c])}.must_raise(Sequel::Error)
+    @ds.db.sqls.must_equal []
+  end
+  
+  it "should handle an expression without a determinable alias" do
+    @ds.select_set{a(t[c])}.must_equal Set[1, 2]
+    @ds.db.sqls.must_equal ['SELECT a(t.c) AS v FROM t']
+  end
+
+  it "should accept a block" do
+    @ds.select_set{a(t[c]).as(b)}.must_equal Set[1, 2]
+    @ds.db.sqls.must_equal ['SELECT a(t.c) AS b FROM t']
+  end
+
+  it "should accept a block with an array of columns" do
+    @ds.select_set{[a(t[c]).as(c), a(t[c]).as(c)]}.must_equal Set[[1, 1], [2, 2]]
+    @ds.db.sqls.must_equal ['SELECT a(t.c) AS c, a(t.c) AS c FROM t']
+  end
+
+  it "should accept a block with a column" do
+    @ds.select_set(:c){a(t[c]).as(c)}.must_equal Set[[1, 1], [2, 2]]
+    @ds.db.sqls.must_equal ['SELECT c, a(t.c) AS c FROM t']
+  end
+
+  it "should accept a block and array of arguments" do
+    @ds.select_set([:c, :c]){[a(t[c]).as(c), a(t[c]).as(c)]}.must_equal Set[[1, 1, 1, 1], [2, 2, 2, 2]]
+    @ds.db.sqls.must_equal ['SELECT c, c, a(t.c) AS c, a(t.c) AS c FROM t']
+  end
+
+  it "should handle an array of columns" do
+    @ds.select_set([:c, :c]).must_equal Set[[1, 1], [2, 2]]
+    @ds.db.sqls.must_equal ['SELECT c, c FROM t']
+    @ds.select_set([Sequel.expr(:d).as(:c), Sequel.qualify(:b, :c), Sequel.identifier(:c), Sequel.identifier(:c).qualify(:b)]).must_equal Set[[1, 1, 1, 1], [2, 2, 2, 2]]
+    @ds.db.sqls.must_equal ['SELECT d AS c, b.c, c, b.c FROM t']
+  end
+
+  with_symbol_splitting "should handle an array of columns with splittable symbols" do
+    @ds.select_set([:a__c, :a__d___c]).must_equal Set[[1, 1], [2, 2]]
+    @ds.db.sqls.must_equal ['SELECT a.c, a.d AS c FROM t']
+  end
+
+  it "should handle an array with a single element" do
+    @ds.select_set([:c]).must_equal Set[[1], [2]]
+    @ds.db.sqls.must_equal ['SELECT c FROM t']
+  end
+end
+
 describe "Dataset#lock_style and for_update" do
   before do
     @ds = Sequel.mock.dataset.from(:t)

spec/integration/dataset_test.rb

@@ -1487,6 +1487,12 @@
     ds.to_hash_groups(:a, :d, :hash => (tmp = {})).must_be_same_as(tmp)
   end
 
+  it "should have working #as_set" do
+    @ds.as_set(:a).must_equal Set[1, 5]
+    @ds.as_set(:b).must_equal Set[2, 6]
+    @ds.as_set([:a, :b]).must_equal Set[[1, 2], [5, 6]]
+  end
+  
   it "should have working #select_map" do
     @ds.select_map(:a).must_equal [1, 5]
     @ds.select_map(:b).must_equal [2, 6]
@@ -1548,6 +1554,22 @@
     ds.select_hash_groups(:a, :d, :hash => (tmp = {})).must_be_same_as(tmp)
     @ds.extension(:null_dataset).nullify.select_hash_groups(:a, :d).must_equal({})
   end
+
+  it "should have working #select_set" do
+    @ds.select_set(:a).must_equal Set[1, 5]
+    @ds.select_set(:b).must_equal Set[2, 6]
+    @ds.select_set([:a]).must_equal Set[[1], [5]]
+    @ds.select_set([:a, :b]).must_equal Set[[1, 2], [5, 6]]
+    @ds.extension(:null_dataset).nullify.select_set([:a, :b]).must_equal Set[]
+
+    @ds.select_set(Sequel[:a].as(:e)).must_equal Set[1, 5]
+    @ds.select_set(Sequel[:b].as(:e)).must_equal Set[2, 6]
+    @ds.select_set([Sequel[:a].as(:e), Sequel[:b].as(:f)]).must_equal Set[[1, 2], [5, 6]]
+    @ds.select_set([Sequel[:a][:a].as(:e), Sequel[:a][:b].as(:f)]).must_equal Set[[1, 2], [5, 6]]
+    @ds.select_set([Sequel.expr(Sequel[:a][:a]).as(:e), Sequel.expr(Sequel[:a][:b]).as(:f)]).must_equal Set[[1, 2], [5, 6]]
+    @ds.select_set([Sequel.qualify(:a, :a).as(:e), Sequel.qualify(:a, :b).as(:f)]).must_equal Set[[1, 2], [5, 6]]
+    @ds.select_set([Sequel.identifier(:a).qualify(:a).as(:e), Sequel.qualify(:a, :b).as(:f)]).must_equal Set[[1, 2], [5, 6]]
+  end
 end
 
 describe "Sequel::Dataset DSL support" do

spec/model/class_dataset_methods_spec.rb

@@ -13,6 +13,8 @@
     @db.sqls.must_equal ["SELECT * FROM items"]
     @c.any?.must_equal true
     @db.sqls.must_equal ["SELECT * FROM items"]
+    @c.as_set(:id).must_equal Set[1]
+    @db.sqls.must_equal ["SELECT * FROM items"]
     @c.avg(:id).must_equal 1
     @db.sqls.must_equal ["SELECT avg(id) AS avg FROM items LIMIT 1"]
     @c.count.must_equal 1
@@ -100,6 +102,8 @@
     @c.select_order_map(:id).must_equal [1]
     @db.sqls.must_equal ["SELECT id FROM items ORDER BY id"]
     @c.select_prepend(:a).sql.must_equal "SELECT a, items.* FROM items"
+    @c.select_set(:id).must_equal Set[1]
+    @db.sqls.must_equal ["SELECT id FROM items"]
     @c.server(:a).opts[:server].must_equal :a
     @c.single_record.must_equal @c.load(:id=>1)
     @db.sqls.must_equal ["SELECT * FROM items LIMIT 1"]