hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with this
 * work for additional information regarding copyright ownership. The ASF
 * 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.apache.hugegraph.backend.query;

import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.function.Function;

import org.apache.tinkerpop.gremlin.structure.util.CloseableIterator;

import org.apache.hugegraph.HugeException;
import org.apache.hugegraph.backend.id.Id;
import org.apache.hugegraph.iterator.CIter;
import org.apache.hugegraph.iterator.FlatMapperIterator;
import org.apache.hugegraph.iterator.ListIterator;
import org.apache.hugegraph.iterator.MapperIterator;
import org.apache.hugegraph.perf.PerfUtil.Watched;
import org.apache.hugegraph.type.Idfiable;
import org.apache.hugegraph.util.E;
import org.apache.hugegraph.util.InsertionOrderUtil;

public class QueryResults<R> {

    private static final Iterator<?> EMPTY_ITERATOR = new EmptyIterator<>();

    private static final QueryResults<?> EMPTY = new QueryResults<>(
                                                 emptyIterator(), Query.NONE);

    private final Iterator<R> results;
    private final List<Query> queries;

    public QueryResults(Iterator<R> results, Query query) {
        this(results);
        this.addQuery(query);
    }

    private QueryResults(Iterator<R> results) {
        this.results = results;
        this.queries = InsertionOrderUtil.newList();
    }

    public void setQuery(Query query) {
        if (this.queries.size() > 0) {
            this.queries.clear();
        }
        this.addQuery(query);
    }

    private void addQuery(Query query) {
        E.checkNotNull(query, "query");
        this.queries.add(query);
    }

    private void addQueries(List<Query> queries) {
        assert !queries.isEmpty();
        for (Query query : queries) {
            this.addQuery(query);
        }
    }

    public Iterator<R> iterator() {
        return this.results;
    }

    public R one() {
        return one(this.results);
    }

    public QueryResults<R> toList() {
        QueryResults<R> fetched = new QueryResults<>(toList(this.results));
        fetched.addQueries(this.queries);
        return fetched;
    }

    public List<Query> queries() {
        return Collections.unmodifiableList(this.queries);
    }

    public <T extends Idfiable> Iterator<T> keepInputOrderIfNeeded(
                                            Iterator<T> origin) {
        if (!origin.hasNext()) {
            // None result found
            return origin;
        }
        Collection<Id> ids;
        if (!this.mustSortByInputIds() || this.paging() ||
            (ids = this.queryIds()).size() <= 1) {
            /*
             * Return the original iterator if it's paging query or if the
             * query input is less than one id, or don't have to do sort.
             * NOTE: queryIds() only return the first batch of index query
             */
            return origin;
        }

        // Fill map with all elements
        Map<Id, T> map = InsertionOrderUtil.newMap();
        QueryResults.fillMap(origin, map);

        if (map.size() > ids.size()) {
            /*
             * This means current query is part of QueryResults. For example,
             * g.V().has('country', 'china').has('city', within('HK', 'BJ'))
             * will be converted to
             * g.V().has('country', 'china').has('city', 'HK') or
             * g.V().has('country', 'china').has('city', 'BJ'),
             * and ids is just first index subquery's id, not all.
             */
            ids = map.keySet();
        }

        return new MapperIterator<>(ids.iterator(), map::get);
    }

    private boolean mustSortByInputIds() {
        assert !this.queries.isEmpty() : this;
        for (Query query : this.queries) {
            if (query instanceof IdQuery &&
                ((IdQuery) query).mustSortByInput()) {
                return true;
            }
        }
        return false;
    }

    private boolean paging() {
        assert !this.queries.isEmpty();
        for (Query query : this.queries) {
            Query origin = query.originQuery();
            if (query.paging() || origin != null && origin.paging()) {
                return true;
            }
        }
        return false;
    }

    @SuppressWarnings("unused")
    private boolean bigCapacity() {
        assert !this.queries.isEmpty();
        for (Query query : this.queries) {
            if (query.bigCapacity()) {
                return true;
            }
        }
        return false;
    }

    private Collection<Id> queryIds() {
        assert !this.queries.isEmpty();
        if (this.queries.size() == 1) {
            return this.queries.get(0).ids();
        }

        Set<Id> ids = InsertionOrderUtil.newSet();
        for (Query query : this.queries) {
            ids.addAll(query.ids());
        }
        return ids;
    }

    @Watched
    public static <T> ListIterator<T> toList(Iterator<T> iterator) {
        try {
            return new ListIterator<>(Query.DEFAULT_CAPACITY, iterator);
        } finally {
            CloseableIterator.closeIterator(iterator);
        }
    }

    @Watched
    public static <T> void fillList(Iterator<T> iterator, List<T> list) {
        try {
            while (iterator.hasNext()) {
                T result = iterator.next();
                list.add(result);
                Query.checkForceCapacity(list.size());
            }
        } finally {
            CloseableIterator.closeIterator(iterator);
        }
    }

    @Watched
    public static <T extends Idfiable> void fillMap(Iterator<T> iterator,
                                                    Map<Id, T> map) {
        try {
            while (iterator.hasNext()) {
                T result = iterator.next();
                assert result.id() != null;
                map.put(result.id(), result);
                Query.checkForceCapacity(map.size());
            }
        } finally {
            CloseableIterator.closeIterator(iterator);
        }
    }

    public static <T, R> QueryResults<R> flatMap(
                  Iterator<T> iterator, Function<T, QueryResults<R>> func) {
        @SuppressWarnings("unchecked")
        QueryResults<R>[] qr = new QueryResults[1];
        qr[0] = new QueryResults<>(new FlatMapperIterator<>(iterator, i -> {
            QueryResults<R> results = func.apply(i);
            if (results == null || !results.iterator().hasNext()) {
                return null;
            }
            /*
             * NOTE: should call results.iterator().hasNext() before
             * results.queries() to collect sub-query with index query
             */
            qr[0].addQueries(results.queries());
            return results.iterator();
        }));
        return qr[0];
    }

    @Watched
    public static <T> T one(Iterator<T> iterator) {
        try {
            if (iterator.hasNext()) {
                T result = iterator.next();
                if (iterator.hasNext()) {
                    throw new HugeException("Expect just one result, " +
                                            "but got at least two: [%s, %s]",
                                            result, iterator.next());
                }
                return result;
            }
        } finally {
            CloseableIterator.closeIterator(iterator);
        }
        return null;
    }

    public static <T> Iterator<T> iterator(T elem) {
        return new OneIterator<>(elem);
    }

    @SuppressWarnings("unchecked")
    public static <T> QueryResults<T> empty() {
        return (QueryResults<T>) EMPTY;
    }

    @SuppressWarnings("unchecked")
    public static <T> Iterator<T> emptyIterator() {
        return (Iterator<T>) EMPTY_ITERATOR;
    }

    public interface Fetcher<R> extends Function<Query, QueryResults<R>> {}

    private static class EmptyIterator<T> implements CIter<T> {

        @Override
        public Object metadata(String meta, Object... args) {
            return null;
        }

        @Override
        public boolean hasNext() {
            return false;
        }

        @Override
        public T next() {
            throw new NoSuchElementException();
        }

        @Override
        public void close() throws Exception {
            // pass
        }
    }

    private static class OneIterator<T> implements CIter<T> {

        private T element;

        public OneIterator(T element) {
            assert element != null;
            this.element = element;
        }

        @Override
        public Object metadata(String meta, Object... args) {
            return null;
        }

        @Override
        public boolean hasNext() {
            return this.element != null;
        }

        @Override
        public T next() {
            if (this.element == null) {
                throw new NoSuchElementException();
            }
            T result = this.element;
            this.element = null;
            return result;
        }

        @Override
        public void close() throws Exception {
            // pass
        }
    }
}