List comprehension in Ruby

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.

(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.

; 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.

-- 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.

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.

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.

8 thoughts on “List comprehension in Ruby”

  1. Hi, cool example!

    To address your comment about it not respecting lazyness, it’s really easy to make it so:

    You’ll see that your code is only slightly modified. Literally just using an Enumerator instead of an Array to accumulate the values :)

    Like

  2. Hi, cool example!

    Addressing your comment about it not respecting lazyness, it’s really easy to make it so:

    You’ll see your code is very slightly modified, literally just using an Enumerator instead of an Array to accumulate the values.

    Like

    1. Yup. This is awesome. Using an enumerator is better than a list. I will change my code in the blog post accordingly. Thanks for the suggestion.

      Like

  3. Hi, I apparently posted the same comment twice (first time I had problems logging in and thought the comment was lost) :-/ Can you please delete one of the copies (along with this comment as well)? Thank you!

    Like

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s