groupoperator.c 47.3 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
  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 已提交
219 220 221 222 223
    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 已提交
224 225
      varDataCopy(pStart, pkey->pData);
      pStart += varDataTLen(pkey->pData);
226
      ASSERT(varDataTLen(pkey->pData) <= pkey->bytes);
H
Haojun Liao 已提交
227 228 229 230 231 232
    } else {
      memcpy(pStart, pkey->pData, pkey->bytes);
      pStart += pkey->bytes;
    }
  }

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

// 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) {
239
    if (pCtx[i].functionId == -1) {  // select count(*),key from t group by key.
H
Haojun Liao 已提交
240 241 242
      SResultRowEntryInfo* pEntryInfo = GET_RES_INFO(&pCtx[i]);

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

wmmhello's avatar
wmmhello 已提交
248 249 250 251
        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)) {
252 253 254 255
          varDataCopy(dest, data);
        } else {
          memcpy(dest, data, pColInfoData->info.bytes);
        }
256
      } else {  // it is a NULL value
H
Haojun Liao 已提交
257
        pEntryInfo->isNullRes = 1;
H
Haojun Liao 已提交
258
      }
H
Haojun Liao 已提交
259 260

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

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

269
  SqlFunctionCtx* pCtx = pOperator->exprSupp.pCtx;
H
Haojun Liao 已提交
270 271 272 273 274 275
  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 已提交
276
  int32_t len = 0;
wmmhello's avatar
wmmhello 已提交
277
  terrno = TSDB_CODE_SUCCESS;
H
Haojun Liao 已提交
278

H
Haojun Liao 已提交
279
  int32_t num = 0;
D
dapan1121 已提交
280
  uint64_t groupId = 0;
H
Haojun Liao 已提交
281 282 283
  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
    int32_t ret = setGroupResultOutputBuf(pOperator, &(pInfo->binfo), pOperator->exprSupp.numOfExprs, pInfo->keyBuf,
311
                                          len, pBlock->info.id.groupId, pInfo->aggSup.pResultBuf, &pInfo->aggSup);
H
Haojun Liao 已提交
312
    if (ret != TSDB_CODE_SUCCESS) {  // null data, too many state code
S
Shengliang Guan 已提交
313
      T_LONG_JMP(pTaskInfo->env, TSDB_CODE_APP_ERROR);
H
Haojun Liao 已提交
314 315 316
    }

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

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

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

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

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

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

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

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

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

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

368 369
  SExecTaskInfo* pTaskInfo = pOperator->pTaskInfo;

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

375 376 377
  int32_t order = TSDB_ORDER_ASC;
  int32_t scanFlag = MAIN_SCAN;

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

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

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

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

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

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

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

409 410 411 412 413 414 415 416 417 418 419 420 421 422
#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
423
  initGroupedResultInfo(&pInfo->groupResInfo, pInfo->aggSup.pResultRowHashTable, 0);
424

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

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

H
Haojun Liao 已提交
438
  SSDataBlock* pResBlock = createDataBlockFromDescNode(pAggNode->node.pOutputDataBlockDesc);
H
Haojun Liao 已提交
439 440 441 442 443 444 445 446 447
  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);
448
  code = initExprSupp(&pInfo->scalarSup, pScalarExprInfo, numOfScalarExpr);
449 450 451
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }
452

H
Haojun Liao 已提交
453
  initResultSizeInfo(&pOperator->resultInfo, 4096);
H
Haojun Liao 已提交
454 455
  blockDataEnsureCapacity(pInfo->binfo.pRes, pOperator->resultInfo.capacity);

H
Haojun Liao 已提交
456
  code = initGroupOptrInfo(&pInfo->pGroupColVals, &pInfo->groupKeyLen, &pInfo->keyBuf, pInfo->pGroupCols);
H
Haojun Liao 已提交
457 458 459 460
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

H
Haojun Liao 已提交
461 462
  int32_t    num = 0;
  SExprInfo* pExprInfo = createExprInfo(pAggNode->pAggFuncs, pAggNode->pGroupKeys, &num);
L
Liu Jicong 已提交
463 464
  code = initAggSup(&pOperator->exprSupp, &pInfo->aggSup, pExprInfo, num, pInfo->groupKeyLen, pTaskInfo->id.str,
                    pTaskInfo->streamInfo.pState);
465 466 467 468
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

469 470 471 472 473
  code = filterInitFromNode((SNode*)pAggNode->node.pConditions, &pOperator->exprSupp.pFilterInfo, 0);
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

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

D
dapan1121 已提交
477 478
  pInfo->binfo.mergeResultBlock = pAggNode->mergeDataBlock;

dengyihao's avatar
dengyihao 已提交
479 480
  pOperator->fpSet = createOperatorFpSet(optrDummyOpenFn, hashGroupbyAggregate, NULL, destroyGroupOperatorInfo,
                                         optrDefaultBufFn, NULL);
H
Haojun Liao 已提交
481
  code = appendDownstream(pOperator, &downstream, 1);
482 483 484 485
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

H
Haojun Liao 已提交
486 487
  return pOperator;

488
_error:
489
  pTaskInfo->code = code;
H
Haojun Liao 已提交
490 491 492
  if (pInfo != NULL) {
    destroyGroupOperatorInfo(pInfo);
  }
H
Haojun Liao 已提交
493 494
  taosMemoryFreeClear(pOperator);
  return NULL;
495 496
}

H
Haojun Liao 已提交
497 498
static void doHashPartition(SOperatorInfo* pOperator, SSDataBlock* pBlock) {
  SPartitionOperatorInfo* pInfo = pOperator->info;
dengyihao's avatar
dengyihao 已提交
499 500
  SExecTaskInfo*          pTaskInfo = pOperator->pTaskInfo;

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

505
    SDataGroupInfo* pGroupInfo = NULL;
506
    void*           pPage = getCurrentDataGroupInfo(pInfo, &pGroupInfo, len);
507 508 509
    if (pPage == NULL) {
      T_LONG_JMP(pTaskInfo->env, terrno);
    }
H
Haojun Liao 已提交
510

511 512 513 514 515
    pGroupInfo->numOfRows += 1;

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

518
    // number of rows
519
    int32_t* rows = (int32_t*)pPage;
H
Haojun Liao 已提交
520

521
    size_t numOfCols = pOperator->exprSupp.numOfExprs;
522
    for (int32_t i = 0; i < numOfCols; ++i) {
523
      SExprInfo* pExpr = &pOperator->exprSupp.pExprInfo[i];
524
      int32_t    slotId = pExpr->base.pParam[0].pCol->slotId;
525 526

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

H
Haojun Liao 已提交
528 529
      int32_t bytes = pColInfoData->info.bytes;
      int32_t startOffset = pInfo->columnOffset[i];
H
Haojun Liao 已提交
530

531
      int32_t* columnLen = NULL;
532
      int32_t  contentLen = 0;
H
Haojun Liao 已提交
533 534

      if (IS_VAR_DATA_TYPE(pColInfoData->info.type)) {
wafwerar's avatar
wafwerar 已提交
535
        int32_t* offset = (int32_t*)((char*)pPage + startOffset);
536 537
        columnLen = (int32_t*)((char*)pPage + startOffset + sizeof(int32_t) * pInfo->rowCapacity);
        char* data = (char*)((char*)columnLen + sizeof(int32_t));
H
Haojun Liao 已提交
538 539 540 541

        if (colDataIsNull_s(pColInfoData, j)) {
          offset[(*rows)] = -1;
          contentLen = 0;
542
        } else if (pColInfoData->info.type == TSDB_DATA_TYPE_JSON) {
wmmhello's avatar
wmmhello 已提交
543
          offset[*rows] = (*columnLen);
544
          char*   src = colDataGetData(pColInfoData, j);
wmmhello's avatar
wmmhello 已提交
545 546 547 548
          int32_t dataLen = getJsonValueLen(src);

          memcpy(data + (*columnLen), src, dataLen);
          int32_t v = (data + (*columnLen) + dataLen - (char*)pPage);
549
          ASSERT(v > 0);
wmmhello's avatar
wmmhello 已提交
550 551

          contentLen = dataLen;
H
Haojun Liao 已提交
552 553 554 555
        } else {
          offset[*rows] = (*columnLen);
          char* src = colDataGetData(pColInfoData, j);
          memcpy(data + (*columnLen), src, varDataTLen(src));
556
          int32_t v = (data + (*columnLen) + varDataTLen(src) - (char*)pPage);
557
          ASSERT(v > 0);
558

H
Haojun Liao 已提交
559 560
          contentLen = varDataTLen(src);
        }
H
Haojun Liao 已提交
561
      } else {
wafwerar's avatar
wafwerar 已提交
562
        char* bitmap = (char*)pPage + startOffset;
563 564
        columnLen = (int32_t*)((char*)pPage + startOffset + BitmapLen(pInfo->rowCapacity));
        char* data = (char*)columnLen + sizeof(int32_t);
H
Haojun Liao 已提交
565 566 567

        bool isNull = colDataIsNull_f(pColInfoData->nullbitmap, j);
        if (isNull) {
H
Haojun Liao 已提交
568
          colDataSetNull_f(bitmap, (*rows));
H
Haojun Liao 已提交
569
        } else {
H
Haojun Liao 已提交
570
          memcpy(data + (*columnLen), colDataGetData(pColInfoData, j), bytes);
571
          ASSERT((data + (*columnLen) + bytes - (char*)pPage) <= getBufPageSize(pInfo->pBuf));
H
Haojun Liao 已提交
572
        }
H
Haojun Liao 已提交
573
        contentLen = bytes;
H
Haojun Liao 已提交
574
      }
H
Haojun Liao 已提交
575 576

      (*columnLen) += contentLen;
H
Haojun Liao 已提交
577 578
    }

H
Haojun Liao 已提交
579 580
    (*rows) += 1;

H
Haojun Liao 已提交
581 582 583
    setBufPageDirty(pPage, true);
    releaseBufPage(pInfo->pBuf, pPage);
  }
H
Haojun Liao 已提交
584 585 586 587 588 589
}

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

  void* pPage = NULL;
590
  if (p == NULL) {  // it is a new group
H
Haojun Liao 已提交
591 592 593 594 595 596 597
    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;
598
    pPage = getNewBufPage(pInfo->pBuf, &pageId);
599 600 601
    if (pPage == NULL) {
      return pPage;
    }
H
Haojun Liao 已提交
602

603
    taosArrayPush(p->pPageList, &pageId);
604
    *(int32_t*)pPage = 0;
H
Haojun Liao 已提交
605 606 607
  } else {
    int32_t* curId = taosArrayGetLast(p->pPageList);
    pPage = getBufPage(pInfo->pBuf, *curId);
608 609 610 611
    if (pPage == NULL) {
      qError("failed to get buffer, code:%s", tstrerror(terrno));
      return pPage;
    }
H
Haojun Liao 已提交
612

613
    int32_t* rows = (int32_t*)pPage;
H
Haojun Liao 已提交
614
    if (*rows >= pInfo->rowCapacity) {
615 616 617
      // release buffer
      releaseBufPage(pInfo->pBuf, pPage);

H
Haojun Liao 已提交
618 619
      // add a new page for current group
      int32_t pageId = 0;
620
      pPage = getNewBufPage(pInfo->pBuf, &pageId);
621 622 623 624 625
      if (pPage == NULL) {
        qError("failed to get new buffer, code:%s", tstrerror(terrno));
        return NULL;
      }

H
Haojun Liao 已提交
626
      taosArrayPush(p->pPageList, &pageId);
627
      memset(pPage, 0, getBufPageSize(pInfo->pBuf));
H
Haojun Liao 已提交
628 629
    }
  }
H
Haojun Liao 已提交
630

H
Haojun Liao 已提交
631 632 633 634 635 636 637 638 639 640 641 642 643 644
  *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 已提交
645 646
}

H
Haojun Liao 已提交
647
int32_t* setupColumnOffset(const SSDataBlock* pBlock, int32_t rowCapacity) {
648
  size_t   numOfCols = taosArrayGetSize(pBlock->pDataBlock);
649
  int32_t* offset = taosMemoryCalloc(numOfCols, sizeof(int32_t));
H
Haojun Liao 已提交
650

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

654
  for (int32_t i = 0; i < numOfCols - 1; ++i) {
H
Haojun Liao 已提交
655 656 657 658
    SColumnInfoData* pColInfoData = taosArrayGet(pBlock->pDataBlock, i);

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

H
Haojun Liao 已提交
660 661 662 663 664 665 666 667 668 669 670 671
    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 已提交
672
static void clearPartitionOperator(SPartitionOperatorInfo* pInfo) {
673 674 675 676
  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 已提交
677
  }
678
  taosArrayClear(pInfo->sortedGroupArray);
5
54liuyao 已提交
679 680 681
  clearDiskbasedBuf(pInfo->pBuf);
}

682 683 684
static int compareDataGroupInfo(const void* group1, const void* group2) {
  const SDataGroupInfo* pGroupInfo1 = group1;
  const SDataGroupInfo* pGroupInfo2 = group2;
685 686 687 688 689

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

690
  return (pGroupInfo1->groupId < pGroupInfo2->groupId) ? -1 : 1;
691 692
}

H
Haojun Liao 已提交
693 694
static SSDataBlock* buildPartitionResult(SOperatorInfo* pOperator) {
  SPartitionOperatorInfo* pInfo = pOperator->info;
dengyihao's avatar
dengyihao 已提交
695 696
  SExecTaskInfo*          pTaskInfo = pOperator->pTaskInfo;

697 698
  SDataGroupInfo* pGroupInfo =
      (pInfo->groupIndex != -1) ? taosArrayGet(pInfo->sortedGroupArray, pInfo->groupIndex) : NULL;
699
  if (pInfo->groupIndex == -1 || pInfo->pageIndex >= taosArrayGetSize(pGroupInfo->pPageList)) {
H
Haojun Liao 已提交
700
    // try next group data
701 702
    ++pInfo->groupIndex;
    if (pInfo->groupIndex >= taosArrayGetSize(pInfo->sortedGroupArray)) {
H
Haojun Liao 已提交
703
      setOperatorCompleted(pOperator);
5
54liuyao 已提交
704
      clearPartitionOperator(pInfo);
H
Haojun Liao 已提交
705 706 707
      return NULL;
    }

708
    pGroupInfo = taosArrayGet(pInfo->sortedGroupArray, pInfo->groupIndex);
H
Haojun Liao 已提交
709 710 711 712
    pInfo->pageIndex = 0;
  }

  int32_t* pageId = taosArrayGet(pGroupInfo->pPageList, pInfo->pageIndex);
713
  void*    page = getBufPage(pInfo->pBuf, *pageId);
714 715 716 717
  if (page == NULL) {
    qError("failed to get buffer, code:%s, %s", tstrerror(terrno), GET_TASKID(pTaskInfo));
    T_LONG_JMP(pTaskInfo->env, terrno);
  }
dengyihao's avatar
dengyihao 已提交
718

719
  blockDataEnsureCapacity(pInfo->binfo.pRes, pInfo->rowCapacity);
H
Haojun Liao 已提交
720
  blockDataFromBuf1(pInfo->binfo.pRes, page, pInfo->rowCapacity);
H
Haojun Liao 已提交
721 722

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

725
  pInfo->binfo.pRes->info.dataLoad = 1;
726
  blockDataUpdateTsWindow(pInfo->binfo.pRes, 0);
H
Haojun Liao 已提交
727
  pInfo->binfo.pRes->info.id.groupId = pGroupInfo->groupId;
728 729

  pOperator->resultInfo.totalRows += pInfo->binfo.pRes->info.rows;
H
Haojun Liao 已提交
730 731 732
  return pInfo->binfo.pRes;
}

733
static SSDataBlock* hashPartition(SOperatorInfo* pOperator) {
H
Haojun Liao 已提交
734 735
  if (pOperator->status == OP_EXEC_DONE) {
    return NULL;
736 737
  }

738 739 740
  SExecTaskInfo* pTaskInfo = pOperator->pTaskInfo;

  SPartitionOperatorInfo* pInfo = pOperator->info;
741
  SSDataBlock*            pRes = pInfo->binfo.pRes;
742

H
Haojun Liao 已提交
743
  if (pOperator->status == OP_RES_TO_RETURN) {
H
Haojun Liao 已提交
744 745
    blockDataCleanup(pRes);
    return buildPartitionResult(pOperator);
H
Haojun Liao 已提交
746 747
  }

748
  int64_t        st = taosGetTimestampUs();
H
Haojun Liao 已提交
749
  SOperatorInfo* downstream = pOperator->pDownstream[0];
H
Haojun Liao 已提交
750

H
Haojun Liao 已提交
751
  while (1) {
752
    SSDataBlock* pBlock = downstream->fpSet.getNextFn(downstream);
H
Haojun Liao 已提交
753 754 755
    if (pBlock == NULL) {
      break;
    }
H
Haojun Liao 已提交
756

757
    // there is an scalar expression that needs to be calculated right before apply the group aggregation.
758
    if (pInfo->scalarSup.pExprInfo != NULL) {
759 760
      pTaskInfo->code = projectApplyFunctions(pInfo->scalarSup.pExprInfo, pBlock, pBlock, pInfo->scalarSup.pCtx,
                                              pInfo->scalarSup.numOfExprs, NULL);
761
      if (pTaskInfo->code != TSDB_CODE_SUCCESS) {
762
        T_LONG_JMP(pTaskInfo->env, pTaskInfo->code);
763 764 765
      }
    }

wmmhello's avatar
wmmhello 已提交
766
    terrno = TSDB_CODE_SUCCESS;
H
Haojun Liao 已提交
767
    doHashPartition(pOperator, pBlock);
wmmhello's avatar
wmmhello 已提交
768
    if (terrno != TSDB_CODE_SUCCESS) {  // group by json error
769
      T_LONG_JMP(pTaskInfo->env, terrno);
wmmhello's avatar
wmmhello 已提交
770
    }
H
Haojun Liao 已提交
771 772
  }

773
  SArray* groupArray = taosArrayInit(taosHashGetSize(pInfo->pGroupSet), sizeof(SDataGroupInfo));
774 775

  void* pGroupIter = taosHashIterate(pInfo->pGroupSet, NULL);
776 777 778 779 780 781 782 783 784 785 786
  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);

787 788
  pOperator->cost.openCost = (taosGetTimestampUs() - st) / 1000.0;

H
Haojun Liao 已提交
789
  pOperator->status = OP_RES_TO_RETURN;
H
Haojun Liao 已提交
790 791 792 793
  blockDataEnsureCapacity(pRes, 4096);
  return buildPartitionResult(pOperator);
}

794
static void destroyPartitionOperatorInfo(void* param) {
H
Haojun Liao 已提交
795
  SPartitionOperatorInfo* pInfo = (SPartitionOperatorInfo*)param;
796
  cleanupBasicInfo(&pInfo->binfo);
H
Haojun Liao 已提交
797
  taosArrayDestroy(pInfo->pGroupCols);
798

799
  for (int i = 0; i < taosArrayGetSize(pInfo->pGroupColVals); i++) {
wmmhello's avatar
wmmhello 已提交
800 801 802
    SGroupKeys key = *(SGroupKeys*)taosArrayGet(pInfo->pGroupColVals, i);
    taosMemoryFree(key.pData);
  }
803

H
Haojun Liao 已提交
804
  taosArrayDestroy(pInfo->pGroupColVals);
H
Haojun Liao 已提交
805
  taosMemoryFree(pInfo->keyBuf);
dengyihao's avatar
dengyihao 已提交
806 807 808 809 810 811

  int32_t size = taosArrayGetSize(pInfo->sortedGroupArray);
  for (int32_t i = 0; i < size; i++) {
    SDataGroupInfo* pGp = taosArrayGet(pInfo->sortedGroupArray, i);
    taosArrayDestroy(pGp->pPageList);
  }
812
  taosArrayDestroy(pInfo->sortedGroupArray);
D
dapan1121 已提交
813 814 815 816 817 818 819 820

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

wmmhello's avatar
wmmhello 已提交
821
  taosHashCleanup(pInfo->pGroupSet);
H
Haojun Liao 已提交
822
  taosMemoryFree(pInfo->columnOffset);
823

824
  cleanupExprSupp(&pInfo->scalarSup);
H
Haojun Liao 已提交
825
  destroyDiskbasedBuf(pInfo->pBuf);
D
dapan1121 已提交
826
  taosMemoryFreeClear(param);
H
Haojun Liao 已提交
827 828
}

829 830
SOperatorInfo* createPartitionOperatorInfo(SOperatorInfo* downstream, SPartitionPhysiNode* pPartNode,
                                           SExecTaskInfo* pTaskInfo) {
H
Haojun Liao 已提交
831
  SPartitionOperatorInfo* pInfo = taosMemoryCalloc(1, sizeof(SPartitionOperatorInfo));
832
  SOperatorInfo*          pOperator = taosMemoryCalloc(1, sizeof(SOperatorInfo));
H
Haojun Liao 已提交
833
  if (pInfo == NULL || pOperator == NULL) {
dengyihao's avatar
dengyihao 已提交
834 835
    terrno = TSDB_CODE_OUT_OF_MEMORY;
    pTaskInfo->code = terrno;
H
Haojun Liao 已提交
836 837
    goto _error;
  }
838

839
  int32_t    numOfCols = 0;
840 841 842 843
  SExprInfo* pExprInfo = createExprInfo(pPartNode->pTargets, NULL, &numOfCols);
  pInfo->pGroupCols = extractPartitionColInfo(pPartNode->pPartitionKeys);

  if (pPartNode->pExprs != NULL) {
844
    int32_t    num = 0;
845
    SExprInfo* pExprInfo1 = createExprInfo(pPartNode->pExprs, NULL, &num);
846
    int32_t    code = initExprSupp(&pInfo->scalarSup, pExprInfo1, num);
847
    if (code != TSDB_CODE_SUCCESS) {
dengyihao's avatar
dengyihao 已提交
848 849
      terrno = code;
      pTaskInfo->code = terrno;
850 851
      goto _error;
    }
852
  }
H
Haojun Liao 已提交
853 854 855 856

  _hash_fn_t hashFn = taosGetDefaultHashFunction(TSDB_DATA_TYPE_BINARY);
  pInfo->pGroupSet = taosHashInit(100, hashFn, false, HASH_NO_LOCK);
  if (pInfo->pGroupSet == NULL) {
dengyihao's avatar
dengyihao 已提交
857 858
    terrno = TSDB_CODE_OUT_OF_MEMORY;
    pTaskInfo->code = terrno;
H
Haojun Liao 已提交
859 860 861
    goto _error;
  }

862
  uint32_t defaultPgsz = 0;
863
  uint32_t defaultBufsz = 0;
H
Haojun Liao 已提交
864

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

wafwerar's avatar
wafwerar 已提交
868 869 870 871 872 873
  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 已提交
874

wafwerar's avatar
wafwerar 已提交
875
  int32_t code = createDiskbasedBuf(&pInfo->pBuf, defaultPgsz, defaultBufsz, pTaskInfo->id.str, tsTempDir);
H
Haojun Liao 已提交
876
  if (code != TSDB_CODE_SUCCESS) {
dengyihao's avatar
dengyihao 已提交
877 878
    terrno = code;
    pTaskInfo->code = code;
H
Haojun Liao 已提交
879 880 881
    goto _error;
  }

H
Haojun Liao 已提交
882 883
  pInfo->rowCapacity = blockDataGetCapacityInRow(pInfo->binfo.pRes, getBufPageSize(pInfo->pBuf));
  pInfo->columnOffset = setupColumnOffset(pInfo->binfo.pRes, pInfo->rowCapacity);
884
  code = initGroupOptrInfo(&pInfo->pGroupColVals, &pInfo->groupKeyLen, &pInfo->keyBuf, pInfo->pGroupCols);
H
Haojun Liao 已提交
885
  if (code != TSDB_CODE_SUCCESS) {
dengyihao's avatar
dengyihao 已提交
886 887
    terrno = code;
    pTaskInfo->code = code;
H
Haojun Liao 已提交
888 889
    goto _error;
  }
H
Haojun Liao 已提交
890

L
Liu Jicong 已提交
891 892
  setOperatorInfo(pOperator, "PartitionOperator", QUERY_NODE_PHYSICAL_PLAN_PARTITION, false, OP_NOT_OPENED, pInfo,
                  pTaskInfo);
893 894
  pOperator->exprSupp.numOfExprs = numOfCols;
  pOperator->exprSupp.pExprInfo = pExprInfo;
895

dengyihao's avatar
dengyihao 已提交
896 897
  pOperator->fpSet =
      createOperatorFpSet(optrDummyOpenFn, hashPartition, NULL, destroyPartitionOperatorInfo, optrDefaultBufFn, NULL);
H
Haojun Liao 已提交
898

H
Haojun Liao 已提交
899
  code = appendDownstream(pOperator, &downstream, 1);
dengyihao's avatar
dengyihao 已提交
900 901 902 903 904 905
  if (code != TSDB_CODE_SUCCESS) {
    terrno = code;
    pTaskInfo->code = code;
    goto _error;
  }

906 907
  return pOperator;

908
_error:
H
Haojun Liao 已提交
909 910 911
  if (pInfo != NULL) {
    destroyPartitionOperatorInfo(pInfo);
  }
H
Haojun Liao 已提交
912
  taosMemoryFreeClear(pOperator);
913
  return NULL;
914 915
}

916 917 918
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;
919
  SResultRowInfo* pResultRowInfo = &binfo->resultRowInfo;
920
  SqlFunctionCtx* pCtx = pOperator->exprSupp.pCtx;
921 922

  SResultRow* pResultRow =
D
dapan1121 已提交
923
      doSetResultOutBufByKey(pBuf, pResultRowInfo, (char*)pData, bytes, true, groupId, pTaskInfo, false, pAggSup, false);
924 925
  assert(pResultRow != NULL);

926
  setResultRowInitCtx(pResultRow, pCtx, numOfCols, pOperator->exprSupp.rowEntryInfoOffset);
927
  return TSDB_CODE_SUCCESS;
D
dapan1121 已提交
928
}
929 930 931

uint64_t calGroupIdByData(SPartitionBySupporter* pParSup, SExprSupp* pExprSup, SSDataBlock* pBlock, int32_t rowId) {
  if (pExprSup->pExprInfo != NULL) {
932 933
    int32_t code =
        projectApplyFunctions(pExprSup->pExprInfo, pBlock, pBlock, pExprSup->pCtx, pExprSup->numOfExprs, NULL);
934 935 936 937 938
    if (code != TSDB_CODE_SUCCESS) {
      qError("calaculate group id error, code:%d", code);
    }
  }
  recordNewGroupKeys(pParSup->pGroupCols, pParSup->pGroupColVals, pBlock, rowId);
939
  int32_t  len = buildGroupKeys(pParSup->keyBuf, pParSup->pGroupColVals);
940 941 942 943
  uint64_t groupId = calcGroupId(pParSup->keyBuf, len);
  return groupId;
}

944
static bool hasRemainPartion(SStreamPartitionOperatorInfo* pInfo) { return pInfo->parIte != NULL; }
5
54liuyao 已提交
945
static bool hasRemainTbName(SStreamPartitionOperatorInfo* pInfo) { return pInfo->pTbNameIte != NULL; }
946 947 948

static SSDataBlock* buildStreamPartitionResult(SOperatorInfo* pOperator) {
  SStreamPartitionOperatorInfo* pInfo = pOperator->info;
949
  SSDataBlock*                  pDest = pInfo->binfo.pRes;
950
  ASSERT(hasRemainPartion(pInfo));
951 952
  SPartitionDataInfo* pParInfo = (SPartitionDataInfo*)pInfo->parIte;
  blockDataCleanup(pDest);
953
  int32_t      rows = taosArrayGetSize(pParInfo->rowIds);
954 955 956 957
  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++) {
958
      int32_t          slotId = pOperator->exprSupp.pExprInfo[j].base.pParam[0].pCol->slotId;
959 960
      SColumnInfoData* pSrcCol = taosArrayGet(pSrc->pDataBlock, slotId);
      SColumnInfoData* pDestCol = taosArrayGet(pDest->pDataBlock, j);
961 962
      bool             isNull = colDataIsNull(pSrcCol, pSrc->info.rows, rowIndex, NULL);
      char*            pSrcData = colDataGetData(pSrcCol, rowIndex);
963
      colDataSetVal(pDestCol, pDest->info.rows, pSrcData, isNull);
964 965
    }
    pDest->info.rows++;
5
54liuyao 已提交
966 967 968 969 970 971 972
  }
  pDest->info.parTbName[0] = 0;
  if (pInfo->tbnameCalSup.numOfExprs > 0) {
    void* tbname = NULL;
    if (streamStateGetParName(pOperator->pTaskInfo->streamInfo.pState, pParInfo->groupId, &tbname) == 0) {
      memcpy(pDest->info.parTbName, tbname, TSDB_TABLE_NAME_LEN);
      tdbFree(tbname);
973
    }
974
  }
975 976
  taosArrayDestroy(pParInfo->rowIds);
  pParInfo->rowIds = NULL;
H
Haojun Liao 已提交
977 978
  pDest->info.dataLoad = 1;

979
  blockDataUpdateTsWindow(pDest, pInfo->tsColIndex);
H
Haojun Liao 已提交
980
  pDest->info.id.groupId = pParInfo->groupId;
981 982
  pOperator->resultInfo.totalRows += pDest->info.rows;
  pInfo->parIte = taosHashIterate(pInfo->pPartitions, pInfo->parIte);
983
  ASSERT(pDest->info.rows > 0);
984 985 986 987
  printDataBlock(pDest, "stream partitionby");
  return pDest;
}

5
54liuyao 已提交
988
void appendCreateTableRow(SStreamState* pState, SExprSupp* pTableSup, SExprSupp* pTagSup, uint64_t groupId,
5
54liuyao 已提交
989 990
                          SSDataBlock* pSrcBlock, int32_t rowId, SSDataBlock* pDestBlock) {
  void* pValue = NULL;
5
54liuyao 已提交
991
  if (streamStateGetParName(pState, groupId, &pValue) != 0) {
5
54liuyao 已提交
992
    SSDataBlock* pTmpBlock = blockCopyOneRow(pSrcBlock, rowId);
5
54liuyao 已提交
993 994
    memset(pTmpBlock->info.parTbName, 0, TSDB_TABLE_NAME_LEN);
    pTmpBlock->info.id.groupId = groupId;
5
54liuyao 已提交
995
    char* tbName = pSrcBlock->info.parTbName;
5
54liuyao 已提交
996 997 998 999
    if (pTableSup->numOfExprs > 0) {
      projectApplyFunctions(pTableSup->pExprInfo, pDestBlock, pTmpBlock, pTableSup->pCtx, pTableSup->numOfExprs, NULL);
      SColumnInfoData* pTbCol = taosArrayGet(pDestBlock->pDataBlock, UD_TABLE_NAME_COLUMN_INDEX);
      memset(tbName, 0, TSDB_TABLE_NAME_LEN);
5
54liuyao 已提交
1000 1001
      int32_t len = 0;
      if (colDataIsNull_s(pTbCol, pDestBlock->info.rows - 1)) {
5
54liuyao 已提交
1002 1003
        len = 1;
        tbName[0] = 0;
5
54liuyao 已提交
1004 1005 1006 1007
      } else {
        void* pData = colDataGetData(pTbCol, pDestBlock->info.rows - 1);
        len = TMIN(varDataLen(pData), TSDB_TABLE_NAME_LEN - 1);
        memcpy(tbName, varDataVal(pData), len);
5
54liuyao 已提交
1008
        streamStatePutParName(pState, groupId, tbName);
5
54liuyao 已提交
1009
      }
5
54liuyao 已提交
1010
      memcpy(pTmpBlock->info.parTbName, tbName, len);
5
54liuyao 已提交
1011 1012 1013
      pDestBlock->info.rows--;
    } else {
      void* pTbNameCol = taosArrayGet(pDestBlock->pDataBlock, UD_TABLE_NAME_COLUMN_INDEX);
X
Xiaoyu Wang 已提交
1014
      colDataSetNULL(pTbNameCol, pDestBlock->info.rows);
5
54liuyao 已提交
1015
      tbName[0] = 0;
5
54liuyao 已提交
1016 1017 1018 1019 1020
    }

    if (pTagSup->numOfExprs > 0) {
      projectApplyFunctions(pTagSup->pExprInfo, pDestBlock, pTmpBlock, pTagSup->pCtx, pTagSup->numOfExprs, NULL);
      pDestBlock->info.rows--;
5
54liuyao 已提交
1021 1022
    } else {
      memcpy(pDestBlock->info.parTbName, pTmpBlock->info.parTbName, TSDB_TABLE_NAME_LEN);
5
54liuyao 已提交
1023 1024 1025 1026 1027 1028
    }

    void* pGpIdCol = taosArrayGet(pDestBlock->pDataBlock, UD_GROUPID_COLUMN_INDEX);
    colDataAppend(pGpIdCol, pDestBlock->info.rows, (const char*)&groupId, false);
    pDestBlock->info.rows++;
    blockDataDestroy(pTmpBlock);
5
54liuyao 已提交
1029 1030
  } else {
    memcpy(pSrcBlock->info.parTbName, pValue, TSDB_TABLE_NAME_LEN);
5
54liuyao 已提交
1031 1032 1033 1034 1035 1036
  }
  streamStateReleaseBuf(pState, NULL, pValue);
}

static SSDataBlock* buildStreamCreateTableResult(SOperatorInfo* pOperator) {
  SStreamPartitionOperatorInfo* pInfo = pOperator->info;
1037 1038
  if ((pInfo->tbnameCalSup.numOfExprs == 0 && pInfo->tagCalSup.numOfExprs == 0) ||
      taosHashGetSize(pInfo->pPartitions) == 0) {
5
54liuyao 已提交
1039 1040 1041 1042 1043 1044
    return NULL;
  }
  blockDataCleanup(pInfo->pCreateTbRes);
  blockDataEnsureCapacity(pInfo->pCreateTbRes, taosHashGetSize(pInfo->pPartitions));
  SSDataBlock* pSrc = pInfo->pInputDataBlock;

5
54liuyao 已提交
1045
  if (pInfo->pTbNameIte != NULL) {
5
54liuyao 已提交
1046
    SPartitionDataInfo* pParInfo = (SPartitionDataInfo*)pInfo->pTbNameIte;
1047
    int32_t             rowId = *(int32_t*)taosArrayGet(pParInfo->rowIds, 0);
5
54liuyao 已提交
1048
    appendCreateTableRow(pOperator->pTaskInfo->streamInfo.pState, &pInfo->tbnameCalSup, &pInfo->tagCalSup,
1049
                         pParInfo->groupId, pSrc, rowId, pInfo->pCreateTbRes);
5
54liuyao 已提交
1050 1051 1052 1053 1054
    pInfo->pTbNameIte = taosHashIterate(pInfo->pPartitions, pInfo->pTbNameIte);
  }
  return pInfo->pCreateTbRes->info.rows > 0 ? pInfo->pCreateTbRes : NULL;
}

1055 1056 1057 1058
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);
1059 1060 1061
    int32_t             keyLen = buildGroupKeys(pInfo->partitionSup.keyBuf, pInfo->partitionSup.pGroupColVals);
    SPartitionDataInfo* pParData =
        (SPartitionDataInfo*)taosHashGet(pInfo->pPartitions, pInfo->partitionSup.keyBuf, keyLen);
1062 1063 1064 1065 1066 1067 1068
    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);
1069
      taosHashPut(pInfo->pPartitions, pInfo->partitionSup.keyBuf, keyLen, &newParData, sizeof(SPartitionDataInfo));
1070 1071 1072 1073 1074 1075 1076 1077 1078
    }
  }
}

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

1079
  SExecTaskInfo*                pTaskInfo = pOperator->pTaskInfo;
1080
  SStreamPartitionOperatorInfo* pInfo = pOperator->info;
5
54liuyao 已提交
1081 1082 1083 1084 1085 1086 1087 1088 1089
  SSDataBlock*                  pCtRes = NULL;

  if (hasRemainTbName(pInfo)) {
    pCtRes = buildStreamCreateTableResult(pOperator);
    if (pCtRes != NULL) {
      return pCtRes;
    }
  }

1090 1091 1092 1093
  if (hasRemainPartion(pInfo)) {
    return buildStreamPartitionResult(pOperator);
  }

1094
  int64_t        st = taosGetTimestampUs();
1095 1096 1097 1098 1099
  SOperatorInfo* downstream = pOperator->pDownstream[0];
  {
    pInfo->pInputDataBlock = NULL;
    SSDataBlock* pBlock = downstream->fpSet.getNextFn(downstream);
    if (pBlock == NULL) {
H
Haojun Liao 已提交
1100
      setOperatorCompleted(pOperator);
1101 1102 1103 1104 1105 1106 1107 1108 1109
      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;
1110 1111 1112
      case STREAM_DELETE_DATA: {
        copyDataBlock(pInfo->pDelRes, pBlock);
        pInfo->pDelRes->info.type = STREAM_DELETE_RESULT;
5
54liuyao 已提交
1113
        printDataBlock(pInfo->pDelRes, "stream partitionby delete");
1114
        return pInfo->pDelRes;
1115
      } break;
1116
      default:
5
54liuyao 已提交
1117
        ASSERTS(pBlock->info.type == STREAM_CREATE_CHILD_TABLE, "invalid SSDataBlock type");
1118 1119 1120 1121 1122
        return pBlock;
    }

    // there is an scalar expression that needs to be calculated right before apply the group aggregation.
    if (pInfo->scalarSup.pExprInfo != NULL) {
5
54liuyao 已提交
1123
      projectApplyFunctions(pInfo->scalarSup.pExprInfo, pBlock, pBlock, pInfo->scalarSup.pCtx,
1124
                                              pInfo->scalarSup.numOfExprs, NULL);
1125 1126 1127 1128 1129
    }
    taosHashClear(pInfo->pPartitions);
    doStreamHashPartitionImpl(pInfo, pBlock);
  }
  pOperator->cost.openCost = (taosGetTimestampUs() - st) / 1000.0;
1130

1131
  pInfo->parIte = taosHashIterate(pInfo->pPartitions, NULL);
5
54liuyao 已提交
1132 1133 1134 1135 1136
  pInfo->pTbNameIte = taosHashIterate(pInfo->pPartitions, NULL);
  pCtRes = buildStreamCreateTableResult(pOperator);
  if (pCtRes != NULL) {
    return pCtRes;
  }
1137 1138 1139 1140 1141 1142 1143 1144
  return buildStreamPartitionResult(pOperator);
}

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

1145
  for (int i = 0; i < taosArrayGetSize(pInfo->partitionSup.pGroupColVals); i++) {
1146 1147 1148 1149 1150 1151 1152
    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 已提交
1153 1154
  cleanupExprSupp(&pInfo->tbnameCalSup);
  cleanupExprSupp(&pInfo->tagCalSup);
1155
  blockDataDestroy(pInfo->pDelRes);
1156
  taosHashCleanup(pInfo->pPartitions);
5
54liuyao 已提交
1157
  blockDataDestroy(pInfo->pCreateTbRes);
1158 1159 1160 1161 1162 1163 1164 1165 1166 1167
  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 已提交
1168
  if (!pScanInfo->igCheckUpdate && !pScanInfo->pUpdateInfo) {
5
54liuyao 已提交
1169 1170
    pScanInfo->pUpdateInfo = updateInfoInit(60000, TSDB_TIME_PRECISION_MILLI, 0);
  }
1171 1172
}

5
54liuyao 已提交
1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212
SSDataBlock* buildCreateTableBlock(SExprSupp* tbName, SExprSupp* tag) {
  SSDataBlock* pBlock = taosMemoryCalloc(1, sizeof(SSDataBlock));
  pBlock->info.hasVarCol = false;
  pBlock->info.id.groupId = 0;
  pBlock->info.rows = 0;
  pBlock->info.type = STREAM_CREATE_CHILD_TABLE;
  pBlock->info.watermark = INT64_MIN;

  pBlock->pDataBlock = taosArrayInit(4, sizeof(SColumnInfoData));
  SColumnInfoData infoData = {0};
  infoData.info.type = TSDB_DATA_TYPE_VARCHAR;
  if (tbName->numOfExprs > 0) {
    infoData.info.bytes = tbName->pExprInfo->base.resSchema.bytes;
  } else {
    infoData.info.bytes = 1;
  }
  pBlock->info.rowSize += infoData.info.bytes;
  // sub table name
  taosArrayPush(pBlock->pDataBlock, &infoData);

  SColumnInfoData gpIdData = {0};
  gpIdData.info.type = TSDB_DATA_TYPE_UBIGINT;
  gpIdData.info.bytes = 8;
  pBlock->info.rowSize += gpIdData.info.bytes;
  // group id
  taosArrayPush(pBlock->pDataBlock, &gpIdData);

  for (int32_t i = 0; i < tag->numOfExprs; i++) {
    SColumnInfoData tagCol = {0};
    tagCol.info.type = tag->pExprInfo[i].base.resSchema.type;
    tagCol.info.bytes = tag->pExprInfo[i].base.resSchema.bytes;
    tagCol.info.precision = tag->pExprInfo[i].base.resSchema.precision;
    // tag info
    taosArrayPush(pBlock->pDataBlock, &tagCol);
    pBlock->info.rowSize += tagCol.info.bytes;
  }

  return pBlock;
}

1213 1214
SOperatorInfo* createStreamPartitionOperatorInfo(SOperatorInfo* downstream, SStreamPartitionPhysiNode* pPartNode,
                                                 SExecTaskInfo* pTaskInfo) {
1215
  int32_t                       code = TSDB_CODE_SUCCESS;
1216 1217 1218
  SStreamPartitionOperatorInfo* pInfo = taosMemoryCalloc(1, sizeof(SStreamPartitionOperatorInfo));
  SOperatorInfo*                pOperator = taosMemoryCalloc(1, sizeof(SOperatorInfo));
  if (pInfo == NULL || pOperator == NULL) {
1219
    code = TSDB_CODE_OUT_OF_MEMORY;
1220 1221
    goto _error;
  }
1222

1223
  pInfo->partitionSup.pGroupCols = extractPartitionColInfo(pPartNode->part.pPartitionKeys);
1224

1225
  if (pPartNode->part.pExprs != NULL) {
1226
    int32_t    num = 0;
1227
    SExprInfo* pCalExprInfo = createExprInfo(pPartNode->part.pExprs, NULL, &num);
1228 1229 1230 1231 1232 1233
    code = initExprSupp(&pInfo->scalarSup, pCalExprInfo, num);
    if (code != TSDB_CODE_SUCCESS) {
      goto _error;
    }
  }

5
54liuyao 已提交
1234
  pInfo->tbnameCalSup.numOfExprs = 0;
1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248
  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;
    }
  }

5
54liuyao 已提交
1249
  pInfo->tagCalSup.numOfExprs = 0;
L
Liu Jicong 已提交
1250 1251
  if (pPartNode->pTags != NULL) {
    int32_t    numOfTags;
5
54liuyao 已提交
1252
    SExprInfo* pTagExpr = createExpr(pPartNode->pTags, &numOfTags);
L
Liu Jicong 已提交
1253
    if (pTagExpr == NULL) {
1254
      code = TSDB_CODE_OUT_OF_MEMORY;
L
Liu Jicong 已提交
1255 1256 1257
      goto _error;
    }
    if (initExprSupp(&pInfo->tagCalSup, pTagExpr, numOfTags) != 0) {
1258
      code = TSDB_CODE_OUT_OF_MEMORY;
L
Liu Jicong 已提交
1259 1260 1261 1262
      goto _error;
    }
  }

5
54liuyao 已提交
1263 1264 1265 1266 1267 1268
  if (pInfo->tbnameCalSup.numOfExprs != 0 || pInfo->tagCalSup.numOfExprs != 0) {
    pInfo->pCreateTbRes = buildCreateTableBlock(&pInfo->tbnameCalSup, &pInfo->tagCalSup);
  } else {
    pInfo->pCreateTbRes = NULL;
  }

1269
  int32_t keyLen = 0;
1270 1271
  code = initGroupOptrInfo(&pInfo->partitionSup.pGroupColVals, &keyLen, &pInfo->partitionSup.keyBuf,
                           pInfo->partitionSup.pGroupCols);
1272 1273 1274 1275 1276
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }
  pInfo->partitionSup.needCalc = true;

H
Haojun Liao 已提交
1277
  pInfo->binfo.pRes = createDataBlockFromDescNode(pPartNode->part.node.pOutputDataBlockDesc);
1278
  if (pInfo->binfo.pRes == NULL) {
1279
    code = TSDB_CODE_OUT_OF_MEMORY;
1280 1281
    goto _error;
  }
1282 1283 1284

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

1285
  pInfo->parIte = NULL;
5
54liuyao 已提交
1286
  pInfo->pTbNameIte = NULL;
1287
  pInfo->pInputDataBlock = NULL;
1288

1289
  _hash_fn_t hashFn = taosGetDefaultHashFunction(TSDB_DATA_TYPE_BINARY);
1290 1291 1292
  pInfo->pPartitions = taosHashInit(1024, hashFn, false, HASH_NO_LOCK);
  pInfo->tsColIndex = 0;
  pInfo->pDelRes = createSpecialDataBlock(STREAM_DELETE_RESULT);
1293

1294
  int32_t    numOfCols = 0;
1295
  SExprInfo* pExprInfo = createExprInfo(pPartNode->part.pTargets, NULL, &numOfCols);
1296

L
Liu Jicong 已提交
1297 1298
  setOperatorInfo(pOperator, "StreamPartitionOperator", QUERY_NODE_PHYSICAL_PLAN_STREAM_PARTITION, false, OP_NOT_OPENED,
                  pInfo, pTaskInfo);
1299 1300
  pOperator->exprSupp.numOfExprs = numOfCols;
  pOperator->exprSupp.pExprInfo = pExprInfo;
dengyihao's avatar
dengyihao 已提交
1301 1302
  pOperator->fpSet = createOperatorFpSet(optrDummyOpenFn, doStreamHashPartition, NULL,
                                         destroyStreamPartitionOperatorInfo, optrDefaultBufFn, NULL);
1303 1304 1305 1306 1307

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

1308
_error:
1309
  pTaskInfo->code = code;
1310
  destroyStreamPartitionOperatorInfo(pInfo);
1311 1312 1313
  taosMemoryFreeClear(pOperator);
  return NULL;
}
H
Haojun Liao 已提交
1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346

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;
}