[Join] Zero-copy storage column check.

Devjiu · Devjiu · commit b42d496d0343 · 2023-08-09T16:09:30.000Z
This commit adds check for valid pointer to column buffer to skip copying. Most effective with enabled `enable-non-lazy-data-import` option. Checks number of chunks in storage. Partially resolves: #574 Signed-off-by: Dmitrii Makarenko <dmitrii.makarenko@intel.com>
diff --git a/omniscidb/ArrowStorage/ArrowStorage.cpp b/omniscidb/ArrowStorage/ArrowStorage.cpp
@@ -187,6 +187,51 @@ std::unique_ptr<AbstractDataToken> ArrowStorage::getZeroCopyBufferMemory(
   return nullptr;
 }
 
+std::unique_ptr<AbstractDataToken> ArrowStorage::getZeroCopyColumnData(
+    const ColumnRef& col_ref) {
+  mapd_shared_lock<mapd_shared_mutex> data_lock(data_mutex_);
+  CHECK_EQ(col_ref.db_id, db_id_);
+  CHECK_EQ(tables_.count(col_ref.table_id), (size_t)1);
+  auto& table = *tables_.at(col_ref.table_id);
+  mapd_shared_lock<mapd_shared_mutex> table_lock(table.mutex);
+  data_lock.unlock();
+
+  auto col_type = getColumnInfo(col_ref.db_id, col_ref.table_id, col_ref.column_id)->type;
+
+  if (col_type->isExtDictionary()) {
+    auto dict_id = col_type->as<hdk::ir::ExtDictionaryType>()->dictId();
+    auto dict_descriptor = getDictMetadata(
+        dict_id);  // this will force materialize the dictionary. it is thread safe
+    CHECK(dict_descriptor);
+  }
+
+  if (!col_type->isVarLen()) {
+    size_t col_idx = columnIndex(col_ref.column_id);
+    if (col_idx >= table.col_data.size()) {
+      return nullptr;
+    }
+    size_t elem_size = col_type->size();
+    const auto* fixed_type =
+        dynamic_cast<const arrow::FixedWidthType*>(table.col_data[col_idx]->type().get());
+    CHECK(fixed_type) << table.col_data[col_idx]->type()->ToString() << " (table "
+                      << col_ref.table_id << ", column " << col_idx << ")";
+    size_t arrow_elem_size = fixed_type->bit_width() / 8;
+    size_t elems = elem_size / arrow_elem_size;
+    CHECK_GT(elems, (size_t)0);
+    auto data_to_fetch = table.col_data[col_idx];
+    if (data_to_fetch->num_chunks() == 1) {
+      auto chunk = data_to_fetch->chunk(0);
+      const int8_t* ptr =
+          chunk->data()->GetValues<int8_t>(1, chunk->data()->offset * arrow_elem_size);
+      size_t chunk_size = chunk->length() * arrow_elem_size;
+      return std::make_unique<ArrowChunkDataToken>(
+          std::move(chunk), col_type, ptr, chunk_size);
+    }
+  }
+
+  return nullptr;
+}
+
 void ArrowStorage::fetchFixedLenData(const TableData& table,
                                      size_t frag_idx,
                                      size_t col_idx,
diff --git a/omniscidb/ArrowStorage/ArrowStorage.h b/omniscidb/ArrowStorage/ArrowStorage.h
@@ -70,6 +70,9 @@ class ArrowStorage : public SimpleSchemaProvider, public AbstractDataProvider {
       const ChunkKey& key,
       size_t num_bytes) override;
 
+  std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override;
+
   TableFragmentsInfo getTableMetadata(int db_id, int table_id) const override;
 
   const DictDescriptor* getDictMetadata(int dict_id, bool load_dict = true) override;
diff --git a/omniscidb/DataMgr/AbstractBufferMgr.h b/omniscidb/DataMgr/AbstractBufferMgr.h
@@ -86,6 +86,8 @@ class AbstractBufferMgr {
   virtual AbstractBuffer* getBuffer(const ChunkKey& key, const size_t numBytes = 0) = 0;
   virtual std::unique_ptr<AbstractDataToken> getZeroCopyBufferMemory(const ChunkKey& key,
                                                                      size_t numBytes) = 0;
+  virtual std::unique_ptr<AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) = 0;
   virtual void fetchBuffer(const ChunkKey& key,
                            AbstractBuffer* destBuffer,
                            const size_t numBytes = 0) = 0;
diff --git a/omniscidb/DataMgr/AbstractDataProvider.h b/omniscidb/DataMgr/AbstractDataProvider.h
@@ -45,6 +45,13 @@ class AbstractDataProvider : public Data_Namespace::AbstractBufferMgr {
     return nullptr;
   }
 
+  // TODO(dmitriim) remove this method after enabling
+  // of hashtable, that takes into a count frag_id and offset
+  std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override {
+    return nullptr;
+  }
+
   void deleteBuffer(const ChunkKey& key, const bool purge = true) override {
     UNREACHABLE();
   }
diff --git a/omniscidb/DataMgr/BufferMgr/BufferMgr.cpp b/omniscidb/DataMgr/BufferMgr/BufferMgr.cpp
@@ -865,6 +865,11 @@ std::unique_ptr<AbstractDataToken> BufferMgr::getZeroCopyBufferMemory(const Chun
   return parent_mgr_->getZeroCopyBufferMemory(key, numBytes);
 }
 
+std::unique_ptr<AbstractDataToken> BufferMgr::getZeroCopyColumnData(
+    const ColumnRef& col_ref) {
+  return parent_mgr_->getZeroCopyColumnData(col_ref);
+}
+
 MemoryInfo BufferMgr::getMemoryInfo() {
   std::unique_lock<std::mutex> sized_segs_lock(sized_segs_mutex_);
   MemoryInfo mi;
diff --git a/omniscidb/DataMgr/BufferMgr/BufferMgr.h b/omniscidb/DataMgr/BufferMgr/BufferMgr.h
@@ -162,6 +162,9 @@ class BufferMgr : public AbstractBufferMgr {  // implements
   std::unique_ptr<AbstractDataToken> getZeroCopyBufferMemory(const ChunkKey& key,
                                                              size_t numBytes) override;
 
+  std::unique_ptr<AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override;
+
   /**
    * @brief Puts the contents of d into the Buffer with ChunkKey key.
    * @param key - Unique identifier for a Chunk.
diff --git a/omniscidb/DataMgr/DataMgr.cpp b/omniscidb/DataMgr/DataMgr.cpp
@@ -458,6 +458,13 @@ AbstractBuffer* DataMgr::getChunkBuffer(const ChunkKey& key,
   return bufferMgrs_[level][deviceId]->getBuffer(key, numBytes);
 }
 
+std::unique_ptr<AbstractDataToken> DataMgr::getZeroCopyColumnData(
+    const ColumnRef& col_ref) {
+  const auto level = static_cast<size_t>(Data_Namespace::CPU_LEVEL);
+  CHECK_LT(level, levelSizes_.size());  // make sure we have a legit buffermgr
+  return bufferMgrs_[level][0]->getZeroCopyColumnData(col_ref);
+}
+
 void DataMgr::deleteChunksWithPrefix(const ChunkKey& keyPrefix) {
   int numLevels = bufferMgrs_.size();
   for (int level = numLevels - 1; level >= 0; --level) {
diff --git a/omniscidb/DataMgr/DataMgr.h b/omniscidb/DataMgr/DataMgr.h
@@ -177,6 +177,8 @@ class DataMgr {
                                  const MemoryLevel memoryLevel,
                                  const int deviceId = 0,
                                  const size_t numBytes = 0);
+  // TODO(dmitriim) remove this method after enabling of hashtable
+  std::unique_ptr<AbstractDataToken> getZeroCopyColumnData(const ColumnRef& col_ref);
   void deleteChunksWithPrefix(const ChunkKey& keyPrefix);
   void deleteChunksWithPrefix(const ChunkKey& keyPrefix, const MemoryLevel memLevel);
   AbstractBuffer* alloc(const MemoryLevel memoryLevel,
diff --git a/omniscidb/DataMgr/DataMgrDataProvider.cpp b/omniscidb/DataMgr/DataMgrDataProvider.cpp
@@ -29,9 +29,16 @@ std::shared_ptr<Chunk_NS::Chunk> DataMgrDataProvider::getChunk(
   return Chunk_NS::Chunk::getChunk(
       col_info, data_mgr_, key, memory_level, device_id, num_bytes, num_elems);
 }
+
+std::unique_ptr<Data_Namespace::AbstractDataToken>
+DataMgrDataProvider::getZeroCopyColumnData(const ColumnRef& col_ref) {
+  return data_mgr_->getZeroCopyColumnData(col_ref);
+}
+
 TableFragmentsInfo DataMgrDataProvider::getTableMetadata(int db_id, int table_id) const {
   return data_mgr_->getTableMetadata(db_id, table_id);
 }
+
 const DictDescriptor* DataMgrDataProvider::getDictMetadata(int dict_id,
                                                            bool load_dict) const {
   return data_mgr_->getDictMetadata(dict_id, load_dict);
diff --git a/omniscidb/DataMgr/DataMgrDataProvider.h b/omniscidb/DataMgr/DataMgrDataProvider.h
@@ -35,6 +35,9 @@ class DataMgrDataProvider : public DataProvider {
       const size_t num_bytes,
       const size_t num_elems) override;
 
+  std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override;
+
   TableFragmentsInfo getTableMetadata(int db_id, int table_id) const override;
 
   const DictDescriptor* getDictMetadata(int dict_id,
diff --git a/omniscidb/DataMgr/PersistentStorageMgr/PersistentStorageMgr.cpp b/omniscidb/DataMgr/PersistentStorageMgr/PersistentStorageMgr.cpp
@@ -55,6 +55,11 @@ std::unique_ptr<AbstractDataToken> PersistentStorageMgr::getZeroCopyBufferMemory
   return getStorageMgrForTableKey(key)->getZeroCopyBufferMemory(key, numBytes);
 }
 
+std::unique_ptr<AbstractDataToken> PersistentStorageMgr::getZeroCopyColumnData(
+    const ColumnRef& col_ref) {
+  return getStorageMgr(col_ref.db_id)->getZeroCopyColumnData(col_ref);
+}
+
 void PersistentStorageMgr::fetchBuffer(const ChunkKey& chunk_key,
                                        AbstractBuffer* destination_buffer,
                                        const size_t num_bytes) {
diff --git a/omniscidb/DataMgr/PersistentStorageMgr/PersistentStorageMgr.h b/omniscidb/DataMgr/PersistentStorageMgr/PersistentStorageMgr.h
@@ -35,6 +35,8 @@ class PersistentStorageMgr : public AbstractBufferMgr {
   AbstractBuffer* getBuffer(const ChunkKey& chunk_key, const size_t num_bytes) override;
   std::unique_ptr<AbstractDataToken> getZeroCopyBufferMemory(const ChunkKey& key,
                                                              size_t numBytes) override;
+  std::unique_ptr<AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override;
   void fetchBuffer(const ChunkKey& chunk_key,
                    AbstractBuffer* destination_buffer,
                    const size_t num_bytes) override;
diff --git a/omniscidb/DataProvider/DataProvider.h b/omniscidb/DataProvider/DataProvider.h
@@ -36,6 +36,11 @@ class DataProvider {
       const size_t num_bytes,
       const size_t num_elems) = 0;
 
+  // CPU only
+  // TODO(dmitriim) remove this method after enabling of hashtable
+  virtual std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) = 0;
+
   virtual TableFragmentsInfo getTableMetadata(int db_id, int table_id) const = 0;
 
   virtual const DictDescriptor* getDictMetadata(int dict_id,
diff --git a/omniscidb/QueryEngine/ColumnFetcher.cpp b/omniscidb/QueryEngine/ColumnFetcher.cpp
@@ -250,6 +250,19 @@ const int8_t* ColumnFetcher::getAllTableColumnFragments(
   const InputDescriptor table_desc(db_id, table_id, int(0));
   {
     std::lock_guard<std::mutex> columnar_conversion_guard(columnar_fetch_mutex_);
+
+    auto col_token = data_provider_->getZeroCopyColumnData(*col_info);
+    if (col_token != nullptr) {
+      size_t num_rows = col_token->getSize() / col_token->getType()->size();
+      auto raw_mem_ptr =
+          executor_->row_set_mem_owner_->saveDataToken(std::move(col_token));
+      ColumnarResults res(
+          {const_cast<int8_t*>(raw_mem_ptr)}, num_rows, col_info->type, thread_idx);
+
+      return ColumnFetcher::transferColumnIfNeeded(
+          &res, 0, memory_level, device_id, device_allocator);
+    }
+
     auto column_it = columnarized_scan_table_cache_.find({table_id, col_id});
     if (column_it == columnarized_scan_table_cache_.end()) {
       for (size_t frag_id = 0; frag_id < frag_count; ++frag_id) {
diff --git a/omniscidb/ResultSet/RowSetMemoryOwner.h b/omniscidb/ResultSet/RowSetMemoryOwner.h
@@ -97,6 +97,16 @@ class RowSetMemoryOwner final : public SimpleAllocator, boost::noncopyable {
     varlen_buffers_.push_back(varlen_buffer);
   }
 
+  const int8_t* saveDataToken(std::unique_ptr<AbstractDataToken>&& token) {
+    std::lock_guard<std::mutex> lock(state_mutex_);
+    auto mem_ptr = token->getMemoryPtr();
+    if (data_tokens_.count(mem_ptr) == 0) {
+      data_tokens_.insert({mem_ptr, std::move(token)});
+    }
+
+    return data_tokens_[mem_ptr]->getMemoryPtr();
+  }
+
   /**
    * Adds a GPU buffer containing a variable length input column. Variable length inputs
    * on GPU are referenced in output projected targets and should not be freed until the
@@ -258,6 +268,7 @@ class RowSetMemoryOwner final : public SimpleAllocator, boost::noncopyable {
   std::vector<void*> col_buffers_;
   std::vector<Data_Namespace::AbstractBuffer*> varlen_input_buffers_;
   std::vector<std::unique_ptr<quantile::TDigest>> t_digests_;
+  std::unordered_map<const int8_t*, std::unique_ptr<AbstractDataToken>> data_tokens_;
 
   DataProvider* data_provider_;  // for metadata lookups
   size_t arena_block_size_;      // for cloning
diff --git a/omniscidb/ResultSetRegistry/ColumnarResults.cpp b/omniscidb/ResultSetRegistry/ColumnarResults.cpp
@@ -145,6 +145,25 @@ ColumnarResults::ColumnarResults(std::shared_ptr<RowSetMemoryOwner> row_set_mem_
   memcpy(((void*)column_buffers_[0]), one_col_buffer, buf_size);
 }
 
+ColumnarResults::ColumnarResults(const std::vector<int8_t*>&& one_col_buffer,
+                                 const size_t num_rows,
+                                 const hdk::ir::Type* target_type,
+                                 const size_t thread_idx)
+    : column_buffers_(1)
+    , num_rows_(num_rows)
+    , target_types_{target_type}
+    , parallel_conversion_(false)
+    , direct_columnar_conversion_(false)
+    , thread_idx_(thread_idx) {
+  auto timer = DEBUG_TIMER(__func__);
+  const bool is_varlen = target_type->isArray() || target_type->isString();
+
+  if (is_varlen) {
+    throw ColumnarConversionNotSupported();
+  }
+  column_buffers_ = std::move(one_col_buffer);
+}
+
 std::unique_ptr<ColumnarResults> ColumnarResults::mergeResults(
     std::shared_ptr<RowSetMemoryOwner> row_set_mem_owner,
     const std::vector<std::unique_ptr<ColumnarResults>>& sub_results) {
diff --git a/omniscidb/ResultSetRegistry/ColumnarResults.h b/omniscidb/ResultSetRegistry/ColumnarResults.h
@@ -73,6 +73,11 @@ class ColumnarResults {
                   const hdk::ir::Type* target_type,
                   const size_t thread_idx);
 
+  ColumnarResults(const std::vector<int8_t*>&& one_col_buffer,
+                  const size_t num_rows,
+                  const hdk::ir::Type* target_type,
+                  const size_t thread_idx);
+
   static std::unique_ptr<ColumnarResults> mergeResults(
       const std::shared_ptr<RowSetMemoryOwner> row_set_mem_owner,
       const std::vector<std::unique_ptr<ColumnarResults>>& sub_results);

Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,13 @@ class AbstractDataProvider : public Data_Namespace::AbstractBufferMgr {`
`45`	`45`	`return nullptr;`
`46`	`46`	`}`
`47`	`47`
	`48`	`+ // TODO(dmitriim) remove this method after enabling`
	`49`	`+ // of hashtable, that takes into a count frag_id and offset`
	`50`	`+ std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(`
	`51`	`+ const ColumnRef& col_ref) override {`
	`52`	`+ return nullptr;`
	`53`	`+ }`
	`54`	`+`
`48`	`55`	`void deleteBuffer(const ChunkKey& key, const bool purge = true) override {`
`49`	`56`	`UNREACHABLE();`
`50`	`57`	`}`