Skip to content

Commit 49b7fa6

Browse files
shizy818shizy818
committed
mock executeForTableModel instead of fetchCteQueryResult
1 parent 140ee02 commit 49b7fa6

File tree

3 files changed

+144
-77
lines changed

3 files changed

+144
-77
lines changed

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializer.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ public class CteMaterializer {
7373
private static final Logger LOGGER = LoggerFactory.getLogger(CteMaterializer.class);
7474

7575
private static final Coordinator coordinator = Coordinator.getInstance();
76+
private static final SessionManager sessionManager = SessionManager.getInstance();
7677

7778
public void materializeCTE(Analysis analysis, MPPQueryContext context) {
7879
analysis
@@ -104,7 +105,7 @@ public void cleanUpCTE(MPPQueryContext context) {
104105

105106
public CteDataStore fetchCteQueryResult(
106107
MPPQueryContext context, Table table, Query query, With with) {
107-
final long queryId = SessionManager.getInstance().requestQueryId();
108+
final long queryId = sessionManager.requestQueryId();
108109
Throwable t = null;
109110
CteDataStore cteDataStore = null;
110111
long startTime = System.nanoTime();
@@ -138,10 +139,9 @@ public CteDataStore fetchCteQueryResult(
138139
coordinator.executeForTableModel(
139140
q,
140141
new SqlParser(),
141-
SessionManager.getInstance().getCurrSession(),
142+
sessionManager.getCurrSession(),
142143
queryId,
143-
SessionManager.getInstance()
144-
.getSessionInfoOfTableModel(SessionManager.getInstance().getCurrSession()),
144+
sessionManager.getSessionInfoOfTableModel(sessionManager.getCurrSession()),
145145
String.format("Materialize query for CTE '%s'", table.getName()),
146146
LocalExecutionPlanner.getInstance().metadata,
147147
context.getCteDataStores(),
@@ -153,6 +153,7 @@ public CteDataStore fetchCteQueryResult(
153153
}
154154
// query execution
155155
QueryExecution execution = (QueryExecution) coordinator.getQueryExecution(queryId);
156+
156157
// get table schema
157158
DatasetHeader datasetHeader = coordinator.getQueryExecution(queryId).getDatasetHeader();
158159
TableSchema tableSchema = getTableSchema(datasetHeader, table.getName().toString());

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerTest.java‎

Lines changed: 127 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -21,44 +21,47 @@
2121

2222
package org.apache.iotdb.db.queryengine.plan.relational.planner;
2323

24-
import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory;
25-
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
24+
import org.apache.iotdb.common.rpc.thrift.TSStatus;
25+
import org.apache.iotdb.commons.exception.IoTDBException;
26+
import org.apache.iotdb.commons.schema.column.ColumnHeader;
27+
import org.apache.iotdb.db.protocol.session.SessionManager;
28+
import org.apache.iotdb.db.queryengine.common.QueryId;
29+
import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
30+
import org.apache.iotdb.db.queryengine.plan.Coordinator;
31+
import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
32+
import org.apache.iotdb.db.queryengine.plan.execution.QueryExecution;
2633
import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan;
27-
import org.apache.iotdb.db.queryengine.plan.relational.metadata.ColumnSchema;
28-
import org.apache.iotdb.db.queryengine.plan.relational.metadata.TableSchema;
2934
import org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern;
3035
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression;
3136
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Expression;
3237
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LongLiteral;
33-
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Query;
34-
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.QuerySpecification;
35-
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SelectItem;
36-
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SingleColumn;
3738
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SymbolReference;
38-
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Table;
39-
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.With;
40-
import org.apache.iotdb.db.utils.cte.CteDataStore;
4139

4240
import com.google.common.collect.ImmutableList;
4341
import com.google.common.collect.ImmutableMap;
4442
import com.google.common.collect.ImmutableSet;
4543
import org.apache.tsfile.enums.TSDataType;
46-
import org.apache.tsfile.read.common.type.Type;
47-
import org.apache.tsfile.read.common.type.TypeFactory;
44+
import org.apache.tsfile.read.common.block.TsBlock;
45+
import org.apache.tsfile.read.common.block.column.LongColumn;
46+
import org.apache.tsfile.read.common.block.column.TimeColumn;
4847
import org.junit.Before;
48+
import org.junit.BeforeClass;
4949
import org.junit.Test;
50+
import org.junit.runner.RunWith;
5051
import org.mockito.Mockito;
51-
import org.mockito.invocation.InvocationOnMock;
52+
import org.powermock.core.classloader.annotations.PowerMockIgnore;
53+
import org.powermock.core.classloader.annotations.PrepareForTest;
54+
import org.powermock.modules.junit4.PowerMockRunner;
5255

5356
import java.util.Collections;
5457
import java.util.List;
5558
import java.util.Optional;
56-
import java.util.stream.Collectors;
5759

5860
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanAssert.assertPlan;
5961
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.aggregation;
6062
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.aggregationFunction;
6163
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.cteScan;
64+
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.explainAnalyze;
6265
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.filter;
6366
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.limit;
6467
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.offset;
@@ -70,26 +73,79 @@
7073
import static org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression.Operator.GREATER_THAN;
7174
import static org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SortItem.NullOrdering.LAST;
7275
import static org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SortItem.Ordering.ASCENDING;
76+
import static org.powermock.api.mockito.PowerMockito.mockStatic;
77+
import static org.powermock.api.mockito.PowerMockito.when;
7378

79+
@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"})
80+
@RunWith(PowerMockRunner.class)
81+
@PrepareForTest({Coordinator.class, SessionManager.class})
7482
public class CteMaterializerTest {
75-
private PlanTester planTester;
83+
private static PlanTester planTester;
84+
private static Coordinator mockCoordinator;
7685

77-
@Before
78-
public void setUp() {
86+
@BeforeClass
87+
public static void prepareEnv() {
7988
planTester = new PlanTester();
80-
mockSubquery();
89+
mockStatic(Coordinator.class);
90+
mockStatic(SessionManager.class);
91+
92+
// Create a mock Coordinator instance
93+
mockCoordinator = Mockito.mock(Coordinator.class);
94+
when(Coordinator.getInstance()).thenReturn(mockCoordinator);
95+
96+
// Create mock SessionManager
97+
SessionManager mockSessionManager = Mockito.mock(SessionManager.class);
98+
when(SessionManager.getInstance()).thenReturn(mockSessionManager);
99+
100+
// Mock TSStatus with success status
101+
TSStatus mockStatus = Mockito.mock(TSStatus.class);
102+
when(mockStatus.getCode()).thenReturn(200); // Success status code
103+
104+
// Create a real ExecutionResult instance
105+
ExecutionResult mockResult = new ExecutionResult(new QueryId("1"), mockStatus);
106+
107+
// Mock the executeForTableModel method
108+
when(mockCoordinator.executeForTableModel(
109+
Mockito.any(), // Statement
110+
Mockito.any(), // SqlParser
111+
Mockito.any(), // IClientSession
112+
Mockito.anyLong(), // queryId
113+
Mockito.any(), // SessionInfo
114+
Mockito.anyString(), // String
115+
Mockito.any(), // Metadata
116+
Mockito.anyMap(), // Map<NodeRef<Table>, CteDataStore>
117+
Mockito.any(), // ExplainType
118+
Mockito.anyLong(), // timeOut
119+
Mockito.anyBoolean())) // userQuery
120+
.thenReturn(mockResult);
81121
}
82122

83-
private Type convertType(String columnName) {
84-
switch (columnName) {
85-
case "time":
86-
return TypeFactory.getType(TSDataType.TIMESTAMP);
87-
case "s1":
88-
case "a1":
89-
return TypeFactory.getType(TSDataType.INT64);
90-
default:
91-
}
92-
return null;
123+
@Before
124+
public void setUp() throws IoTDBException {
125+
// Create QueryExecution mock
126+
QueryExecution mockQueryExecution = Mockito.mock(QueryExecution.class);
127+
when(mockQueryExecution.hasNextResult())
128+
.thenReturn(true) // First call returns true
129+
.thenReturn(false); // Subsequent calls return false
130+
131+
// Create a real DatasetHeader with time and s1 columns
132+
List<ColumnHeader> columnHeaders =
133+
ImmutableList.of(
134+
new ColumnHeader("time", TSDataType.TIMESTAMP),
135+
new ColumnHeader("s1", TSDataType.INT64));
136+
DatasetHeader mockDatasetHeader = new DatasetHeader(columnHeaders, false);
137+
when(mockQueryExecution.getDatasetHeader()).thenReturn(mockDatasetHeader);
138+
139+
// Create a TSBlock with sample data for getBatchResult
140+
long[] timestamps = {1000L, 2000L, 3000L};
141+
long[] values = {10L, 20L, 30L};
142+
TimeColumn timeColumn = new TimeColumn(3, timestamps);
143+
LongColumn valueColumn = new LongColumn(3, Optional.empty(), values);
144+
TsBlock sampleTsBlock = new TsBlock(timeColumn, valueColumn);
145+
when(mockQueryExecution.getBatchResult()).thenReturn(Optional.of(sampleTsBlock));
146+
147+
// Mock coordinator methods
148+
when(mockCoordinator.getQueryExecution(Mockito.anyLong())).thenReturn(mockQueryExecution);
93149
}
94150

95151
private void mockException() {
@@ -100,44 +156,6 @@ private void mockException() {
100156
CteMaterializer.setInstance(cteMaterializer);
101157
}
102158

103-
private void mockSubquery() {
104-
CteMaterializer cteMaterializer = Mockito.spy(new CteMaterializer());
105-
106-
Mockito.doAnswer(
107-
(InvocationOnMock invocation) -> {
108-
Table table = invocation.getArgument(1);
109-
Query query = invocation.getArgument(2);
110-
List<SelectItem> selectItems =
111-
((QuerySpecification) query.getQueryBody()).getSelect().getSelectItems();
112-
List<ColumnSchema> columnsSchemas =
113-
selectItems.stream()
114-
.map(
115-
selectItem -> {
116-
SingleColumn column = ((SingleColumn) selectItem);
117-
String columnName =
118-
column.getAlias().isPresent()
119-
? column.getAlias().get().toString()
120-
: column.getExpression().toString();
121-
return new ColumnSchema(
122-
columnName,
123-
convertType(columnName),
124-
false,
125-
TsTableColumnCategory.FIELD);
126-
})
127-
.collect(Collectors.toList());
128-
129-
TableSchema tableSchema = new TableSchema(table.getName().toString(), columnsSchemas);
130-
return new CteDataStore(query, tableSchema, ImmutableList.of(0, 1));
131-
})
132-
.when(cteMaterializer)
133-
.fetchCteQueryResult(
134-
Mockito.any(MPPQueryContext.class),
135-
Mockito.any(Table.class),
136-
Mockito.any(Query.class),
137-
Mockito.any(With.class));
138-
CteMaterializer.setInstance(cteMaterializer);
139-
}
140-
141159
@Test
142160
public void testSimpleCte() {
143161
String sql = "with cte1 as materialized (SELECT time, s1 FROM table1) select * from cte1";
@@ -157,13 +175,13 @@ public void testSimpleCte() {
157175
@Test
158176
public void testFieldFilterCte() {
159177
String sql =
160-
"with cte1 as materialized (SELECT time, s1 as a1 FROM table1) select * from cte1 where a1 > 10";
178+
"with cte1 as materialized (SELECT time, s1 FROM table1) select * from cte1 where s1 > 10";
161179

162180
LogicalQueryPlan logicalQueryPlan = planTester.createPlan(sql);
163181

164182
Expression filterPredicate =
165-
new ComparisonExpression(GREATER_THAN, new SymbolReference("a1"), new LongLiteral("10"));
166-
PlanMatchPattern cteScan = cteScan("cte1", ImmutableList.of("time", "a1"));
183+
new ComparisonExpression(GREATER_THAN, new SymbolReference("s1"), new LongLiteral("10"));
184+
PlanMatchPattern cteScan = cteScan("cte1", ImmutableList.of("time", "s1"));
167185

168186
// Verify full LogicalPlan
169187
/*
@@ -263,6 +281,7 @@ public void testAggCte() {
263281

264282
@Test
265283
public void testCteQueryException() {
284+
CteMaterializer originalCteMaterializer = CteMaterializer.getInstance();
266285
mockException();
267286

268287
String sql = "with cte1 as materialized (SELECT time, s1 FROM table1) select * from cte1";
@@ -278,5 +297,43 @@ public void testCteQueryException() {
278297
* └──TableScanNode
279298
*/
280299
assertPlan(logicalQueryPlan, output(tableScan));
300+
301+
// reset original CteMaterializer
302+
CteMaterializer.setInstance(originalCteMaterializer);
303+
}
304+
305+
@Test
306+
public void testExplainAnalyze() {
307+
String sql =
308+
"explain analyze with cte1 as materialized (SELECT time, s1 FROM table1) select * from cte1";
309+
310+
LogicalQueryPlan logicalQueryPlan = planTester.createPlan(sql);
311+
312+
PlanMatchPattern cteScan = cteScan("cte1", ImmutableList.of("time", "s1"));
313+
314+
// Verify full LogicalPlan
315+
/*
316+
* └──OutputNode
317+
* └──ExplainAnalyzeNode
318+
* └──CteScanNode
319+
*/
320+
assertPlan(logicalQueryPlan, output(explainAnalyze(cteScan)));
321+
}
322+
323+
/** This test primarily ensures code coverage for handleCteExplainResults method. */
324+
@Test
325+
public void testExplain() {
326+
String sql = "with cte1 as materialized (SELECT time, s1 FROM table1) select * from cte1";
327+
328+
LogicalQueryPlan logicalQueryPlan = planTester.createPlan(sql, true);
329+
330+
PlanMatchPattern cteScan = cteScan("cte1", ImmutableList.of("time", "s1"));
331+
332+
// Verify full LogicalPlan
333+
/*
334+
* └──OutputNode
335+
* └──CteScanNode
336+
*/
337+
assertPlan(logicalQueryPlan, output(cteScan));
281338
}
282339
}

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/PlanTester.java‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import org.apache.iotdb.commons.conf.IoTDBConstant;
2424
import org.apache.iotdb.db.conf.IoTDBDescriptor;
2525
import org.apache.iotdb.db.protocol.session.IClientSession;
26+
import org.apache.iotdb.db.queryengine.common.ExplainType;
2627
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
2728
import org.apache.iotdb.db.queryengine.common.QueryId;
2829
import org.apache.iotdb.db.queryengine.common.SessionInfo;
@@ -104,20 +105,28 @@ public PlanTester(Metadata metadata) {
104105
}
105106

106107
public LogicalQueryPlan createPlan(String sql) {
107-
return createPlan(sessionInfo, sql, NOOP, createPlanOptimizersStatsCollector());
108+
return createPlan(sessionInfo, sql, NOOP, createPlanOptimizersStatsCollector(), false);
109+
}
110+
111+
public LogicalQueryPlan createPlan(String sql, boolean explain) {
112+
return createPlan(sessionInfo, sql, NOOP, createPlanOptimizersStatsCollector(), explain);
108113
}
109114

110115
public LogicalQueryPlan createPlan(SessionInfo sessionInfo, String sql) {
111-
return createPlan(sessionInfo, sql, NOOP, createPlanOptimizersStatsCollector());
116+
return createPlan(sessionInfo, sql, NOOP, createPlanOptimizersStatsCollector(), false);
112117
}
113118

114119
public LogicalQueryPlan createPlan(
115120
SessionInfo sessionInfo,
116121
String sql,
117122
WarningCollector warningCollector,
118-
PlanOptimizersStatsCollector planOptimizersStatsCollector) {
123+
PlanOptimizersStatsCollector planOptimizersStatsCollector,
124+
boolean explain) {
119125
distributedQueryPlan = null;
120126
MPPQueryContext context = new MPPQueryContext(sql, queryId, sessionInfo, null, null);
127+
if (explain) {
128+
context.setExplainType(ExplainType.EXPLAIN);
129+
}
121130

122131
Analysis analysis = analyze(sql, metadata, context);
123132
this.analysis = analysis;

0 commit comments

Comments
 (0)