io.trino.benchmark.driver.Suite Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of trino-benchmark-driver Show documentation
There is a newer version: 363
Show newest version
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.trino.benchmark.driver;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.regex.Pattern;

import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.Streams.stream;
import static io.airlift.json.JsonCodec.mapJsonCodec;
import static java.util.Objects.requireNonNull;

public class Suite
{
    private final String name;
    private final Map sessionProperties;
    private final List schemaNameTemplates;
    private final List queryNamePatterns;

    public Suite(String name, Map sessionProperties, Iterable schemaNameTemplates, Iterable queryNamePatterns)
    {
        this.name = requireNonNull(name, "name is null");
        this.sessionProperties = sessionProperties == null ? ImmutableMap.of() : ImmutableMap.copyOf(sessionProperties);
        this.schemaNameTemplates = ImmutableList.copyOf(requireNonNull(schemaNameTemplates, "schemaNameTemplates is null"));
        this.queryNamePatterns = ImmutableList.copyOf(requireNonNull(queryNamePatterns, "queryNamePatterns is null"));
    }

    public String getName()
    {
        return name;
    }

    public Map getSessionProperties()
    {
        return sessionProperties;
    }

    public List getSchemaNameTemplates()
    {
        return schemaNameTemplates;
    }

    public List getQueryNamePatterns()
    {
        return queryNamePatterns;
    }

    public List selectSchemas(Iterable schemas)
    {
        if (schemaNameTemplates.isEmpty()) {
            return ImmutableList.of();
        }

        ImmutableList.Builder benchmarkSchemas = ImmutableList.builder();
        for (RegexTemplate schemaNameTemplate : schemaNameTemplates) {
            for (String schema : schemas) {
                Optional> tags = schemaNameTemplate.parse(schema);
                if (tags.isPresent()) {
                    benchmarkSchemas.add(new BenchmarkSchema(schema, tags.get()));
                }
            }
        }
        return benchmarkSchemas.build();
    }

    public List selectQueries(Iterable queries)
    {
        if (getQueryNamePatterns().isEmpty()) {
            return ImmutableList.copyOf(queries);
        }

        List filteredQueries = stream(queries)
                .filter(query -> getQueryNamePatterns().stream().anyMatch(pattern -> pattern.matcher(query.getName()).matches()))
                .collect(toImmutableList());

        return filteredQueries;
    }

    @Override
    public String toString()
    {
        return toStringHelper(this)
                .add("name", name)
                .add("sessionProperties", sessionProperties)
                .add("queryNamePatterns", queryNamePatterns)
                .toString();
    }

    public static List readSuites(File file)
            throws IOException
    {
        requireNonNull(file, "file is null");
        checkArgument(file.canRead(), "Cannot read file: %s", file);
        byte[] json = Files.readAllBytes(file.toPath());
        Map options = mapJsonCodec(String.class, OptionsJson.class).fromJson(json);
        ImmutableList.Builder runOptions = ImmutableList.builder();
        for (Entry entry : options.entrySet()) {
            runOptions.add(entry.getValue().toSuite(entry.getKey()));
        }
        return runOptions.build();
    }

    public static class OptionsJson
    {
        private final List schema;
        private final Map session;
        private final List query;

        @JsonCreator
        public OptionsJson(
                @JsonProperty("schema") List schema,
                @JsonProperty("session") Map session,
                @JsonProperty("query") List query)
        {
            this.schema = ImmutableList.copyOf(requireNonNull(schema, "schema is null"));
            this.session = ImmutableMap.copyOf(requireNonNull(session, "session is null"));
            this.query = requireNonNull(query, "query is null");
        }

        public Suite toSuite(String name)
        {
            ImmutableList.Builder queryNameTemplates = ImmutableList.builder();
            for (String q : query) {
                queryNameTemplates.add(Pattern.compile(q));
            }
            ImmutableList.Builder schemaNameTemplates = ImmutableList.builder();
            for (String s : schema) {
                schemaNameTemplates.add(new RegexTemplate(s));
            }
            return new Suite(name, session, schemaNameTemplates.build(), queryNameTemplates.build());
        }
    }
}