From 7cc21d120fcbeb44b1989a1858937f943aec22a8 Mon Sep 17 00:00:00 2001 From: TJ Banghart Date: Sat, 25 Oct 2025 12:56:15 -0700 Subject: [PATCH] [CALCITE-7427] Connection config for enabling RuleMatchVisualize --- .../config/CalciteConnectionConfig.java | 3 + .../config/CalciteConnectionConfigImpl.java | 5 + .../config/CalciteConnectionProperty.java | 7 +- .../calcite/jdbc/CalciteConnectionImpl.java | 13 ++ .../plan/visualizer/RuleMatchVisualizer.java | 8 +- .../calcite/util/RuleMatchVisualizerHook.java | 155 ++++++++++++++ .../org/apache/calcite/test/JdbcTest.java | 1 + .../util/RuleMatchVisualizerHookTest.java | 200 ++++++++++++++++++ 8 files changed, 388 insertions(+), 4 deletions(-) create mode 100644 core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java create mode 100644 core/src/test/java/org/apache/calcite/util/RuleMatchVisualizerHookTest.java diff --git a/core/src/main/java/org/apache/calcite/config/CalciteConnectionConfig.java b/core/src/main/java/org/apache/calcite/config/CalciteConnectionConfig.java index 3999a6ce9b64..d7bd82d96f99 100644 --- a/core/src/main/java/org/apache/calcite/config/CalciteConnectionConfig.java +++ b/core/src/main/java/org/apache/calcite/config/CalciteConnectionConfig.java @@ -111,6 +111,9 @@ public interface CalciteConnectionConfig extends ConnectionConfig { /** Returns the value of {@link CalciteConnectionProperty#TOPDOWN_OPT}. */ boolean topDownOpt(); + /** Returns the value of {@link CalciteConnectionProperty#RULE_VISUALIZER_DIR}. */ + @Nullable String ruleVisualizerDir(); + /** Returns the value of {@link CalciteConnectionProperty#META_TABLE_FACTORY}, * or a default meta table factory if not set. If * {@code defaultMetaTableFactory} is not null, the result is never null. */ diff --git a/core/src/main/java/org/apache/calcite/config/CalciteConnectionConfigImpl.java b/core/src/main/java/org/apache/calcite/config/CalciteConnectionConfigImpl.java index 7d9f9f76d1cc..2272c14b12a1 100644 --- a/core/src/main/java/org/apache/calcite/config/CalciteConnectionConfigImpl.java +++ b/core/src/main/java/org/apache/calcite/config/CalciteConnectionConfigImpl.java @@ -215,6 +215,11 @@ public boolean isSet(CalciteConnectionProperty property) { .getBoolean(); } + @Override public @Nullable String ruleVisualizerDir() { + return CalciteConnectionProperty.RULE_VISUALIZER_DIR.wrap(properties) + .getString(); + } + @Override public @PolyNull T metaTableFactory( Class metaTableFactoryClass, @PolyNull T defaultMetaTableFactory) { diff --git a/core/src/main/java/org/apache/calcite/config/CalciteConnectionProperty.java b/core/src/main/java/org/apache/calcite/config/CalciteConnectionProperty.java index a9c7627025cd..0969834a9c8d 100644 --- a/core/src/main/java/org/apache/calcite/config/CalciteConnectionProperty.java +++ b/core/src/main/java/org/apache/calcite/config/CalciteConnectionProperty.java @@ -186,7 +186,12 @@ public enum CalciteConnectionProperty implements ConnectionProperty { LENIENT_OPERATOR_LOOKUP("lenientOperatorLookup", Type.BOOLEAN, false, false), /** Whether to enable top-down optimization in Volcano planner. */ - TOPDOWN_OPT("topDownOpt", Type.BOOLEAN, CalciteSystemProperty.TOPDOWN_OPT.value(), false); + TOPDOWN_OPT("topDownOpt", Type.BOOLEAN, CalciteSystemProperty.TOPDOWN_OPT.value(), false), + + /** Directory path for RuleMatchVisualizer output. + * If set, enables visualization of the rule matching process during query optimization. + * The visualizer will create HTML and JSON files in the specified directory. */ + RULE_VISUALIZER_DIR("ruleVisualizerDir", Type.STRING, null, false); private final String camelName; private final Type type; diff --git a/core/src/main/java/org/apache/calcite/jdbc/CalciteConnectionImpl.java b/core/src/main/java/org/apache/calcite/jdbc/CalciteConnectionImpl.java index ac2fbc127f51..d6cf006421df 100644 --- a/core/src/main/java/org/apache/calcite/jdbc/CalciteConnectionImpl.java +++ b/core/src/main/java/org/apache/calcite/jdbc/CalciteConnectionImpl.java @@ -67,6 +67,7 @@ import org.apache.calcite.tools.RelRunner; import org.apache.calcite.util.BuiltInMethod; import org.apache.calcite.util.Holder; +import org.apache.calcite.util.RuleMatchVisualizerHook; import org.apache.calcite.util.Util; import com.google.common.collect.ImmutableList; @@ -188,6 +189,18 @@ void init() { true, true); } } + + // Enable RuleMatchVisualizer if configured + CalciteConnectionConfig cfg = config(); + String vizDir = cfg.ruleVisualizerDir(); + if (vizDir != null && !vizDir.isEmpty()) { + try { + RuleMatchVisualizerHook.INSTANCE.enableFromConnection(this); + } catch (Exception e) { + // Log but don't fail connection if visualizer setup fails + System.err.println("Warning: Failed to enable RuleMatchVisualizer: " + e.getMessage()); + } + } } @Override public T unwrap(Class iface) throws SQLException { diff --git a/core/src/main/java/org/apache/calcite/plan/visualizer/RuleMatchVisualizer.java b/core/src/main/java/org/apache/calcite/plan/visualizer/RuleMatchVisualizer.java index 57bf1a990683..63ef7999eed4 100644 --- a/core/src/main/java/org/apache/calcite/plan/visualizer/RuleMatchVisualizer.java +++ b/core/src/main/java/org/apache/calcite/plan/visualizer/RuleMatchVisualizer.java @@ -80,6 +80,8 @@ public class RuleMatchVisualizer implements RelOptListener { private static final String INITIAL = "INITIAL"; private static final String FINAL = "FINAL"; public static final String DEFAULT_SET = "default"; + public static final String HTML_FILE_PREFIX = "planner-viz"; + public static final String DATA_FILE_PREFIX = "planner-viz-data"; // default HTML template can be edited at // core/src/main/resources/org/apache/calcite/plan/visualizer/viz-template.html @@ -393,10 +395,10 @@ public void writeToFile() { requireNonNull(cl.getResourceAsStream(templatePath)); String htmlTemplate = IOUtils.toString(resourceAsStream, UTF_8); - String htmlFileName = "planner-viz" + outputSuffix + ".html"; - String dataFileName = "planner-viz-data" + outputSuffix + ".js"; + String htmlFileName = HTML_FILE_PREFIX + outputSuffix + ".html"; + String dataFileName = DATA_FILE_PREFIX + outputSuffix + ".js"; - String replaceString = "src=\"planner-viz-data.js\""; + String replaceString = "src=\"" + DATA_FILE_PREFIX + ".js\""; int replaceIndex = htmlTemplate.indexOf(replaceString); String htmlContent = htmlTemplate.substring(0, replaceIndex) + "src=\"" + dataFileName + "\"" diff --git a/core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java b/core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java new file mode 100644 index 000000000000..953e8b422897 --- /dev/null +++ b/core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java @@ -0,0 +1,155 @@ +/* + * 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.calcite.util; + +import org.apache.calcite.config.CalciteConnectionConfig; +import org.apache.calcite.jdbc.CalciteConnection; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.visualizer.RuleMatchVisualizer; +import org.apache.calcite.runtime.Hook; + +import java.io.File; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +/** + * Utility class to enable RuleMatchVisualizer for Calcite connections. + * + *

This class provides hooks to automatically attach a RuleMatchVisualizer + * to planners when a connection specifies the ruleVisualizerDir property. + * + *

Usage in JDBC URL: + * 

+ * jdbc:calcite:ruleVisualizerDir=/tmp/calcite-viz
+ *
+ * + *

Or programmatically: + * 

+ * RuleMatchVisualizerHook.enable("/tmp/calcite-viz");
+ *
+ */ +public class RuleMatchVisualizerHook { + public static final RuleMatchVisualizerHook INSTANCE = new RuleMatchVisualizerHook(); + + private final Map visualizerMap = new HashMap<>(); + private final AtomicInteger queryCounter = new AtomicInteger(0); + + private Hook.Closeable hookCloseable = Hook.Closeable.EMPTY; + + /** Private constructor to prevent instantiation. */ + private RuleMatchVisualizerHook() {} + + /** + * Enables the visualizer for all subsequent queries with the specified output directory. + * + * @param outputDir Directory where visualization files will be created + */ + public synchronized void enable(String outputDir) { + hookCloseable.close(); + + // Ensure the output directory exists + File dir = new File(outputDir); + if (!dir.exists()) { + boolean madeDir = dir.mkdirs(); + assert madeDir : "Failed to create directory: " + outputDir; + } + + // Install the hook + hookCloseable = Hook.PLANNER.addThread((Consumer) planner -> { + attachVisualizer(planner, outputDir); + }); + } + + /** + * Enables the visualizer using the connection's configuration. + * This method checks if the connection has the ruleVisualizerDir property set. + * + * @param connection The Calcite connection + */ + public synchronized void enableFromConnection(CalciteConnection connection) { + CalciteConnectionConfig config = connection.config(); + String vizDir = config.ruleVisualizerDir(); + + if (vizDir != null && !vizDir.isEmpty()) { + enable(vizDir); + } + } + + /** + * Disables the visualizer. + */ + public synchronized void disable() { + hookCloseable.close(); + + // Write any pending visualizations + for (RuleMatchVisualizer viz : visualizerMap.values()) { + viz.writeToFile(); + } + visualizerMap.clear(); + } + + /** + * Attaches a visualizer to the given planner. + */ + private void attachVisualizer(RelOptPlanner planner, String outputDir) { + + // Check if we've already attached a visualizer to this planner + if (visualizerMap.containsKey(planner)) { + return; + } + + int queryNum = queryCounter.incrementAndGet(); + int queryStart = (int) System.currentTimeMillis() / 1000; + String suffix = String.format(Locale.ROOT, "query_%d_%d", queryNum, queryStart); + + // Create and attach the visualizer + RuleMatchVisualizer visualizer = new RuleMatchVisualizer(outputDir, suffix); + visualizer.attachTo(planner); + visualizerMap.put(planner, visualizer); + + // For HepPlanner, we need to manually write the output + if (planner instanceof HepPlanner) { + // Add a hook to write the visualization after the planner finishes + Hook.PLAN_BEFORE_IMPLEMENTATION.addThread(relRoot -> { + RuleMatchVisualizer viz = visualizerMap.get(planner); + if (viz != null) { + viz.writeToFile(); + visualizerMap.remove(planner); + } + }); + } + // VolcanoPlanner automatically calls writeToFile() when done + + System.out.println("RuleMatchVisualizer enabled: Output will be written to " + + outputDir + File.separator + suffix + "*"); + } + + /** + * Checks the system property and enables visualization if set. + * This can be called at application startup. + */ + public void checkSystemProperty() { + String vizDir = System.getProperty("calcite.visualizer.dir"); + if (vizDir != null && !vizDir.isEmpty()) { + enable(vizDir); + } + } +} diff --git a/core/src/test/java/org/apache/calcite/test/JdbcTest.java b/core/src/test/java/org/apache/calcite/test/JdbcTest.java index de66b636f59c..93088450422f 100644 --- a/core/src/test/java/org/apache/calcite/test/JdbcTest.java +++ b/core/src/test/java/org/apache/calcite/test/JdbcTest.java @@ -975,6 +975,7 @@ static void checkMockDdl(AtomicInteger counter, boolean hasCommit, assertTrue(names.contains("SCHEMA")); assertTrue(names.contains("TIME_ZONE")); assertTrue(names.contains("MATERIALIZATIONS_ENABLED")); + assertTrue(names.contains("RULE_VISUALIZER_DIR")); } /** diff --git a/core/src/test/java/org/apache/calcite/util/RuleMatchVisualizerHookTest.java b/core/src/test/java/org/apache/calcite/util/RuleMatchVisualizerHookTest.java new file mode 100644 index 000000000000..4616521cdfda --- /dev/null +++ b/core/src/test/java/org/apache/calcite/util/RuleMatchVisualizerHookTest.java @@ -0,0 +1,200 @@ +/* + * 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.calcite.util; + +import org.apache.calcite.jdbc.CalciteConnection; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.nio.file.Path; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.apache.calcite.plan.visualizer.RuleMatchVisualizer.DATA_FILE_PREFIX; +import static org.apache.calcite.plan.visualizer.RuleMatchVisualizer.HTML_FILE_PREFIX; + +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.hasSize; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import static java.util.Objects.requireNonNull; + +/** + * Tests for {@link RuleMatchVisualizerHook}. + */ +class RuleMatchVisualizerHookTest { + + @TempDir + Path tempDir; + + @AfterEach + void cleanup() { + RuleMatchVisualizerHook.INSTANCE.disable(); + } + + @Test void testEnableDisable() { + String vizDir = tempDir.toString(); + + // Enable visualizer + RuleMatchVisualizerHook.INSTANCE.enable(vizDir); + + // Directory should be created + File dir = new File(vizDir); + assertTrue(dir.exists()); + assertTrue(dir.isDirectory()); + + // Disable visualizer + RuleMatchVisualizerHook.INSTANCE.disable(); + } + + @Test void testEnableFromConnection() throws Exception { + String vizDir = tempDir.resolve("viz").toString(); + + // Create connection with visualizer property + String url = "jdbc:calcite:ruleVisualizerDir=" + vizDir; + try (Connection conn = DriverManager.getConnection(url)) { + assertThat(conn, instanceOf(CalciteConnection.class)); + CalciteConnection calciteConn = (CalciteConnection) conn; + + // Check that the property is set + String configuredDir = calciteConn.config().ruleVisualizerDir(); + assertThat(configuredDir, notNullValue()); + assertThat(configuredDir, is(vizDir)); + + // The hook should be enabled automatically by the connection + // Let's run a simple query to trigger the planner + try (Statement stmt = conn.createStatement()) { + String sql = "SELECT 1 FROM (VALUES (1))"; + try (ResultSet rs = stmt.executeQuery(sql)) { + assertTrue(rs.next()); + assertThat(rs.getInt(1), is(1)); + } + } + } + + File dir = new File(vizDir); + assertTrue(dir.exists()); + assertTrue(dir.isDirectory()); + Map> matched = + Arrays.stream(requireNonNull(dir.listFiles())) + .collect(Collectors.partitioningBy(f -> f.getName().contains(DATA_FILE_PREFIX))); + + List dataFiles = matched.get(true); + List htmlFiles = matched.get(false); + + assertThat(dataFiles, hasSize(1)); + assertThat(htmlFiles, hasSize(1)); + assertThat(dataFiles.get(0).getName(), containsString(DATA_FILE_PREFIX)); + assertThat(htmlFiles.get(0).getName(), containsString(HTML_FILE_PREFIX)); + } + + @Test void testSystemProperty() { + String vizDir = tempDir.resolve("sysprop").toString(); + + // Set system property + System.setProperty("calcite.visualizer.dir", vizDir); + + try { + // Check system property + RuleMatchVisualizerHook.INSTANCE.checkSystemProperty(); + + // Directory should be created + File dir = new File(vizDir); + assertTrue(dir.exists()); + assertTrue(dir.isDirectory()); + } finally { + // Clean up system property + System.clearProperty("calcite.visualizer.dir"); + RuleMatchVisualizerHook.INSTANCE.disable(); + } + } + + @Test void testMultipleQueries() throws Exception { + String vizDir = tempDir.resolve("multi").toString(); + + // Create connection with visualizer property + String url = "jdbc:calcite:ruleVisualizerDir=" + vizDir; + try (Connection conn = DriverManager.getConnection(url)) { + try (Statement stmt = conn.createStatement()) { + // Execute multiple queries + for (int i = 1; i <= 3; i++) { + String sql = "SELECT " + i; + try (ResultSet rs = stmt.executeQuery(sql)) { + assertTrue(rs.next()); + assertTrue(rs.getInt(1) == i); + } + } + + String sql = "SELECT * FROM (VALUES (1, 'a'), (2, 'b')) AS t(x, y) WHERE x > 0"; + try (ResultSet rs = stmt.executeQuery(sql)) { + assertTrue(rs.next()); + } + } + } + + // Directory should exist + File dir = new File(vizDir); + assertTrue(dir.exists()); + assertTrue(dir.isDirectory()); + + Map> matched = + Arrays.stream(requireNonNull(dir.listFiles())) + .collect(Collectors.partitioningBy(f -> f.getName().contains(DATA_FILE_PREFIX))); + + List dataFiles = matched.get(true); + List htmlFiles = matched.get(false); + + assertThat(dataFiles, hasSize(3)); + assertThat(htmlFiles, hasSize(3)); + } + + @Test void testConnectionWithoutVisualizer() throws Exception { + // Create connection without visualizer property + String url = "jdbc:calcite:"; + try (Connection conn = DriverManager.getConnection(url)) { + assertThat(conn, instanceOf(CalciteConnection.class)); + CalciteConnection calciteConn = (CalciteConnection) conn; + + // Check that the property is not set + String configuredDir = calciteConn.config().ruleVisualizerDir(); + assertNull(configuredDir); + + // Query should still work + try (Statement stmt = conn.createStatement()) { + String sql = "SELECT 1 FROM (VALUES (1))"; + try (ResultSet rs = stmt.executeQuery(sql)) { + assertTrue(rs.next()); + assertThat(rs.getInt(1), is(1)); + } + } + } + } +}