001/*
002 * (C) Copyright 2014 Nuxeo SA (http://nuxeo.com/) and others.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 *
016 * Contributors:
017 *     bdelbosc
018 */
019package org.nuxeo.elasticsearch.aggregate;
020
021import java.util.ArrayList;
022import java.util.Collection;
023import java.util.Collections;
024import java.util.Comparator;
025import java.util.List;
026
027import org.codehaus.jackson.annotate.JsonIgnore;
028import org.elasticsearch.index.query.FilterBuilders;
029import org.elasticsearch.index.query.OrFilterBuilder;
030import org.elasticsearch.index.query.RangeFilterBuilder;
031import org.elasticsearch.search.aggregations.AggregationBuilders;
032import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
033import org.elasticsearch.search.aggregations.bucket.range.Range;
034import org.elasticsearch.search.aggregations.bucket.range.RangeBuilder;
035import org.nuxeo.ecm.core.api.DocumentModel;
036import org.nuxeo.ecm.platform.query.api.AggregateDefinition;
037import org.nuxeo.ecm.platform.query.api.AggregateRangeDefinition;
038import org.nuxeo.ecm.platform.query.core.BucketRange;
039
040/**
041 * @since 6.0
042 */
043public class RangeAggregate extends AggregateEsBase<BucketRange> {
044
045    public RangeAggregate(AggregateDefinition definition, DocumentModel searchDocument) {
046        super(definition, searchDocument);
047    }
048
049    @JsonIgnore
050    @Override
051    public RangeBuilder getEsAggregate() {
052        RangeBuilder ret = AggregationBuilders.range(getId()).field(getField());
053        for (AggregateRangeDefinition range : getRanges()) {
054            if (range.getFrom() != null) {
055                if (range.getTo() != null) {
056                    ret.addRange(range.getKey(), range.getFrom(), range.getTo());
057                } else {
058                    ret.addUnboundedFrom(range.getKey(), range.getFrom());
059                }
060            } else if (range.getTo() != null) {
061                ret.addUnboundedTo(range.getKey(), range.getTo());
062            }
063        }
064        return ret;
065    }
066
067    @JsonIgnore
068    @Override
069    public OrFilterBuilder getEsFilter() {
070        if (getSelection().isEmpty()) {
071            return null;
072        }
073        OrFilterBuilder ret = FilterBuilders.orFilter();
074        for (AggregateRangeDefinition range : getRanges()) {
075            if (getSelection().contains(range.getKey())) {
076                RangeFilterBuilder rangeFilter = FilterBuilders.rangeFilter(getField());
077                if (range.getFrom() != null) {
078                    rangeFilter.gte(range.getFrom());
079                }
080                if (range.getTo() != null) {
081                    rangeFilter.lt(range.getTo());
082                }
083                ret.add(rangeFilter);
084            }
085        }
086        return ret;
087    }
088
089    @JsonIgnore
090    @Override
091    public void parseEsBuckets(Collection<? extends MultiBucketsAggregation.Bucket> buckets) {
092        List<BucketRange> nxBuckets = new ArrayList<>(buckets.size());
093        for (MultiBucketsAggregation.Bucket bucket : buckets) {
094            Range.Bucket rangeBucket = (Range.Bucket) bucket;
095            nxBuckets.add(new BucketRange(bucket.getKey(), rangeBucket.getFrom(), rangeBucket.getTo(),
096                    rangeBucket.getDocCount()));
097        }
098        Collections.sort(nxBuckets, new BucketRangeComparator());
099        this.buckets = nxBuckets;
100    }
101
102    protected class BucketRangeComparator implements Comparator<BucketRange> {
103        @Override
104        public int compare(BucketRange arg0, BucketRange arg1) {
105            return definition.getAggregateRangeDefinitionOrderMap().get(arg0.getKey()).compareTo(
106                    definition.getAggregateRangeDefinitionOrderMap().get(arg1.getKey()));
107        }
108    }
109
110}