org.elasticsearch.xpack.core.ml.dataframe.analyses.Types Maven / Gradle / Ivy

/*
 * 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; you may not use this file except in compliance with the Elastic License
 * 2.0.
 */
package org.elasticsearch.xpack.core.ml.dataframe.analyses;

import org.elasticsearch.index.mapper.BooleanFieldMapper;
import org.elasticsearch.index.mapper.IpFieldMapper;
import org.elasticsearch.index.mapper.KeywordFieldMapper;
import org.elasticsearch.index.mapper.NumberFieldMapper.NumberType;
import org.elasticsearch.index.mapper.TextFieldMapper;

import java.util.Collections;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * Helper class that defines groups of types
 */
public final class Types {

    private Types() {}

    private static final Set CATEGORICAL_TYPES =
        Collections.unmodifiableSet(
            Stream.of(TextFieldMapper.CONTENT_TYPE, KeywordFieldMapper.CONTENT_TYPE, IpFieldMapper.CONTENT_TYPE)
                .collect(Collectors.toSet()));

    private static final Set NUMERICAL_TYPES =
        Collections.unmodifiableSet(
            Stream.concat(Stream.of(NumberType.values()).map(NumberType::typeName), Stream.of("scaled_float", "unsigned_long"))
                .collect(Collectors.toSet()));

    private static final Set DISCRETE_NUMERICAL_TYPES =
        Collections.unmodifiableSet(
            Stream.concat(
                Stream.of(NumberType.BYTE, NumberType.SHORT, NumberType.INTEGER, NumberType.LONG).map(NumberType::typeName),
                Stream.of("unsigned_long"))
                .collect(Collectors.toSet()));

    private static final Set BOOL_TYPES = Collections.singleton(BooleanFieldMapper.CONTENT_TYPE);

    public static Set categorical() {
        return CATEGORICAL_TYPES;
    }

    public static Set numerical() {
        return NUMERICAL_TYPES;
    }

    public static Set discreteNumerical() {
        return DISCRETE_NUMERICAL_TYPES;
    }

    public static Set bool() {
        return BOOL_TYPES;
    }
}