List comprehension in Ruby

By Vagmi on 21 Apr 2015

Having partitioned a large postgres table into multiple partitions, I wanted a quick way to dump the partitioned data into CSV files using the /copy command. As the tables are named tablename_yyyymm. I wanted an easy way to generate the yyyymm sequences. One way to do it is inner loops.

1
(2011..2015).map { |y| (1..12).map {|m| (y*100+m).to_s } }

But that looks quite inelegant. Most languages have a way to do list comprehension. For example, if you have a set [1,2,3] and another set [4,5,6], we can produce a list like [[1,4],[1,5],[1,6],[2,4],[2,5],[2,6],[3,4],[3,5],[3,6]] using list comprehension. We can consider map to be a degenerate case of list comprehension.

Here are a few code samples from some of my favorite languages.

1
2
3
4
5
; Clojure
user=> (for [x (range 1 4) y (range 4 7)] [x y])
([1 4] [1 5] [1 6]
 [2 4] [2 5] [2 6]
 [3 4] [3 5] [3 6])

This is how it is done in Haskell.

1
2
3
4
5
-- Haskell
Prelude> [[x,y] | x <- [1..3], y <- [4..6]]
[[1,4],[1,5],[1,6],
 [2,4],[2,5],[2,6],
 [3,4],[3,5],[3,6]]

But Ruby does not support it. So, I decided to add them in Ruby. Well, it is not quite the same but it is enough to get my job done. True array comprehension will respect laziness and support filtering operations. But for my use case, this is sufficient. update - I have updated the code as per @sgporras comment. To make the comprehension lazy and return a enumerator instead of a list.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def comprehend(*enums)
  cycles = enums.map(&:cycle)
 
  Enumerator.new do |comprehended|
    loop do
      value = cycles.map(&:peek)
      value = yield value if block_given?
      comprehended << value
      (cycles.length - 1).downto(0) do |index|
cycles[index].next
break unless cycles[index].peek == enums[index].first
raise StopIteration if index == 0
      end
    end
  end.lazy # requires Ruby > 2.0
end

You can call it with a block that will perform a map inline or call it just as a method to return just the elements that are generated.

1
2
3
4
5
6
irb> comprehend(1..3,4..6).to_a
=> [[1, 4], [1, 5], [1, 6],
    [2, 4], [2, 5], [2, 6],
    [3, 4], [3, 5], [3, 6]]
irb> comprehend(1..3,4..6) { |(x,y)| "#{x}-#{y}" }.to_a
=> ["1-4", "1-5", "1-6", "2-4", "2-5", "2-6", "3-4", "3-5", "3-6"]

I thought of refining the array and adding a comprehend method on the Array but I prefer a function implementation in a module rather than monkey patching the Array. Notice that the arguments x and y are pattern matched on the block supplied to the comprehend method.

comments powered by Disqus