001/*
002 * (C) Copyright 2014-2015 Nuxeo SA (http://nuxeo.com/) and contributors.
003 *
004 * All rights reserved. This program and the accompanying materials
005 * are made available under the terms of the GNU Lesser General Public License
006 * (LGPL) version 2.1 which accompanies this distribution, and is available at
007 * http://www.gnu.org/licenses/lgpl-2.1.html
008 *
009 * This library is distributed in the hope that it will be useful,
010 * but WITHOUT ANY WARRANTY; without even the implied warranty of
011 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
012 * Lesser General Public License for more details.
013 *
014 * Contributors:
015 *     Benoit Delbosc
016 *     Florent Guillaume
017 */
018package org.nuxeo.elasticsearch.core;
019
020import java.io.Serializable;
021import java.util.ArrayList;
022import java.util.HashMap;
023import java.util.Iterator;
024import java.util.List;
025import java.util.Map;
026import java.util.NoSuchElementException;
027
028import org.elasticsearch.action.search.SearchResponse;
029import org.elasticsearch.search.SearchHit;
030import org.elasticsearch.search.SearchHitField;
031import org.nuxeo.ecm.core.api.IterableQueryResult;
032import org.nuxeo.ecm.core.query.sql.NXQL;
033import org.nuxeo.ecm.core.schema.types.Type;
034import org.nuxeo.ecm.core.schema.types.primitives.DateType;
035
036/**
037 * Iterable query result of the results of an Elasticsearch query.
038 * <p>
039 * Loads all results in memory.
040 *
041 * @since 7.2
042 */
043public class EsResultSetImpl implements IterableQueryResult, Iterator<Map<String, Serializable>> {
044
045    private final SearchResponse response;
046
047    private final Map<String, Type> selectFieldsAndTypes;
048
049    boolean closed;
050
051    protected List<Map<String, Serializable>> maps;
052
053    protected long size;
054
055    private long pos;
056
057    public EsResultSetImpl(SearchResponse response, Map<String, Type> selectFieldsAndTypes) {
058        this.response = response;
059        this.selectFieldsAndTypes = selectFieldsAndTypes;
060        maps = buildMaps();
061        size = maps.size();
062    }
063
064    protected List<Map<String, Serializable>> buildMaps() {
065        List<Map<String, Serializable>> rows = new ArrayList<>(response.getHits().getHits().length);
066        Map<String, Serializable> emptyRow = new HashMap<>(selectFieldsAndTypes.size());
067        for (String fieldName : selectFieldsAndTypes.keySet()) {
068            emptyRow.put(fieldName, null);
069        }
070        for (SearchHit hit : response.getHits().getHits()) {
071            Map<String, Serializable> row = new HashMap<>(emptyRow);
072            for (SearchHitField field : hit.getFields().values()) {
073                String name = field.getName();
074                Serializable value = field.<Serializable> getValue();
075                // type conversion
076                Type type;
077                if (value instanceof String && (type = selectFieldsAndTypes.get(name)) instanceof DateType) {
078                    // convert back to calendar
079                    value = (Serializable) type.decode(((String) value));
080                }
081                row.put(name, value);
082            }
083            if (selectFieldsAndTypes.containsKey(NXQL.ECM_FULLTEXT_SCORE)) {
084                row.put(NXQL.ECM_FULLTEXT_SCORE, Double.valueOf(hit.getScore()));
085            }
086            rows.add(row);
087        }
088        return rows;
089    }
090
091    @Override
092    public void close() {
093        closed = true;
094        pos = -1;
095    }
096
097    @Override
098    public boolean isLife() {
099        return !closed;
100    }
101
102    @Override
103    public long size() {
104        return response.getHits().getTotalHits();
105    }
106
107    @Override
108    public long pos() {
109        return pos;
110    }
111
112    @Override
113    public void skipTo(long pos) {
114        if (pos < 0) {
115            pos = 0;
116        } else if (pos > size) {
117            pos = size;
118        }
119        this.pos = pos;
120    }
121
122    @Override
123    public Iterator<Map<String, Serializable>> iterator() {
124        return this;
125    }
126
127    @Override
128    public boolean hasNext() {
129        return pos < size;
130    }
131
132    @Override
133    public Map<String, Serializable> next() {
134        if (closed || pos == size) {
135            throw new NoSuchElementException();
136        }
137        Map<String, Serializable> map = maps.get((int) pos);
138        pos++;
139        return map;
140    }
141
142}