Get API: Allow to specify which fields to load, close #65.

modules/elasticsearch/src/main/java/org/elasticsearch/action/get/GetField.java

@@ -0,0 +1,129 @@
+/*
+ * Licensed to Elastic Search and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Elastic Search licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.action.get;
+
+import org.elasticsearch.util.io.Streamable;
+
+import java.io.DataInput;
+import java.io.DataOutput;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+/**
+ * @author kimchy (shay.banon)
+ */
+public class GetField implements Streamable, Iterable<Object> {
+
+    private String name;
+
+    private List<Object> values;
+
+    private GetField() {
+
+    }
+
+    public GetField(String name, List<Object> values) {
+        this.name = name;
+        this.values = values;
+    }
+
+    public String name() {
+        return name;
+    }
+
+    public List<Object> values() {
+        return values;
+    }
+
+    @Override public Iterator<Object> iterator() {
+        return values.iterator();
+    }
+
+    public static GetField readGetField(DataInput in) throws IOException, ClassNotFoundException {
+        GetField result = new GetField();
+        result.readFrom(in);
+        return result;
+    }
+
+    @Override public void readFrom(DataInput in) throws IOException, ClassNotFoundException {
+        name = in.readUTF();
+        int size = in.readInt();
+        values = new ArrayList<Object>(size);
+        for (int i = 0; i < size; i++) {
+            Object value;
+            byte type = in.readByte();
+            if (type == 0) {
+                value = in.readUTF();
+            } else if (type == 1) {
+                value = in.readInt();
+            } else if (type == 2) {
+                value = in.readLong();
+            } else if (type == 3) {
+                value = in.readFloat();
+            } else if (type == 4) {
+                value = in.readDouble();
+            } else if (type == 5) {
+                value = in.readBoolean();
+            } else if (type == 6) {
+                int bytesSize = in.readInt();
+                value = new byte[bytesSize];
+                in.readFully(((byte[]) value));
+            } else {
+                throw new IOException("Can't read unknown type [" + type + "]");
+            }
+            values.add(value);
+        }
+    }
+
+    @Override public void writeTo(DataOutput out) throws IOException {
+        out.writeUTF(name);
+        out.writeInt(values.size());
+        for (Object obj : values) {
+            Class type = obj.getClass();
+            if (type == String.class) {
+                out.write(0);
+                out.writeUTF((String) obj);
+            } else if (type == Integer.class) {
+                out.write(1);
+                out.writeInt((Integer) obj);
+            } else if (type == Long.class) {
+                out.write(2);
+                out.writeLong((Long) obj);
+            } else if (type == Float.class) {
+                out.write(3);
+                out.writeFloat((Float) obj);
+            } else if (type == Double.class) {
+                out.write(4);
+                out.writeDouble((Double) obj);
+            } else if (type == Boolean.class) {
+                out.write(5);
+                out.writeBoolean((Boolean) obj);
+            } else if (type == byte[].class) {
+                out.write(6);
+                out.writeInt(((byte[]) obj).length);
+                out.write(((byte[]) obj));
+            } else {
+                throw new IOException("Can't write type [" + type + "]");
+            }
+        }
+    }
+}

modules/elasticsearch/src/main/java/org/elasticsearch/action/get/GetRequest.java

@@ -40,6 +40,8 @@ import java.io.IOException;
  */
 public class GetRequest extends SingleOperationRequest {
 
+    private String[] fields;
+
     GetRequest() {
     }
 
@@ -78,6 +80,23 @@ public class GetRequest extends SingleOperationRequest {
         return this;
     }
 
+    /**
+     * Explicitly specify the fields that will be returned. By default, the <tt>_source</tt>
+     * field will be returned.
+     */
+    public GetRequest fields(String... fields) {
+        this.fields = fields;
+        return this;
+    }
+
+    /**
+     * Explicitly specify the fields that will be returned. By default, the <tt>_source</tt>
+     * field will be returned.
+     */
+    public String[] fields() {
+        return this.fields;
+    }
+
     /**
      * Should the listener be called on a separate thread if needed.
      */
@@ -96,10 +115,25 @@ public class GetRequest extends SingleOperationRequest {
 
     @Override public void readFrom(DataInput in) throws IOException, ClassNotFoundException {
         super.readFrom(in);
+        int size = in.readInt();
+        if (size >= 0) {
+            fields = new String[size];
+            for (int i = 0; i < size; i++) {
+                fields[i] = in.readUTF();
+            }
+        }
     }
 
     @Override public void writeTo(DataOutput out) throws IOException {
         super.writeTo(out);
+        if (fields == null) {
+            out.writeInt(-1);
+        } else {
+            out.writeInt(fields.length);
+            for (String field : fields) {
+                out.writeUTF(field);
+            }
+        }
     }
 
     @Override public String toString() {

modules/elasticsearch/src/main/java/org/elasticsearch/action/get/GetResponse.java

@@ -27,8 +27,12 @@ import org.elasticsearch.util.io.Streamable;
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
+import java.util.Iterator;
 import java.util.Map;
 
+import static com.google.common.collect.Iterators.*;
+import static com.google.common.collect.Maps.*;
+import static org.elasticsearch.action.get.GetField.*;
 import static org.elasticsearch.util.json.Jackson.*;
 
 /**
@@ -38,7 +42,7 @@ import static org.elasticsearch.util.json.Jackson.*;
  * @see GetRequest
  * @see org.elasticsearch.client.Client#get(GetRequest)
  */
-public class GetResponse implements ActionResponse, Streamable {
+public class GetResponse implements ActionResponse, Streamable, Iterable<GetField> {
 
     private String index;
 
@@ -46,23 +50,29 @@ public class GetResponse implements ActionResponse, Streamable {
 
     private String id;
 
+    private boolean exists;
+
+    private Map<String, GetField> fields;
+
     private byte[] source;
 
     GetResponse() {
     }
 
-    GetResponse(String index, String type, String id, byte[] source) {
+    GetResponse(String index, String type, String id, boolean exists, byte[] source, Map<String, GetField> fields) {
         this.index = index;
         this.type = type;
         this.id = id;
+        this.exists = exists;
         this.source = source;
+        this.fields = fields;
     }
 
     /**
      * Does the document exists.
      */
     public boolean exists() {
-        return source != null && source.length > 0;
+        return exists;
     }
 
     /**
@@ -103,8 +113,9 @@ public class GetResponse implements ActionResponse, Streamable {
     /**
      * The source of the document (As a map).
      */
+    @SuppressWarnings({"unchecked"})
     public Map<String, Object> sourceAsMap() throws ElasticSearchParseException {
-        if (!exists()) {
+        if (source == null) {
             return null;
         }
         try {
@@ -114,26 +125,63 @@ public class GetResponse implements ActionResponse, Streamable {
         }
     }
 
+    public Map<String, GetField> fields() {
+        return this.fields;
+    }
+
+    public GetField field(String name) {
+        return fields.get(name);
+    }
+
+    @Override public Iterator<GetField> iterator() {
+        if (fields == null) {
+            return emptyIterator();
+        }
+        return fields.values().iterator();
+    }
+
     @Override public void readFrom(DataInput in) throws IOException, ClassNotFoundException {
         index = in.readUTF();
         type = in.readUTF();
         id = in.readUTF();
+        exists = in.readBoolean();
+        if (exists) {
             int size = in.readInt();
             if (size > 0) {
                 source = new byte[size];
                 in.readFully(source);
             }
+            size = in.readInt();
+            if (size > 0) {
+                fields = newHashMapWithExpectedSize(size);
+                for (int i = 0; i < size; i++) {
+                    GetField field = readGetField(in);
+                    fields.put(field.name(), field);
+                }
+            }
+        }
     }
 
     @Override public void writeTo(DataOutput out) throws IOException {
         out.writeUTF(index);
         out.writeUTF(type);
         out.writeUTF(id);
+        out.writeBoolean(exists);
+        if (exists) {
             if (source == null) {
                 out.writeInt(0);
             } else {
                 out.writeInt(source.length);
                 out.write(source);
             }
+            if (fields == null) {
+                out.writeInt(0);
+            } else {
+                out.writeInt(fields.size());
+                for (GetField field : fields.values()) {
+                    field.writeTo(out);
+                }
+            }
+        }
     }
 }

modules/elasticsearch/src/main/java/org/elasticsearch/action/get/TransportGetAction.java

@@ -20,16 +20,29 @@
 package org.elasticsearch.action.get;
 
 import com.google.inject.Inject;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.FieldSelector;
+import org.apache.lucene.document.Fieldable;
 import org.elasticsearch.ElasticSearchException;
 import org.elasticsearch.action.TransportActions;
 import org.elasticsearch.action.support.single.TransportSingleOperationAction;
 import org.elasticsearch.cluster.ClusterService;
+import org.elasticsearch.index.engine.Engine;
+import org.elasticsearch.index.mapper.*;
+import org.elasticsearch.index.service.IndexService;
 import org.elasticsearch.index.shard.service.IndexShard;
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
+import org.elasticsearch.util.lucene.Lucene;
 import org.elasticsearch.util.settings.Settings;
 
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Map;
+
+import static com.google.common.collect.Maps.*;
+
 /**
  * Performs the get operation.
  *
@@ -51,8 +64,97 @@ public class TransportGetAction extends TransportSingleOperationAction<GetReques
     }
 
     @Override protected GetResponse shardOperation(GetRequest request, int shardId) throws ElasticSearchException {
-        IndexShard indexShard = indicesService.indexServiceSafe(request.index()).shardSafe(shardId);
+        IndexService indexService = indicesService.indexServiceSafe(request.index());
-        return new GetResponse(request.index(), request.type(), request.id(), indexShard.get(request.type(), request.id()));
+        IndexShard indexShard = indexService.shardSafe(shardId);
+
+        DocumentMapper docMapper = indexService.mapperService().type(request.type());
+        if (docMapper == null) {
+            throw new DocumentMapperNotFoundException("No mapper found for type [" + request.type() + "]");
+        }
+
+        Engine.Searcher searcher = indexShard.searcher();
+        boolean exists = false;
+        byte[] source = null;
+        Map<String, GetField> fields = null;
+        try {
+            int docId = Lucene.docId(searcher.reader(), docMapper.uidMapper().term(request.type(), request.id()));
+            if (docId != Lucene.NO_DOC) {
+                exists = true;
+                FieldSelector fieldSelector = buildFieldSelectors(docMapper, request.fields());
+                if (fieldSelector != null) {
+                    Document doc = searcher.reader().document(docId, fieldSelector);
+                    source = extractSource(doc, docMapper);
+
+                    for (Object oField : doc.getFields()) {
+                        Fieldable field = (Fieldable) oField;
+                        String name = field.name();
+                        Object value = null;
+                        FieldMappers fieldMappers = docMapper.mappers().indexName(field.name());
+                        if (fieldMappers != null) {
+                            FieldMapper mapper = fieldMappers.mapper();
+                            if (mapper != null) {
+                                name = mapper.names().fullName();
+                                value = mapper.valueForSearch(field);
+                            }
+                        }
+                        if (value == null) {
+                            if (field.isBinary()) {
+                                value = field.getBinaryValue();
+                            } else {
+                                value = field.stringValue();
+                            }
+                        }
+
+                        if (fields == null) {
+                            fields = newHashMapWithExpectedSize(2);
+                        }
+
+                        GetField getField = fields.get(name);
+                        if (getField == null) {
+                            getField = new GetField(name, new ArrayList<Object>(2));
+                            fields.put(name, getField);
+                        }
+                        getField.values().add(value);
+                    }
+                }
+            }
+        } catch (IOException e) {
+            throw new ElasticSearchException("Failed to get type [" + request.type() + "] and id [" + request.id() + "]", e);
+        } finally {
+            searcher.release();
+        }
+        return new GetResponse(request.index(), request.type(), request.id(), exists, source, fields);
+    }
+
+    private FieldSelector buildFieldSelectors(DocumentMapper docMapper, String... fields) {
+        if (fields == null) {
+            return docMapper.sourceMapper().fieldSelector();
+        }
+
+        // don't load anything
+        if (fields.length == 0) {
+            return null;
+        }
+
+        FieldMappersFieldSelector fieldSelector = new FieldMappersFieldSelector();
+        for (String fieldName : fields) {
+            FieldMappers x = docMapper.mappers().smartName(fieldName);
+            if (x == null) {
+                throw new ElasticSearchException("No mapping for field [" + fieldName + "] in type [" + docMapper.type() + "]");
+            }
+            fieldSelector.add(x);
+        }
+        return fieldSelector;
+    }
+
+    private byte[] extractSource(Document doc, DocumentMapper documentMapper) {
+        byte[] source = null;
+        Fieldable sourceField = doc.getFieldable(documentMapper.sourceMapper().names().indexName());
+        if (sourceField != null) {
+            source = documentMapper.sourceMapper().value(sourceField);
+            doc.removeField(documentMapper.sourceMapper().names().indexName());
+        }
+        return source;
     }
 
     @Override protected GetRequest newRequest() {

modules/elasticsearch/src/main/java/org/elasticsearch/rest/action/get/RestGetAction.java

@@ -21,6 +21,7 @@ package org.elasticsearch.rest.action.get;
 
 import com.google.inject.Inject;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.get.GetField;
 import org.elasticsearch.action.get.GetRequest;
 import org.elasticsearch.action.get.GetResponse;
 import org.elasticsearch.client.Client;
@@ -29,7 +30,10 @@ import org.elasticsearch.util.json.JsonBuilder;
 import org.elasticsearch.util.settings.Settings;
 
 import java.io.IOException;
+import java.util.List;
+import java.util.regex.Pattern;
 
+import static com.google.common.collect.Lists.*;
 import static org.elasticsearch.rest.RestRequest.Method.*;
 import static org.elasticsearch.rest.RestResponse.Status.*;
 import static org.elasticsearch.rest.action.support.RestJsonBuilder.*;
@@ -39,6 +43,12 @@ import static org.elasticsearch.rest.action.support.RestJsonBuilder.*;
  */
 public class RestGetAction extends BaseRestHandler {
 
+    private final static Pattern fieldsPattern;
+
+    static {
+        fieldsPattern = Pattern.compile(",");
+    }
+
     @Inject public RestGetAction(Settings settings, Client client, RestController controller) {
         super(settings, client);
         controller.registerHandler(GET, "/{index}/{type}/{id}", this);
@@ -50,25 +60,69 @@ public class RestGetAction extends BaseRestHandler {
         getRequest.listenerThreaded(false);
         // if we have a local operation, execute it on a thread since we don't spawn
         getRequest.threadedOperation(true);
+
+
+        List<String> fields = request.params("field");
+        String sField = request.param("fields");
+        if (sField != null) {
+            String[] sFields = fieldsPattern.split(sField);
+            if (sFields != null) {
+                if (fields == null) {
+                    fields = newArrayListWithExpectedSize(sField.length());
+                }
+                for (String field : sFields) {
+                    fields.add(field);
+                }
+            }
+        }
+        if (fields != null) {
+            getRequest.fields(fields.toArray(new String[fields.size()]));
+        }
+
+
         client.get(getRequest, new ActionListener<GetResponse>() {
-            @Override public void onResponse(GetResponse result) {
+            @Override public void onResponse(GetResponse response) {
                 try {
-                    if (!result.exists()) {
+                    if (!response.exists()) {
                         JsonBuilder builder = restJsonBuilder(request);
                         builder.startObject();
-                        builder.field("_index", result.index());
+                        builder.field("_index", response.index());
-                        builder.field("_type", result.type());
+                        builder.field("_type", response.type());
-                        builder.field("_id", result.id());
+                        builder.field("_id", response.id());
                         builder.endObject();
                         channel.sendResponse(new JsonRestResponse(request, NOT_FOUND, builder));
                     } else {
                         JsonBuilder builder = restJsonBuilder(request);
                         builder.startObject();
-                        builder.field("_index", result.index());
+                        builder.field("_index", response.index());
-                        builder.field("_type", result.type());
+                        builder.field("_type", response.type());
-                        builder.field("_id", result.id());
+                        builder.field("_id", response.id());
+                        if (response.source() != null) {
                             builder.raw(", \"_source\" : ");
-                        builder.raw(result.source());
+                            builder.raw(response.source());
+                        }
+
+                        if (response.fields() != null && !response.fields().isEmpty()) {
+                            builder.startObject("fields");
+                            for (GetField field : response.fields().values()) {
+                                if (field.values().isEmpty()) {
+                                    continue;
+                                }
+                                if (field.values().size() == 1) {
+                                    builder.field(field.name(), field.values().get(0));
+                                } else {
+                                    builder.field(field.name());
+                                    builder.startArray();
+                                    for (Object value : field.values()) {
+                                        builder.value(value);
+                                    }
+                                    builder.endArray();
+                                }
+                            }
+                            builder.endObject();
+                        }
+
+
                         builder.endObject();
                         channel.sendResponse(new JsonRestResponse(request, OK, builder));
                     }