groupoperator.c 42.0 KB
Newer Older
H
Haojun Liao 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/*
 * Copyright (c) 2019 TAOS Data, Inc. <jhtao@taosdata.com>
 *
 * This program is free software: you can use, redistribute, and/or modify
 * it under the terms of the GNU Affero General Public License, version 3
 * or later ("AGPL"), as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

16
#include "filter.h"
H
Haojun Liao 已提交
17
#include "function.h"
18
#include "os.h"
H
Haojun Liao 已提交
19 20 21 22 23
#include "tname.h"

#include "tdatablock.h"
#include "tmsg.h"

24
#include "executorInt.h"
H
Haojun Liao 已提交
25 26 27 28 29
#include "executorimpl.h"
#include "tcompare.h"
#include "thash.h"
#include "ttypes.h"

H
Haojun Liao 已提交
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59
typedef struct SGroupbyOperatorInfo {
  SOptrBasicInfo binfo;
  SAggSupporter  aggSup;
  SArray*        pGroupCols;     // group by columns, SArray<SColumn>
  SArray*        pGroupColVals;  // current group column values, SArray<SGroupKeys>
  bool           isInit;         // denote if current val is initialized or not
  char*          keyBuf;         // group by keys for hash
  int32_t        groupKeyLen;    // total group by column width
  SGroupResInfo  groupResInfo;
  SExprSupp      scalarSup;
} SGroupbyOperatorInfo;

// The sort in partition may be needed later.
typedef struct SPartitionOperatorInfo {
  SOptrBasicInfo binfo;
  SArray*        pGroupCols;
  SArray*        pGroupColVals;  // current group column values, SArray<SGroupKeys>
  char*          keyBuf;         // group by keys for hash
  int32_t        groupKeyLen;    // total group by column width
  SHashObj*      pGroupSet;      // quick locate the window object for each result

  SDiskbasedBuf* pBuf;              // query result buffer based on blocked-wised disk file
  int32_t        rowCapacity;       // maximum number of rows for each buffer page
  int32_t*       columnOffset;      // start position for each column data
  SArray*        sortedGroupArray;  // SDataGroupInfo sorted by group id
  int32_t        groupIndex;        // group index
  int32_t        pageIndex;         // page index of current group
  SExprSupp      scalarSup;
} SPartitionOperatorInfo;

60
static void*    getCurrentDataGroupInfo(const SPartitionOperatorInfo* pInfo, SDataGroupInfo** pGroupInfo, int32_t len);
H
Haojun Liao 已提交
61
static int32_t* setupColumnOffset(const SSDataBlock* pBlock, int32_t rowCapacity);
62 63
static int32_t  setGroupResultOutputBuf(SOperatorInfo* pOperator, SOptrBasicInfo* binfo, int32_t numOfCols, char* pData,
                                        int16_t bytes, uint64_t groupId, SDiskbasedBuf* pBuf, SAggSupporter* pAggSup);
H
Haojun Liao 已提交
64
static SArray*  extractColumnInfo(SNodeList* pNodeList);
H
Haojun Liao 已提交
65

H
Haojun Liao 已提交
66
static void freeGroupKey(void* param) {
67
  SGroupKeys* pKey = (SGroupKeys*)param;
H
Haojun Liao 已提交
68 69 70
  taosMemoryFree(pKey->pData);
}

71
static void destroyGroupOperatorInfo(void* param) {
H
Haojun Liao 已提交
72
  SGroupbyOperatorInfo* pInfo = (SGroupbyOperatorInfo*)param;
73 74 75 76
  if (pInfo == NULL) {
    return;
  }

77
  cleanupBasicInfo(&pInfo->binfo);
H
Haojun Liao 已提交
78 79
  taosMemoryFreeClear(pInfo->keyBuf);
  taosArrayDestroy(pInfo->pGroupCols);
H
Haojun Liao 已提交
80
  taosArrayDestroyEx(pInfo->pGroupColVals, freeGroupKey);
81
  cleanupExprSupp(&pInfo->scalarSup);
H
Haojun Liao 已提交
82 83 84

  cleanupGroupResInfo(&pInfo->groupResInfo);
  cleanupAggSup(&pInfo->aggSup);
D
dapan1121 已提交
85
  taosMemoryFreeClear(param);
H
Haojun Liao 已提交
86 87
}

wmmhello's avatar
wmmhello 已提交
88
static int32_t initGroupOptrInfo(SArray** pGroupColVals, int32_t* keyLen, char** keyBuf, const SArray* pGroupColList) {
H
Haojun Liao 已提交
89 90
  *pGroupColVals = taosArrayInit(4, sizeof(SGroupKeys));
  if ((*pGroupColVals) == NULL) {
H
Haojun Liao 已提交
91 92 93 94 95
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  int32_t numOfGroupCols = taosArrayGetSize(pGroupColList);
  for (int32_t i = 0; i < numOfGroupCols; ++i) {
5
54liuyao 已提交
96
    SColumn* pCol = (SColumn*)taosArrayGet(pGroupColList, i);
97
    (*keyLen) += pCol->bytes;  // actual data + null_flag
H
Haojun Liao 已提交
98

99
    SGroupKeys key = {0};
100 101
    key.bytes = pCol->bytes;
    key.type = pCol->type;
H
Haojun Liao 已提交
102
    key.isNull = false;
103
    key.pData = taosMemoryCalloc(1, pCol->bytes);
H
Haojun Liao 已提交
104 105 106 107
    if (key.pData == NULL) {
      return TSDB_CODE_OUT_OF_MEMORY;
    }

H
Haojun Liao 已提交
108
    taosArrayPush((*pGroupColVals), &key);
H
Haojun Liao 已提交
109 110 111
  }

  int32_t nullFlagSize = sizeof(int8_t) * numOfGroupCols;
112
  (*keyLen) += nullFlagSize;
H
Haojun Liao 已提交
113

114
  (*keyBuf) = taosMemoryCalloc(1, (*keyLen));
H
Haojun Liao 已提交
115
  if ((*keyBuf) == NULL) {
H
Haojun Liao 已提交
116 117 118 119 120 121
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  return TSDB_CODE_SUCCESS;
}

122 123
static bool groupKeyCompare(SArray* pGroupCols, SArray* pGroupColVals, SSDataBlock* pBlock, int32_t rowIndex,
                            int32_t numOfGroupCols) {
H
Haojun Liao 已提交
124 125
  SColumnDataAgg* pColAgg = NULL;
  for (int32_t i = 0; i < numOfGroupCols; ++i) {
H
Haojun Liao 已提交
126
    SColumn*         pCol = taosArrayGet(pGroupCols, i);
H
Haojun Liao 已提交
127 128
    SColumnInfoData* pColInfoData = taosArrayGet(pBlock->pDataBlock, pCol->slotId);
    if (pBlock->pBlockAgg != NULL) {
129
      pColAgg = pBlock->pBlockAgg[pCol->slotId];  // TODO is agg data matched?
H
Haojun Liao 已提交
130 131 132 133
    }

    bool isNull = colDataIsNull(pColInfoData, pBlock->info.rows, rowIndex, pColAgg);

H
Haojun Liao 已提交
134
    SGroupKeys* pkey = taosArrayGet(pGroupColVals, i);
H
Haojun Liao 已提交
135 136 137 138 139 140 141 142 143 144
    if (pkey->isNull && isNull) {
      continue;
    }

    if (isNull || pkey->isNull) {
      return false;
    }

    char* val = colDataGetData(pColInfoData, rowIndex);

wmmhello's avatar
wmmhello 已提交
145 146 147
    if (pkey->type == TSDB_DATA_TYPE_JSON) {
      int32_t dataLen = getJsonValueLen(val);

148
      if (memcmp(pkey->pData, val, dataLen) == 0) {
wmmhello's avatar
wmmhello 已提交
149 150 151 152 153
        continue;
      } else {
        return false;
      }
    } else if (IS_VAR_DATA_TYPE(pkey->type)) {
H
Haojun Liao 已提交
154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169
      int32_t len = varDataLen(val);
      if (len == varDataLen(pkey->pData) && memcmp(varDataVal(pkey->pData), varDataVal(val), len) == 0) {
        continue;
      } else {
        return false;
      }
    } else {
      if (memcmp(pkey->pData, val, pkey->bytes) != 0) {
        return false;
      }
    }
  }

  return true;
}

wmmhello's avatar
wmmhello 已提交
170
static void recordNewGroupKeys(SArray* pGroupCols, SArray* pGroupColVals, SSDataBlock* pBlock, int32_t rowIndex) {
H
Haojun Liao 已提交
171 172
  SColumnDataAgg* pColAgg = NULL;

173 174
  size_t numOfGroupCols = taosArrayGetSize(pGroupCols);

H
Haojun Liao 已提交
175
  for (int32_t i = 0; i < numOfGroupCols; ++i) {
H
Haojun Liao 已提交
176
    SColumn*         pCol = taosArrayGet(pGroupCols, i);
H
Haojun Liao 已提交
177 178 179
    SColumnInfoData* pColInfoData = taosArrayGet(pBlock->pDataBlock, pCol->slotId);

    if (pBlock->pBlockAgg != NULL) {
180
      pColAgg = pBlock->pBlockAgg[pCol->slotId];  // TODO is agg data matched?
H
Haojun Liao 已提交
181 182
    }

H
Haojun Liao 已提交
183
    SGroupKeys* pkey = taosArrayGet(pGroupColVals, i);
H
Haojun Liao 已提交
184 185 186
    if (colDataIsNull(pColInfoData, pBlock->info.rows, rowIndex, pColAgg)) {
      pkey->isNull = true;
    } else {
187
      pkey->isNull = false;
H
Haojun Liao 已提交
188
      char* val = colDataGetData(pColInfoData, rowIndex);
wmmhello's avatar
wmmhello 已提交
189
      if (pkey->type == TSDB_DATA_TYPE_JSON) {
190
        if (tTagIsJson(val)) {
wmmhello's avatar
wmmhello 已提交
191 192 193
          terrno = TSDB_CODE_QRY_JSON_IN_GROUP_ERROR;
          return;
        }
wmmhello's avatar
wmmhello 已提交
194 195 196
        int32_t dataLen = getJsonValueLen(val);
        memcpy(pkey->pData, val, dataLen);
      } else if (IS_VAR_DATA_TYPE(pkey->type)) {
H
Haojun Liao 已提交
197
        memcpy(pkey->pData, val, varDataTLen(val));
198
        ASSERT(varDataTLen(val) <= pkey->bytes);
H
Haojun Liao 已提交
199 200 201 202 203 204 205
      } else {
        memcpy(pkey->pData, val, pkey->bytes);
      }
    }
  }
}

wmmhello's avatar
wmmhello 已提交
206
static int32_t buildGroupKeys(void* pKey, const SArray* pGroupColVals) {
H
Haojun Liao 已提交
207 208 209 210 211 212 213 214 215 216 217 218 219
  ASSERT(pKey != NULL);
  size_t numOfGroupCols = taosArrayGetSize(pGroupColVals);

  char* isNull = (char*)pKey;
  char* pStart = (char*)pKey + sizeof(int8_t) * numOfGroupCols;
  for (int32_t i = 0; i < numOfGroupCols; ++i) {
    SGroupKeys* pkey = taosArrayGet(pGroupColVals, i);
    if (pkey->isNull) {
      isNull[i] = 1;
      continue;
    }

    isNull[i] = 0;
wmmhello's avatar
wmmhello 已提交
220 221 222 223 224
    if (pkey->type == TSDB_DATA_TYPE_JSON) {
      int32_t dataLen = getJsonValueLen(pkey->pData);
      memcpy(pStart, (pkey->pData), dataLen);
      pStart += dataLen;
    } else if (IS_VAR_DATA_TYPE(pkey->type)) {
H
Haojun Liao 已提交
225 226 227 228 229 230 231 232 233
      varDataCopy(pStart, pkey->pData);
      pStart += varDataTLen(pkey->pData);
      ASSERT(varDataTLen(pkey->pData) <= pkey->bytes);
    } else {
      memcpy(pStart, pkey->pData, pkey->bytes);
      pStart += pkey->bytes;
    }
  }

234
  return (int32_t)(pStart - (char*)pKey);
H
Haojun Liao 已提交
235 236 237 238 239
}

// assign the group keys or user input constant values if required
static void doAssignGroupKeys(SqlFunctionCtx* pCtx, int32_t numOfOutput, int32_t totalRows, int32_t rowIndex) {
  for (int32_t i = 0; i < numOfOutput; ++i) {
240
    if (pCtx[i].functionId == -1) {  // select count(*),key from t group by key.
H
Haojun Liao 已提交
241 242 243
      SResultRowEntryInfo* pEntryInfo = GET_RES_INFO(&pCtx[i]);

      SColumnInfoData* pColInfoData = pCtx[i].input.pData[0];
244
      // todo OPT all/all not NULL
H
Haojun Liao 已提交
245 246 247 248
      if (!colDataIsNull(pColInfoData, totalRows, rowIndex, NULL)) {
        char* dest = GET_ROWCELL_INTERBUF(pEntryInfo);
        char* data = colDataGetData(pColInfoData, rowIndex);

wmmhello's avatar
wmmhello 已提交
249 250 251 252
        if (pColInfoData->info.type == TSDB_DATA_TYPE_JSON) {
          int32_t dataLen = getJsonValueLen(data);
          memcpy(dest, data, dataLen);
        } else if (IS_VAR_DATA_TYPE(pColInfoData->info.type)) {
253 254 255 256
          varDataCopy(dest, data);
        } else {
          memcpy(dest, data, pColInfoData->info.bytes);
        }
257
      } else {  // it is a NULL value
H
Haojun Liao 已提交
258
        pEntryInfo->isNullRes = 1;
H
Haojun Liao 已提交
259
      }
H
Haojun Liao 已提交
260 261

      pEntryInfo->numOfRes = 1;
H
Haojun Liao 已提交
262 263 264 265 266 267 268 269
    }
  }
}

static void doHashGroupbyAgg(SOperatorInfo* pOperator, SSDataBlock* pBlock) {
  SExecTaskInfo*        pTaskInfo = pOperator->pTaskInfo;
  SGroupbyOperatorInfo* pInfo = pOperator->info;

270
  SqlFunctionCtx* pCtx = pOperator->exprSupp.pCtx;
H
Haojun Liao 已提交
271 272 273 274 275 276
  int32_t         numOfGroupCols = taosArrayGetSize(pInfo->pGroupCols);
  //  if (type == TSDB_DATA_TYPE_FLOAT || type == TSDB_DATA_TYPE_DOUBLE) {
  // qError("QInfo:0x%"PRIx64" group by not supported on double/float columns, abort", GET_TASKID(pRuntimeEnv));
  //    return;
  //  }

H
Haojun Liao 已提交
277
  int32_t len = 0;
wmmhello's avatar
wmmhello 已提交
278
  terrno = TSDB_CODE_SUCCESS;
H
Haojun Liao 已提交
279

H
Haojun Liao 已提交
280 281 282 283
  int32_t num = 0;
  for (int32_t j = 0; j < pBlock->info.rows; ++j) {
    // Compare with the previous row of this column, and do not set the output buffer again if they are identical.
    if (!pInfo->isInit) {
284
      recordNewGroupKeys(pInfo->pGroupCols, pInfo->pGroupColVals, pBlock, j);
wmmhello's avatar
wmmhello 已提交
285
      if (terrno != TSDB_CODE_SUCCESS) {  // group by json error
286
        T_LONG_JMP(pTaskInfo->env, terrno);
wmmhello's avatar
wmmhello 已提交
287
      }
H
Haojun Liao 已提交
288 289 290 291 292
      pInfo->isInit = true;
      num++;
      continue;
    }

H
Haojun Liao 已提交
293
    bool equal = groupKeyCompare(pInfo->pGroupCols, pInfo->pGroupColVals, pBlock, j, numOfGroupCols);
H
Haojun Liao 已提交
294 295 296 297 298
    if (equal) {
      num++;
      continue;
    }

H
Haojun Liao 已提交
299
    // The first row of a new block does not belongs to the previous existed group
300
    if (j == 0) {
H
Haojun Liao 已提交
301
      num++;
302
      recordNewGroupKeys(pInfo->pGroupCols, pInfo->pGroupColVals, pBlock, j);
wmmhello's avatar
wmmhello 已提交
303
      if (terrno != TSDB_CODE_SUCCESS) {  // group by json error
304
        T_LONG_JMP(pTaskInfo->env, terrno);
wmmhello's avatar
wmmhello 已提交
305
      }
H
Haojun Liao 已提交
306 307 308
      continue;
    }

H
Haojun Liao 已提交
309
    len = buildGroupKeys(pInfo->keyBuf, pInfo->pGroupColVals);
310 311
    int32_t ret = setGroupResultOutputBuf(pOperator, &(pInfo->binfo), pOperator->exprSupp.numOfExprs, pInfo->keyBuf,
                                          len, pBlock->info.groupId, pInfo->aggSup.pResultBuf, &pInfo->aggSup);
H
Haojun Liao 已提交
312
    if (ret != TSDB_CODE_SUCCESS) {  // null data, too many state code
313
      T_LONG_JMP(pTaskInfo->env, TSDB_CODE_QRY_APP_ERROR);
H
Haojun Liao 已提交
314 315 316
    }

    int32_t rowIndex = j - num;
H
Haojun Liao 已提交
317
    applyAggFunctionOnPartialTuples(pTaskInfo, pCtx, NULL, rowIndex, num, pBlock->info.rows, pOperator->exprSupp.numOfExprs);
H
Haojun Liao 已提交
318 319

    // assign the group keys or user input constant values if required
320
    doAssignGroupKeys(pCtx, pOperator->exprSupp.numOfExprs, pBlock->info.rows, rowIndex);
321
    recordNewGroupKeys(pInfo->pGroupCols, pInfo->pGroupColVals, pBlock, j);
H
Haojun Liao 已提交
322 323 324 325
    num = 1;
  }

  if (num > 0) {
H
Haojun Liao 已提交
326
    len = buildGroupKeys(pInfo->keyBuf, pInfo->pGroupColVals);
327 328
    int32_t ret = setGroupResultOutputBuf(pOperator, &(pInfo->binfo), pOperator->exprSupp.numOfExprs, pInfo->keyBuf,
                                          len, pBlock->info.groupId, pInfo->aggSup.pResultBuf, &pInfo->aggSup);
H
Haojun Liao 已提交
329
    if (ret != TSDB_CODE_SUCCESS) {
330
      T_LONG_JMP(pTaskInfo->env, TSDB_CODE_QRY_APP_ERROR);
H
Haojun Liao 已提交
331 332 333
    }

    int32_t rowIndex = pBlock->info.rows - num;
H
Haojun Liao 已提交
334
    applyAggFunctionOnPartialTuples(pTaskInfo, pCtx, NULL, rowIndex, num, pBlock->info.rows, pOperator->exprSupp.numOfExprs);
335
    doAssignGroupKeys(pCtx, pOperator->exprSupp.numOfExprs, pBlock->info.rows, rowIndex);
H
Haojun Liao 已提交
336 337 338
  }
}

339 340 341 342
static SSDataBlock* buildGroupResultDataBlock(SOperatorInfo* pOperator) {
  SGroupbyOperatorInfo* pInfo = pOperator->info;

  SSDataBlock* pRes = pInfo->binfo.pRes;
343
  while (1) {
344
    doBuildResultDatablock(pOperator, &pInfo->binfo, &pInfo->groupResInfo, pInfo->aggSup.pResultBuf);
H
Haojun Liao 已提交
345
    doFilter(pRes, pOperator->exprSupp.pFilterInfo, NULL);
346

347
    if (!hasRemainResults(&pInfo->groupResInfo)) {
H
Haojun Liao 已提交
348
      setOperatorCompleted(pOperator);
349 350 351 352 353 354 355 356 357
      break;
    }

    if (pRes->info.rows > 0) {
      break;
    }
  }

  pOperator->resultInfo.totalRows += pRes->info.rows;
358
  return (pRes->info.rows == 0) ? NULL : pRes;
359 360
}

361
static SSDataBlock* hashGroupbyAggregate(SOperatorInfo* pOperator) {
H
Haojun Liao 已提交
362 363 364 365
  if (pOperator->status == OP_EXEC_DONE) {
    return NULL;
  }

366 367
  SExecTaskInfo* pTaskInfo = pOperator->pTaskInfo;

H
Haojun Liao 已提交
368 369
  SGroupbyOperatorInfo* pInfo = pOperator->info;
  if (pOperator->status == OP_RES_TO_RETURN) {
370
    return buildGroupResultDataBlock(pOperator);
H
Haojun Liao 已提交
371 372
  }

373 374 375
  int32_t order = TSDB_ORDER_ASC;
  int32_t scanFlag = MAIN_SCAN;

376
  int64_t        st = taosGetTimestampUs();
H
Haojun Liao 已提交
377 378 379
  SOperatorInfo* downstream = pOperator->pDownstream[0];

  while (1) {
380
    SSDataBlock* pBlock = downstream->fpSet.getNextFn(downstream);
H
Haojun Liao 已提交
381 382 383 384
    if (pBlock == NULL) {
      break;
    }

385 386
    int32_t code = getTableScanInfo(pOperator, &order, &scanFlag);
    if (code != TSDB_CODE_SUCCESS) {
387
      T_LONG_JMP(pTaskInfo->env, code);
388 389
    }

H
Haojun Liao 已提交
390
    // the pDataBlock are always the same one, no need to call this again
391
    setInputDataBlock(&pOperator->exprSupp, pBlock, order, scanFlag, true);
392

393
    // there is an scalar expression that needs to be calculated right before apply the group aggregation.
394
    if (pInfo->scalarSup.pExprInfo != NULL) {
395 396
      pTaskInfo->code = projectApplyFunctions(pInfo->scalarSup.pExprInfo, pBlock, pBlock, pInfo->scalarSup.pCtx,
                                              pInfo->scalarSup.numOfExprs, NULL);
397
      if (pTaskInfo->code != TSDB_CODE_SUCCESS) {
398
        T_LONG_JMP(pTaskInfo->env, pTaskInfo->code);
399
      }
400 401
    }

H
Haojun Liao 已提交
402 403 404 405
    doHashGroupbyAgg(pOperator, pBlock);
  }

  pOperator->status = OP_RES_TO_RETURN;
H
Haojun Liao 已提交
406

407 408 409 410 411 412 413 414 415 416 417 418 419 420
#if 0
  if(pOperator->fpSet.encodeResultRow){
    char *result = NULL;
    int32_t length = 0;
    pOperator->fpSet.encodeResultRow(pOperator, &result, &length);
    SAggSupporter* pSup = &pInfo->aggSup;
    taosHashClear(pSup->pResultRowHashTable);
    pInfo->binfo.resultRowInfo.size = 0;
    pOperator->fpSet.decodeResultRow(pOperator, result);
    if(result){
      taosMemoryFree(result);
    }
  }
#endif
421
  initGroupedResultInfo(&pInfo->groupResInfo, pInfo->aggSup.pResultRowHashTable, 0);
422

423
  pOperator->cost.openCost = (taosGetTimestampUs() - st) / 1000.0;
424
  return buildGroupResultDataBlock(pOperator);
H
Haojun Liao 已提交
425 426
}

5
54liuyao 已提交
427
SOperatorInfo* createGroupOperatorInfo(SOperatorInfo* downstream, SAggPhysiNode* pAggNode, SExecTaskInfo* pTaskInfo) {
H
Haojun Liao 已提交
428 429 430 431 432 433
  SGroupbyOperatorInfo* pInfo = taosMemoryCalloc(1, sizeof(SGroupbyOperatorInfo));
  SOperatorInfo*        pOperator = taosMemoryCalloc(1, sizeof(SOperatorInfo));
  if (pInfo == NULL || pOperator == NULL) {
    goto _error;
  }

H
Haojun Liao 已提交
434
  SSDataBlock* pResBlock = createDataBlockFromDescNode(pAggNode->node.pOutputDataBlockDesc);
H
Haojun Liao 已提交
435 436 437 438 439 440 441 442 443
  initBasicInfo(&pInfo->binfo, pResBlock);

  int32_t    numOfScalarExpr = 0;
  SExprInfo* pScalarExprInfo = NULL;
  if (pAggNode->pExprs != NULL) {
    pScalarExprInfo = createExprInfo(pAggNode->pExprs, NULL, &numOfScalarExpr);
  }

  pInfo->pGroupCols = extractColumnInfo(pAggNode->pGroupKeys);
444 445 446 447
  int32_t code = initExprSupp(&pInfo->scalarSup, pScalarExprInfo, numOfScalarExpr);
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }
448

H
Haojun Liao 已提交
449
  initResultSizeInfo(&pOperator->resultInfo, 4096);
H
Haojun Liao 已提交
450 451
  blockDataEnsureCapacity(pInfo->binfo.pRes, pOperator->resultInfo.capacity);

H
Haojun Liao 已提交
452
  code = initGroupOptrInfo(&pInfo->pGroupColVals, &pInfo->groupKeyLen, &pInfo->keyBuf, pInfo->pGroupCols);
H
Haojun Liao 已提交
453 454 455 456
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

H
Haojun Liao 已提交
457 458
  int32_t    num = 0;
  SExprInfo* pExprInfo = createExprInfo(pAggNode->pAggFuncs, pAggNode->pGroupKeys, &num);
H
Haojun Liao 已提交
459
  code = initAggInfo(&pOperator->exprSupp, &pInfo->aggSup, pExprInfo, num, pInfo->groupKeyLen, pTaskInfo->id.str);
460 461 462 463
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

464 465 466 467 468
  code = filterInitFromNode((SNode*)pAggNode->node.pConditions, &pOperator->exprSupp.pFilterInfo, 0);
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

469
  initResultRowInfo(&pInfo->binfo.resultRowInfo);
470
  setOperatorInfo(pOperator, "GroupbyAggOperator", 0, true, OP_NOT_OPENED, pInfo, pTaskInfo);
H
Haojun Liao 已提交
471

5
54liuyao 已提交
472
  pOperator->fpSet =
H
Haojun Liao 已提交
473
      createOperatorFpSet(operatorDummyOpenFn, hashGroupbyAggregate, NULL, destroyGroupOperatorInfo, NULL);
H
Haojun Liao 已提交
474
  code = appendDownstream(pOperator, &downstream, 1);
475 476 477 478
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

H
Haojun Liao 已提交
479 480
  return pOperator;

481
_error:
H
Haojun Liao 已提交
482
  pTaskInfo->code = TSDB_CODE_OUT_OF_MEMORY;
H
Haojun Liao 已提交
483 484 485
  if (pInfo != NULL) {
    destroyGroupOperatorInfo(pInfo);
  }
H
Haojun Liao 已提交
486 487
  taosMemoryFreeClear(pOperator);
  return NULL;
488 489
}

H
Haojun Liao 已提交
490 491 492 493
static void doHashPartition(SOperatorInfo* pOperator, SSDataBlock* pBlock) {
  SPartitionOperatorInfo* pInfo = pOperator->info;

  for (int32_t j = 0; j < pBlock->info.rows; ++j) {
494
    recordNewGroupKeys(pInfo->pGroupCols, pInfo->pGroupColVals, pBlock, j);
H
Haojun Liao 已提交
495 496
    int32_t len = buildGroupKeys(pInfo->keyBuf, pInfo->pGroupColVals);

497
    SDataGroupInfo* pGroupInfo = NULL;
498
    void*           pPage = getCurrentDataGroupInfo(pInfo, &pGroupInfo, len);
H
Haojun Liao 已提交
499

500 501 502 503 504
    pGroupInfo->numOfRows += 1;

    // group id
    if (pGroupInfo->groupId == 0) {
      pGroupInfo->groupId = calcGroupId(pInfo->keyBuf, len);
H
Haojun Liao 已提交
505 506
    }

507
    // number of rows
508
    int32_t* rows = (int32_t*)pPage;
H
Haojun Liao 已提交
509

510
    size_t numOfCols = pOperator->exprSupp.numOfExprs;
511
    for (int32_t i = 0; i < numOfCols; ++i) {
512
      SExprInfo* pExpr = &pOperator->exprSupp.pExprInfo[i];
513
      int32_t    slotId = pExpr->base.pParam[0].pCol->slotId;
514 515

      SColumnInfoData* pColInfoData = taosArrayGet(pBlock->pDataBlock, slotId);
H
Haojun Liao 已提交
516

H
Haojun Liao 已提交
517 518
      int32_t bytes = pColInfoData->info.bytes;
      int32_t startOffset = pInfo->columnOffset[i];
H
Haojun Liao 已提交
519

520
      int32_t* columnLen = NULL;
521
      int32_t  contentLen = 0;
H
Haojun Liao 已提交
522 523

      if (IS_VAR_DATA_TYPE(pColInfoData->info.type)) {
wafwerar's avatar
wafwerar 已提交
524
        int32_t* offset = (int32_t*)((char*)pPage + startOffset);
525 526
        columnLen = (int32_t*)((char*)pPage + startOffset + sizeof(int32_t) * pInfo->rowCapacity);
        char* data = (char*)((char*)columnLen + sizeof(int32_t));
H
Haojun Liao 已提交
527 528 529 530

        if (colDataIsNull_s(pColInfoData, j)) {
          offset[(*rows)] = -1;
          contentLen = 0;
531
        } else if (pColInfoData->info.type == TSDB_DATA_TYPE_JSON) {
wmmhello's avatar
wmmhello 已提交
532
          offset[*rows] = (*columnLen);
533
          char*   src = colDataGetData(pColInfoData, j);
wmmhello's avatar
wmmhello 已提交
534 535 536 537 538 539 540
          int32_t dataLen = getJsonValueLen(src);

          memcpy(data + (*columnLen), src, dataLen);
          int32_t v = (data + (*columnLen) + dataLen - (char*)pPage);
          ASSERT(v > 0);

          contentLen = dataLen;
H
Haojun Liao 已提交
541 542 543 544
        } else {
          offset[*rows] = (*columnLen);
          char* src = colDataGetData(pColInfoData, j);
          memcpy(data + (*columnLen), src, varDataTLen(src));
545 546 547
          int32_t v = (data + (*columnLen) + varDataTLen(src) - (char*)pPage);
          ASSERT(v > 0);

H
Haojun Liao 已提交
548 549
          contentLen = varDataTLen(src);
        }
H
Haojun Liao 已提交
550
      } else {
wafwerar's avatar
wafwerar 已提交
551
        char* bitmap = (char*)pPage + startOffset;
552 553
        columnLen = (int32_t*)((char*)pPage + startOffset + BitmapLen(pInfo->rowCapacity));
        char* data = (char*)columnLen + sizeof(int32_t);
H
Haojun Liao 已提交
554 555 556

        bool isNull = colDataIsNull_f(pColInfoData->nullbitmap, j);
        if (isNull) {
H
Haojun Liao 已提交
557
          colDataSetNull_f(bitmap, (*rows));
H
Haojun Liao 已提交
558
        } else {
H
Haojun Liao 已提交
559
          memcpy(data + (*columnLen), colDataGetData(pColInfoData, j), bytes);
H
Haojun Liao 已提交
560
          ASSERT((data + (*columnLen) + bytes - (char*)pPage) <= getBufPageSize(pInfo->pBuf));
H
Haojun Liao 已提交
561
        }
H
Haojun Liao 已提交
562
        contentLen = bytes;
H
Haojun Liao 已提交
563
      }
H
Haojun Liao 已提交
564 565

      (*columnLen) += contentLen;
566
      ASSERT(*columnLen >= 0);
H
Haojun Liao 已提交
567 568
    }

H
Haojun Liao 已提交
569 570
    (*rows) += 1;

H
Haojun Liao 已提交
571 572 573
    setBufPageDirty(pPage, true);
    releaseBufPage(pInfo->pBuf, pPage);
  }
H
Haojun Liao 已提交
574 575 576 577 578 579
}

void* getCurrentDataGroupInfo(const SPartitionOperatorInfo* pInfo, SDataGroupInfo** pGroupInfo, int32_t len) {
  SDataGroupInfo* p = taosHashGet(pInfo->pGroupSet, pInfo->keyBuf, len);

  void* pPage = NULL;
580
  if (p == NULL) {  // it is a new group
H
Haojun Liao 已提交
581 582 583 584 585 586 587
    SDataGroupInfo gi = {0};
    gi.pPageList = taosArrayInit(100, sizeof(int32_t));
    taosHashPut(pInfo->pGroupSet, pInfo->keyBuf, len, &gi, sizeof(SDataGroupInfo));

    p = taosHashGet(pInfo->pGroupSet, pInfo->keyBuf, len);

    int32_t pageId = 0;
588
    pPage = getNewBufPage(pInfo->pBuf, &pageId);
H
Haojun Liao 已提交
589 590
    taosArrayPush(p->pPageList, &pageId);

591
    *(int32_t*)pPage = 0;
H
Haojun Liao 已提交
592 593 594 595
  } else {
    int32_t* curId = taosArrayGetLast(p->pPageList);
    pPage = getBufPage(pInfo->pBuf, *curId);

596
    int32_t* rows = (int32_t*)pPage;
H
Haojun Liao 已提交
597
    if (*rows >= pInfo->rowCapacity) {
598 599 600
      // release buffer
      releaseBufPage(pInfo->pBuf, pPage);

H
Haojun Liao 已提交
601 602
      // add a new page for current group
      int32_t pageId = 0;
603
      pPage = getNewBufPage(pInfo->pBuf, &pageId);
H
Haojun Liao 已提交
604
      taosArrayPush(p->pPageList, &pageId);
605
      memset(pPage, 0, getBufPageSize(pInfo->pBuf));
H
Haojun Liao 已提交
606 607
    }
  }
H
Haojun Liao 已提交
608

H
Haojun Liao 已提交
609 610 611 612 613 614 615 616 617 618 619 620 621 622
  *pGroupInfo = p;
  return pPage;
}

uint64_t calcGroupId(char* pData, int32_t len) {
  T_MD5_CTX context;
  tMD5Init(&context);
  tMD5Update(&context, (uint8_t*)pData, len);
  tMD5Final(&context);

  // NOTE: only extract the initial 8 bytes of the final MD5 digest
  uint64_t id = 0;
  memcpy(&id, context.digest, sizeof(uint64_t));
  return id;
H
Haojun Liao 已提交
623 624
}

H
Haojun Liao 已提交
625
int32_t* setupColumnOffset(const SSDataBlock* pBlock, int32_t rowCapacity) {
626
  size_t   numOfCols = taosArrayGetSize(pBlock->pDataBlock);
627
  int32_t* offset = taosMemoryCalloc(numOfCols, sizeof(int32_t));
H
Haojun Liao 已提交
628

629 630
  offset[0] = sizeof(int32_t) +
              sizeof(uint64_t);  // the number of rows in current page, ref to SSDataBlock paged serialization format
H
Haojun Liao 已提交
631

632
  for (int32_t i = 0; i < numOfCols - 1; ++i) {
H
Haojun Liao 已提交
633 634 635 636
    SColumnInfoData* pColInfoData = taosArrayGet(pBlock->pDataBlock, i);

    int32_t bytes = pColInfoData->info.bytes;
    int32_t payloadLen = bytes * rowCapacity;
637

H
Haojun Liao 已提交
638 639 640 641 642 643 644 645 646 647 648 649
    if (IS_VAR_DATA_TYPE(pColInfoData->info.type)) {
      // offset segment + content length + payload
      offset[i + 1] = rowCapacity * sizeof(int32_t) + sizeof(int32_t) + payloadLen + offset[i];
    } else {
      // bitmap + content length + payload
      offset[i + 1] = BitmapLen(rowCapacity) + sizeof(int32_t) + payloadLen + offset[i];
    }
  }

  return offset;
}

5
54liuyao 已提交
650
static void clearPartitionOperator(SPartitionOperatorInfo* pInfo) {
651 652 653 654
  int32_t size = taosArrayGetSize(pInfo->sortedGroupArray);
  for (int32_t i = 0; i < size; i++) {
    SDataGroupInfo* pGp = taosArrayGet(pInfo->sortedGroupArray, i);
    taosArrayDestroy(pGp->pPageList);
5
54liuyao 已提交
655
  }
656
  taosArrayClear(pInfo->sortedGroupArray);
5
54liuyao 已提交
657 658 659
  clearDiskbasedBuf(pInfo->pBuf);
}

660 661 662
static int compareDataGroupInfo(const void* group1, const void* group2) {
  const SDataGroupInfo* pGroupInfo1 = group1;
  const SDataGroupInfo* pGroupInfo2 = group2;
663 664 665 666 667 668

  if (pGroupInfo1->groupId == pGroupInfo2->groupId) {
    ASSERT(0);
    return 0;
  }

669
  return (pGroupInfo1->groupId < pGroupInfo2->groupId) ? -1 : 1;
670 671
}

H
Haojun Liao 已提交
672 673 674
static SSDataBlock* buildPartitionResult(SOperatorInfo* pOperator) {
  SPartitionOperatorInfo* pInfo = pOperator->info;

675 676
  SDataGroupInfo* pGroupInfo =
      (pInfo->groupIndex != -1) ? taosArrayGet(pInfo->sortedGroupArray, pInfo->groupIndex) : NULL;
677
  if (pInfo->groupIndex == -1 || pInfo->pageIndex >= taosArrayGetSize(pGroupInfo->pPageList)) {
H
Haojun Liao 已提交
678
    // try next group data
679 680
    ++pInfo->groupIndex;
    if (pInfo->groupIndex >= taosArrayGetSize(pInfo->sortedGroupArray)) {
H
Haojun Liao 已提交
681
      setOperatorCompleted(pOperator);
5
54liuyao 已提交
682
      clearPartitionOperator(pInfo);
H
Haojun Liao 已提交
683 684 685
      return NULL;
    }

686
    pGroupInfo = taosArrayGet(pInfo->sortedGroupArray, pInfo->groupIndex);
H
Haojun Liao 已提交
687 688 689 690
    pInfo->pageIndex = 0;
  }

  int32_t* pageId = taosArrayGet(pGroupInfo->pPageList, pInfo->pageIndex);
691
  void*    page = getBufPage(pInfo->pBuf, *pageId);
H
Haojun Liao 已提交
692

693
  blockDataEnsureCapacity(pInfo->binfo.pRes, pInfo->rowCapacity);
H
Haojun Liao 已提交
694
  blockDataFromBuf1(pInfo->binfo.pRes, page, pInfo->rowCapacity);
H
Haojun Liao 已提交
695 696

  pInfo->pageIndex += 1;
697
  releaseBufPage(pInfo->pBuf, page);
H
Haojun Liao 已提交
698

699
  blockDataUpdateTsWindow(pInfo->binfo.pRes, 0);
H
Haojun Liao 已提交
700
  pInfo->binfo.pRes->info.groupId = pGroupInfo->groupId;
701 702

  pOperator->resultInfo.totalRows += pInfo->binfo.pRes->info.rows;
H
Haojun Liao 已提交
703 704 705
  return pInfo->binfo.pRes;
}

706
static SSDataBlock* hashPartition(SOperatorInfo* pOperator) {
H
Haojun Liao 已提交
707 708
  if (pOperator->status == OP_EXEC_DONE) {
    return NULL;
709 710
  }

711 712 713
  SExecTaskInfo* pTaskInfo = pOperator->pTaskInfo;

  SPartitionOperatorInfo* pInfo = pOperator->info;
714
  SSDataBlock*            pRes = pInfo->binfo.pRes;
715

H
Haojun Liao 已提交
716
  if (pOperator->status == OP_RES_TO_RETURN) {
H
Haojun Liao 已提交
717 718
    blockDataCleanup(pRes);
    return buildPartitionResult(pOperator);
H
Haojun Liao 已提交
719 720
  }

721
  int64_t        st = taosGetTimestampUs();
H
Haojun Liao 已提交
722
  SOperatorInfo* downstream = pOperator->pDownstream[0];
H
Haojun Liao 已提交
723

H
Haojun Liao 已提交
724
  while (1) {
725
    SSDataBlock* pBlock = downstream->fpSet.getNextFn(downstream);
H
Haojun Liao 已提交
726 727 728
    if (pBlock == NULL) {
      break;
    }
H
Haojun Liao 已提交
729

730
    // there is an scalar expression that needs to be calculated right before apply the group aggregation.
731
    if (pInfo->scalarSup.pExprInfo != NULL) {
732 733
      pTaskInfo->code = projectApplyFunctions(pInfo->scalarSup.pExprInfo, pBlock, pBlock, pInfo->scalarSup.pCtx,
                                              pInfo->scalarSup.numOfExprs, NULL);
734
      if (pTaskInfo->code != TSDB_CODE_SUCCESS) {
735
        T_LONG_JMP(pTaskInfo->env, pTaskInfo->code);
736 737 738
      }
    }

wmmhello's avatar
wmmhello 已提交
739
    terrno = TSDB_CODE_SUCCESS;
H
Haojun Liao 已提交
740
    doHashPartition(pOperator, pBlock);
wmmhello's avatar
wmmhello 已提交
741
    if (terrno != TSDB_CODE_SUCCESS) {  // group by json error
742
      T_LONG_JMP(pTaskInfo->env, terrno);
wmmhello's avatar
wmmhello 已提交
743
    }
H
Haojun Liao 已提交
744 745
  }

746
  SArray* groupArray = taosArrayInit(taosHashGetSize(pInfo->pGroupSet), sizeof(SDataGroupInfo));
747 748

  void* pGroupIter = taosHashIterate(pInfo->pGroupSet, NULL);
749 750 751 752 753 754 755 756 757 758 759
  while (pGroupIter != NULL) {
    SDataGroupInfo* pGroupInfo = pGroupIter;
    taosArrayPush(groupArray, pGroupInfo);
    pGroupIter = taosHashIterate(pInfo->pGroupSet, pGroupIter);
  }

  taosArraySort(groupArray, compareDataGroupInfo);
  pInfo->sortedGroupArray = groupArray;
  pInfo->groupIndex = -1;
  taosHashClear(pInfo->pGroupSet);

760 761
  pOperator->cost.openCost = (taosGetTimestampUs() - st) / 1000.0;

H
Haojun Liao 已提交
762
  pOperator->status = OP_RES_TO_RETURN;
H
Haojun Liao 已提交
763 764 765 766
  blockDataEnsureCapacity(pRes, 4096);
  return buildPartitionResult(pOperator);
}

767
static void destroyPartitionOperatorInfo(void* param) {
H
Haojun Liao 已提交
768
  SPartitionOperatorInfo* pInfo = (SPartitionOperatorInfo*)param;
769
  cleanupBasicInfo(&pInfo->binfo);
H
Haojun Liao 已提交
770
  taosArrayDestroy(pInfo->pGroupCols);
771

772
  for (int i = 0; i < taosArrayGetSize(pInfo->pGroupColVals); i++) {
wmmhello's avatar
wmmhello 已提交
773 774 775
    SGroupKeys key = *(SGroupKeys*)taosArrayGet(pInfo->pGroupColVals, i);
    taosMemoryFree(key.pData);
  }
776

H
Haojun Liao 已提交
777
  taosArrayDestroy(pInfo->pGroupColVals);
H
Haojun Liao 已提交
778
  taosMemoryFree(pInfo->keyBuf);
779
  taosArrayDestroy(pInfo->sortedGroupArray);
D
dapan1121 已提交
780 781 782 783 784 785 786 787

  void* pGroupIter = taosHashIterate(pInfo->pGroupSet, NULL);
  while (pGroupIter != NULL) {
    SDataGroupInfo* pGroupInfo = pGroupIter;
    taosArrayDestroy(pGroupInfo->pPageList);
    pGroupIter = taosHashIterate(pInfo->pGroupSet, pGroupIter);
  }

wmmhello's avatar
wmmhello 已提交
788
  taosHashCleanup(pInfo->pGroupSet);
H
Haojun Liao 已提交
789
  taosMemoryFree(pInfo->columnOffset);
790

791
  cleanupExprSupp(&pInfo->scalarSup);
H
Haojun Liao 已提交
792
  destroyDiskbasedBuf(pInfo->pBuf);
D
dapan1121 已提交
793
  taosMemoryFreeClear(param);
H
Haojun Liao 已提交
794 795
}

796 797
SOperatorInfo* createPartitionOperatorInfo(SOperatorInfo* downstream, SPartitionPhysiNode* pPartNode,
                                           SExecTaskInfo* pTaskInfo) {
H
Haojun Liao 已提交
798
  SPartitionOperatorInfo* pInfo = taosMemoryCalloc(1, sizeof(SPartitionOperatorInfo));
799
  SOperatorInfo*          pOperator = taosMemoryCalloc(1, sizeof(SOperatorInfo));
H
Haojun Liao 已提交
800
  if (pInfo == NULL || pOperator == NULL) {
H
Haojun Liao 已提交
801 802
    goto _error;
  }
803

804
  int32_t    numOfCols = 0;
805 806 807 808
  SExprInfo* pExprInfo = createExprInfo(pPartNode->pTargets, NULL, &numOfCols);
  pInfo->pGroupCols = extractPartitionColInfo(pPartNode->pPartitionKeys);

  if (pPartNode->pExprs != NULL) {
809
    int32_t    num = 0;
810
    SExprInfo* pExprInfo1 = createExprInfo(pPartNode->pExprs, NULL, &num);
811
    int32_t    code = initExprSupp(&pInfo->scalarSup, pExprInfo1, num);
812 813 814
    if (code != TSDB_CODE_SUCCESS) {
      goto _error;
    }
815
  }
H
Haojun Liao 已提交
816 817 818 819 820 821 822

  _hash_fn_t hashFn = taosGetDefaultHashFunction(TSDB_DATA_TYPE_BINARY);
  pInfo->pGroupSet = taosHashInit(100, hashFn, false, HASH_NO_LOCK);
  if (pInfo->pGroupSet == NULL) {
    goto _error;
  }

823
  uint32_t defaultPgsz = 0;
824
  uint32_t defaultBufsz = 0;
H
Haojun Liao 已提交
825

H
Haojun Liao 已提交
826
  pInfo->binfo.pRes = createDataBlockFromDescNode(pPartNode->node.pOutputDataBlockDesc);
H
Haojun Liao 已提交
827
  getBufferPgSize(pInfo->binfo.pRes->info.rowSize, &defaultPgsz, &defaultBufsz);
828

wafwerar's avatar
wafwerar 已提交
829 830 831 832 833 834
  if (!osTempSpaceAvailable()) {
    terrno = TSDB_CODE_NO_AVAIL_DISK;
    pTaskInfo->code = terrno;
    qError("Create partition operator info failed since %s", terrstr(terrno));
    goto _error;
  }
H
Haojun Liao 已提交
835

wafwerar's avatar
wafwerar 已提交
836
  int32_t code = createDiskbasedBuf(&pInfo->pBuf, defaultPgsz, defaultBufsz, pTaskInfo->id.str, tsTempDir);
H
Haojun Liao 已提交
837 838 839 840
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

H
Haojun Liao 已提交
841 842
  pInfo->rowCapacity = blockDataGetCapacityInRow(pInfo->binfo.pRes, getBufPageSize(pInfo->pBuf));
  pInfo->columnOffset = setupColumnOffset(pInfo->binfo.pRes, pInfo->rowCapacity);
843
  code = initGroupOptrInfo(&pInfo->pGroupColVals, &pInfo->groupKeyLen, &pInfo->keyBuf, pInfo->pGroupCols);
H
Haojun Liao 已提交
844 845 846
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }
H
Haojun Liao 已提交
847

L
Liu Jicong 已提交
848 849
  setOperatorInfo(pOperator, "PartitionOperator", QUERY_NODE_PHYSICAL_PLAN_PARTITION, false, OP_NOT_OPENED, pInfo,
                  pTaskInfo);
850 851
  pOperator->exprSupp.numOfExprs = numOfCols;
  pOperator->exprSupp.pExprInfo = pExprInfo;
852

L
Liu Jicong 已提交
853
  pOperator->fpSet = createOperatorFpSet(operatorDummyOpenFn, hashPartition, NULL, destroyPartitionOperatorInfo, NULL);
H
Haojun Liao 已提交
854

H
Haojun Liao 已提交
855
  code = appendDownstream(pOperator, &downstream, 1);
856 857
  return pOperator;

858
_error:
H
Haojun Liao 已提交
859
  pTaskInfo->code = TSDB_CODE_OUT_OF_MEMORY;
H
Haojun Liao 已提交
860 861 862
  if (pInfo != NULL) {
    destroyPartitionOperatorInfo(pInfo);
  }
H
Haojun Liao 已提交
863
  taosMemoryFreeClear(pOperator);
864
  return NULL;
865 866
}

867 868 869
int32_t setGroupResultOutputBuf(SOperatorInfo* pOperator, SOptrBasicInfo* binfo, int32_t numOfCols, char* pData,
                                int16_t bytes, uint64_t groupId, SDiskbasedBuf* pBuf, SAggSupporter* pAggSup) {
  SExecTaskInfo*  pTaskInfo = pOperator->pTaskInfo;
870
  SResultRowInfo* pResultRowInfo = &binfo->resultRowInfo;
871
  SqlFunctionCtx* pCtx = pOperator->exprSupp.pCtx;
872 873 874 875 876

  SResultRow* pResultRow =
      doSetResultOutBufByKey(pBuf, pResultRowInfo, (char*)pData, bytes, true, groupId, pTaskInfo, false, pAggSup);
  assert(pResultRow != NULL);

877
  setResultRowInitCtx(pResultRow, pCtx, numOfCols, pOperator->exprSupp.rowEntryInfoOffset);
878
  return TSDB_CODE_SUCCESS;
D
dapan1121 已提交
879
}
880 881 882

uint64_t calGroupIdByData(SPartitionBySupporter* pParSup, SExprSupp* pExprSup, SSDataBlock* pBlock, int32_t rowId) {
  if (pExprSup->pExprInfo != NULL) {
883 884
    int32_t code =
        projectApplyFunctions(pExprSup->pExprInfo, pBlock, pBlock, pExprSup->pCtx, pExprSup->numOfExprs, NULL);
885 886 887 888 889
    if (code != TSDB_CODE_SUCCESS) {
      qError("calaculate group id error, code:%d", code);
    }
  }
  recordNewGroupKeys(pParSup->pGroupCols, pParSup->pGroupColVals, pBlock, rowId);
890
  int32_t  len = buildGroupKeys(pParSup->keyBuf, pParSup->pGroupColVals);
891 892 893 894
  uint64_t groupId = calcGroupId(pParSup->keyBuf, len);
  return groupId;
}

895
static bool hasRemainPartion(SStreamPartitionOperatorInfo* pInfo) { return pInfo->parIte != NULL; }
896 897 898

static SSDataBlock* buildStreamPartitionResult(SOperatorInfo* pOperator) {
  SStreamPartitionOperatorInfo* pInfo = pOperator->info;
899
  SSDataBlock*                  pDest = pInfo->binfo.pRes;
900 901 902
  ASSERT(hasRemainPartion(pInfo));
  SPartitionDataInfo* pParInfo = (SPartitionDataInfo*)pInfo->parIte;
  blockDataCleanup(pDest);
903
  int32_t      rows = taosArrayGetSize(pParInfo->rowIds);
904 905 906 907
  SSDataBlock* pSrc = pInfo->pInputDataBlock;
  for (int32_t i = 0; i < rows; i++) {
    int32_t rowIndex = *(int32_t*)taosArrayGet(pParInfo->rowIds, i);
    for (int32_t j = 0; j < pOperator->exprSupp.numOfExprs; j++) {
908
      int32_t          slotId = pOperator->exprSupp.pExprInfo[j].base.pParam[0].pCol->slotId;
909 910
      SColumnInfoData* pSrcCol = taosArrayGet(pSrc->pDataBlock, slotId);
      SColumnInfoData* pDestCol = taosArrayGet(pDest->pDataBlock, j);
911 912
      bool             isNull = colDataIsNull(pSrcCol, pSrc->info.rows, rowIndex, NULL);
      char*            pSrcData = colDataGetData(pSrcCol, rowIndex);
913 914 915
      colDataAppend(pDestCol, pDest->info.rows, pSrcData, isNull);
    }
    pDest->info.rows++;
L
Liu Jicong 已提交
916
    if (pInfo->tbnameCalSup.numOfExprs > 0 && i == 0) {
L
Liu Jicong 已提交
917 918 919 920
      void* tbname = NULL;
      if (streamStateGetParName(pOperator->pTaskInfo->streamInfo.pState, pParInfo->groupId, &tbname) == 0) {
        memcpy(pDest->info.parTbName, tbname, TSDB_TABLE_NAME_LEN);
        tdbFree(tbname);
921
      } else {
L
Liu Jicong 已提交
922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948
        SSDataBlock* pTmpBlock = blockCopyOneRow(pSrc, rowIndex);
        SSDataBlock* pResBlock = createDataBlock();
        pResBlock->info.rowSize = TSDB_TABLE_NAME_LEN;
        SColumnInfoData data = createColumnInfoData(TSDB_DATA_TYPE_VARCHAR, TSDB_TABLE_NAME_LEN, 0);
        taosArrayPush(pResBlock->pDataBlock, &data);
        blockDataEnsureCapacity(pResBlock, 1);
        projectApplyFunctions(pInfo->tbnameCalSup.pExprInfo, pResBlock, pTmpBlock, pInfo->tbnameCalSup.pCtx, 1, NULL);
        ASSERT(pResBlock->info.rows == 1);
        ASSERT(taosArrayGetSize(pResBlock->pDataBlock) == 1);
        SColumnInfoData* pCol = taosArrayGet(pResBlock->pDataBlock, 0);
        ASSERT(pCol->info.type == TSDB_DATA_TYPE_VARCHAR);
        void* pData = colDataGetVarData(pCol, 0);
        // TODO check tbname validity
        if (pData != (void*)-1) {
          memset(pDest->info.parTbName, 0, TSDB_TABLE_NAME_LEN);
          int32_t len = TMIN(varDataLen(pData), TSDB_TABLE_NAME_LEN - 1);
          memcpy(pDest->info.parTbName, varDataVal(pData), len);
          /*pDest->info.parTbName[len + 1] = 0;*/
        } else {
          pDest->info.parTbName[0] = 0;
        }
        if (pParInfo->groupId && pDest->info.parTbName[0]) {
          streamStatePutParName(pOperator->pTaskInfo->streamInfo.pState, pParInfo->groupId, pDest->info.parTbName);
        }
        /*printf("\n\n set name %s\n\n", pDest->info.parTbName);*/
        blockDataDestroy(pTmpBlock);
        blockDataDestroy(pResBlock);
L
Liu Jicong 已提交
949
      }
950
    }
951
  }
952 953
  taosArrayDestroy(pParInfo->rowIds);
  pParInfo->rowIds = NULL;
954 955 956 957 958 959 960 961 962 963 964 965 966
  blockDataUpdateTsWindow(pDest, pInfo->tsColIndex);
  pDest->info.groupId = pParInfo->groupId;
  pOperator->resultInfo.totalRows += pDest->info.rows;
  pInfo->parIte = taosHashIterate(pInfo->pPartitions, pInfo->parIte);
  ASSERT(pDest->info.rows > 0);
  printDataBlock(pDest, "stream partitionby");
  return pDest;
}

static void doStreamHashPartitionImpl(SStreamPartitionOperatorInfo* pInfo, SSDataBlock* pBlock) {
  pInfo->pInputDataBlock = pBlock;
  for (int32_t i = 0; i < pBlock->info.rows; ++i) {
    recordNewGroupKeys(pInfo->partitionSup.pGroupCols, pInfo->partitionSup.pGroupColVals, pBlock, i);
967 968 969
    int32_t             keyLen = buildGroupKeys(pInfo->partitionSup.keyBuf, pInfo->partitionSup.pGroupColVals);
    SPartitionDataInfo* pParData =
        (SPartitionDataInfo*)taosHashGet(pInfo->pPartitions, pInfo->partitionSup.keyBuf, keyLen);
970 971 972 973 974 975 976
    if (pParData) {
      taosArrayPush(pParData->rowIds, &i);
    } else {
      SPartitionDataInfo newParData = {0};
      newParData.groupId = calcGroupId(pInfo->partitionSup.keyBuf, keyLen);
      newParData.rowIds = taosArrayInit(64, sizeof(int32_t));
      taosArrayPush(newParData.rowIds, &i);
977
      taosHashPut(pInfo->pPartitions, pInfo->partitionSup.keyBuf, keyLen, &newParData, sizeof(SPartitionDataInfo));
978 979 980 981 982 983 984 985 986
    }
  }
}

static SSDataBlock* doStreamHashPartition(SOperatorInfo* pOperator) {
  if (pOperator->status == OP_EXEC_DONE) {
    return NULL;
  }

987
  SExecTaskInfo*                pTaskInfo = pOperator->pTaskInfo;
988 989 990 991 992
  SStreamPartitionOperatorInfo* pInfo = pOperator->info;
  if (hasRemainPartion(pInfo)) {
    return buildStreamPartitionResult(pOperator);
  }

993
  int64_t        st = taosGetTimestampUs();
994 995 996 997 998
  SOperatorInfo* downstream = pOperator->pDownstream[0];
  {
    pInfo->pInputDataBlock = NULL;
    SSDataBlock* pBlock = downstream->fpSet.getNextFn(downstream);
    if (pBlock == NULL) {
H
Haojun Liao 已提交
999
      setOperatorCompleted(pOperator);
1000 1001 1002 1003 1004 1005 1006 1007 1008
      return NULL;
    }
    printDataBlock(pBlock, "stream partitionby recv");
    switch (pBlock->info.type) {
      case STREAM_NORMAL:
      case STREAM_PULL_DATA:
      case STREAM_INVALID:
        pInfo->binfo.pRes->info.type = pBlock->info.type;
        break;
1009 1010 1011
      case STREAM_DELETE_DATA: {
        copyDataBlock(pInfo->pDelRes, pBlock);
        pInfo->pDelRes->info.type = STREAM_DELETE_RESULT;
5
54liuyao 已提交
1012
        printDataBlock(pInfo->pDelRes, "stream partitionby delete");
1013
        return pInfo->pDelRes;
1014
      } break;
1015 1016 1017 1018 1019 1020
      default:
        return pBlock;
    }

    // there is an scalar expression that needs to be calculated right before apply the group aggregation.
    if (pInfo->scalarSup.pExprInfo != NULL) {
1021 1022
      pTaskInfo->code = projectApplyFunctions(pInfo->scalarSup.pExprInfo, pBlock, pBlock, pInfo->scalarSup.pCtx,
                                              pInfo->scalarSup.numOfExprs, NULL);
1023 1024 1025 1026 1027 1028 1029 1030
      if (pTaskInfo->code != TSDB_CODE_SUCCESS) {
        longjmp(pTaskInfo->env, pTaskInfo->code);
      }
    }
    taosHashClear(pInfo->pPartitions);
    doStreamHashPartitionImpl(pInfo, pBlock);
  }
  pOperator->cost.openCost = (taosGetTimestampUs() - st) / 1000.0;
1031

1032 1033 1034 1035 1036 1037 1038 1039 1040
  pInfo->parIte = taosHashIterate(pInfo->pPartitions, NULL);
  return buildStreamPartitionResult(pOperator);
}

static void destroyStreamPartitionOperatorInfo(void* param) {
  SStreamPartitionOperatorInfo* pInfo = (SStreamPartitionOperatorInfo*)param;
  cleanupBasicInfo(&pInfo->binfo);
  taosArrayDestroy(pInfo->partitionSup.pGroupCols);

1041
  for (int i = 0; i < taosArrayGetSize(pInfo->partitionSup.pGroupColVals); i++) {
1042 1043 1044 1045 1046 1047 1048
    SGroupKeys key = *(SGroupKeys*)taosArrayGet(pInfo->partitionSup.pGroupColVals, i);
    taosMemoryFree(key.pData);
  }
  taosArrayDestroy(pInfo->partitionSup.pGroupColVals);

  taosMemoryFree(pInfo->partitionSup.keyBuf);
  cleanupExprSupp(&pInfo->scalarSup);
L
Liu Jicong 已提交
1049 1050
  cleanupExprSupp(&pInfo->tbnameCalSup);
  cleanupExprSupp(&pInfo->tagCalSup);
1051
  blockDataDestroy(pInfo->pDelRes);
1052
  taosHashCleanup(pInfo->pPartitions);
1053 1054 1055 1056 1057 1058 1059 1060 1061 1062
  taosMemoryFreeClear(param);
}

void initParDownStream(SOperatorInfo* downstream, SPartitionBySupporter* pParSup, SExprSupp* pExpr) {
  if (downstream->operatorType != QUERY_NODE_PHYSICAL_PLAN_STREAM_SCAN) {
    return;
  }
  SStreamScanInfo* pScanInfo = downstream->info;
  pScanInfo->partitionSup = *pParSup;
  pScanInfo->pPartScalarSup = pExpr;
5
54liuyao 已提交
1063 1064 1065
  if (!pScanInfo->pUpdateInfo) {
    pScanInfo->pUpdateInfo = updateInfoInit(60000, TSDB_TIME_PRECISION_MILLI, 0);
  }
1066 1067
}

1068 1069
SOperatorInfo* createStreamPartitionOperatorInfo(SOperatorInfo* downstream, SStreamPartitionPhysiNode* pPartNode,
                                                 SExecTaskInfo* pTaskInfo) {
1070 1071 1072 1073 1074 1075
  SStreamPartitionOperatorInfo* pInfo = taosMemoryCalloc(1, sizeof(SStreamPartitionOperatorInfo));
  SOperatorInfo*                pOperator = taosMemoryCalloc(1, sizeof(SOperatorInfo));
  if (pInfo == NULL || pOperator == NULL) {
    goto _error;
  }
  int32_t code = TSDB_CODE_SUCCESS;
1076
  pInfo->partitionSup.pGroupCols = extractPartitionColInfo(pPartNode->part.pPartitionKeys);
1077

1078
  if (pPartNode->part.pExprs != NULL) {
1079
    int32_t    num = 0;
1080
    SExprInfo* pCalExprInfo = createExprInfo(pPartNode->part.pExprs, NULL, &num);
1081 1082 1083 1084 1085 1086
    code = initExprSupp(&pInfo->scalarSup, pCalExprInfo, num);
    if (code != TSDB_CODE_SUCCESS) {
      goto _error;
    }
  }

1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100
  if (pPartNode->pSubtable != NULL) {
    SExprInfo* pSubTableExpr = taosMemoryCalloc(1, sizeof(SExprInfo));
    if (pSubTableExpr == NULL) {
      code = TSDB_CODE_OUT_OF_MEMORY;
      goto _error;
    }
    pInfo->tbnameCalSup.pExprInfo = pSubTableExpr;
    createExprFromOneNode(pSubTableExpr, pPartNode->pSubtable, 0);
    code = initExprSupp(&pInfo->tbnameCalSup, pSubTableExpr, 1);
    if (code != TSDB_CODE_SUCCESS) {
      goto _error;
    }
  }

L
Liu Jicong 已提交
1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113
  if (pPartNode->pTags != NULL) {
    int32_t    numOfTags;
    SExprInfo* pTagExpr = createExprInfo(pPartNode->pTags, NULL, &numOfTags);
    if (pTagExpr == NULL) {
      terrno = TSDB_CODE_OUT_OF_MEMORY;
      goto _error;
    }
    if (initExprSupp(&pInfo->tagCalSup, pTagExpr, numOfTags) != 0) {
      terrno = TSDB_CODE_OUT_OF_MEMORY;
      goto _error;
    }
  }

1114
  int32_t keyLen = 0;
1115 1116
  code = initGroupOptrInfo(&pInfo->partitionSup.pGroupColVals, &keyLen, &pInfo->partitionSup.keyBuf,
                           pInfo->partitionSup.pGroupCols);
1117 1118 1119 1120 1121
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }
  pInfo->partitionSup.needCalc = true;

H
Haojun Liao 已提交
1122
  pInfo->binfo.pRes = createDataBlockFromDescNode(pPartNode->part.node.pOutputDataBlockDesc);
1123
  if (pInfo->binfo.pRes == NULL) {
1124 1125
    goto _error;
  }
1126 1127 1128

  blockDataEnsureCapacity(pInfo->binfo.pRes, 4096);

1129 1130
  pInfo->parIte = NULL;
  pInfo->pInputDataBlock = NULL;
1131

1132
  _hash_fn_t hashFn = taosGetDefaultHashFunction(TSDB_DATA_TYPE_BINARY);
1133 1134 1135
  pInfo->pPartitions = taosHashInit(1024, hashFn, false, HASH_NO_LOCK);
  pInfo->tsColIndex = 0;
  pInfo->pDelRes = createSpecialDataBlock(STREAM_DELETE_RESULT);
1136

1137
  int32_t    numOfCols = 0;
1138
  SExprInfo* pExprInfo = createExprInfo(pPartNode->part.pTargets, NULL, &numOfCols);
1139

L
Liu Jicong 已提交
1140 1141
  setOperatorInfo(pOperator, "StreamPartitionOperator", QUERY_NODE_PHYSICAL_PLAN_STREAM_PARTITION, false, OP_NOT_OPENED,
                  pInfo, pTaskInfo);
1142 1143
  pOperator->exprSupp.numOfExprs = numOfCols;
  pOperator->exprSupp.pExprInfo = pExprInfo;
L
Liu Jicong 已提交
1144 1145
  pOperator->fpSet =
      createOperatorFpSet(operatorDummyOpenFn, doStreamHashPartition, NULL, destroyStreamPartitionOperatorInfo, NULL);
1146 1147 1148 1149 1150

  initParDownStream(downstream, &pInfo->partitionSup, &pInfo->scalarSup);
  code = appendDownstream(pOperator, &downstream, 1);
  return pOperator;

1151
_error:
1152
  pTaskInfo->code = TSDB_CODE_OUT_OF_MEMORY;
1153
  destroyStreamPartitionOperatorInfo(pInfo);
1154 1155 1156
  taosMemoryFreeClear(pOperator);
  return NULL;
}
H
Haojun Liao 已提交
1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189

SArray* extractColumnInfo(SNodeList* pNodeList) {
  size_t  numOfCols = LIST_LENGTH(pNodeList);
  SArray* pList = taosArrayInit(numOfCols, sizeof(SColumn));
  if (pList == NULL) {
    terrno = TSDB_CODE_OUT_OF_MEMORY;
    return NULL;
  }

  for (int32_t i = 0; i < numOfCols; ++i) {
    STargetNode* pNode = (STargetNode*)nodesListGetNode(pNodeList, i);

    if (nodeType(pNode->pExpr) == QUERY_NODE_COLUMN) {
      SColumnNode* pColNode = (SColumnNode*)pNode->pExpr;

      SColumn c = extractColumnFromColumnNode(pColNode);
      taosArrayPush(pList, &c);
    } else if (nodeType(pNode->pExpr) == QUERY_NODE_VALUE) {
      SValueNode* pValNode = (SValueNode*)pNode->pExpr;
      SColumn     c = {0};
      c.slotId = pNode->slotId;
      c.colId = pNode->slotId;
      c.type = pValNode->node.type;
      c.bytes = pValNode->node.resType.bytes;
      c.scale = pValNode->node.resType.scale;
      c.precision = pValNode->node.resType.precision;

      taosArrayPush(pList, &c);
    }
  }

  return pList;
}