From 77c8c11ec40fdbfd7fe6c4e87b7e30d29d3f7ba1 Mon Sep 17 00:00:00 2001
From: Dan Buch <daniel.buch@gmail.com>
Date: Wed, 29 Feb 2012 09:32:48 -0500
Subject: [PATCH] Trying to play with sorted sets for querying in range of
 magnitudes

---
 redis-with-ruby/earthquakes.rb | 25 +++++++++++++++++++++----
 redis-with-ruby/eq-info.rb     |  8 +++++---
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/redis-with-ruby/earthquakes.rb b/redis-with-ruby/earthquakes.rb
index 40588d6..b1fc41f 100644
--- a/redis-with-ruby/earthquakes.rb
+++ b/redis-with-ruby/earthquakes.rb
@@ -7,14 +7,17 @@ require 'open-uri'
 
 
 class Earthquakes
-  attr_accessor :id_set, :region_set, :magnitude_set, :data_url
+  attr_accessor :id_set, :region_set, :magnitude_set
+  attr_accessor :magnitude_scores_set, :data_url
 
   def initialize
     @redis = Redis.new
     @id_set = 'earthquake-ids'
     @region_set = 'earthquake-regions'
     @magnitude_set = 'earthquake-magnitudes'
-    @data_url = 'http://earthquake.usgs.gov/earthquakes/catalogs/eqs7day-M1.txt'
+    @magnitude_scores_set = 'earthquakes-magnitudes-scores'
+    @data_url =
+      'http://earthquake.usgs.gov/earthquakes/catalogs/eqs7day-M1.txt'
 
     @logger = Logger.new(STDOUT)
     @logger.formatter = proc do |severity, datetime, progname, msg|
@@ -30,8 +33,12 @@ class Earthquakes
     @redis.smembers(with_prefix(region, 'eqregion'))
   end
 
-  def with_magnitude(magnitude)
-    @redis.smembers(with_prefix(magnitude, 'eqmagnitude'))
+  def with_magnitude_in_range(min_magnitude, max_magnitude)
+    min_mag, max_mag =
+      magnitude_as_int(min_magnitude), magnitude_as_int(max_magnitude)
+    @logger.debug(
+      "Fetching magnitudes with scores in range #{min_mag}-#{max_mag}")
+    @redis.zrange(@magnitude_scores_set, min_mag, max_mag)
   end
 
   def all_ids
@@ -61,6 +68,9 @@ class Earthquakes
 
     @logger.info('Removing earthquake magnitude set')
     @redis.del(@magnitude_set)
+
+    @logger.info('Removing earthquake magnitude scores set')
+    @redis.del(@magnitude_scores_set)
   end
 
   def load_data
@@ -101,6 +111,9 @@ class Earthquakes
     magnitude_key = with_prefix(row[:magnitude], 'eqmagnitude')
     redis_multi.sadd(magnitude_key, id_key)
     redis_multi.sadd(@magnitude_set, magnitude_key)
+    score = magnitude_as_int(row[:magnitude])
+    @logger.info("Setting score of #{score} for row #{row.inspect}")
+    redis_multi.zadd(@magnitude_scores_set, score, id_key)
   end
 
   def with_prefix(key, prefix)
@@ -114,6 +127,10 @@ class Earthquakes
   def sans_prefix(key, prefix)
     key.gsub(/^#{prefix}:/, '')
   end
+
+  def magnitude_as_int(magnitude)
+    magnitude.to_s.gsub(/\./, '').to_i
+  end
 end
 
 
diff --git a/redis-with-ruby/eq-info.rb b/redis-with-ruby/eq-info.rb
index c1795bf..21281c0 100755
--- a/redis-with-ruby/eq-info.rb
+++ b/redis-with-ruby/eq-info.rb
@@ -13,7 +13,7 @@ def main
     on :R, :listall_regions, 'List available earthquake regions.'
     on :M, :listall_magnitudes, 'List available earthquakes magnitudes.'
     on :r, :region=, 'Show earthquakes for region.'
-    on :m, :magnitude=, 'Show earthquakes of magnitude.'
+    on :m, :magnitude_range=, 'Show earthquakes of magnitude in range, e.g. "1.1-2.9".'
     on :i, :eqid=, 'Show all available info for earthquake with id.'
   end
 
@@ -51,11 +51,13 @@ def main
     return 0
   end
 
-  if opts[:magnitude]
+  if opts[:magnitude_range]
     puts "/* Earthquakes with magnitude #{opts[:magnitude]} */"
 
+    min_mag, max_mag = opts[:magnitude_range].split('-')
+
     results = []
-    earthquakes.with_magnitude(opts[:magnitude]).each do |i|
+    earthquakes.with_magnitude_in_range(min_mag, max_mag).each do |i|
       results << earthquakes.get_info(i)
     end