wonderdog 0.0.1

data/notes/pigstorefunc.pig

@@ -0,0 +1,45 @@
+--
+-- Doesn't work at the moment, just some notes on how the storefunc might look.
+--
+
+
+--
+-- Right now the ElasticSearchOutputFormat gets all its options from the
+-- Job object. We can use the call to setStoreLocation in the storefunc
+-- to set the required parameters. Need to make sure the following are
+-- set:
+-- 
+-- wonderdog.index.name  - should be set by the storefunc constructor
+-- wonderdog.bulk.size   - should be set by the storefunc constructor
+-- wonderdog.field.names - should be set by the call to checkSchema
+-- wonderdog.id.field    - should be set by the storefunc constructor
+-- wonderdog.object.type - should be set by the storefunc constructor
+-- wonderdog.plugins.dir - should be set by call to setStoreLocation
+-- wonderdog.config      - should be set by call to setStoreLocation
+--
+-- FIXME: options used in the ElasticSearchOutputFormat should NOT be
+-- namespaced with 'wonderdog'
+
+%default INDEX 'es_index'
+%default OBJ   'text_obj'
+
+        
+records         = LOAD '$DATA'   AS (text_field:chararray);
+records_with_id = LOAD '$IDDATA' AS (id_field:int, text_field:chararray);
+
+-- Here we would use the elasticsearch index name as the uri, pass in a
+-- comma separated list of field names as the first arg, the id field
+-- as the second arg and the bulk size as the third. 
+-- 
+-- and so on.
+STORE records INTO '$INDEX/$OBJ' USING ElasticSearchStorage('my_text_field', '-1', '1000');
+
+
+-- but it would be really nice to duplicate what's in WonderDog.java in that,
+-- should a bulk request fail, the failed records are written to hdfs. The
+-- user should have some control of this. Also, it should be possible to generate
+-- the field names directly from the pig schema? (We'd have to be VERY explicit in the
+-- docs about this as it would be a point of headscratching/swearing...) In this
+-- case we might have something like:
+named_records = FOREACH records GENERATE text_field AS text_field_name;
+STORE records INTO '/path/to/failed_requests' USING ElasticSearchStorage('$INDEX/$OBJ', '-1', '1000');

data/pom.xml

@@ -0,0 +1,80 @@
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+  xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+
+  <groupId>com.infochimps.elasticsearch</groupId>
+  <artifactId>wonderdog</artifactId>
+  <version>1.0-SNAPSHOT</version>
+  <packaging>jar</packaging>
+
+  <name>wonderdog</name>
+  <url>http://maven.apache.org</url>
+
+  <properties>
+    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+  </properties>
+
+  <dependencies>
+    <dependency>
+      <groupId>org.elasticsearch</groupId>
+      <artifactId>elasticsearch</artifactId>
+      <version>0.19.8</version>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.hadoop</groupId>
+      <artifactId>hadoop-core</artifactId>
+      <version>0.20.2</version>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.pig</groupId>
+      <artifactId>pig</artifactId>
+      <version>0.8.0</version>
+    </dependency>
+
+    <dependency>
+      <groupId>org.codehaus.jackson</groupId>
+      <artifactId>jackson-mapper-asl</artifactId>
+      <version>1.5.2</version>
+    </dependency>    
+  </dependencies>
+
+  <build>
+    <plugins>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-compiler-plugin</artifactId>
+        <configuration>
+          <source>1.5</source>
+          <target>1.5</target>
+        </configuration>
+      </plugin>
+    </plugins>
+  </build>
+  
+  <repositories>
+    <repository>
+      <id>com.cloudera</id>
+      <url>https://repository.cloudera.com/content/repositories/releases</url>
+    </repository>
+
+    <repository>
+      <id>sonatype-releases</id>
+      <url>http://oss.sonatype.org/content/repositories/releases</url>
+    </repository>
+
+    <repository>
+      <id>codehaus</id>
+      <url>http://repository.codehaus.org/org/codehaus</url>
+      <snapshots>
+        <enabled>true</enabled>
+      </snapshots>
+      <releases>
+        <enabled>true</enabled>
+      </releases>
+    </repository>
+
+  </repositories>
+  
+</project>

data/spec/spec_helper.rb

@@ -0,0 +1,22 @@
+require 'wonderdog'
+require 'wukong/spec_helpers'
+require_relative('support/integration_helper')
+require_relative('support/driver_helper')
+
+
+RSpec.configure do |config|
+
+  config.before(:each) do
+    @orig_reg = Wukong.registry.show
+  end
+
+  config.after(:each) do
+    Wukong.registry.clear!
+    Wukong.registry.merge!(@orig_reg)
+  end
+    
+  include Wukong::SpecHelpers
+  include Wukong::Elasticsearch::IntegrationHelper
+  include Wukong::Elasticsearch::DriverHelper
+end
+

data/spec/support/driver_helper.rb

@@ -0,0 +1,15 @@
+module Wukong
+  module Elasticsearch
+    module DriverHelper
+
+      def driver *args
+        params   = Elasticsearch.configure(Hadoop.configure(Configliere::Param.new))
+        params.resolve!
+        params.merge!(args.pop) if args.last.is_a?(Hash)
+        Hadoop::Driver.new(params, *args)
+      end
+      
+    end
+  end
+end
+

data/spec/support/integration_helper.rb

@@ -0,0 +1,30 @@
+module Wukong
+  module Elasticsearch
+    module IntegrationHelper
+
+      def root
+        @root ||= Pathname.new(File.expand_path('../../..', __FILE__))
+      end
+
+      def lib_dir
+        root.join('lib')
+      end
+
+      def bin_dir
+        root.join('bin')
+      end
+      
+      def integration_env
+        {
+          "RUBYLIB" => [lib_dir.to_s, ENV["RUBYLIB"]].compact.join(':')
+        }
+      end
+
+      def integration_cwd
+        root.to_s
+      end
+
+    end
+  end
+end
+

data/spec/wonderdog/hadoop_invocation_override_spec.rb

@@ -0,0 +1,81 @@
+require 'spec_helper'
+
+describe Wukong::Elasticsearch::HadoopInvocationOverride do
+
+  let(:no_es)      { driver('regexp',  'count', input: '/tmp/input_file',        output: '/tmp/output_file')         }
+  let(:es_reader)  { driver('regexp',  'count', input: 'es://the_index/the_map', output: '/tmp/output_file')         }
+  let(:es_writer)  { driver('regexp',  'count', input: '/tmp/input_file',        output: 'es:///the_index/the_map')  }
+  let(:es_complex) { driver('regexp',  'count', input: 'es://the_index/the_map', output: 'es:///the_index/the_map', es_query: '{"hi": "there"}', es_request_size: 1000, es_index_field: 'ID') }
+
+  context "not interacting with Elasticsearch" do
+    subject                  { no_es                                            }
+    # input
+    its(:input_paths)        { should == '/tmp/input_file'                      }
+    its(:hadoop_commandline) { should     match(%r{-input.*/tmp/input_file}i)   }
+
+    # output
+    its(:output_path)        { should == '/tmp/output_file'                     }
+    its(:hadoop_commandline) { should     match(%r{-output.*/tmp/output_file}i) }
+
+    # no elasticsearch anything
+    its(:hadoop_commandline) { should_not match(/elasticsearch/i)               }
+  end
+
+  context "reading from Elasticsearch" do
+    subject                  { es_reader                                                          }
+    
+    # input
+    its(:input_paths)        { should     match(%r{/user.*wukong.*the_index.*the_map})            }
+    its(:hadoop_commandline) { should     match(/-inputformat.*elasticsearch/i)                   }
+    its(:hadoop_commandline) { should     match(%r{-input.*/user.*wukong.*the_index.*the_map}i)   }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.input\.index.*the_index/i)   }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.input\.map.*the_map/i)       }
+    
+    # output
+    its(:output_path)        { should == '/tmp/output_file'                                       }
+    its(:hadoop_commandline) { should_not match(/-outputformat/i)                                 }
+    its(:hadoop_commandline) { should     match(%r{-output.*/tmp/output_file}i)                   }
+    its(:hadoop_commandline) { should_not match(/-D\s+elasticsearch\.output/i)                    }
+  end
+
+  context "writing to Elasticsearch" do
+    subject                  { es_writer                                                          }
+
+    # input
+    its(:input_paths)        { should == '/tmp/input_file'                                        }
+    its(:hadoop_commandline) { should_not match(/-inputformat/i)                                  }
+    its(:hadoop_commandline) { should     match(%r{-input.*/tmp/input_file}i)                     }
+    its(:hadoop_commandline) { should_not match(/-D\s+elasticsearch\.input/i)                     }
+
+    # output
+    its(:output_path)        { should     match(%r{/user.*wukong.*the_index.*the_map})            }
+    its(:hadoop_commandline) { should     match(/-outputformat.*elasticsearch/i)                  }
+    its(:hadoop_commandline) { should     match(%r{-output.*/user.*wukong.*the_index.*the_map}i)  }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.output\.index.*the_index/i)  }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.output\.map.*the_map/i)      }
+  end
+
+  context "reading and writing with many options" do
+    subject                  { es_complex                                                         }
+
+    # input
+    its(:input_paths)        { should     match(%r{/user.*wukong.*the_index.*the_map})            }
+    its(:hadoop_commandline) { should     match(/-inputformat.*elasticsearch/i)                   }
+    its(:hadoop_commandline) { should     match(%r{-input.*/user.*wukong.*the_index.*the_map}i)   }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.input\.index.*the_index/i)   }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.input\.map.*the_map/i)       }
+
+    # output
+    its(:output_path)        { should     match(%r{/user.*wukong.*the_index.*the_map})            }
+    its(:hadoop_commandline) { should     match(/-outputformat.*elasticsearch/i)                  }
+    its(:hadoop_commandline) { should     match(%r{-output.*/user.*wukong.*the_index.*the_map}i)  }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.output\.index.*the_index/i)  }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.output\.map.*the_map/i)      }
+
+    # options
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.input\.query.*hi.*there/i)   }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.input\.request_size.*1000/i) }
+    its(:hadoop_commandline) { should     match(/-D\s+elasticsearch\.output\.index\.field.*ID/i)  }
+  end
+  
+end

data/spec/wonderdog/index_and_type_spec.rb

@@ -0,0 +1,73 @@
+require 'spec_helper'
+
+describe Wukong::Elasticsearch::IndexAndMapping do
+
+  subject { Wukong::Elasticsearch::IndexAndMapping }
+
+  let(:filesystem_path)      { '/some/path'                                        }
+  let(:filesystem_paths)     { '/some/path,/some/other/path'                       }
+  
+  let(:hdfs_path)            { 'hdfs://some/hdfs/path'                             }
+  let(:hdfs_paths)           { 'hdfs://some/hdfs/path,hdfs://some/other/hdfs/path' }
+  
+  let(:es_index_and_mapping)    { 'es://index/mapping'                                   }
+  let(:es_indices_and_mapping)  { 'es://index1,index2/mapping'                           }
+  let(:es_index_and_mappings)   { 'es://index/mapping1,mapping2'                         }
+  let(:es_indices_and_mappings) { 'es://index1,index2/mapping1,mapping2'                 }
+
+  fails  = %w[filesystem_path filesystem_paths hdfs_path hdfs_paths]
+  passes = %w[es_index_and_mapping es_indices_and_mapping es_index_and_mappings es_indices_and_mappings]
+  
+  context 'recognizing possible es://index/mapping specifications' do
+    fails.each do |name|
+      it "doesn't recognize a #{name}" do
+        subject.matches?(self.send(name)).should be_false
+      end
+    end
+    passes.each do |name|
+      it "recognizes a #{name}" do
+        subject.matches?(self.send(name)).should be_true
+      end
+    end
+  end
+
+  context "parsing es://index/mapping specifications" do
+    fails.each do |name|
+      it "raises an error on a #{name}" do
+        lambda { subject.new(self.send(name)) }.should raise_error(Wukong::Error, /not an elasticsearch.*index\/mapping/i)
+      end
+    end
+
+    it "raises an error on a specification with too many parts" do
+      lambda { subject.new('es://index/mapping/extra') }.should raise_error(Wukong::Error, /not an elasticsearch.*index\/mapping/i)
+    end
+
+    it "raises an error on a specification with too few parts" do
+      lambda { subject.new('es://') }.should raise_error(Wukong::Error, /not an elasticsearch.*index\/mapping/i)
+    end
+    
+    context "on an index and mapping" do
+      subject { Wukong::Elasticsearch::IndexAndMapping.new(es_index_and_mapping) }
+      its(:index)   { should == 'index'   }
+      its(:mapping) { should == 'mapping' }
+    end
+    context "on indices and a mapping" do
+      subject { Wukong::Elasticsearch::IndexAndMapping.new(es_indices_and_mapping) }
+      its(:index)   { should == 'index1,index2'   }
+      its(:mapping) { should == 'mapping'         }
+    end
+    context "on an index and mappings" do
+      subject { Wukong::Elasticsearch::IndexAndMapping.new(es_index_and_mappings) }
+      its(:index)   { should == 'index'             }
+      its(:mapping) { should == 'mapping1,mapping2' }
+    end
+    context "on indices and mappings" do
+      subject { Wukong::Elasticsearch::IndexAndMapping.new(es_indices_and_mappings) }
+      its(:index)   { should == 'index1,index2'      }
+      its(:mapping) { should == 'mapping1,mapping2'  }
+    end
+    
+    
+  end
+  
+end

data/src/main/java/com/infochimps/elasticsearch/ElasticSearchInputFormat.java

@@ -0,0 +1,268 @@
+package com.infochimps.elasticsearch;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.List;
+import java.util.ArrayList;
+import java.util.Iterator;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.apache.hadoop.io.*;
+import org.apache.hadoop.conf.Configurable;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.mapreduce.InputFormat;
+import org.apache.hadoop.mapreduce.RecordReader;
+import org.apache.hadoop.mapreduce.InputSplit;
+import org.apache.hadoop.mapreduce.JobContext;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+
+import com.infochimps.elasticsearch.hadoop.util.HadoopUtils;
+
+import org.elasticsearch.node.Node;
+import org.elasticsearch.node.NodeBuilder;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.client.Requests;
+import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.search.SearchType;
+import org.elasticsearch.index.query.FilterBuilders.*;
+import org.elasticsearch.index.query.QueryBuilders;
+
+/**
+   
+   A Hadoop InputFormat to read data from an Elasticsearch index. The RecordReader
+   divulges records where the key is the record id in elasticsearch and the value
+   is a json string of the (source) record contents.
+   
+ */
+public class ElasticSearchInputFormat extends InputFormat<Text, Text> implements Configurable {
+
+    static Log LOG = LogFactory.getLog(ElasticSearchInputFormat.class);
+    private Configuration conf = null;
+
+    private Node node;
+    private Client client;
+
+    private Integer requestSize;
+    private Long numHits;
+    private Long numSplits;
+    private Long numSplitRecords;
+    private String indexName;
+    private String objType;
+    private String queryString;
+    
+    private static final String ES_REQUEST_SIZE = "elasticsearch.request.size";           // number of records to fetch at one time
+    private static final String ES_NUM_SPLITS = "elasticsearch.num.input.splits";       // number of hadoop map tasks to launch
+    private static final String ES_QUERY_STRING = "elasticsearch.query.string";
+    
+    private static final String ES_CONFIG_NAME = "elasticsearch.yml";
+    private static final String ES_PLUGINS_NAME = "plugins";
+    private static final String ES_INDEX_NAME = "elasticsearch.index.name";
+    private static final String ES_OBJECT_TYPE = "elasticsearch.object.type";
+    private static final String ES_CONFIG = "es.config";
+    private static final String ES_PLUGINS = "es.path.plugins";
+    private static final String SLASH = "/";
+    
+    public RecordReader<Text,Text> createRecordReader(InputSplit inputSplit,
+                                                              TaskAttemptContext context) {
+        return new ElasticSearchRecordReader();
+    }
+
+    /**
+       The number of splits is specified in the Hadoop configuration object. 
+     */
+    public List<InputSplit> getSplits(JobContext context) {
+        setConf(context.getConfiguration());
+        List<InputSplit> splits = new ArrayList<InputSplit>(numSplits.intValue());
+        for(int i = 0; i < numSplits; i++) {
+            Long size = (numSplitRecords == 1) ? 1 : numSplitRecords-1; 
+            splits.add(new ElasticSearchSplit(queryString, i*numSplitRecords, size));
+        }
+        if (numHits % numSplits > 0) splits.add(new ElasticSearchSplit(queryString, numSplits*numSplitRecords, numHits % numSplits - 1));
+        LOG.info("Created ["+splits.size()+"] splits for ["+numHits+"] hits");
+        return splits;
+    }
+    
+    /**
+       Sets the configuration object, opens a connection to elasticsearch, and
+       initiates the initial search request.
+     */
+    @Override
+    public void setConf(Configuration configuration) {
+        this.conf = configuration;
+        this.indexName = conf.get(ES_INDEX_NAME);
+        this.objType = conf.get(ES_OBJECT_TYPE);
+        this.requestSize = Integer.parseInt(conf.get(ES_REQUEST_SIZE));
+        this.numSplits = Long.parseLong(conf.get(ES_NUM_SPLITS));
+        this.queryString = conf.get(ES_QUERY_STRING);
+
+        //
+        // Need to ensure that this is set in the hadoop configuration so we can
+        // instantiate a local client. The reason is that no files are in the
+        // distributed cache when this is called.
+        //
+        System.setProperty(ES_CONFIG, conf.get(ES_CONFIG));
+        System.setProperty(ES_PLUGINS, conf.get(ES_PLUGINS));
+        
+        start_embedded_client();
+
+        initiate_search();
+    }
+
+    @Override
+    public Configuration getConf() {
+        return conf;
+    }
+
+    /**
+       Starts an embedded elasticsearch client (ie. data = false)
+    */
+    private void start_embedded_client() {
+        LOG.info("Starting embedded elasticsearch client ...");
+        this.node   = NodeBuilder.nodeBuilder().client(true).node();
+        this.client = node.client();
+    }
+
+    private void initiate_search() {
+        SearchResponse response = client.prepareSearch(indexName)
+            .setTypes(objType)
+            .setSearchType(SearchType.COUNT)
+            .setQuery(QueryBuilders.queryString(queryString))
+            .setSize(requestSize)
+            .execute()
+            .actionGet();
+        this.numHits = response.hits().totalHits();
+        if(numSplits > numHits) numSplits = numHits; // This could be bad
+        this.numSplitRecords = (numHits/numSplits);
+    }
+
+    protected class ElasticSearchRecordReader extends RecordReader<Text, Text> {
+
+        private Node node;
+        private Client client;
+    
+        private String indexName;
+        private String objType;
+        private Long numSplitRecords;
+        private Integer requestSize;
+        private Text currentKey;
+        private Text currentValue;
+        private Integer recordsRead;
+        private Iterator<SearchHit> hitsItr = null;        
+
+
+        private String queryString;
+        private Long from;
+        private Long recsToRead;
+        
+        public ElasticSearchRecordReader() {
+        }
+
+        public void initialize(InputSplit split, TaskAttemptContext context) throws IOException {
+            Configuration conf = context.getConfiguration();
+            this.indexName = conf.get(ES_INDEX_NAME);
+            this.objType    = conf.get(ES_OBJECT_TYPE);
+            LOG.info("Initializing elasticsearch record reader on index ["+indexName+"] and object type ["+objType+"]");
+
+            //
+            // Fetches elasticsearch.yml and the plugins directory from the distributed cache
+            //
+            try {
+                String taskConfigPath = HadoopUtils.fetchFileFromCache(ES_CONFIG_NAME, conf);
+                LOG.info("Using ["+taskConfigPath+"] as es.config");
+                String taskPluginsPath = HadoopUtils.fetchArchiveFromCache(ES_PLUGINS_NAME, conf);
+                LOG.info("Using ["+taskPluginsPath+"] as es.plugins.dir");
+                System.setProperty(ES_CONFIG, taskConfigPath);
+                System.setProperty(ES_PLUGINS, taskPluginsPath+SLASH+ES_PLUGINS_NAME);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+
+            queryString = ((ElasticSearchSplit)split).getQueryString();
+            from = ((ElasticSearchSplit)split).getFrom();
+            recsToRead = ((ElasticSearchSplit)split).getSize();
+
+            LOG.info("elasticsearch record reader: query ["+queryString+"], from ["+from+"], size ["+recsToRead+"]");
+            start_embedded_client();
+            recordsRead = 0;
+        }
+
+        /**
+           Starts an embedded elasticsearch client (ie. data = false)
+        */
+        private void start_embedded_client() {
+            LOG.info("Starting embedded elasticsearch client ...");
+            this.node   = NodeBuilder.nodeBuilder().client(true).node();
+            this.client = node.client();
+        }
+
+        private Iterator<SearchHit> fetchNextHits() {
+            SearchResponse response = client.prepareSearch(indexName)
+                .setTypes(objType)                
+                .setFrom(from.intValue())
+                .setSize(recsToRead.intValue())
+                .setQuery(QueryBuilders.queryString(queryString))                
+                .execute()
+                .actionGet();
+            return response.hits().iterator();
+        }
+        
+        @Override
+        public boolean nextKeyValue() throws IOException {
+            if (hitsItr!=null) {
+                if (recordsRead < recsToRead) {
+                    if (hitsItr.hasNext()) {
+                        SearchHit hit = hitsItr.next();
+                        currentKey = new Text(hit.id());
+                        currentValue = new Text(hit.sourceAsString());
+                        recordsRead += 1;
+                        return true;
+                    }
+                } else {
+                    hitsItr = null;
+                }
+            } else {
+                if (recordsRead < recsToRead) {
+                    hitsItr = fetchNextHits();
+                    if (hitsItr.hasNext()) {
+                        SearchHit hit = hitsItr.next();
+                        currentKey = new Text(hit.id());
+                        currentValue = new Text(hit.sourceAsString());
+                        recordsRead += 1;
+                        return true;
+                    }
+                }
+            }
+            return false;
+        }
+    
+        @Override
+        public Text getCurrentKey() {
+            return currentKey;
+        }
+    
+        @Override
+        public Text getCurrentValue() {
+            return currentValue;
+        }
+        
+        @Override
+        public float getProgress() throws IOException {
+            return 0;
+        }
+
+        @Override
+        public void close() throws IOException {
+            LOG.info("Closing record reader");
+            client.close();
+            LOG.info("Client is closed");
+            if (node != null) {
+                 node.close();
+            }
+            LOG.info("Record reader closed.");
+        }
+
+    }
+}