InternalScriptFacet.java

/*
 * Copyright 2012 Igor Motov
 *
 * Licensed 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.search.facet.script;

import org.elasticsearch.client.Client;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.bytes.HashedBytesArray;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentBuilderString;
import org.elasticsearch.script.ExecutableScript;
import org.elasticsearch.script.ScriptService;
import org.elasticsearch.search.facet.Facet;
import org.elasticsearch.search.facet.InternalFacet;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.elasticsearch.common.collect.Lists.newArrayList;

/**
 *
 */
public class InternalScriptFacet extends InternalFacet implements ScriptFacet {
    private static final BytesReference STREAM_TYPE = new HashedBytesArray(Strings.toUTF8Bytes("script"));

    private Object facet;
    private String scriptLang;
    private String reduceScript;
    private Map<String, Object> reduceParams;
    private ScriptService scriptService;
    private Client client;


    public static void registerStreams(ScriptService scriptService, Client client) {
        Streams.registerStream(new ScriptFacetStream(scriptService, client), STREAM_TYPE);
    }

    private InternalScriptFacet(ScriptService scriptService, Client client) {
        this.scriptService = scriptService;
        this.client = client;
    }

    private InternalScriptFacet(String name, ScriptService scriptService, Client client) {
        super(name);
        this.scriptService = scriptService;
        this.client = client;
    }

    public InternalScriptFacet(String name, Object facet, String scriptLang, String reduceScript, Map<String, Object> reduceParams, ScriptService scriptService, Client client) {
        this(name, scriptService, client);
        this.facet = facet;
        this.reduceScript = reduceScript;
        this.reduceParams = reduceParams;
        this.scriptLang = scriptLang;
    }

    @Override
    public BytesReference streamType() {
        return STREAM_TYPE;
    }

    @Override
    public Facet reduce(ReduceContext reduceContext) {
        List<Object> facetObjects = newArrayList();
        for (Facet facet : reduceContext.facets()) {
            InternalScriptFacet mapReduceFacet = (InternalScriptFacet) facet;
            facetObjects.add(mapReduceFacet.facet());
        }
        InternalScriptFacet firstFacet = ((InternalScriptFacet) reduceContext.facets().get(0));
        Object facet;
        if (firstFacet.reduceScript() != null) {
            Map<String, Object> params;
            if (firstFacet.reduceParams() != null) {
                params = new HashMap<String, Object>(firstFacet.reduceParams());
            } else {
                params = new HashMap<String, Object>();
            }
            params.put("facets", facetObjects);
            params.put("_client", client);
            ExecutableScript script = scriptService.executable(firstFacet.scriptLang(), firstFacet.reduceScript(), params);
            facet = script.run();
        } else {
            facet = facetObjects;
        }
        return new InternalScriptFacet(firstFacet.getName(), facet, firstFacet.scriptLang(), firstFacet.reduceScript(), firstFacet.reduceParams(), scriptService, client);
    }

    @Override
    public String getType() {
        return ScriptFacet.TYPE;
    }

    @Override
    public void readFrom(StreamInput in) throws IOException {
        super.readFrom(in);
        scriptLang = in.readOptionalString();
        reduceScript = in.readOptionalString();
        reduceParams = in.readMap();
        facet = in.readGenericValue();
    }

    @Override
    public void writeTo(StreamOutput out) throws IOException {
        super.writeTo(out);
        out.writeOptionalString(scriptLang);
        out.writeOptionalString(reduceScript);
        out.writeMap(reduceParams);
        out.writeGenericValue(facet);
    }

    @Override
    public Object facet() {
        return facet;
    }

    @Override
    public Object getFacet() {
        return facet();
    }

    public String scriptLang() {
        return scriptLang;
    }

    public String reduceScript() {
        return reduceScript;
    }

    public Map<String, Object> reduceParams() {
        return reduceParams;
    }

    static final class Fields {
        static final XContentBuilderString _TYPE = new XContentBuilderString("_type");
        static final XContentBuilderString FACET = new XContentBuilderString("facet");
    }

    @Override
    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
        builder.startObject(getName());
        builder.field(Fields._TYPE, ScriptFacet.TYPE);
        builder.field(Fields.FACET, facet);
        builder.endObject();
        return builder;
    }

    public static InternalScriptFacet readMapReduceFacet(StreamInput in, ScriptService scriptService, Client client) throws IOException {
        InternalScriptFacet facet = new InternalScriptFacet(scriptService, client);
        facet.readFrom(in);
        return facet;
    }

    private static class ScriptFacetStream implements InternalFacet.Stream {

        private ScriptService scriptService;
        private Client client;


        public ScriptFacetStream(ScriptService scriptService, Client client) {
            this.scriptService = scriptService;
            this.client = client;
        }

        @Override
        public Facet readFacet(StreamInput in) throws IOException {
            return InternalScriptFacet.readMapReduceFacet(in, scriptService, client);
        }

    }
}