Plan compare tests

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/tpc/AbstractTpcQueryPlannerTest.java

@@ -0,0 +1,182 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.ignite.internal.processors.query.calcite.planner.tpc;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.sql.SqlExplainLevel;
+import org.apache.ignite.cache.query.annotations.QuerySqlFunction;
+import org.apache.ignite.calcite.CalciteQueryEngineConfiguration;
+import org.apache.ignite.configuration.CacheConfiguration;
+import org.apache.ignite.configuration.IgniteConfiguration;
+import org.apache.ignite.internal.IgniteEx;
+import org.apache.ignite.internal.IgnitionEx;
+import org.apache.ignite.internal.processors.query.calcite.CalciteQueryProcessor;
+import org.apache.ignite.internal.processors.query.calcite.planner.AbstractPlannerTest;
+import org.apache.ignite.internal.processors.query.calcite.planner.tpc.PlanChecker.AfterPlansTest;
+import org.apache.ignite.internal.processors.query.calcite.planner.tpc.PlanChecker.BeforePlansTest;
+import org.apache.ignite.internal.processors.query.calcite.schema.IgniteSchema;
+import org.apache.ignite.testframework.GridTestUtils;
+import org.junit.Test;
+import org.junit.runners.Parameterized;
+
+import static org.apache.ignite.internal.processors.query.calcite.planner.tpc.TpchHelper.loadFromResource;
+import static org.apache.ignite.internal.processors.query.calcite.planner.tpc.TpchHelper.scriptToQueries;
+import static org.apache.ignite.internal.processors.query.calcite.planner.tpc.TpchHelper.sql;
+import static org.apache.ignite.internal.processors.query.calcite.planner.tpc.TpchHelper.sqlTestName;
+
+/**
+ * Abstract test class to ensure a planner generates expected plan for TPC queries.
+ */
+public class AbstractTpcQueryPlannerTest extends AbstractPlannerTest {
+    /** Set to {@code true} to write plan files, instead of checking. */
+    private static final boolean UPDATE_PLAN = false;
+
+    /** Server node to run queries on. */
+    private static IgniteEx srv;
+
+    /** Query id. Set by {@link PlanChecker}. */
+    @Parameterized.Parameter
+    public String queryId;
+
+    /** Run once, before all queries check. */
+    @BeforePlansTest
+    public static void startAll(Class<?> testClass) throws Exception {
+        AbstractTpcQueryPlannerTest mock = new AbstractTpcQueryPlannerTest();
+
+        mock.beforeFirstTest();
+
+        IgniteConfiguration cfg = mock.getConfiguration("server");
+
+        cfg.getSqlConfiguration()
+            .setQueryEnginesConfiguration(new CalciteQueryEngineConfiguration().setDefault(true));
+
+        srv = (IgniteEx)IgnitionEx.start(cfg);
+
+        srv.getOrCreateCache(new CacheConfiguration<>("mock")
+            .setSqlFunctionClasses(AbstractTpcQueryPlannerTest.TpchUDF.class)
+            .setSqlSchema("PUBLIC"));
+
+        TpchHelper.testFiles(testClass, "ddl")
+            .filter(p -> p.toString().endsWith(".sql"))
+            .map(p -> loadFromResource(p.toString()))
+            .forEach(ddl -> scriptToQueries(ddl).forEach(q -> sql(srv, q)));
+    }
+
+    /** Run once, after all queries check. */
+    @AfterPlansTest
+    public static void stopAll(Class<?> testClass) {
+        if (srv != null) {
+            srv.close();
+
+            srv = null;
+        }
+    }
+
+    /** Test single query. */
+    @Test
+    public void testQuery() {
+        String actualPlan = queryPlan();
+
+        if (UPDATE_PLAN) {
+            updatePlan(actualPlan);
+
+            return;
+        }
+
+        boolean match = false;
+
+        String expPlan = TpchHelper.replaceIdAndHash(loadFromResource(String.format("%s/%s.plan", sqlTestName(getClass()), queryId)));
+
+        for (String possiblePlan : TpchHelper.expandTemplates(expPlan)) {
+            if (possiblePlan.equals(actualPlan)) {
+                match = true;
+
+                break;
+            }
+        }
+
+        if (!match) {
+            // This assertion will print nice diff in IDE that will help to investigate.
+            // Test will fail anyway.
+            assertEquals(expPlan, actualPlan);
+
+            assert false : "Should not happen";
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override protected boolean isSafeTopology() {
+        return false;
+    }
+
+    /** */
+    private String queryPlan() {
+        CalciteQueryProcessor engine = (CalciteQueryProcessor)srv.context().query().defaultQueryEngine();
+
+        Map<String, IgniteSchema> schemas = GridTestUtils.getFieldValue(engine.schemaHolder(), "igniteSchemas");
+
+        List<String> res = scriptToQueries(loadFromResource(sqlTestName(getClass()) + "/" + queryId + ".sql")).stream().map(qry -> {
+            try {
+                return RelOptUtil.toString(physicalPlan(plannerCtx(qry, schemas.values(), null)), SqlExplainLevel.ALL_ATTRIBUTES);
+            }
+            catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        })
+            .map(TpchHelper::replaceIdAndHash)
+            .collect(Collectors.toList());
+
+        assertEquals(1, res.size());
+
+        return res.get(0);
+    }
+
+    /** */
+    private void updatePlan(String newPlan) {
+        Path targetDir = Path.of("./src/test/resources/" + sqlTestName(getClass()));
+
+        // A targetDirectory must be specified by hand when expected plans are generated.
+        if (targetDir == null) {
+            throw new RuntimeException("Please provide target directory to where save generated plans."
+                + " Usually plans are kept in resource folder of tests within the same module.");
+        }
+
+        try {
+            Files.createDirectories(targetDir);
+
+            Files.writeString(targetDir.resolve(String.format("%s.plan", queryId)), newPlan);
+        }
+        catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /** */
+    public static class TpchUDF {
+        /** */
+        @QuerySqlFunction(alias = "SUBSTR")
+        public static String substr(String str, int from, int cnt) {
+            return str.substring(from, from + cnt);
+        }
+    }
+}

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/tpc/PlanChecker.java

@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.ignite.internal.processors.query.calcite.planner.tpc;
+
+import java.io.IOException;
+import java.lang.annotation.Annotation;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.junit.runner.Runner;
+import org.junit.runner.notification.RunNotifier;
+import org.junit.runners.Suite;
+import org.junit.runners.model.InitializationError;
+import org.junit.runners.model.TestClass;
+import org.junit.runners.parameterized.BlockJUnit4ClassRunnerWithParameters;
+import org.junit.runners.parameterized.TestWithParameters;
+
+/**
+ * Suite to check queries plans.
+ * Makes pretty test names to simplify CI tracking.
+ */
+public class PlanChecker extends Suite {
+    /**
+     * Only called reflectively. Do not use programmatically.
+     */
+    public PlanChecker(Class<?> klass) throws Throwable {
+        super(klass, createRunnersForParameters(new TestClass(klass)).collect(Collectors.toList()));
+    }
+
+    /** */
+    private static Stream<Runner> createRunnersForParameters(TestClass testClass) throws IOException {
+        Stream<String> queries = TpchHelper.testFiles(testClass.getJavaClass())
+            .filter(p -> p.toString().endsWith(".sql"))
+            .sorted()
+            .map(p -> p.getFileName().toString().replace(".sql", ""));
+
+        return queries
+            .map(qryId -> new TestWithParameters("[queryId=" + qryId + "]", testClass, Collections.singletonList(qryId)))
+            .map(test -> {
+                try {
+                    return new BlockJUnit4ClassRunnerWithParameters(test);
+                }
+                catch (InitializationError e) {
+                    throw new RuntimeException(e);
+                }
+            });
+    }
+
+    /** {@inheritDoc} */
+    @Override public void run(RunNotifier notifier) {
+        runAnnotated(BeforePlansTest.class);
+
+        try {
+            super.run(notifier);
+        }
+        finally {
+            runAnnotated(AfterPlansTest.class);
+        }
+    }
+
+    /** Runs static void method of test class annotated with the specific annotation. */
+    private void runAnnotated(Class<? extends Annotation> annotation) {
+        getTestClass().getAnnotatedMethods(annotation).forEach(m -> {
+            List<Throwable> errors = new ArrayList<>();
+
+            m.validatePublicVoid(true, errors);
+
+            try {
+                if (!errors.isEmpty())
+                    throw errors.get(0);
+
+                m.invokeExplosively(getTestClass().getJavaClass(), getTestClass().getJavaClass());
+            }
+            catch (Throwable e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+    /** */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.TYPE)
+    public @interface PlansTest {
+        /** @return Test name. */
+        String name() default "";
+    }
+
+    /** */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.METHOD)
+    public @interface BeforePlansTest {
+        // No-op.
+    }
+
+    /** */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.METHOD)
+    public @interface AfterPlansTest {
+        // No-op.
+    }
+}