Examples with AccumulatorStateFactory io.trino.spi.function.AccumulatorStateFactory used on opensource projects

Example 1 with AccumulatorStateFactory

use of io.trino.spi.function.AccumulatorStateFactory in project trino by trinodb.

the class StateCompiler method generateStateFactory.

public static <T extends AccumulatorState> AccumulatorStateFactory<T> generateStateFactory(Class<T> clazz, Map<String, Type> fieldTypes) {
    AccumulatorStateMetadata metadata = getMetadataAnnotation(clazz);
    if (metadata != null && metadata.stateFactoryClass() != AccumulatorStateFactory.class) {
        try {
            // noinspection unchecked
            return (AccumulatorStateFactory<T>) metadata.stateFactoryClass().getConstructor().newInstance();
        } catch (InstantiationException | IllegalAccessException | NoSuchMethodException | InvocationTargetException e) {
            throw new RuntimeException(e);
        }
    }
    // grouped aggregation state fields use engine classes, so generated class must be able to see both plugin and system classes
    DynamicClassLoader classLoader = new DynamicClassLoader(clazz.getClassLoader(), StateCompiler.class.getClassLoader());
    Class<? extends T> singleStateClass = generateSingleStateClass(clazz, fieldTypes, classLoader);
    Class<? extends T> groupedStateClass = generateGroupedStateClass(clazz, fieldTypes, classLoader);
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName(clazz.getSimpleName() + "Factory"), type(Object.class), type(AccumulatorStateFactory.class));
    // Generate constructor
    definition.declareDefaultConstructor(a(PUBLIC));
    // Generate single state creation method
    definition.declareMethod(a(PUBLIC), "createSingleState", type(AccumulatorState.class)).getBody().newObject(singleStateClass).dup().invokeConstructor(singleStateClass).retObject();
    // Generate grouped state creation method
    definition.declareMethod(a(PUBLIC), "createGroupedState", type(AccumulatorState.class)).getBody().newObject(groupedStateClass).dup().invokeConstructor(groupedStateClass).retObject();
    Class<?> factoryClass = defineClass(definition, AccumulatorStateFactory.class, classLoader);
    try {
        // noinspection unchecked
        return (AccumulatorStateFactory<T>) factoryClass.getConstructor().newInstance();
    } catch (ReflectiveOperationException e) {
        throw new RuntimeException(e);
    }
}