KYLIN-4260 When using server side PreparedStatement cache, the query …

core-metadata/src/main/java/org/apache/kylin/measure/topn/TopNMeasureType.java

@@ -395,6 +395,11 @@ public boolean needRewrite() {
 
     @Override
     public void adjustSqlDigest(List<MeasureDesc> measureDescs, SQLDigest sqlDigest) {
+        // If sqlDiegest is already adjusted, then not to adjust it again.
+        if (sqlDigest.isBorrowedContext) {
+            return;
+        }
+
         if (sqlDigest.aggregations.size() > 1) {
             return;
         }

core-metadata/src/main/java/org/apache/kylin/metadata/realization/SQLDigest.java

@@ -80,6 +80,7 @@ public SQLCall(String function, Iterable<Object> args) {
     public List<TblColRef> sortColumns;
     public List<OrderEnum> sortOrders;
     public boolean isRawQuery;
+    public boolean isBorrowedContext;
     public boolean limitPrecedesAggr;
     public boolean hasLimit;
 
@@ -92,7 +93,7 @@ public SQLDigest(String factTable, Set<TblColRef> allColumns, List<JoinDesc> joi
             List<DynamicFunctionDesc> dynAggregations, //
             Set<TblColRef> rtDimensionColumns, Set<TblColRef> rtMetricColumns, // dynamic col related columns
             Set<TblColRef> filterColumns, TupleFilter filter, TupleFilter havingFilter, // filter
-            List<TblColRef> sortColumns, List<OrderEnum> sortOrders, boolean limitPrecedesAggr, boolean hasLimit, // sort & limit
+            List<TblColRef> sortColumns, List<OrderEnum> sortOrders, boolean limitPrecedesAggr, boolean hasLimit, boolean isBorrowedContext, // sort & limit
             Set<MeasureDesc> involvedMeasure
     ) {
         this.factTable = factTable;
@@ -121,6 +122,7 @@ public SQLDigest(String factTable, Set<TblColRef> allColumns, List<JoinDesc> joi
         this.sortColumns = sortColumns;
         this.sortOrders = sortOrders;
         this.isRawQuery = isRawQuery();
+        this.isBorrowedContext = isBorrowedContext;
         this.limitPrecedesAggr = limitPrecedesAggr;
         this.hasLimit = hasLimit;
 

core-storage/src/test/java/org/apache/kylin/storage/hybrid/HybridInstanceTest.java

@@ -74,6 +74,7 @@ public void testIsCapable() {
                 new LinkedList<>(),
                 true,
                 true,
+                false,
                 new LinkedHashSet<>());
         CapabilityResult capabilityResult = hybridInstance.isCapable(sQLDigest);
 

kylin-it/src/test/java/org/apache/kylin/storage/hbase/ITStorageTest.java

@@ -148,7 +148,7 @@ private int search(List<TblColRef> groups, List<FunctionDesc> aggregations, Tupl
                     /*runtimeDimensionColumns*/ Collections.<TblColRef> emptySet(), //
                     /*runtimeMetricColumns*/ Collections.<TblColRef> emptySet(), //
                     /*filter col*/ Collections.<TblColRef> emptySet(), filter, null, //
-                    /*sortCol*/ new ArrayList<TblColRef>(), new ArrayList<SQLDigest.OrderEnum>(), false, false, new HashSet<MeasureDesc>());
+                    /*sortCol*/ new ArrayList<TblColRef>(), new ArrayList<SQLDigest.OrderEnum>(), false, false, false, new HashSet<MeasureDesc>());
             iterator = storageEngine.search(context, sqlDigest, mockup.newTupleInfo(groups, aggregations));
             while (iterator.hasNext()) {
                 ITuple tuple = iterator.next();

query/src/main/java/org/apache/kylin/query/enumerator/OLAPEnumerator.java

@@ -106,7 +106,10 @@ private ITupleIterator queryStorage() {
         // bind dynamic variables
         olapContext.bindVariable(optiqContext);
 
-        olapContext.resetSQLDigest();
+        // If olapContext is cached, then inherit it.
+        if (!olapContext.isBorrowedContext) {
+            olapContext.resetSQLDigest();
+        }
         SQLDigest sqlDigest = olapContext.getSQLDigest();
 
         // query storage engine

query/src/main/java/org/apache/kylin/query/relnode/OLAPContext.java

@@ -157,6 +157,7 @@ public OLAPContext(int seq) {
     public TupleFilter havingFilter;
     public List<JoinDesc> joins = new LinkedList<>();
     public JoinsTree joinsTree;
+    public boolean isBorrowedContext = false; // Whether preparedContext is borrowed from cache
     List<TblColRef> sortColumns;
     List<SQLDigest.OrderEnum> sortOrders;
 
@@ -200,7 +201,7 @@ public SQLDigest getSQLDigest() {
                     metricsColumns, aggregations, aggrSqlCalls, dynFuncs, // aggregation
                     rtDimColumns, rtMetricColumns, // runtime related columns
                     filterColumns, filter, havingFilter, // filter
-                    sortColumns, sortOrders, limitPrecedesAggr, hasLimit, // sort & limit
+                    sortColumns, sortOrders, limitPrecedesAggr, hasLimit, isBorrowedContext, // sort & limit
                     involvedMeasure);
         }
         return sqlDigest;

server-base/src/main/java/org/apache/kylin/rest/service/QueryService.java

@@ -696,6 +696,13 @@ private SQLResponse queryWithSqlMassage(SQLRequest sqlRequest) throws Exception
             DBUtils.closeQuietly(conn);
             if (preparedContext != null) {
                 if (borrowPrepareContext) {
+                    // Set tag isBorrowedContext true, when return preparedContext back
+                    for (OLAPContext olapContext : preparedContext.olapContexts) {
+                        if (borrowPrepareContext) {
+                            olapContext.isBorrowedContext = true;
+                        }
+                    }
+
                     preparedContextPool.returnObject(preparedContextKey, preparedContext);
                 } else {
                     preparedContext.close();
@@ -1253,6 +1260,10 @@ private static PreparedContext createPreparedContext(String project, String sql)
         Connection conn = QueryConnection.getConnection(project);
         PreparedStatement preparedStatement = conn.prepareStatement(sql);
         Collection<OLAPContext> olapContexts = OLAPContext.getThreadLocalContexts();
+        // If the preparedContext is first initialized, then set the borrowed tag to false
+        for (OLAPContext olapContext : olapContexts) {
+            olapContext.isBorrowedContext = false;
+        }
         return new PreparedContext(conn, preparedStatement, olapContexts);
     }