TDengine/source/dnode/vnode/src/tsdb/tsdbCacheRead.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 "taoserror.h"
#include "tarray.h"
#include "tcommon.h"
#include "tsdb.h"

#define HASTYPE(_type, _t) (((_type) & (_t)) == (_t))

static void saveOneRow(SArray* pRow, SSDataBlock* pBlock, SCacheRowsReader* pReader, const int32_t* slotIds,
                       void** pRes) {
  ASSERT(pReader->numOfCols <= taosArrayGetSize(pBlock->pDataBlock));
  int32_t numOfRows = pBlock->info.rows;

  if (HASTYPE(pReader->type, CACHESCAN_RETRIEVE_LAST)) {
    bool allNullRow = true;

    for (int32_t i = 0; i < pReader->numOfCols; ++i) {
      SColumnInfoData* pColInfoData = taosArrayGet(pBlock->pDataBlock, i);
      SFirstLastRes*   p = (SFirstLastRes*)varDataVal(pRes[i]);

      if (slotIds[i] == -1) {  // the primary timestamp
        SLastCol* pColVal = (SLastCol*)taosArrayGet(pRow, 0);
        p->ts = pColVal->ts;
        p->bytes = TSDB_KEYSIZE;
        *(int64_t*)p->buf = pColVal->ts;
        allNullRow = false;
      } else {
        int32_t   slotId = slotIds[i];
        SLastCol* pColVal = (SLastCol*)taosArrayGet(pRow, slotId);

        p->ts = pColVal->ts;
        p->isNull = !COL_VAL_IS_VALUE(&pColVal->colVal);
        allNullRow = p->isNull & allNullRow;

        if (!p->isNull) {
          if (IS_VAR_DATA_TYPE(pColVal->colVal.type)) {
            varDataSetLen(p->buf, pColVal->colVal.value.nData);
            memcpy(varDataVal(p->buf), pColVal->colVal.value.pData, pColVal->colVal.value.nData);
            p->bytes = pColVal->colVal.value.nData + VARSTR_HEADER_SIZE;  // binary needs to plus the header size
          } else {
            memcpy(p->buf, &pColVal->colVal.value, pReader->pSchema->columns[slotId].bytes);
            p->bytes = pReader->pSchema->columns[slotId].bytes;
          }
        }
      }

      // pColInfoData->info.bytes includes the VARSTR_HEADER_SIZE, need to substruct it
      p->hasResult = true;
      varDataSetLen(pRes[i], pColInfoData->info.bytes - VARSTR_HEADER_SIZE);
      colDataAppend(pColInfoData, numOfRows, (const char*)pRes[i], false);
    }

    pBlock->info.rows += allNullRow? 0:1;
  } else {
    ASSERT(HASTYPE(pReader->type, CACHESCAN_RETRIEVE_LAST_ROW));

    for (int32_t i = 0; i < pReader->numOfCols; ++i) {
      SColumnInfoData* pColInfoData = taosArrayGet(pBlock->pDataBlock, i);

      if (slotIds[i] == -1) {
        SLastCol* pColVal = (SLastCol*)taosArrayGet(pRow, 0);
        colDataAppend(pColInfoData, numOfRows, (const char*)&pColVal->ts, false);
      } else {
        int32_t   slotId = slotIds[i];
        SLastCol* pColVal = (SLastCol*)taosArrayGet(pRow, slotId);
        SColVal*  pVal = &pColVal->colVal;

        if (IS_VAR_DATA_TYPE(pColVal->colVal.type)) {
          if (!COL_VAL_IS_VALUE(&pColVal->colVal)) {
            colDataAppendNULL(pColInfoData, numOfRows);
          } else {
            varDataSetLen(pReader->transferBuf[slotId], pVal->value.nData);
            memcpy(varDataVal(pReader->transferBuf[slotId]), pVal->value.pData, pVal->value.nData);
            colDataAppend(pColInfoData, numOfRows, pReader->transferBuf[slotId], false);
          }
        } else {
          colDataAppend(pColInfoData, numOfRows, (const char*)&pVal->value.val, !COL_VAL_IS_VALUE(pVal));
        }
      }
    }

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

int32_t tsdbCacherowsReaderOpen(void* pVnode, int32_t type, SArray* pTableIdList, int32_t numOfCols, uint64_t suid,
                                void** pReader) {
  *pReader = NULL;

  SCacheRowsReader* p = taosMemoryCalloc(1, sizeof(SCacheRowsReader));
  if (p == NULL) {
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  p->type = type;
  p->pVnode = pVnode;
  p->numOfCols = numOfCols;
  p->suid = suid;

  if (taosArrayGetSize(pTableIdList) == 0) {
    *pReader = p;
    return TSDB_CODE_SUCCESS;
  }

  STableKeyInfo* pKeyInfo = taosArrayGet(pTableIdList, 0);
  p->pSchema = metaGetTbTSchema(p->pVnode->pMeta, pKeyInfo->uid, -1, 1);
  p->pTableList = pTableIdList;

  p->transferBuf = taosMemoryCalloc(p->pSchema->numOfCols, POINTER_BYTES);
  if (p->transferBuf == NULL) {
    tsdbCacherowsReaderClose(p);
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  for (int32_t i = 0; i < p->pSchema->numOfCols; ++i) {
    if (IS_VAR_DATA_TYPE(p->pSchema->columns[i].type)) {
      p->transferBuf[i] = taosMemoryMalloc(p->pSchema->columns[i].bytes);
      if (p->transferBuf[i] == NULL) {
        tsdbCacherowsReaderClose(p);
        return TSDB_CODE_OUT_OF_MEMORY;
      }
    }
  }

  p->pLoadInfo = tCreateLastBlockLoadInfo(p->pSchema, NULL, 0);
  if (p->pLoadInfo == NULL) {
    tsdbCacherowsReaderClose(p);
    return TSDB_CODE_OUT_OF_MEMORY;
  }

  *pReader = p;
  return TSDB_CODE_SUCCESS;
}

void* tsdbCacherowsReaderClose(void* pReader) {
  SCacheRowsReader* p = pReader;

  if (p->pSchema != NULL) {
    for (int32_t i = 0; i < p->pSchema->numOfCols; ++i) {
      taosMemoryFreeClear(p->transferBuf[i]);
    }

    taosMemoryFree(p->transferBuf);
    taosMemoryFree(p->pSchema);
  }

  destroyLastBlockLoadInfo(p->pLoadInfo);

  taosMemoryFree(pReader);
  return NULL;
}

static int32_t doExtractCacheRow(SCacheRowsReader* pr, SLRUCache* lruCache, uint64_t uid, SArray** pRow,
                                 LRUHandle** h) {
  int32_t code = TSDB_CODE_SUCCESS;
  *pRow = NULL;

  if (HASTYPE(pr->type, CACHESCAN_RETRIEVE_LAST_ROW)) {
    code = tsdbCacheGetLastrowH(lruCache, uid, pr, h);
  } else {
    code = tsdbCacheGetLastH(lruCache, uid, pr, h);
  }

  if (code != TSDB_CODE_SUCCESS) {
    return code;
  }

  // no data in the table of Uid
  if (*h != NULL) {
    *pRow = (SArray*)taosLRUCacheValue(lruCache, *h);
  }

  return code;
}

static void freeItem(void* pItem) {
  SLastCol* pCol = (SLastCol*)pItem;
  if (IS_VAR_DATA_TYPE(pCol->colVal.type)) {
    taosMemoryFree(pCol->colVal.value.pData);
  }
}

int32_t tsdbRetrieveCacheRows(void* pReader, SSDataBlock* pResBlock, const int32_t* slotIds, SArray* pTableUidList) {
  if (pReader == NULL || pResBlock == NULL) {
    return TSDB_CODE_INVALID_PARA;
  }

  SCacheRowsReader* pr = pReader;

  int32_t    code = TSDB_CODE_SUCCESS;
  SLRUCache* lruCache = pr->pVnode->pTsdb->lruCache;
  LRUHandle* h = NULL;
  SArray*    pRow = NULL;
  size_t     numOfTables = taosArrayGetSize(pr->pTableList);
  bool       hasRes = false;
  SArray*    pLastCols = NULL;

  void** pRes = taosMemoryCalloc(pr->numOfCols, POINTER_BYTES);
  if (pRes == NULL) {
    code = TSDB_CODE_OUT_OF_MEMORY;
    goto _end;
  }

  for (int32_t j = 0; j < pr->numOfCols; ++j) {
    pRes[j] = taosMemoryCalloc(1, sizeof(SFirstLastRes) + pr->pSchema->columns[slotIds[j]].bytes + VARSTR_HEADER_SIZE);
    SFirstLastRes* p = (SFirstLastRes*)varDataVal(pRes[j]);
    p->ts = INT64_MIN;
  }

  pLastCols = taosArrayInit(pr->pSchema->numOfCols, sizeof(SLastCol));
  if (pLastCols == NULL) {
    code = TSDB_CODE_OUT_OF_MEMORY;
    goto _end;
  }

  for (int32_t i = 0; i < pr->pSchema->numOfCols; ++i) {
    struct STColumn* pCol = &pr->pSchema->columns[i];
    SLastCol         p = {.ts = INT64_MIN, .colVal.type = pCol->type};

    if (IS_VAR_DATA_TYPE(pCol->type)) {
      p.colVal.value.pData = taosMemoryCalloc(pCol->bytes, sizeof(char));
    }
    taosArrayPush(pLastCols, &p);
  }

  tsdbTakeReadSnap(pr->pVnode->pTsdb, &pr->pReadSnap, "cache-l");
  pr->pDataFReader = NULL;

  // retrieve the only one last row of all tables in the uid list.
  if (HASTYPE(pr->type, CACHESCAN_RETRIEVE_TYPE_SINGLE)) {
    for (int32_t i = 0; i < numOfTables; ++i) {
      STableKeyInfo* pKeyInfo = taosArrayGet(pr->pTableList, i);

      code = doExtractCacheRow(pr, lruCache, pKeyInfo->uid, &pRow, &h);
      if (code != TSDB_CODE_SUCCESS) {
        return code;
      }

      if (h == NULL) {
        continue;
      }

      {
        for (int32_t k = 0; k < pr->numOfCols; ++k) {
          int32_t slotId = slotIds[k];

          if (slotId == -1) {  // the primary timestamp
            SLastCol* p = taosArrayGet(pLastCols, 0);
            SLastCol* pCol = (SLastCol*)taosArrayGet(pRow, 0);
            if (pCol->ts > p->ts) {
              hasRes = true;
              p->ts = pCol->ts;
              p->colVal = pCol->colVal;

              // only set value for last row query
              if (HASTYPE(pr->type, CACHESCAN_RETRIEVE_LAST_ROW)) {
                if (taosArrayGetSize(pTableUidList) == 0) {
                  taosArrayPush(pTableUidList, &pKeyInfo->uid);
                } else {
                  taosArraySet(pTableUidList, 0, &pKeyInfo->uid);
                }
              }
            }
          } else {
            SLastCol* p = taosArrayGet(pLastCols, slotId);
            SLastCol* pColVal = (SLastCol*)taosArrayGet(pRow, slotId);

            if (pColVal->ts > p->ts) {
              if (!COL_VAL_IS_VALUE(&pColVal->colVal) && HASTYPE(pr->type, CACHESCAN_RETRIEVE_LAST)) {
                continue;
              }

              hasRes = true;
              p->ts = pColVal->ts;

              uint8_t* px = p->colVal.value.pData;
              p->colVal = pColVal->colVal;

              if (COL_VAL_IS_VALUE(&pColVal->colVal) && IS_VAR_DATA_TYPE(pColVal->colVal.type)) {
                p->colVal.value.pData = px;
                memcpy(px, pColVal->colVal.value.pData, pColVal->colVal.value.nData);
              }
            }
          }
        }
      }

      tsdbCacheRelease(lruCache, h);
    }

    if (hasRes) {
      saveOneRow(pLastCols, pResBlock, pr, slotIds, pRes);
    }

  } else if (HASTYPE(pr->type, CACHESCAN_RETRIEVE_TYPE_ALL)) {
    for (int32_t i = pr->tableIndex; i < numOfTables; ++i) {
      STableKeyInfo* pKeyInfo = (STableKeyInfo*)taosArrayGet(pr->pTableList, i);
      code = doExtractCacheRow(pr, lruCache, pKeyInfo->uid, &pRow, &h);
      if (code != TSDB_CODE_SUCCESS) {
        return code;
      }

      if (h == NULL) {
        continue;
      }

      saveOneRow(pRow, pResBlock, pr, slotIds, pRes);
      // TODO reset the pRes

      taosArrayPush(pTableUidList, &pKeyInfo->uid);
      tsdbCacheRelease(lruCache, h);

      pr->tableIndex += 1;
      if (pResBlock->info.rows >= pResBlock->info.capacity) {
        goto _end;
      }
    }
  } else {
    code = TSDB_CODE_INVALID_PARA;
  }

_end:
  tsdbDataFReaderClose(&pr->pDataFReader);

  tsdbUntakeReadSnap(pr->pVnode->pTsdb, pr->pReadSnap, "cache-l");

  for (int32_t j = 0; j < pr->numOfCols; ++j) {
    taosMemoryFree(pRes[j]);
  }

  taosMemoryFree(pRes);
  taosArrayDestroyEx(pLastCols, freeItem);
  return code;
}