Skip to content

Commit

Permalink
Add tests that verify timestamp indexes can be used in MSE
Browse files Browse the repository at this point in the history
  • Loading branch information
gortiz committed Dec 20, 2024
1 parent 1338c15 commit 4689d87
Show file tree
Hide file tree
Showing 6 changed files with 268 additions and 57 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,6 @@
import org.apache.pinot.spi.utils.CommonConstants.Broker;
import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
import org.apache.pinot.spi.utils.DataSizeUtils;
import org.apache.pinot.spi.utils.TimestampIndexUtils;
import org.apache.pinot.spi.utils.builder.TableNameBuilder;
import org.apache.pinot.sql.FilterKind;
import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
package org.apache.pinot.integration.tests;

import com.fasterxml.jackson.databind.JsonNode;
import com.google.common.collect.Lists;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
Expand All @@ -28,7 +29,6 @@
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -91,9 +91,10 @@ public abstract class BaseClusterIntegrationTest extends ClusterTest {
protected static final List<String> DEFAULT_NO_DICTIONARY_COLUMNS =
Arrays.asList("ActualElapsedTime", "ArrDelay", "DepDelay", "CRSDepTime");
protected static final String DEFAULT_SORTED_COLUMN = "Carrier";
protected static final List<String> DEFAULT_INVERTED_INDEX_COLUMNS = Arrays.asList("FlightNum", "Origin", "Quarter");
private static final List<String> DEFAULT_BLOOM_FILTER_COLUMNS = Arrays.asList("FlightNum", "Origin");
private static final List<String> DEFAULT_RANGE_INDEX_COLUMNS = Collections.singletonList("Origin");
protected static final List<String> DEFAULT_INVERTED_INDEX_COLUMNS
= Lists.newArrayList("FlightNum", "Origin", "Quarter");
private static final List<String> DEFAULT_BLOOM_FILTER_COLUMNS = Lists.newArrayList("FlightNum", "Origin");
private static final List<String> DEFAULT_RANGE_INDEX_COLUMNS = Lists.newArrayList("Origin");
protected static final int DEFAULT_NUM_REPLICAS = 1;
protected static final boolean DEFAULT_NULL_HANDLING_ENABLED = false;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -509,7 +509,7 @@ protected JsonNode getDebugInfo(final String uri)
/**
* Queries the broker's sql query endpoint (/query/sql)
*/
protected JsonNode postQuery(String query)
public JsonNode postQuery(String query)
throws Exception {
return postQuery(query, getBrokerQueryApiUrl(getBrokerBaseApiUrl(), useMultiStageQueryEngine()), null,
getExtraQueryProperties());
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/**
* 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.pinot.integration.tests;

import com.fasterxml.jackson.databind.JsonNode;
import org.intellij.lang.annotations.Language;
import org.testng.Assert;


public interface ExplainIntegrationTestTrait {

JsonNode postQuery(@Language("sql") String query)
throws Exception;

default void explainLogical(@Language("sql") String query, String expected) {
try {
JsonNode jsonNode = postQuery("explain plan without implementation for " + query);
JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);

Assert.assertEquals(plan.asText(), expected);
} catch (RuntimeException e) {
throw e;
} catch (Exception e) {
throw new RuntimeException(e);
}
}

default void explain(@Language("sql") String query, String expected) {
try {
JsonNode jsonNode = postQuery("explain plan for " + query);
JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);

Assert.assertEquals(plan.asText(), expected);
} catch (RuntimeException e) {
throw e;
} catch (Exception e) {
throw new RuntimeException(e);
}
}

default void explainVerbose(@Language("sql") String query, String expected) {
try {
JsonNode jsonNode = postQuery("set explainPlanVerbose=true; explain plan for " + query);
JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);

String actual = plan.asText()
.replaceAll("numDocs=\\[[^\\]]*]", "numDocs=[any]")
.replaceAll("segment=\\[[^\\]]*]", "segment=[any]")
.replaceAll("totalDocs=\\[[^\\]]*]", "totalDocs=[any]");


Assert.assertEquals(actual, expected);
} catch (RuntimeException e) {
throw e;
} catch (Exception e) {
throw new RuntimeException(e);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,24 +18,22 @@
*/
package org.apache.pinot.integration.tests;

import com.fasterxml.jackson.databind.JsonNode;
import java.io.File;
import java.util.List;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.util.TestUtils;
import org.intellij.lang.annotations.Language;
import org.testcontainers.shaded.org.apache.commons.io.FileUtils;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;


public class MultiStageEngineExplainIntegrationTest extends BaseClusterIntegrationTest {
public class MultiStageEngineExplainIntegrationTest extends BaseClusterIntegrationTest
implements ExplainIntegrationTestTrait {

@BeforeClass
public void setUp()
Expand Down Expand Up @@ -161,7 +159,7 @@ public void simpleQueryVerbose() {
+ " Project(columns=[[]])\n"
+ " DocIdSet(maxDocs=[10000])\n"
+ " FilterMatchEntireSegment(numDocs=[any])\n");
//@formatter:on
//@formatter:on
}

@Test
Expand All @@ -171,7 +169,7 @@ public void simpleQueryLogical() {
"Execution Plan\n"
+ "LogicalProject(EXPR$0=[1])\n"
+ " LogicalTableScan(table=[[default, mytable]])\n");
//@formatter:on
//@formatter:on
}

@AfterClass
Expand All @@ -186,49 +184,4 @@ public void tearDown()

FileUtils.deleteDirectory(_tempDir);
}

private void explainVerbose(@Language("sql") String query, String expected) {
try {
JsonNode jsonNode = postQuery("set explainPlanVerbose=true; explain plan for " + query);
JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);

String actual = plan.asText()
.replaceAll("numDocs=\\[[^\\]]*]", "numDocs=[any]")
.replaceAll("segment=\\[[^\\]]*]", "segment=[any]")
.replaceAll("totalDocs=\\[[^\\]]*]", "totalDocs=[any]");


Assert.assertEquals(actual, expected);
} catch (RuntimeException e) {
throw e;
} catch (Exception e) {
throw new RuntimeException(e);
}
}

private void explain(@Language("sql") String query, String expected) {
try {
JsonNode jsonNode = postQuery("explain plan for " + query);
JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);

Assert.assertEquals(plan.asText(), expected);
} catch (RuntimeException e) {
throw e;
} catch (Exception e) {
throw new RuntimeException(e);
}
}

private void explainLogical(@Language("sql") String query, String expected) {
try {
JsonNode jsonNode = postQuery("set explainAskingServers=false; explain plan for " + query);
JsonNode plan = jsonNode.get("resultTable").get("rows").get(0).get(1);

Assert.assertEquals(plan.asText(), expected);
} catch (RuntimeException e) {
throw e;
} catch (Exception e) {
throw new RuntimeException(e);
}
}
}
Loading

0 comments on commit 4689d87

Please sign in to comment.