tsort.c 22.1 KB
Newer Older
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/>.
 */

S
common  
Shengliang Guan 已提交
16
#include "tcommon.h"
17 18
#include "query.h"

H
Haojun Liao 已提交
19
#include "tdatablock.h"
S
Shengliang Guan 已提交
20
#include "tdef.h"
21 22
#include "tlosertree.h"
#include "tpagedbuf.h"
H
Haojun Liao 已提交
23
#include "tsort.h"
24
#include "tutil.h"
25
#include "tcompare.h"
26

27 28 29 30 31
struct STupleHandle {
  SSDataBlock* pBlock;
  int32_t      rowIndex;
};

H
Haojun Liao 已提交
32
struct SSortHandle {
33 34 35 36 37
  int32_t           type;
  int32_t           pageSize;
  int32_t           numOfPages;
  SDiskbasedBuf    *pBuf;

H
Haojun Liao 已提交
38
  SArray           *pSortInfo;
39 40
  SArray           *pOrderedSource;

41
  int32_t           loops;
42
  uint64_t          sortElapsed;
43
  int64_t           startTs;
44 45 46 47 48 49 50 51 52 53 54
  uint64_t          totalElapsed;

  int32_t           sourceId;
  SSDataBlock      *pDataBlock;
  SMsortComparParam cmpParam;
  int32_t           numOfCompletedSources;
  bool              opened;
  const char       *idStr;
  bool              inMemSort;
  bool              needAdjust;
  STupleHandle      tupleHandle;
55 56
  void             *param;
  void (*beforeFp)(SSDataBlock* pBlock, void* param);
57 58 59 60

  _sort_fetch_block_fn_t  fetchfp;
  _sort_merge_compar_fn_t comparFn;
  SMultiwayMergeTreeInfo *pMergeTree;
H
Haojun Liao 已提交
61
};
62

63 64
static int32_t msortComparFn(const void *pLeft, const void *pRight, void *param);

65 66
SSDataBlock* tsortGetSortedDataBlock(const SSortHandle* pSortHandle) {
  return createOneDataBlock(pSortHandle->pDataBlock, false);
67 68 69 70 71 72 73
}

/**
 *
 * @param type
 * @return
 */
74
SSortHandle* tsortCreateSortHandle(SArray* pSortInfo, SArray* pIndexMap, int32_t type, int32_t pageSize, int32_t numOfPages, SSDataBlock* pBlock, const char* idstr) {
wafwerar's avatar
wafwerar 已提交
75
  SSortHandle* pSortHandle = taosMemoryCalloc(1, sizeof(SSortHandle));
76 77 78 79

  pSortHandle->type       = type;
  pSortHandle->pageSize   = pageSize;
  pSortHandle->numOfPages = numOfPages;
H
Haojun Liao 已提交
80
  pSortHandle->pSortInfo  = pSortInfo;
81
  pSortHandle->loops      = 0;
82 83 84 85

  if (pBlock != NULL) {
    pSortHandle->pDataBlock = createOneDataBlock(pBlock, false);
  }
H
Haojun Liao 已提交
86 87 88

  pSortHandle->pOrderedSource     = taosArrayInit(4, POINTER_BYTES);
  pSortHandle->cmpParam.orderInfo = pSortInfo;
89

H
Haojun Liao 已提交
90
  tsortSetComparFp(pSortHandle, msortComparFn);
91 92 93 94 95 96 97 98

  if (idstr != NULL) {
    pSortHandle->idStr    = strdup(idstr);
  }

  return pSortHandle;
}

wmmhello's avatar
wmmhello 已提交
99 100
static int32_t sortComparClearup(SMsortComparParam* cmpParam) {
  for(int32_t i = 0; i < cmpParam->numOfSources; ++i) {
wmmhello's avatar
wmmhello 已提交
101
    SSortSource* pSource = cmpParam->pSources[i];    // NOTICE: pSource may be SGenericSource *, if it is SORT_MULTISOURCE_MERGE
wmmhello's avatar
wmmhello 已提交
102 103 104 105 106 107 108 109
    blockDataDestroy(pSource->src.pBlock);
    taosMemoryFreeClear(pSource);
  }

  cmpParam->numOfSources = 0;
  return TSDB_CODE_SUCCESS;
}

H
Haojun Liao 已提交
110 111
void tsortDestroySortHandle(SSortHandle* pSortHandle) {
  tsortClose(pSortHandle);
112 113 114 115
  if (pSortHandle->pMergeTree != NULL) {
    tMergeTreeDestroy(pSortHandle->pMergeTree);
  }

H
Haojun Liao 已提交
116
  destroyDiskbasedBuf(pSortHandle->pBuf);
wafwerar's avatar
wafwerar 已提交
117
  taosMemoryFreeClear(pSortHandle->idStr);
wmmhello's avatar
wmmhello 已提交
118
  blockDataDestroy(pSortHandle->pDataBlock);
wmmhello's avatar
wmmhello 已提交
119
  for (size_t i = 0; i < taosArrayGetSize(pSortHandle->pOrderedSource); i++){
wmmhello's avatar
wmmhello 已提交
120
    SSortSource** pSource = taosArrayGet(pSortHandle->pOrderedSource, i);
wmmhello's avatar
wmmhello 已提交
121 122
    blockDataDestroy((*pSource)->src.pBlock);
    taosMemoryFreeClear(*pSource);
wmmhello's avatar
wmmhello 已提交
123 124
  }
  taosArrayDestroy(pSortHandle->pOrderedSource);
wafwerar's avatar
wafwerar 已提交
125
  taosMemoryFreeClear(pSortHandle);
126 127
}

H
Haojun Liao 已提交
128
int32_t tsortAddSource(SSortHandle* pSortHandle, void* pSource) {
H
Haojun Liao 已提交
129
  taosArrayPush(pSortHandle->pOrderedSource, &pSource);
130
  return TSDB_CODE_SUCCESS;
131 132 133
}

static int32_t doAddNewExternalMemSource(SDiskbasedBuf *pBuf, SArray* pAllSources, SSDataBlock* pBlock, int32_t* sourceId) {
wmmhello's avatar
wmmhello 已提交
134
  SSortSource* pSource = taosMemoryCalloc(1, sizeof(SSortSource));
135 136 137 138 139 140 141 142 143 144 145 146
  if (pSource == NULL) {
    return TSDB_CODE_QRY_OUT_OF_MEMORY;
  }

  pSource->pageIdList = getDataBufPagesIdList(pBuf, (*sourceId));
  pSource->src.pBlock = pBlock;

  taosArrayPush(pAllSources, &pSource);

  (*sourceId) += 1;

  int32_t rowSize = blockDataGetSerialRowSize(pSource->src.pBlock);
wmmhello's avatar
wmmhello 已提交
147 148
  int32_t numOfRows = (getBufPageSize(pBuf) - blockDataGetSerialMetaSize(pBlock))/rowSize;  // The value of numOfRows must be greater than 0, which is guaranteed by the previous memory allocation
  ASSERT(numOfRows > 0);
149 150 151 152 153 154 155
  return blockDataEnsureCapacity(pSource->src.pBlock, numOfRows);
}

static int32_t doAddToBuf(SSDataBlock* pDataBlock, SSortHandle* pHandle) {
  int32_t start = 0;

  if (pHandle->pBuf == NULL) {
156
    int32_t code = createDiskbasedBuf(&pHandle->pBuf, pHandle->pageSize, pHandle->numOfPages * pHandle->pageSize, "doAddToBuf", TD_TMP_DIR_PATH);
H
Haojun Liao 已提交
157
    dBufSetPrintInfo(pHandle->pBuf);
158 159 160 161 162 163 164
    if (code != TSDB_CODE_SUCCESS) {
      return code;
    }
  }

  while(start < pDataBlock->info.rows) {
    int32_t stop = 0;
165
    blockDataSplitRows(pDataBlock, pDataBlock->info.hasVarCol, start, &stop, pHandle->pageSize);
166 167 168 169 170 171
    SSDataBlock* p = blockDataExtractBlock(pDataBlock, start, stop - start + 1);
    if (p == NULL) {
      return terrno;
    }

    int32_t pageId = -1;
172
    void* pPage = getNewBufPage(pHandle->pBuf, pHandle->sourceId, &pageId);
173
    if (pPage == NULL) {
wmmhello's avatar
wmmhello 已提交
174
      blockDataDestroy(p);
175 176 177 178 179 180
      return terrno;
    }

    int32_t size = blockDataGetSize(p) + sizeof(int32_t)  + p->info.numOfCols * sizeof(int32_t);
    assert(size <= getBufPageSize(pHandle->pBuf));

181
    blockDataToBuf(pPage, p);
182 183 184 185 186 187 188 189

    setBufPageDirty(pPage, true);
    releaseBufPage(pHandle->pBuf, pPage);

    blockDataDestroy(p);
    start = stop + 1;
  }

190
  blockDataCleanup(pDataBlock);
191

192
  SSDataBlock* pBlock = createOneDataBlock(pDataBlock, false);
193
  return doAddNewExternalMemSource(pHandle->pBuf, pHandle->pOrderedSource, pBlock, &pHandle->sourceId);
194 195 196 197 198 199 200 201
}

static int32_t sortComparInit(SMsortComparParam* cmpParam, SArray* pSources, int32_t startIndex, int32_t endIndex, SSortHandle* pHandle) {
  cmpParam->pSources  = taosArrayGet(pSources, startIndex);
  cmpParam->numOfSources = (endIndex - startIndex + 1);

  int32_t code = 0;

H
Haojun Liao 已提交
202
  if (pHandle->type == SORT_SINGLESOURCE_SORT) {
203
    for (int32_t i = 0; i < cmpParam->numOfSources; ++i) {
wmmhello's avatar
wmmhello 已提交
204
      SSortSource* pSource = cmpParam->pSources[i];
205 206
      SPageInfo*          pPgInfo = *(SPageInfo**)taosArrayGet(pSource->pageIdList, pSource->pageIndex);

207 208
      void* pPage = getBufPage(pHandle->pBuf, getPageId(pPgInfo));
      code = blockDataFromBuf(pSource->src.pBlock, pPage);
209 210 211 212 213 214 215 216 217
      if (code != TSDB_CODE_SUCCESS) {
        return code;
      }

      releaseBufPage(pHandle->pBuf, pPage);
    }
  } else {
    // multi-pass internal merge sort is required
    if (pHandle->pBuf == NULL) {
218
      code = createDiskbasedBuf(&pHandle->pBuf, pHandle->pageSize, pHandle->numOfPages * pHandle->pageSize, "sortComparInit", TD_TMP_DIR_PATH);
H
Haojun Liao 已提交
219
      dBufSetPrintInfo(pHandle->pBuf);
220 221 222 223 224 225
      if (code != TSDB_CODE_SUCCESS) {
        return code;
      }
    }

    for (int32_t i = 0; i < cmpParam->numOfSources; ++i) {
wmmhello's avatar
wmmhello 已提交
226
      SSortSource* pSource = cmpParam->pSources[i];
227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243
      pSource->src.pBlock = pHandle->fetchfp(pSource->param);
    }
  }

  return code;
}

static void appendOneRowToDataBlock(SSDataBlock *pBlock, const SSDataBlock* pSource, int32_t* rowIndex) {
  for (int32_t i = 0; i < pBlock->info.numOfCols; ++i) {
    SColumnInfoData* pColInfo = taosArrayGet(pBlock->pDataBlock, i);

    SColumnInfoData* pSrcColInfo = taosArrayGet(pSource->pDataBlock, i);
    bool isNull = colDataIsNull(pSrcColInfo, pSource->info.rows, *rowIndex, NULL);

    if (isNull) {
      colDataAppend(pColInfo, pBlock->info.rows, NULL, true);
    } else {
244
      char* pData = colDataGetData(pSrcColInfo, *rowIndex);
245 246 247 248 249 250 251 252
      colDataAppend(pColInfo, pBlock->info.rows, pData, false);
    }
  }

  pBlock->info.rows += 1;
  *rowIndex += 1;
}

wmmhello's avatar
wmmhello 已提交
253
static int32_t adjustMergeTreeForNextTuple(SSortSource *pSource, SMultiwayMergeTreeInfo *pTree, SSortHandle *pHandle, int32_t* numOfCompleted) {
254 255 256 257 258 259 260
  /*
   * load a new SDataBlock into memory of a given intermediate data-set source,
   * since it's last record in buffer has been chosen to be processed, as the winner of loser-tree
   */
  if (pSource->src.rowIndex >= pSource->src.pBlock->info.rows) {
    pSource->src.rowIndex = 0;

261
    if (pHandle->type == SORT_SINGLESOURCE_SORT) {
wmmhello's avatar
wmmhello 已提交
262
      pSource->pageIndex ++;
263 264 265 266 267 268
      if (pSource->pageIndex >= taosArrayGetSize(pSource->pageIdList)) {
        (*numOfCompleted) += 1;
        pSource->src.rowIndex = -1;
        pSource->pageIndex = -1;
        pSource->src.pBlock = blockDataDestroy(pSource->src.pBlock);
      } else {
269 270
        SPageInfo* pPgInfo = *(SPageInfo**)taosArrayGet(pSource->pageIdList, pSource->pageIndex);

wmmhello's avatar
wmmhello 已提交
271 272
        void* pPage = getBufPage(pHandle->pBuf, getPageId(pPgInfo));
        int32_t    code = blockDataFromBuf(pSource->src.pBlock, pPage);
273 274 275 276 277
        if (code != TSDB_CODE_SUCCESS) {
          return code;
        }

        releaseBufPage(pHandle->pBuf, pPage);
278 279
      }
    } else {
wmmhello's avatar
wmmhello 已提交
280
      pSource->src.pBlock = pHandle->fetchfp(((SSortSource*)pSource)->param);
281 282 283
      if (pSource->src.pBlock == NULL) {
        (*numOfCompleted) += 1;
        pSource->src.rowIndex = -1;
284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304
      }
    }
  }

  /*
   * Adjust loser tree otherwise, according to new candidate data
   * if the loser tree is rebuild completed, we do not need to adjust
   */
  int32_t leafNodeIndex = tMergeTreeGetAdjustIndex(pTree);

#ifdef _DEBUG_VIEW
  printf("before adjust:\t");
  tMergeTreePrint(pTree);
#endif

  tMergeTreeAdjust(pTree, leafNodeIndex);

#ifdef _DEBUG_VIEW
  printf("\nafter adjust:\t");
  tMergeTreePrint(pTree);
#endif
305
  return TSDB_CODE_SUCCESS;
306 307
}

wmmhello's avatar
wmmhello 已提交
308
static SSDataBlock* getSortedBlockDataInner(SSortHandle* pHandle, SMsortComparParam* cmpParam, int32_t capacity) {
309
  blockDataCleanup(pHandle->pDataBlock);
310 311 312 313 314 315 316 317

  while(1) {
    if (cmpParam->numOfSources == pHandle->numOfCompletedSources) {
      break;
    }

    int32_t index = tMergeTreeGetChosenIndex(pHandle->pMergeTree);

wmmhello's avatar
wmmhello 已提交
318
    SSortSource *pSource = (*cmpParam).pSources[index];
319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334
    appendOneRowToDataBlock(pHandle->pDataBlock, pSource->src.pBlock, &pSource->src.rowIndex);

    int32_t code = adjustMergeTreeForNextTuple(pSource, pHandle->pMergeTree, pHandle, &pHandle->numOfCompletedSources);
    if (code != TSDB_CODE_SUCCESS) {
      terrno = code;
      return NULL;
    }

    if (pHandle->pDataBlock->info.rows >= capacity) {
      return pHandle->pDataBlock;
    }
  }

  return (pHandle->pDataBlock->info.rows > 0)? pHandle->pDataBlock:NULL;
}

335
int32_t msortComparFn(const void *pLeft, const void *pRight, void *param) {
336 337 338 339 340 341 342
  int32_t pLeftIdx  = *(int32_t *)pLeft;
  int32_t pRightIdx = *(int32_t *)pRight;

  SMsortComparParam *pParam = (SMsortComparParam *)param;

  SArray *pInfo = pParam->orderInfo;

wmmhello's avatar
wmmhello 已提交
343 344
  SSortSource* pLeftSource  = pParam->pSources[pLeftIdx];
  SSortSource* pRightSource = pParam->pSources[pRightIdx];
345 346 347 348 349 350 351 352 353 354 355 356 357 358 359

  // this input is exhausted, set the special value to denote this
  if (pLeftSource->src.rowIndex == -1) {
    return 1;
  }

  if (pRightSource->src.rowIndex == -1) {
    return -1;
  }

  SSDataBlock* pLeftBlock = pLeftSource->src.pBlock;
  SSDataBlock* pRightBlock = pRightSource->src.pBlock;

  for(int32_t i = 0; i < pInfo->size; ++i) {
    SBlockOrderInfo* pOrder = TARRAY_GET_ELEM(pInfo, i);
H
Haojun Liao 已提交
360
    SColumnInfoData* pLeftColInfoData = TARRAY_GET_ELEM(pLeftBlock->pDataBlock, pOrder->slotId);
361 362 363

    bool leftNull  = false;
    if (pLeftColInfoData->hasNull) {
364
      leftNull = colDataIsNull(pLeftColInfoData, pLeftBlock->info.rows, pLeftSource->src.rowIndex, pLeftBlock->pBlockAgg[pOrder->slotId]);
365 366
    }

H
Haojun Liao 已提交
367
    SColumnInfoData* pRightColInfoData = TARRAY_GET_ELEM(pRightBlock->pDataBlock, pOrder->slotId);
368 369
    bool rightNull = false;
    if (pRightColInfoData->hasNull) {
370
      rightNull = colDataIsNull(pRightColInfoData, pRightBlock->info.rows, pRightSource->src.rowIndex, pRightBlock->pBlockAgg[pOrder->slotId]);
371 372 373 374 375 376 377
    }

    if (leftNull && rightNull) {
      continue; // continue to next slot
    }

    if (rightNull) {
wmmhello's avatar
wmmhello 已提交
378
      return pOrder->nullFirst? 1:-1;
379 380 381
    }

    if (leftNull) {
wmmhello's avatar
wmmhello 已提交
382
      return pOrder->nullFirst? -1:1;
383 384
    }

385 386
    void* left1  = colDataGetData(pLeftColInfoData, pLeftSource->src.rowIndex);
    void* right1 = colDataGetData(pRightColInfoData, pRightSource->src.rowIndex);
387

388
    __compar_fn_t fn = getKeyComparFunc(pLeftColInfoData->info.type, pOrder->order);
389

390 391 392 393 394
    int ret = fn(left1, right1);
    if (ret == 0) {
      continue;
    } else {
      return ret;
395 396
    }
  }
397
  return 0;
398 399 400 401
}

static int32_t doInternalMergeSort(SSortHandle* pHandle) {
  size_t numOfSources = taosArrayGetSize(pHandle->pOrderedSource);
H
Haojun Liao 已提交
402 403 404
  if (numOfSources == 0) {
    return 0;
  }
405 406 407 408 409

  // Calculate the I/O counts to complete the data sort.
  double sortPass = floorl(log2(numOfSources) / log2(pHandle->numOfPages));

  pHandle->totalElapsed = taosGetTimestampUs() - pHandle->startTs;
410
  qDebug("%s %d rounds mergesort required to complete the sort, first-round sorted data size:%"PRIzu", sort elapsed:%"PRId64", total elapsed:%"PRId64,
411 412
         pHandle->idStr, (int32_t) (sortPass + 1), getTotalBufSize(pHandle->pBuf), pHandle->sortElapsed, pHandle->totalElapsed);

H
Haojun Liao 已提交
413
  int32_t numOfRows = blockDataGetCapacityInRow(pHandle->pDataBlock, pHandle->pageSize);
wmmhello's avatar
wmmhello 已提交
414
  blockDataEnsureCapacity(pHandle->pDataBlock, numOfRows);
415

416 417 418
  // the initial pass + sortPass + final mergePass
  pHandle->loops = sortPass + 2;

419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449
  size_t numOfSorted = taosArrayGetSize(pHandle->pOrderedSource);
  for(int32_t t = 0; t < sortPass; ++t) {
    int64_t st = taosGetTimestampUs();

    SArray* pResList = taosArrayInit(4, POINTER_BYTES);

    int32_t numOfInputSources = pHandle->numOfPages;
    int32_t sortGroup = (numOfSorted + numOfInputSources - 1) / numOfInputSources;

    // Only *numOfInputSources* can be loaded into buffer to perform the external sort.
    for(int32_t i = 0; i < sortGroup; ++i) {
      pHandle->sourceId += 1;

      int32_t end = (i + 1) * numOfInputSources - 1;
      if (end > numOfSorted - 1) {
        end = numOfSorted - 1;
      }

      pHandle->cmpParam.numOfSources = end - i * numOfInputSources + 1;

      int32_t code = sortComparInit(&pHandle->cmpParam, pHandle->pOrderedSource, i * numOfInputSources, end, pHandle);
      if (code != TSDB_CODE_SUCCESS) {
        return code;
      }

      code = tMergeTreeCreate(&pHandle->pMergeTree, pHandle->cmpParam.numOfSources, &pHandle->cmpParam, pHandle->comparFn);
      if (code != TSDB_CODE_SUCCESS) {
        return code;
      }

      while (1) {
wmmhello's avatar
wmmhello 已提交
450
        SSDataBlock* pDataBlock = getSortedBlockDataInner(pHandle, &pHandle->cmpParam, numOfRows);
451 452 453 454 455
        if (pDataBlock == NULL) {
          break;
        }

        int32_t pageId = -1;
456
        void* pPage = getNewBufPage(pHandle->pBuf, pHandle->sourceId, &pageId);
457 458 459 460 461 462 463
        if (pPage == NULL) {
          return terrno;
        }

        int32_t size = blockDataGetSize(pDataBlock) + sizeof(int32_t) + pDataBlock->info.numOfCols * sizeof(int32_t);
        assert(size <= getBufPageSize(pHandle->pBuf));

464
        blockDataToBuf(pPage, pDataBlock);
465 466 467 468

        setBufPageDirty(pPage, true);
        releaseBufPage(pHandle->pBuf, pPage);

469
        blockDataCleanup(pDataBlock);
470 471
      }

wmmhello's avatar
wmmhello 已提交
472
      sortComparClearup(&pHandle->cmpParam);
473 474 475
      tMergeTreeDestroy(pHandle->pMergeTree);
      pHandle->numOfCompletedSources = 0;

476
      SSDataBlock* pBlock = createOneDataBlock(pHandle->pDataBlock, false);
477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495
      code = doAddNewExternalMemSource(pHandle->pBuf, pResList, pBlock, &pHandle->sourceId);
      if (code != 0) {
        return code;
      }
    }

    taosArrayClear(pHandle->pOrderedSource);
    taosArrayAddAll(pHandle->pOrderedSource, pResList);
    taosArrayDestroy(pResList);

    numOfSorted = taosArrayGetSize(pHandle->pOrderedSource);

    int64_t el = taosGetTimestampUs() - st;
    pHandle->totalElapsed += el;

    SDiskbasedBufStatis statis = getDBufStatis(pHandle->pBuf);
    qDebug("%s %d round mergesort, elapsed:%"PRId64" readDisk:%.2f Kb, flushDisk:%.2f Kb", pHandle->idStr, t + 1, el, statis.loadBytes/1024.0,
           statis.flushBytes/1024.0);

H
Haojun Liao 已提交
496 497
    if (pHandle->type == SORT_MULTISOURCE_MERGE) {
      pHandle->type = SORT_SINGLESOURCE_SORT;
498 499 500 501 502 503 504 505
      pHandle->comparFn = msortComparFn;
    }
  }

  pHandle->cmpParam.numOfSources = taosArrayGetSize(pHandle->pOrderedSource);
  return 0;
}

506
static int32_t createInitialSources(SSortHandle* pHandle) {
507
  size_t sortBufSize = pHandle->numOfPages * pHandle->pageSize;
H
Haojun Liao 已提交
508 509

  if (pHandle->type == SORT_SINGLESOURCE_SORT) {
wmmhello's avatar
wmmhello 已提交
510
    SSortSource* source = taosArrayGetP(pHandle->pOrderedSource, 0);
H
Haojun Liao 已提交
511
    taosArrayClear(pHandle->pOrderedSource);
512
    
513
    while (1) {
H
Haojun Liao 已提交
514
      SSDataBlock* pBlock = pHandle->fetchfp(source->param);
515 516 517 518 519
      if (pBlock == NULL) {
        break;
      }

      if (pHandle->pDataBlock == NULL) {
520
        pHandle->pDataBlock = createOneDataBlock(pBlock, false);
521 522 523 524 525 526 527 528

        // calculate the buffer pages according to the total available buffers.
        int32_t rowSize = blockDataGetRowSize(pBlock);
        if (rowSize * 4 > 4096) {
          pHandle->pageSize = rowSize * 4;
        } else {
          pHandle->pageSize = 4096;
        }
529
        
530 531 532
        // todo!!
        pHandle->numOfPages = 1024;
        sortBufSize = pHandle->numOfPages * pHandle->pageSize;
533 534
      }

535 536 537 538 539 540
      // perform the scalar function calculation before apply the sort
      if (pHandle->beforeFp != NULL) {
        pHandle->beforeFp(pBlock, pHandle->param);
      }

      // todo relocate the columns
541
      int32_t code = blockDataMerge(pHandle->pDataBlock, pBlock);
542 543 544 545 546 547 548 549
      if (code != 0) {
        return code;
      }

      size_t size = blockDataGetSize(pHandle->pDataBlock);
      if (size > sortBufSize) {
        // Perform the in-memory sort and then flush data in the buffer into disk.
        int64_t p = taosGetTimestampUs();
H
Haojun Liao 已提交
550
        blockDataSort(pHandle->pDataBlock, pHandle->pSortInfo);
551 552 553 554 555 556 557 558

        int64_t el = taosGetTimestampUs() - p;
        pHandle->sortElapsed += el;

        doAddToBuf(pHandle->pDataBlock, pHandle);
      }
    }

559
    if (pHandle->pDataBlock != NULL && pHandle->pDataBlock->info.rows > 0) {
560 561 562
      size_t size = blockDataGetSize(pHandle->pDataBlock);

      // Perform the in-memory sort and then flush data in the buffer into disk.
563 564
      int64_t p = taosGetTimestampUs();

H
Haojun Liao 已提交
565
      blockDataSort(pHandle->pDataBlock, pHandle->pSortInfo);
566

567 568 569
      int64_t el = taosGetTimestampUs() - p;
      pHandle->sortElapsed += el;

570 571 572 573 574
      // All sorted data can fit in memory, external memory sort is not needed. Return to directly
      if (size <= sortBufSize) {
        pHandle->cmpParam.numOfSources = 1;
        pHandle->inMemSort = true;

575
        pHandle->loops = 1;
576 577 578 579 580 581 582
        pHandle->tupleHandle.rowIndex  = -1;
        pHandle->tupleHandle.pBlock = pHandle->pDataBlock;
        return 0;
      } else {
        doAddToBuf(pHandle->pDataBlock, pHandle);
      }
    }
H
Haojun Liao 已提交
583 584 585 586 587
  }

  return TSDB_CODE_SUCCESS;
}

H
Haojun Liao 已提交
588
int32_t tsortOpen(SSortHandle* pHandle) {
H
Haojun Liao 已提交
589 590 591 592 593 594 595 596 597 598
  if (pHandle->opened) {
    return 0;
  }

  if (pHandle->fetchfp == NULL || pHandle->comparFn == NULL) {
    return -1;
  }

  pHandle->opened = true;

599
  int32_t code = createInitialSources(pHandle);
H
Haojun Liao 已提交
600 601
  if (code != TSDB_CODE_SUCCESS) {
    return code;
602 603 604
  }

  // do internal sort
H
Haojun Liao 已提交
605
  code = doInternalMergeSort(pHandle);
606 607 608 609 610
  if (code != TSDB_CODE_SUCCESS) {
    return code;
  }

  int32_t numOfSources = taosArrayGetSize(pHandle->pOrderedSource);
H
Haojun Liao 已提交
611 612 613 614
  if (pHandle->pBuf != NULL) {
    ASSERT(numOfSources <= getNumOfInMemBufPages(pHandle->pBuf));
  }

H
Haojun Liao 已提交
615 616 617 618
  if (numOfSources == 0) {
    return 0;
  }

619 620 621 622 623
  code = sortComparInit(&pHandle->cmpParam, pHandle->pOrderedSource, 0, numOfSources - 1, pHandle);
  if (code != TSDB_CODE_SUCCESS) {
    return code;
  }

624
  return tMergeTreeCreate(&pHandle->pMergeTree, pHandle->cmpParam.numOfSources, &pHandle->cmpParam, pHandle->comparFn);
625 626
}

H
Haojun Liao 已提交
627
int32_t tsortClose(SSortHandle* pHandle) {
628
  // do nothing
629
  return TSDB_CODE_SUCCESS;
630 631
}

632 633 634 635
int32_t tsortSetFetchRawDataFp(SSortHandle* pHandle, _sort_fetch_block_fn_t fetchFp, void (*fp)(SSDataBlock*, void*), void* param) {
  pHandle->fetchfp = fetchFp;
  pHandle->beforeFp = fp;
  pHandle->param = param;
636
  return TSDB_CODE_SUCCESS;
637 638
}

H
Haojun Liao 已提交
639
int32_t tsortSetComparFp(SSortHandle* pHandle, _sort_merge_compar_fn_t fp) {
640
  pHandle->comparFn = fp;
641
  return TSDB_CODE_SUCCESS;
642 643
}

H
Haojun Liao 已提交
644
STupleHandle* tsortNextTuple(SSortHandle* pHandle) {
645 646 647 648
  if (pHandle->cmpParam.numOfSources == pHandle->numOfCompletedSources) {
    return NULL;
  }

649
  // All the data are hold in the buffer, no external sort is invoked.
650 651 652 653 654 655 656 657 658 659 660
  if (pHandle->inMemSort) {
    pHandle->tupleHandle.rowIndex += 1;
    if (pHandle->tupleHandle.rowIndex == pHandle->pDataBlock->info.rows) {
      pHandle->numOfCompletedSources = 1;
      return NULL;
    }

    return &pHandle->tupleHandle;
  }

  int32_t index = tMergeTreeGetChosenIndex(pHandle->pMergeTree);
wmmhello's avatar
wmmhello 已提交
661
  SSortSource *pSource = pHandle->cmpParam.pSources[index];
662 663 664 665 666 667 668 669 670 671 672 673 674

  if (pHandle->needAdjust) {
    int32_t code = adjustMergeTreeForNextTuple(pSource, pHandle->pMergeTree, pHandle, &pHandle->numOfCompletedSources);
    if (code != TSDB_CODE_SUCCESS) {
      terrno = code;
      return NULL;
    }
  }

  if (pHandle->cmpParam.numOfSources == pHandle->numOfCompletedSources) {
    return NULL;
  }

675
  // Get the adjusted value after the loser tree is updated.
676 677 678 679 680 681 682 683 684 685 686 687 688 689
  index = tMergeTreeGetChosenIndex(pHandle->pMergeTree);
  pSource = pHandle->cmpParam.pSources[index];

  assert(pSource->src.pBlock != NULL);

  pHandle->tupleHandle.rowIndex = pSource->src.rowIndex;
  pHandle->tupleHandle.pBlock = pSource->src.pBlock;

  pHandle->needAdjust = true;
  pSource->src.rowIndex += 1;

  return &pHandle->tupleHandle;
}

H
Haojun Liao 已提交
690
bool tsortIsNullVal(STupleHandle* pVHandle, int32_t colIndex) {
691
  SColumnInfoData* pColInfoSrc = taosArrayGet(pVHandle->pBlock->pDataBlock, colIndex);
692
  return colDataIsNull_s(pColInfoSrc, pVHandle->rowIndex);
693 694
}

H
Haojun Liao 已提交
695
void* tsortGetValue(STupleHandle* pVHandle, int32_t colIndex) {
696
  SColumnInfoData* pColInfo = TARRAY_GET_ELEM(pVHandle->pBlock->pDataBlock, colIndex);
697
  return colDataGetData(pColInfo, pVHandle->rowIndex);
698
}
699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715

SSortExecInfo tsortGetSortExecInfo(SSortHandle* pHandle) {
  SSortExecInfo info = {0};

  info.sortBuffer = pHandle->pageSize * pHandle->numOfPages;
  info.sortMethod = pHandle->inMemSort? SORT_QSORT_T:SORT_SPILLED_MERGE_SORT_T;
  info.loops  = pHandle->loops;

  if (pHandle->pBuf != NULL) {
    SDiskbasedBufStatis st = getDBufStatis(pHandle->pBuf);
    info.writeBytes = st.flushBytes;
    info.readBytes = st.loadBytes;
  }

  return info;
}