001/*
002 * (C) Copyright 2014 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 *     bdelbosc
016 */
017package org.nuxeo.elasticsearch.aggregate;
018
019import java.util.ArrayList;
020import java.util.Collection;
021import java.util.Collections;
022import java.util.Comparator;
023import java.util.List;
024
025import org.codehaus.jackson.annotate.JsonIgnore;
026import org.elasticsearch.index.query.FilterBuilders;
027import org.elasticsearch.index.query.OrFilterBuilder;
028import org.elasticsearch.index.query.RangeFilterBuilder;
029import org.elasticsearch.search.aggregations.AggregationBuilders;
030import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
031import org.elasticsearch.search.aggregations.bucket.range.Range;
032import org.elasticsearch.search.aggregations.bucket.range.RangeBuilder;
033import org.nuxeo.ecm.core.api.DocumentModel;
034import org.nuxeo.ecm.platform.query.api.AggregateDefinition;
035import org.nuxeo.ecm.platform.query.api.AggregateRangeDefinition;
036import org.nuxeo.ecm.platform.query.core.BucketRange;
037
038/**
039 * @since 6.0
040 */
041public class RangeAggregate extends AggregateEsBase<BucketRange> {
042
043    public RangeAggregate(AggregateDefinition definition, DocumentModel searchDocument) {
044        super(definition, searchDocument);
045    }
046
047    @JsonIgnore
048    @Override
049    public RangeBuilder getEsAggregate() {
050        RangeBuilder ret = AggregationBuilders.range(getId()).field(getField());
051        for (AggregateRangeDefinition range : getRanges()) {
052            if (range.getFrom() != null) {
053                if (range.getTo() != null) {
054                    ret.addRange(range.getKey(), range.getFrom(), range.getTo());
055                } else {
056                    ret.addUnboundedFrom(range.getKey(), range.getFrom());
057                }
058            } else if (range.getTo() != null) {
059                ret.addUnboundedTo(range.getKey(), range.getTo());
060            }
061        }
062        return ret;
063    }
064
065    @JsonIgnore
066    @Override
067    public OrFilterBuilder getEsFilter() {
068        if (getSelection().isEmpty()) {
069            return null;
070        }
071        OrFilterBuilder ret = FilterBuilders.orFilter();
072        for (AggregateRangeDefinition range : getRanges()) {
073            if (getSelection().contains(range.getKey())) {
074                RangeFilterBuilder rangeFilter = FilterBuilders.rangeFilter(getField());
075                if (range.getFrom() != null) {
076                    rangeFilter.gte(range.getFrom());
077                }
078                if (range.getTo() != null) {
079                    rangeFilter.lt(range.getTo());
080                }
081                ret.add(rangeFilter);
082            }
083        }
084        return ret;
085    }
086
087    @JsonIgnore
088    @Override
089    public void parseEsBuckets(Collection<? extends MultiBucketsAggregation.Bucket> buckets) {
090        List<BucketRange> nxBuckets = new ArrayList<>(buckets.size());
091        for (MultiBucketsAggregation.Bucket bucket : buckets) {
092            Range.Bucket rangeBucket = (Range.Bucket) bucket;
093            nxBuckets.add(new BucketRange(bucket.getKey(), rangeBucket.getFrom(), rangeBucket.getTo(),
094                    rangeBucket.getDocCount()));
095        }
096        Collections.sort(nxBuckets, new BucketRangeComparator());
097        this.buckets = nxBuckets;
098    }
099
100    protected class BucketRangeComparator implements Comparator<BucketRange> {
101        @Override
102        public int compare(BucketRange arg0, BucketRange arg1) {
103            return definition.getAggregateRangeDefinitionOrderMap().get(arg0.getKey()).compareTo(
104                    definition.getAggregateRangeDefinitionOrderMap().get(arg1.getKey()));
105        }
106    }
107
108}