org.elasticsearch.search.aggregations.support.ValuesSourceRegistry Maven / Gradle / Ivy

Go to download
/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0 and the Server Side Public License, v 1; you may not use this file except
 * in compliance with, at your election, the Elastic License 2.0 or the Server
 * Side Public License, v 1.
 */
package org.elasticsearch.search.aggregations.support;

import org.elasticsearch.index.query.SearchExecutionContext;
import org.elasticsearch.search.SearchModule;
import org.elasticsearch.search.aggregations.AggregationExecutionException;

import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * {@link ValuesSourceRegistry} holds the mapping from {@link ValuesSourceType}s to functions for building aggregation components.  DO NOT
 * directly instantiate this class, instead get an already-configured copy from {@link SearchExecutionContext#getValuesSourceRegistry()},
 * or (in the case of some test scenarios only) directly from {@link SearchModule#getValuesSourceRegistry()}
 *
 */
public class ValuesSourceRegistry {

    public static final class RegistryKey {
        private final String name;
        private final Class supplierType;

        public RegistryKey(String name, Class supplierType) {
            this.name = Objects.requireNonNull(name);
            this.supplierType = Objects.requireNonNull(supplierType);
        }

        public String getName() {
            return name;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            RegistryKey that = (RegistryKey) o;
            return name.equals(that.name) && supplierType.equals(that.supplierType);
        }

        @Override
        public int hashCode() {
            return Objects.hash(name, supplierType);
        }
    }

    @SuppressWarnings("rawtypes")
    public static final RegistryKey UNREGISTERED_KEY = new RegistryKey<>("unregistered", RegistryKey.class);

    public static class Builder {
        private final AggregationUsageService.Builder usageServiceBuilder;
        private Map, List>> aggregatorRegistry = new HashMap<>();

        public Builder() {
            this.usageServiceBuilder = new AggregationUsageService.Builder();
        }

        /**
         * Register a ValuesSource to Aggregator mapping. This method registers mappings that only apply to a
         * single {@link ValuesSourceType}
         * @param registryKey The name of the family of aggregations paired with the expected component supplier type for this
         *                    family of aggregations.  Generally, the aggregation builder is expected to define a constant for use as the
         *                    registryKey
         * @param valuesSourceType The ValuesSourceType this mapping applies to.
         * @param aggregatorSupplier An Aggregation-specific specialization of ComponentSupplier which will construct the mapped aggregator
         * @param registerUsage Flag to indicate if this aggregation values source combo should be added to the usage registry.
         *                      Aggregations that set this to false should register with the usage registry through some other path.
         */
        public  void register(
            RegistryKey registryKey,
            ValuesSourceType valuesSourceType,
            T aggregatorSupplier,
            boolean registerUsage
        ) {
            if (aggregatorRegistry.containsKey(registryKey) == false) {
                aggregatorRegistry.put(registryKey, new ArrayList<>());
            }
            aggregatorRegistry.get(registryKey).add(new AbstractMap.SimpleEntry<>(valuesSourceType, aggregatorSupplier));
            if (registerUsage) {
                registerUsage(registryKey.getName(), valuesSourceType);
            }
        }

        /**
         * Register a ValuesSource to Aggregator mapping. This version provides a convenience method for mappings that apply to a
         * known list of {@link ValuesSourceType}
         * @param registryKey The name of the family of aggregations paired with the expected component supplier type for this
         *                    family of aggregations.  Generally, the aggregation builder is expected to define a constant for use as the
         *                    registryKey
         * @param valuesSourceTypes The ValuesSourceTypes this mapping applies to.
         * @param aggregatorSupplier An Aggregation-specific specialization of ComponentSupplier which will construct the mapped aggregator
         * @param registerUsage Flag to indicate if this aggregation values source combo should be added to the usage registry.
         *                      Aggregations that set this to false should register with the usage registry through some other path.
         */
        public  void register(
            RegistryKey registryKey,
            List valuesSourceTypes,
            T aggregatorSupplier,
            boolean registerUsage
        ) {
            for (ValuesSourceType valuesSourceType : valuesSourceTypes) {
                register(registryKey, valuesSourceType, aggregatorSupplier, registerUsage);
            }
        }

        public void registerUsage(String aggregationName, ValuesSourceType valuesSourceType) {
            usageServiceBuilder.registerAggregationUsage(aggregationName, valuesSourceType.typeName());
        }

        public void registerUsage(String aggregationName) {
            usageServiceBuilder.registerAggregationUsage(aggregationName);
        }

        public ValuesSourceRegistry build() {
            return new ValuesSourceRegistry(aggregatorRegistry, usageServiceBuilder.build());
        }
    }

    private static Map, Map> copyMap(
        Map, List>> mutableMap
    ) {
        /*
         Make an immutable copy of our input map. Since this is write once, read many, we'll spend a bit of extra time to shape this
         into a Map.of(), which is more read optimized than just using a hash map.
         */
        Map, Map> tmp = new HashMap<>();
        mutableMap.forEach((key, value) -> tmp.put(key, value.stream().collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))));
        return Collections.unmodifiableMap(tmp);
    }

    /** Maps Aggregation names to (ValuesSourceType, Supplier) pairs, keyed by ValuesSourceType */
    private final AggregationUsageService usageService;
    private final Map, Map> aggregatorRegistry;

    public ValuesSourceRegistry(
        Map, List>> aggregatorRegistry,
        AggregationUsageService usageService
    ) {
        this.usageService = usageService;
        this.aggregatorRegistry = copyMap(aggregatorRegistry);
    }

    public boolean isRegistered(RegistryKey registryKey) {
        return aggregatorRegistry.containsKey(registryKey);
    }

    public  T getAggregator(RegistryKey registryKey, ValuesSourceConfig valuesSourceConfig) {
        if (registryKey != null && aggregatorRegistry.containsKey(registryKey)) {
            @SuppressWarnings("unchecked")
            T supplier = (T) aggregatorRegistry.get(registryKey).get(valuesSourceConfig.valueSourceType());
            if (supplier == null) {
                throw new IllegalArgumentException(
                    valuesSourceConfig.getDescription() + " is not supported for aggregation [" + registryKey.getName() + "]"
                );
            }
            return supplier;
        }
        throw new AggregationExecutionException("Unregistered Aggregation [" + registryKey.getName() + "]");
    }

    public AggregationUsageService getUsageService() {
        return usageService;
    }
}