#5044

Committed 06 May 2026 02:35AM UTC coverage: 73.169% (+0.06%) from 73.107%

Build # #5044

Build Type

push

travis-ci

Committed by

web-flow

Commit Message

feat: [6659794715] cpu limit (#35153)

Coverage Stats

244 of 275 new or added lines in 23 files covered. (88.73%)

526 existing lines in 141 files now uncovered.

277745 of 379596 relevant lines covered (73.17%)

133740972.66 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

83.04

/source/libs/executor/src/sortoperator.c

/*
 * 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/>.
 */

#include "executorInt.h"
#include "filter.h"
#include "operator.h"
#include "querytask.h"
#include "tdatablock.h"

typedef struct SSortOpGroupIdCalc {
  STupleHandle* pSavedTuple;
  SArray*       pSortColsArr;
  char*         keyBuf;
  int32_t       lastKeysLen; // default to be 0
  uint64_t      lastGroupId;
  bool          excludePKCol;
} SSortOpGroupIdCalc;

typedef struct SSortOperatorInfo {
  SOptrBasicInfo      binfo;
  uint32_t            sortBufSize;  // max buffer size for in-memory sort
  SArray*             pSortInfo;
  SSortHandle*        pSortHandle;
  SColMatchInfo       matchInfo;
  int32_t             bufPageSize;
  int64_t             startTs;      // sort start time
  uint64_t            sortElapsed;  // sort elapsed time, time to flush to disk not included.
  SLimitInfo          limitInfo;
  uint64_t            maxTupleLength;
  int64_t             maxRows;
  SSortOpGroupIdCalc* pGroupIdCalc;
} SSortOperatorInfo;

static int32_t doSort(SOperatorInfo* pOperator, SSDataBlock** pResBlock);
static int32_t doOpenSortOperator(SOperatorInfo* pOperator);
static int32_t getExplainExecInfo(SOperatorInfo* pOptr, void** pOptrExplain, uint32_t* len);
static int32_t doGroupSort(SOperatorInfo* pOperator, SSDataBlock** pResBlock);

static void destroySortOperatorInfo(void* param);
static void calcSortOperMaxTupleLength(SSortOperatorInfo* pSortOperInfo, SNodeList* pSortKeys);

// Check whether the given slotId is the output of a Sort scalar pre-calculation expression.
// The planner's rewritePrecalcExprs pushes expression results (e.g. ts+1000) into extra slots
// in the child's output descriptor. This helper identifies such slots.
static bool sortIsExprResultSlot(const SOperatorInfo* pOperator, int32_t slotId) {
  if (pOperator == NULL || pOperator->exprSupp.pExprInfo == NULL || slotId < 0) {
    return false;
  }
  for (int32_t idx = 0; idx < pOperator->exprSupp.numOfExprs; ++idx) {
    if (pOperator->exprSupp.pExprInfo[idx].base.resSchema.slotId == slotId) {
      return true;
    }
  }
  return false;
}

// Find the slot index of the original (non-expression) primary timestamp column in the
// Sort's internal data block.  The internal block may contain both the original ts column
// and an expression-derived ts column (e.g. ts+1000).  We return the first TIMESTAMP
// column whose slot is NOT an expression result.
static int32_t sortFindOrigTsSlot(const SOperatorInfo* pOperator, const SSDataBlock* pBlock) {
  if (pBlock == NULL || pBlock->pDataBlock == NULL) {
    return -1;
  }
  int32_t colCount = (int32_t)taosArrayGetSize(pBlock->pDataBlock);
  for (int32_t idx = 0; idx < colCount; ++idx) {
    SColumnInfoData* pCol = taosArrayGet(pBlock->pDataBlock, idx);
    if (pCol != NULL && pCol->info.type == TSDB_DATA_TYPE_TIMESTAMP &&
        !sortIsExprResultSlot(pOperator, idx)) {
      return idx;
    }
  }
  return -1;
}

static void destroySortOpGroupIdCalc(SSortOpGroupIdCalc* pCalc);

static int32_t resetSortOperState(SOperatorInfo* pOper) {
  SSortOperatorInfo* pInfo = pOper->info;
  SExecTaskInfo*           pTaskInfo = pOper->pTaskInfo;
  pOper->status = OP_NOT_OPENED;

  resetBasicOperatorState(&pInfo->binfo);
  destroySqlFunctionCtx(pOper->exprSupp.pCtx, pOper->exprSupp.pExprInfo, pOper->exprSupp.numOfExprs);
  taosMemoryFreeClear(pOper->exprSupp.rowEntryInfoOffset);
  pOper->exprSupp.pCtx =
      createSqlFunctionCtx(pOper->exprSupp.pExprInfo, pOper->exprSupp.numOfExprs, &pOper->exprSupp.rowEntryInfoOffset, &pTaskInfo->storageAPI.functionStore);

  tsortDestroySortHandle(pInfo->pSortHandle);
  pInfo->pSortHandle = NULL;

  if (pInfo->pGroupIdCalc) {
    pInfo->pGroupIdCalc->lastGroupId = 0;
    pInfo->pGroupIdCalc->lastKeysLen = 0;
  }

  return 0;
}

// todo add limit/offset impl
int32_t createSortOperatorInfo(SOperatorInfo* downstream, SSortPhysiNode* pSortNode, SExecTaskInfo* pTaskInfo, SOperatorInfo** pOptrInfo) {
  QRY_PARAM_CHECK(pOptrInfo);

  int32_t code = 0;
  int32_t lino = 0;

  SSortOperatorInfo* pInfo = taosMemoryCalloc(1, sizeof(SSortOperatorInfo));
  SOperatorInfo*     pOperator = taosMemoryCalloc(1, sizeof(SOperatorInfo));
  if (pInfo == NULL || pOperator == NULL) {
    code = terrno;
    goto _error;
  }
  initOperatorCostInfo(pOperator);

  pOperator->pTaskInfo = pTaskInfo;
  SDataBlockDescNode* pDescNode = pSortNode->node.pOutputDataBlockDesc;

  int32_t numOfCols = 0;
  code = createExprInfo(pSortNode->pExprs, NULL, &pOperator->exprSupp.pExprInfo, &numOfCols);
  QUERY_CHECK_CODE(code, lino, _error);

  pOperator->exprSupp.numOfExprs = numOfCols;
  int32_t numOfOutputCols = 0;
  code =
      extractColMatchInfo(pSortNode->pTargets, pDescNode, &numOfOutputCols, COL_MATCH_FROM_SLOT_ID, &pInfo->matchInfo);
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }
  
  calcSortOperMaxTupleLength(pInfo, pSortNode->pSortKeys);
  pInfo->maxRows = -1;
  if (pSortNode->node.pLimit && ((SLimitNode*)pSortNode->node.pLimit)->limit) {
    SLimitNode* pLimit = (SLimitNode*)pSortNode->node.pLimit;
    if (pLimit->limit->datum.i > 0) {
      pInfo->maxRows = pLimit->limit->datum.i + (pLimit->offset ? pLimit->offset->datum.i : 0);
    }
  }

  pOperator->exprSupp.pCtx =
      createSqlFunctionCtx(pOperator->exprSupp.pExprInfo, numOfCols, &pOperator->exprSupp.rowEntryInfoOffset, &pTaskInfo->storageAPI.functionStore);
  QUERY_CHECK_NULL(pOperator->exprSupp.pCtx, code, lino, _error, terrno);
  initResultSizeInfo(&pOperator->resultInfo, 1024);
  code = filterInitFromNode((SNode*)pSortNode->node.pConditions, &pOperator->exprSupp.pFilterInfo, 0,
                            pTaskInfo->pStreamRuntimeInfo);
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

  pInfo->binfo.pRes = createDataBlockFromDescNode(pDescNode);
  QUERY_CHECK_NULL(pInfo->binfo.pRes, code, lino, _error, terrno);

  pInfo->pSortInfo = createSortInfo(pSortNode->pSortKeys);
  TSDB_CHECK_NULL(pInfo->pSortInfo, code, lino, _error, terrno);

  if (pSortNode->calcGroupId) {
    int32_t keyLen;
    SSortOpGroupIdCalc* pGroupIdCalc = pInfo->pGroupIdCalc = taosMemoryCalloc(1, sizeof(SSortOpGroupIdCalc));
    if (!pGroupIdCalc) {
      code = terrno;
      goto _error;
    }
    SNodeList* pSortColsNodeArr = makeColsNodeArrFromSortKeys(pSortNode->pSortKeys);
    if (!pSortColsNodeArr) code = terrno;
    if (TSDB_CODE_SUCCESS == code) {
      pGroupIdCalc->pSortColsArr = makeColumnArrayFromList(pSortColsNodeArr);
      if (!pGroupIdCalc->pSortColsArr) code = terrno;
      nodesClearList(pSortColsNodeArr);
    }
    if (TSDB_CODE_SUCCESS == code) {
      // PK ts col should always at last, see partColOptCreateSort
      if (pSortNode->excludePkCol) taosArrayPop(pGroupIdCalc->pSortColsArr);
      code = extractKeysLen(pGroupIdCalc->pSortColsArr, &keyLen);
      QUERY_CHECK_CODE(code, lino, _error);
    }
    if (TSDB_CODE_SUCCESS == code) {
      pGroupIdCalc->lastKeysLen = 0;
      pGroupIdCalc->keyBuf = taosMemoryCalloc(1, keyLen);
      if (!pGroupIdCalc->keyBuf) {
        code = terrno;
      }
    }
  }
  if (code != TSDB_CODE_SUCCESS) goto _error;

  pInfo->binfo.inputTsOrder = pSortNode->node.inputTsOrder;
  pInfo->binfo.outputTsOrder = pSortNode->node.outputTsOrder;
  initLimitInfo(pSortNode->node.pLimit, pSortNode->node.pSlimit, &pInfo->limitInfo);

  setOperatorInfo(pOperator, "SortOperator", QUERY_NODE_PHYSICAL_PLAN_SORT, true, OP_NOT_OPENED, pInfo, pTaskInfo);


  // lazy evaluation for the following parameter since the input datablock is not known till now.
  //  pInfo->bufPageSize  = rowSize < 1024 ? 1024 * 2 : rowSize * 2;
  //  there are headers, so pageSize = rowSize + header pInfo->sortBufSize  = pInfo->bufPageSize * 16;
  // TODO dynamic set the available sort buffer

  pOperator->fpSet =
      createOperatorFpSet(doOpenSortOperator, doSort, NULL, destroySortOperatorInfo, optrDefaultBufFn, getExplainExecInfo, optrDefaultGetNextExtFn, NULL);

  setOperatorResetStateFn(pOperator, resetSortOperState);
  code = appendDownstream(pOperator, &downstream, 1);
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

  *pOptrInfo = pOperator;
  return TSDB_CODE_SUCCESS;

_error:
  if (pInfo != NULL) {
    destroySortOperatorInfo(pInfo);
  }
  destroyOperatorAndDownstreams(pOperator, &downstream, 1);
  pTaskInfo->code = code;
  return code;
}

int32_t appendOneRowToDataBlock(SSDataBlock* pBlock, STupleHandle* pTupleHandle) {
  int32_t code = 0;
  for (int32_t i = 0; i < taosArrayGetSize(pBlock->pDataBlock); ++i) {
    SColumnInfoData* pColInfo = taosArrayGet(pBlock->pDataBlock, i);
    if (pColInfo == NULL) {
      return terrno;
    }

    bool isNull = tsortIsNullVal(pTupleHandle, i);
    if (isNull) {
      colDataSetNULL(pColInfo, pBlock->info.rows);
    } else {
      char* pData = NULL;
      tsortGetValue(pTupleHandle, i, (void**) &pData);

      if (pData != NULL) {
        code = colDataSetVal(pColInfo, pBlock->info.rows, pData, false);
        if (code) {
          return code;
        }
      }
    }
  }

  pBlock->info.dataLoad = 1;

  SDataBlockInfo info = {0};
  tsortGetBlockInfo(pTupleHandle, &info);

  pBlock->info.scanFlag = info.scanFlag;
  pBlock->info.rows += 1;
  return code;
}

/**
 * @brief get next tuple with group id attached, here assume that all tuples are sorted by group keys
 * @param [in, out] pBlock the output block, the group id will be saved in it
 * @retval NULL if next group tuple arrived and this new group tuple will be saved in pInfo.pSavedTuple
 */
static int32_t nextTupleWithGroupId(SSortHandle* pHandle, SSortOperatorInfo* pInfo, SSDataBlock* pBlock,
                                    STupleHandle** pTupleHandle) {
  QRY_PARAM_CHECK(pTupleHandle);

  int32_t       code = 0;
  STupleHandle* retTuple = pInfo->pGroupIdCalc->pSavedTuple;
  if (!retTuple) {
    code = tsortNextTuple(pHandle, &retTuple);
    if (code) {
      qError("failed to get next tuple, code:%s", tstrerror(code));
      return code;
    }
  }

  if (retTuple) {
    int32_t newGroup;
    if (pInfo->pGroupIdCalc->pSavedTuple) {
      newGroup = true;
      pInfo->pGroupIdCalc->pSavedTuple = NULL;
    } else {
      newGroup = tsortCompAndBuildKeys(pInfo->pGroupIdCalc->pSortColsArr, pInfo->pGroupIdCalc->keyBuf,
                                       &pInfo->pGroupIdCalc->lastKeysLen, retTuple);
    }

    bool emptyBlock = (pBlock->info.rows == 0);
    if (newGroup) {
      if (!emptyBlock) {
        // new group arrived, and we have already copied some tuples for cur group, save the new group tuple, return
        // NULL. Note that the keyBuf and lastKeysLen has been updated to new value
        pInfo->pGroupIdCalc->pSavedTuple = retTuple;
        retTuple = NULL;
      } else {
        // new group with empty block
        pInfo->pGroupIdCalc->lastGroupId = pBlock->info.id.groupId =
            calcGroupId(pInfo->pGroupIdCalc->keyBuf, pInfo->pGroupIdCalc->lastKeysLen);
      }
    } else {
      if (emptyBlock) {
        // new block but not new group, assign last group id to it
        pBlock->info.id.groupId = pInfo->pGroupIdCalc->lastGroupId;
      } else {
        // not new group and not empty block and ret NOT NULL, just return the tuple
      }
    }
  }

  *pTupleHandle = retTuple;
  return code;
}

static int32_t getSortedBlockData(SSortHandle* pHandle, SSDataBlock* pDataBlock, int32_t capacity, SArray* pColMatchInfo,
                                  const SOperatorInfo* pOperator, SSortOperatorInfo* pInfo, SSDataBlock** pResBlock) {
  QRY_PARAM_CHECK(pResBlock);
  blockDataCleanup(pDataBlock);

  int32_t       lino = 0;
  int32_t       code = 0;
  STupleHandle* pTupleHandle = NULL;
  SSDataBlock*  p = NULL;

  code = tsortGetSortedDataBlock(pHandle, &p);
  if (p == NULL || (code != 0)) {
    return code;
  }

  code = blockDataEnsureCapacity(p, capacity);
  QUERY_CHECK_CODE(code, lino, _error);

  while (1) {
    if (pInfo->pGroupIdCalc) {
      code = nextTupleWithGroupId(pHandle, pInfo, p, &pTupleHandle);
    } else {
      code = tsortNextTuple(pHandle, &pTupleHandle);
    }

    TSDB_CHECK_CODE(code, lino, _error);
    if (pTupleHandle == NULL) {
      break;
    }

    code = appendOneRowToDataBlock(p, pTupleHandle);
    QUERY_CHECK_CODE(code, lino, _error);

    if (p->info.rows >= capacity) {
      break;
    }
  }

  QUERY_CHECK_CODE(code, lino, _error);

  if (p->info.rows > 0) {
    code = blockDataEnsureCapacity(pDataBlock, capacity);
    QUERY_CHECK_CODE(code, lino, _error);

    // todo extract function to handle this
    int32_t numOfCols = taosArrayGetSize(pColMatchInfo);
    for (int32_t i = 0; i < numOfCols; ++i) {
      SColMatchItem* pmInfo = taosArrayGet(pColMatchInfo, i);
      QUERY_CHECK_NULL(pmInfo, code, lino, _error, terrno);

      int32_t srcSlotId = pmInfo->srcSlotId;

      // Fix: when the planner's setListSlotId resolves Sort pTargets by name, it may
      // mistakenly bind the primary timestamp column to the scalar expression result slot
      // (e.g. slot for ts+1000) instead of the original ts slot, because pushdownDataBlockSlots
      // added the expression slot with the same column name.  This causes the downstream
      // Project operator to apply the expression again (ts+1000 becomes ts+2000).
      // Detect this case and fall back to the original timestamp slot.
      if (pmInfo->colId == PRIMARYKEY_TIMESTAMP_COL_ID &&
          pmInfo->dataType.type == TSDB_DATA_TYPE_TIMESTAMP &&
          sortIsExprResultSlot(pOperator, srcSlotId)) {
        int32_t origSlot = sortFindOrigTsSlot(pOperator, p);
        if (origSlot >= 0 && origSlot != srcSlotId) {
          srcSlotId = origSlot;
        }
      }

      SColumnInfoData* pSrc = taosArrayGet(p->pDataBlock, srcSlotId);
      QUERY_CHECK_NULL(pSrc, code, lino, _error, terrno);

      SColumnInfoData* pDst = taosArrayGet(pDataBlock->pDataBlock, pmInfo->dstSlotId);
      QUERY_CHECK_NULL(pDst, code, lino, _error, terrno);

      code = colDataAssign(pDst, pSrc, p->info.rows, &pDataBlock->info);
      QUERY_CHECK_CODE(code, lino, _error);
    }

    pDataBlock->info.dataLoad = 1;
    pDataBlock->info.rows = p->info.rows;
    pDataBlock->info.scanFlag = p->info.scanFlag;
    // propagate both C-group id and baseGId from upstream
    pDataBlock->info.id.groupId = p->info.id.groupId;
    pDataBlock->info.id.baseGId = p->info.id.baseGId;
  }

  blockDataDestroy(p);
  *pResBlock = (pDataBlock->info.rows > 0) ? pDataBlock : NULL;
  return code;

  _error:
  qError("%s failed at line %d since %s", __func__, lino, tstrerror(code));

  blockDataDestroy(p);
  return code;
}

int32_t loadNextDataBlock(void* param, SSDataBlock** ppBlock) {
  SOperatorInfo* pOperator = (SOperatorInfo*)param;
  int32_t        code = pOperator->fpSet.getNextFn(pOperator, ppBlock);
  if (code) {
    qError("failed to get next data block from upstream, %s code:%s", __func__, tstrerror(code));
  } else {
    code = blockDataCheck(*ppBlock);
    if (code) {
      qError("failed to check block data, %s code:%s", __func__, tstrerror(code));
    }
  }
  return code;
}

// todo refactor: merged with fetch fp
void applyScalarFunction(SSDataBlock* pBlock, void* param) {
  SOperatorInfo*     pOperator = param;
  SSortOperatorInfo* pSort = pOperator->info;
  if (pOperator->exprSupp.pExprInfo != NULL && pOperator->exprSupp.numOfExprs > 0) {
    int32_t code = projectApplyFunctions(pOperator->exprSupp.pExprInfo, pBlock, pBlock, pOperator->exprSupp.pCtx,
                                         pOperator->exprSupp.numOfExprs, NULL,
                                         GET_STM_RTINFO(pOperator->pTaskInfo), pOperator->pTaskInfo);
    if (code != TSDB_CODE_SUCCESS) {
      T_LONG_JMP(pOperator->pTaskInfo->env, code);
    }
  }
}

int32_t doOpenSortOperator(SOperatorInfo* pOperator) {
  SSortOperatorInfo* pInfo = pOperator->info;
  SExecTaskInfo*     pTaskInfo = pOperator->pTaskInfo;
  int32_t            code = TSDB_CODE_SUCCESS;
  int32_t            lino = 0;
  SSortSource* pSource =NULL;

  if (pOperator->exprSupp.pFilterInfo != NULL) {
    filterSetExecContext(pOperator->exprSupp.pFilterInfo, pTaskInfo, isTaskKilled);
  }

  if (OPTR_IS_OPENED(pOperator)) {
    return code;
  }

  //  pInfo->binfo.pRes is not equalled to the input datablock.
  pInfo->pSortHandle = NULL;
  code =
      tsortCreateSortHandle(pInfo->pSortInfo, SORT_SINGLESOURCE_SORT, -1, -1, NULL, pTaskInfo->id.str, pInfo->maxRows,
                            pInfo->maxTupleLength, tsPQSortMemThreshold * 1024 * 1024, &pInfo->pSortHandle);
  QUERY_CHECK_CODE(code, lino, _end);

  tsortSetFetchRawDataFp(pInfo->pSortHandle, loadNextDataBlock, applyScalarFunction, pOperator);

  pSource = taosMemoryCalloc(1, sizeof(SSortSource));
  QUERY_CHECK_NULL(pSource, code, lino, _end, terrno);

  pSource->param = pOperator->pDownstream[0];
  pSource->onlyRef = true;

  code = tsortAddSource(pInfo->pSortHandle, pSource);
  QUERY_CHECK_CODE(code, lino, _end);
  pSource = NULL;

  code = tsortOpen(pInfo->pSortHandle);
  QUERY_CHECK_CODE(code, lino, _end);
  pOperator->status = OP_RES_TO_RETURN;
  OPTR_SET_OPENED(pOperator);

_end:
  if (pSource) {
    taosMemoryFree(pSource);
  }
  if (code != TSDB_CODE_SUCCESS) {
    qError("%s failed at line %d since %s", __func__, lino, tstrerror(code));
    pTaskInfo->code = code;
    T_LONG_JMP(pTaskInfo->env, code);
  }
  return code;
}

int32_t doSort(SOperatorInfo* pOperator, SSDataBlock** pResBlock) {
  QRY_PARAM_CHECK(pResBlock);
  int32_t code = TSDB_CODE_SUCCESS;
  int32_t lino = 0;
  if (pOperator->status == OP_EXEC_DONE) {
    return code;
  }

  SExecTaskInfo*     pTaskInfo = pOperator->pTaskInfo;
  SSortOperatorInfo* pInfo = pOperator->info;

  code = pOperator->fpSet._openFn(pOperator);
  QUERY_CHECK_CODE(code, lino, _end);

  // multi-group case not handle here
  SSDataBlock* pBlock = NULL;
  while (1) {
    if (tsortIsClosed(pInfo->pSortHandle)) {
      code = TSDB_CODE_TSC_QUERY_CANCELLED;
      QUERY_CHECK_CODE(code, lino, _end);
    }

    recordOpExecBeforeDownstream(pOperator);
    code = getSortedBlockData(pInfo->pSortHandle, pInfo->binfo.pRes, pOperator->resultInfo.capacity,
                                pInfo->matchInfo.pList, pOperator, pInfo, &pBlock);
    recordOpExecAfterDownstream(pOperator, pBlock ? pBlock->info.rows : 0);
    QUERY_CHECK_CODE(code, lino, _end);
    if (pBlock == NULL) {
      setOperatorCompleted(pOperator);
      return code;
    }

    code = doFilter(pBlock, pOperator->exprSupp.pFilterInfo, &pInfo->matchInfo, NULL);
    QUERY_CHECK_CODE(code, lino, _end);

    if (blockDataGetNumOfRows(pBlock) == 0) {
      continue;
    }

    // there are bugs?
    bool limitReached = applyLimitOffset(&pInfo->limitInfo, pBlock, pTaskInfo);
    if (limitReached) {
      resetLimitInfoForNextGroup(&pInfo->limitInfo);
    }

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

  *pResBlock = blockDataGetNumOfRows(pBlock) > 0 ? pBlock : NULL;
_end:
  if (code != TSDB_CODE_SUCCESS) {
    qError("%s failed at line %d since %s", __func__, lino, tstrerror(code));
    pTaskInfo->code = code;
    T_LONG_JMP(pTaskInfo->env, code);
  }
  return code;
}

void destroySortOperatorInfo(void* param) {
  SSortOperatorInfo* pInfo = (SSortOperatorInfo*)param;
  blockDataDestroy(pInfo->binfo.pRes);
  pInfo->binfo.pRes = NULL;

  tsortDestroySortHandle(pInfo->pSortHandle);
  taosArrayDestroy(pInfo->pSortInfo);
  taosArrayDestroy(pInfo->matchInfo.pList);
  destroySortOpGroupIdCalc(pInfo->pGroupIdCalc);
  taosMemoryFreeClear(param);
}

int32_t getExplainExecInfo(SOperatorInfo* pOptr, void** pOptrExplain, uint32_t* len) {
  SSortExecInfo* pInfo = taosMemoryCalloc(1, sizeof(SSortExecInfo));
  if (pInfo == NULL) {
    return terrno;
  }

  SSortOperatorInfo* pOperatorInfo = (SSortOperatorInfo*)pOptr->info;

  *pInfo = tsortGetSortExecInfo(pOperatorInfo->pSortHandle);
  *pOptrExplain = pInfo;
  *len = sizeof(SSortExecInfo);
  return TSDB_CODE_SUCCESS;
}

static void calcSortOperMaxTupleLength(SSortOperatorInfo* pSortOperInfo, SNodeList* pSortKeys) {
  SColMatchInfo* pColItem = &pSortOperInfo->matchInfo;
  size_t         size = taosArrayGetSize(pColItem->pList);
  for (size_t i = 0; i < size; ++i) {
    SColMatchItem* pInfo = taosArrayGet(pColItem->pList, i);
    if (pInfo == NULL) {
      continue;
    }

    pSortOperInfo->maxTupleLength += pInfo->dataType.bytes;
  }

  size = LIST_LENGTH(pSortKeys);
  for (size_t i = 0; i < size; ++i) {
    SOrderByExprNode* pOrderExprNode = (SOrderByExprNode*)nodesListGetNode(pSortKeys, i);
    pSortOperInfo->maxTupleLength += ((SColumnNode*)pOrderExprNode->pExpr)->node.resType.bytes;
  }
}

static void destroySortOpGroupIdCalc(SSortOpGroupIdCalc* pCalc) {
  if (pCalc) {
    taosArrayDestroy(pCalc->pSortColsArr);
    taosMemoryFree(pCalc->keyBuf);
    taosMemoryFree(pCalc);
  }
}

//=====================================================================================
// Group Sort Operator
typedef enum EChildOperatorStatus { CHILD_OP_NEW_GROUP, CHILD_OP_SAME_GROUP, CHILD_OP_FINISHED } EChildOperatorStatus;

typedef struct SGroupSortOperatorInfo {
  SOptrBasicInfo       binfo;
  SArray*              pSortInfo;
  SColMatchInfo        matchInfo;
  int64_t              startTs;
  uint64_t             sortElapsed;
  bool                 hasGroupId;
  uint64_t             currGroupId;
  uint64_t             currBaseGId;
  SSDataBlock*         prefetchedSortInput;
  SSortHandle*         pCurrSortHandle;
  EChildOperatorStatus childOpStatus;
  SSortExecInfo        sortExecInfo;
} SGroupSortOperatorInfo;

int32_t getGroupSortedBlockData(SSortHandle* pHandle, SSDataBlock* pDataBlock, int32_t capacity, SArray* pColMatchInfo,
                                SGroupSortOperatorInfo* pInfo, SSDataBlock** pResBlock) {
  QRY_PARAM_CHECK(pResBlock);

  blockDataCleanup(pDataBlock);
  int32_t code = blockDataEnsureCapacity(pDataBlock, capacity);
  if (code) {
    return code;
  }

  SSDataBlock* p = NULL;
  code = tsortGetSortedDataBlock(pHandle, &p);
  if (p == NULL || (code != 0)) {
    return code;
  }

  code = blockDataEnsureCapacity(p, capacity);
  if (code) {
    return code;
  }

  while (1) {
    STupleHandle* pTupleHandle = NULL;
    code = tsortNextTuple(pHandle, &pTupleHandle);
    if (pTupleHandle == NULL || code != 0) {
      break;
    }

    code = appendOneRowToDataBlock(p, pTupleHandle);
    if (code) {
      break;
    }

    if (p->info.rows >= capacity) {
      break;
    }
  }

  if (p->info.rows > 0) {
    int32_t numOfCols = taosArrayGetSize(pColMatchInfo);
    for (int32_t i = 0; i < numOfCols; ++i) {
      SColMatchItem* pmInfo = taosArrayGet(pColMatchInfo, i);
      if (pmInfo == NULL) {
        return terrno;
      }

      SColumnInfoData* pSrc = taosArrayGet(p->pDataBlock, pmInfo->srcSlotId);
      if (pSrc == NULL) {
        return terrno;
      }

      SColumnInfoData* pDst = taosArrayGet(pDataBlock->pDataBlock, pmInfo->dstSlotId);
      if (pDst == NULL) {
        return terrno;
      }

      code = colDataAssign(pDst, pSrc, p->info.rows, &pDataBlock->info);
      if (code) {
        return code;
      }
    }

    pDataBlock->info.rows = p->info.rows;
    pDataBlock->info.capacity = p->info.rows;
    pDataBlock->info.scanFlag = p->info.scanFlag;
    // propagate ids for the current group
    pDataBlock->info.id.groupId = pInfo->currGroupId;
    pDataBlock->info.id.baseGId = pInfo->currBaseGId;
  }

  blockDataDestroy(p);
  *pResBlock = (pDataBlock->info.rows > 0) ? pDataBlock : NULL;
  return code;
}

typedef struct SGroupSortSourceParam {
  SOperatorInfo*          childOpInfo;
  SGroupSortOperatorInfo* grpSortOpInfo;
} SGroupSortSourceParam;

int32_t fetchNextGroupSortDataBlock(void* param, SSDataBlock** ppBlock) {
  int32_t                 code = 0;
  int32_t                 lino = 0;
  SGroupSortSourceParam*  source = param;
  SGroupSortOperatorInfo* grpSortOpInfo = source->grpSortOpInfo;
  SSDataBlock*            block = NULL;

  QRY_PARAM_CHECK(ppBlock);

  if (grpSortOpInfo->prefetchedSortInput) {
    block = grpSortOpInfo->prefetchedSortInput;
    grpSortOpInfo->prefetchedSortInput = NULL;
    *ppBlock = block;
  } else {
    SOperatorInfo* childOp = source->childOpInfo;
    code = childOp->fpSet.getNextFn(childOp, &block);
    QUERY_CHECK_CODE(code, lino, _end);

    if (block != NULL) {
      code = blockDataCheck(block);
      QUERY_CHECK_CODE(code, lino, _end);
      if (block->info.id.groupId == grpSortOpInfo->currGroupId) {
        grpSortOpInfo->childOpStatus = CHILD_OP_SAME_GROUP;
        *ppBlock = block;
      } else {
        grpSortOpInfo->childOpStatus = CHILD_OP_NEW_GROUP;
        grpSortOpInfo->prefetchedSortInput = block;
      }
    } else {
      grpSortOpInfo->childOpStatus = CHILD_OP_FINISHED;
    }
  }

  return code;
_end:
  if (code != 0) {
    qError("%s failed at line %d since %s", __func__, lino, tstrerror(code));
  }
  return code;
}

int32_t beginSortGroup(SOperatorInfo* pOperator) {
  SGroupSortOperatorInfo* pInfo = pOperator->info;
  SExecTaskInfo*          pTaskInfo = pOperator->pTaskInfo;

  //  pInfo->binfo.pRes is not equalled to the input datablock.
  pInfo->pCurrSortHandle = NULL;

  int32_t code = tsortCreateSortHandle(pInfo->pSortInfo, SORT_SINGLESOURCE_SORT, -1, -1, NULL, pTaskInfo->id.str, 0, 0,
                                       0, &pInfo->pCurrSortHandle);
  if (code) {
    return code;
  }

  tsortSetFetchRawDataFp(pInfo->pCurrSortHandle, fetchNextGroupSortDataBlock, applyScalarFunction, pOperator);

  SSortSource*           ps = taosMemoryCalloc(1, sizeof(SSortSource));
  SGroupSortSourceParam* param = taosMemoryCalloc(1, sizeof(SGroupSortSourceParam));
  if (ps == NULL || param == NULL) {
    taosMemoryFree(ps);
    taosMemoryFree(param);
    return terrno;
  }

  param->childOpInfo = pOperator->pDownstream[0];
  param->grpSortOpInfo = pInfo;

  ps->param = param;
  ps->onlyRef = false;
  code = tsortAddSource(pInfo->pCurrSortHandle, ps);
  if (code != 0) {
    return code;
  }

  code = tsortOpen(pInfo->pCurrSortHandle);
  return code;
}

int32_t finishSortGroup(SOperatorInfo* pOperator) {
  SGroupSortOperatorInfo* pInfo = pOperator->info;

  SSortExecInfo sortExecInfo = tsortGetSortExecInfo(pInfo->pCurrSortHandle);

  pInfo->sortExecInfo.sortMethod = sortExecInfo.sortMethod;
  pInfo->sortExecInfo.sortBuffer = sortExecInfo.sortBuffer;
  pInfo->sortExecInfo.loops += sortExecInfo.loops;
  pInfo->sortExecInfo.readBytes += sortExecInfo.readBytes;
  pInfo->sortExecInfo.writeBytes += sortExecInfo.writeBytes;

  tsortDestroySortHandle(pInfo->pCurrSortHandle);
  pInfo->pCurrSortHandle = NULL;

  return TSDB_CODE_SUCCESS;
}

int32_t doGroupSort(SOperatorInfo* pOperator, SSDataBlock** pResBlock) {
  QRY_PARAM_CHECK(pResBlock);
  SExecTaskInfo*          pTaskInfo = pOperator->pTaskInfo;
  SGroupSortOperatorInfo* pInfo = pOperator->info;
  int32_t                 code = TSDB_CODE_SUCCESS;
  int32_t                 lino = 0;

  if (pOperator->status == OP_EXEC_DONE) {
    return code;
  }

  code = pOperator->fpSet._openFn(pOperator);
  QUERY_CHECK_CODE(code, lino, _end);

  if (!pInfo->hasGroupId) {
    pInfo->hasGroupId = true;

    pInfo->prefetchedSortInput = getNextBlockFromDownstream(pOperator, 0);
    if (pInfo->prefetchedSortInput == NULL) {
      setOperatorCompleted(pOperator);
      return code;
    }

    pInfo->currGroupId = pInfo->prefetchedSortInput->info.id.groupId;
    pInfo->currBaseGId = pInfo->prefetchedSortInput->info.id.baseGId;
    pInfo->childOpStatus = CHILD_OP_NEW_GROUP;
    code = beginSortGroup(pOperator);
    QUERY_CHECK_CODE(code, lino, _end);
  }

  SSDataBlock* pBlock = NULL;
  while (pInfo->pCurrSortHandle != NULL) {
    if (tsortIsClosed(pInfo->pCurrSortHandle)) {
      code = TSDB_CODE_TSC_QUERY_CANCELLED;
      QUERY_CHECK_CODE(code, lino, _end);
    }

    // beginSortGroup would fetch all child blocks of pInfo->currGroupId;
    if (pInfo->childOpStatus == CHILD_OP_SAME_GROUP) {
      code = TSDB_CODE_QRY_EXECUTOR_INTERNAL_ERROR;
      QUERY_CHECK_CODE(code, lino, _end);
    }

    recordOpExecBeforeDownstream(pOperator);
    code = getGroupSortedBlockData(pInfo->pCurrSortHandle, pInfo->binfo.pRes, pOperator->resultInfo.capacity,
                                     pInfo->matchInfo.pList, pInfo, &pBlock);
    recordOpExecAfterDownstream(pOperator, pBlock ? pBlock->info.rows : 0);
    QUERY_CHECK_CODE(code, lino, _end);
    if (pBlock != NULL) {
      // keep both ids aligned with current group
      pBlock->info.id.groupId = pInfo->currGroupId;
      pBlock->info.id.baseGId = pInfo->currBaseGId;
      // baseGId follows upstream; if upstream is empty here, preserve current
      // (no-op if not set). We cannot reconstruct baseGId here; rely on upstream propagation.
      pOperator->resultInfo.totalRows += pBlock->info.rows;
      *pResBlock = pBlock;
      return code;
    } else {
      if (pInfo->childOpStatus == CHILD_OP_NEW_GROUP) {
        (void) finishSortGroup(pOperator);
        pInfo->currGroupId = pInfo->prefetchedSortInput->info.id.groupId;
        pInfo->currBaseGId = pInfo->prefetchedSortInput->info.id.baseGId;
        code = beginSortGroup(pOperator);
        QUERY_CHECK_CODE(code, lino, _end);
      } else if (pInfo->childOpStatus == CHILD_OP_FINISHED) {
        (void) finishSortGroup(pOperator);
        setOperatorCompleted(pOperator);
        return code;
      }
    }
  }

_end:
  if (code != TSDB_CODE_SUCCESS) {
    qError("%s failed at line %d since %s", __func__, lino, tstrerror(code));
    pTaskInfo->code = code;
    T_LONG_JMP(pTaskInfo->env, code);
  }
  return code;
}

int32_t getGroupSortExplainExecInfo(SOperatorInfo* pOptr, void** pOptrExplain, uint32_t* len) {
  SGroupSortOperatorInfo* pInfo = (SGroupSortOperatorInfo*)pOptr->info;
  *pOptrExplain = &pInfo->sortExecInfo;
  *len = sizeof(SSortExecInfo);
  return TSDB_CODE_SUCCESS;
}

void destroyGroupSortOperatorInfo(void* param) {
  SGroupSortOperatorInfo* pInfo = (SGroupSortOperatorInfo*)param;
  blockDataDestroy(pInfo->binfo.pRes);
  pInfo->binfo.pRes = NULL;

  taosArrayDestroy(pInfo->pSortInfo);
  taosArrayDestroy(pInfo->matchInfo.pList);

  tsortDestroySortHandle(pInfo->pCurrSortHandle);
  pInfo->pCurrSortHandle = NULL;

  taosMemoryFreeClear(param);
}

static int32_t resetGroupSortOperState(SOperatorInfo* pOper) {
  SGroupSortOperatorInfo* pInfo = pOper->info;
  SExecTaskInfo*           pTaskInfo = pOper->pTaskInfo;
  pOper->status = OP_NOT_OPENED;

  pInfo->currGroupId = 0;
  pInfo->hasGroupId = false;
  pInfo->prefetchedSortInput = NULL;
  pInfo->childOpStatus = CHILD_OP_NEW_GROUP;
  pInfo->sortExecInfo = (SSortExecInfo){0};
  
  resetBasicOperatorState(&pInfo->binfo);
  destroySqlFunctionCtx(pOper->exprSupp.pCtx, pOper->exprSupp.pExprInfo, pOper->exprSupp.numOfExprs);
  taosMemoryFreeClear(pOper->exprSupp.rowEntryInfoOffset);
  pOper->exprSupp.pCtx =
      createSqlFunctionCtx(pOper->exprSupp.pExprInfo, pOper->exprSupp.numOfExprs, &pOper->exprSupp.rowEntryInfoOffset, &pTaskInfo->storageAPI.functionStore);

  tsortDestroySortHandle(pInfo->pCurrSortHandle);
  pInfo->pCurrSortHandle = NULL;

  return 0;
}

int32_t createGroupSortOperatorInfo(SOperatorInfo* downstream, SGroupSortPhysiNode* pSortPhyNode,
                                    SExecTaskInfo* pTaskInfo, SOperatorInfo** pOptrInfo) {
  QRY_PARAM_CHECK(pOptrInfo);
  int32_t code = 0;
  int32_t lino = 0;

  SGroupSortOperatorInfo* pInfo = taosMemoryCalloc(1, sizeof(SGroupSortOperatorInfo));
  SOperatorInfo*          pOperator = taosMemoryCalloc(1, sizeof(SOperatorInfo));
  if (pInfo == NULL || pOperator == NULL) {
    code = terrno;
    goto _error;
  }
  initOperatorCostInfo(pOperator);

  SExprSupp*          pSup = &pOperator->exprSupp;
  SDataBlockDescNode* pDescNode = pSortPhyNode->node.pOutputDataBlockDesc;

  int32_t    numOfCols = 0;
  SExprInfo* pExprInfo = NULL;
  code = createExprInfo(pSortPhyNode->pExprs, NULL, &pExprInfo, &numOfCols);
  QUERY_CHECK_CODE(code, lino, _error);

  pSup->pExprInfo = pExprInfo;
  pSup->numOfExprs = numOfCols;

  initResultSizeInfo(&pOperator->resultInfo, 1024);
  pOperator->exprSupp.pCtx = createSqlFunctionCtx(pExprInfo, numOfCols, &pOperator->exprSupp.rowEntryInfoOffset,
                                                  &pTaskInfo->storageAPI.functionStore);
  QUERY_CHECK_NULL(pOperator->exprSupp.pCtx, code, lino, _error, terrno);

  pInfo->binfo.pRes = createDataBlockFromDescNode(pDescNode);
  QUERY_CHECK_NULL(pInfo->binfo.pRes, code, lino, _error, terrno);

  code = blockDataEnsureCapacity(pInfo->binfo.pRes, pOperator->resultInfo.capacity);
  TSDB_CHECK_CODE(code, lino, _error);

  pInfo->binfo.inputTsOrder = pSortPhyNode->node.inputTsOrder;
  pInfo->binfo.outputTsOrder = pSortPhyNode->node.outputTsOrder;

  int32_t numOfOutputCols = 0;
  code = extractColMatchInfo(pSortPhyNode->pTargets, pDescNode, &numOfOutputCols, COL_MATCH_FROM_SLOT_ID,
                             &pInfo->matchInfo);
  TSDB_CHECK_CODE(code, lino, _error);

  pInfo->pSortInfo = createSortInfo(pSortPhyNode->pSortKeys);
  setOperatorInfo(pOperator, "GroupSortOperator", QUERY_NODE_PHYSICAL_PLAN_GROUP_SORT, false, OP_NOT_OPENED, pInfo,
                  pTaskInfo);
  pOperator->fpSet = createOperatorFpSet(optrDummyOpenFn, doGroupSort, NULL, destroyGroupSortOperatorInfo,
                                         optrDefaultBufFn, getGroupSortExplainExecInfo, optrDefaultGetNextExtFn, NULL);

  setOperatorResetStateFn(pOperator, resetGroupSortOperState);
                                         
  code = appendDownstream(pOperator, &downstream, 1);
  if (code != TSDB_CODE_SUCCESS) {
    goto _error;
  }

  *pOptrInfo = pOperator;
  return TSDB_CODE_SUCCESS;

_error:
  pTaskInfo->code = code;
  if (pInfo != NULL) {
    destroyGroupSortOperatorInfo(pInfo);
  }
  destroyOperatorAndDownstreams(pOperator, &downstream, 1);
  return code;
}

taosdata / TDengine / #5044

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous