[Join] Remove redundant copies.

Devjiu · Devjiu · commit 40fd5486c6e7 · 2023-08-07T14:10:19.000Z
This commit removes useless copying(memcpy) in `getAllTableColumnFragments`. Also some parallelization added. 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);
+    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;
+    // For fixed size arrays we simply use elem type in arrow and therefore have to scale
+    // to get a proper slice.
+    size_t elems = elem_size / arrow_elem_size;
+    CHECK_GT(elems, (size_t)0);
+    auto data_to_fetch = table.col_data[col_idx];
+    LOG(ERROR) << "getZeroCopyColumnData num_chunks: " << data_to_fetch->num_chunks();
+    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
@@ -84,8 +84,11 @@ class AbstractBufferMgr {
   virtual void deleteBuffersWithPrefix(const ChunkKey& keyPrefix,
                                        const bool purge = true) = 0;
   virtual AbstractBuffer* getBuffer(const ChunkKey& key, const size_t numBytes = 0) = 0;
+  // virtual AbstractBuffer* getZeroCopyColumnBuffer(const ColumnRef &col_ref) = 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
@@ -42,6 +42,13 @@ class AbstractDataProvider : public Data_Namespace::AbstractBufferMgr {
   std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyBufferMemory(
       const ChunkKey& key,
       size_t numBytes) override {
+    UNREACHABLE();
+    return nullptr;
+  }
+
+  std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override {
+    UNREACHABLE();
     return nullptr;
   }
 
diff --git a/omniscidb/DataMgr/BufferMgr/BufferMgr.cpp b/omniscidb/DataMgr/BufferMgr/BufferMgr.cpp
@@ -865,6 +865,20 @@ 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);
+}
+
+// AbstractBuffer* BufferMgr::getZeroCopyColumnBuffer(const ColumnRef& col_ref) {
+//   AbstractBuffer* res = nullptr;
+
+//   if (auto token = getZeroCopyColumnData(col_ref)) {
+//     res = createZeroCopyBuffer({}, std::move(token));
+//   }
+//   return res;
+// }
+
 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
@@ -159,8 +159,12 @@ class BufferMgr : public AbstractBufferMgr {  // implements
   /// Returns the a pointer to the chunk with the specified key.
   AbstractBuffer* getBuffer(const ChunkKey& key, const size_t num_bytes = 0) override;
 
+  // AbstractBuffer* BufferMgr::getZeroCopyColumnBuffer(const ColumnRef &col_ref);
+
   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.
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,7 @@ class DataMgr {
                                  const MemoryLevel memoryLevel,
                                  const int deviceId = 0,
                                  const size_t numBytes = 0);
+  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,10 @@ class DataProvider {
       const size_t num_bytes,
       const size_t num_elems) = 0;
 
+  // CPU only
+  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
@@ -16,6 +16,7 @@
 
 #include "QueryEngine/ColumnFetcher.h"
 
+#include <tbb/parallel_for.h>
 #include <memory>
 
 #include "DataMgr/ArrayNoneEncoder.h"
@@ -238,6 +239,14 @@ const int8_t* ColumnFetcher::getAllTableColumnFragments(
     const int device_id,
     DeviceAllocator* device_allocator,
     const size_t thread_idx) const {
+  const auto& col_token = data_provider_->getZeroCopyColumnData(*col_info);
+
+  if (col_token != nullptr) {
+    // executor_->row_set_mem_owner_->addColBuffer(col_buffers);
+    return col_token->getMemoryPtr();
+  }
+  // throw std::runtime_error("Didn't worked");
+
   int db_id = col_info->db_id;
   int table_id = col_info->table_id;
   int col_id = col_info->column_id;
@@ -251,44 +260,93 @@ const int8_t* ColumnFetcher::getAllTableColumnFragments(
   {
     std::lock_guard<std::mutex> columnar_conversion_guard(columnar_fetch_mutex_);
     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) {
-        if (executor_->getConfig()
-                .exec.interrupt.enable_non_kernel_time_query_interrupt &&
-            executor_->checkNonKernelTimeInterrupted()) {
-          throw QueryExecutionError(Executor::ERR_INTERRUPTED);
-        }
-        std::list<std::shared_ptr<Chunk_NS::Chunk>> chunk_holder;
-        std::list<ChunkIter> chunk_iter_holder;
-        const auto& fragment = (*fragments)[frag_id];
-        if (fragment.isEmptyPhysicalFragment()) {
-          continue;
+    if (column_it != columnarized_scan_table_cache_.end()) {
+      table_column = column_it->second.get();
+      return ColumnFetcher::transferColumnIfNeeded(
+          table_column, 0, memory_level, device_id, device_allocator);
+    }
+
+    size_t total_row_count = 0;
+    for (size_t frag_id = 0; frag_id < frag_count; ++frag_id) {
+      if (executor_->getConfig().exec.interrupt.enable_non_kernel_time_query_interrupt &&
+          executor_->checkNonKernelTimeInterrupted()) {
+        throw QueryExecutionError(Executor::ERR_INTERRUPTED);
+      }
+      const auto& fragment = (*fragments)[frag_id];
+      const auto& rows_in_frag = fragment.getNumTuples();
+      total_row_count += rows_in_frag;
+    }
+
+    const auto& type_width = col_info->type->size();
+    auto write_ptr =
+        executor_->row_set_mem_owner_->allocate(type_width * total_row_count);
+    tbb::concurrent_vector<std::pair<int8_t*, size_t>> write_ptrs;
+    for (size_t frag_id = 0; frag_id < frag_count; ++frag_id) {
+      const auto& fragment = (*fragments)[frag_id];
+      if (!fragment.getNumTuples()) {
+        continue;
+      }
+      CHECK_EQ(type_width, fragment.getChunkMetadataMap().at(col_id)->type()->size());
+      write_ptrs.push_back({write_ptr, fragment.getNumTuples() * type_width});
+      write_ptr += fragment.getNumTuples() * type_width;
+    }
+
+    CHECK_EQ(frag_count, write_ptrs.size());
+    for (size_t frag_id = 0; frag_id < frag_count; ++frag_id) {
+      std::list<std::shared_ptr<Chunk_NS::Chunk>> chunk_holder;
+      std::list<ChunkIter> chunk_iter_holder;
+      const auto& fragment = (*fragments)[frag_id];
+      if (fragment.isEmptyPhysicalFragment()) {
+        continue;
+      }
+      auto chunk_meta_it = fragment.getChunkMetadataMap().find(col_id);
+      CHECK(chunk_meta_it != fragment.getChunkMetadataMap().end());
+      std::shared_ptr<Chunk_NS::Chunk> chunk;
+      // Fixed length arrays are also included here.
+      const bool is_varlen = col_info->type->isString() || col_info->type->isArray();
+      int8_t* col_buffer;
+      {
+        ChunkKey chunk_key{db_id, fragment.physicalTableId, col_id, fragment.fragmentId};
+        std::unique_ptr<std::lock_guard<std::mutex>> varlen_chunk_lock;
+        if (is_varlen) {
+          varlen_chunk_lock.reset(
+              new std::lock_guard<std::mutex>(varlen_chunk_fetch_mutex_));
         }
-        auto chunk_meta_it = fragment.getChunkMetadataMap().find(col_id);
+        chunk = data_provider_->getChunk(col_info,
+                                         chunk_key,
+                                         Data_Namespace::CPU_LEVEL,
+                                         0,
+                                         chunk_meta_it->second->numBytes(),
+                                         chunk_meta_it->second->numElements());
+        std::lock_guard<std::mutex> chunk_list_lock(chunk_list_mutex_);
+        chunk_holder.push_back(chunk);
+      }
+      if (is_varlen) {
+        CHECK_GT(table_id, 0);
         CHECK(chunk_meta_it != fragment.getChunkMetadataMap().end());
-        auto col_buffer = getOneTableColumnFragment(col_info,
-                                                    static_cast<int>(frag_id),
-                                                    all_tables_fragments,
-                                                    chunk_holder,
-                                                    chunk_iter_holder,
-                                                    Data_Namespace::CPU_LEVEL,
-                                                    int(0),
-                                                    device_allocator);
-        column_frags.push_back(
-            std::make_unique<ColumnarResults>(executor_->row_set_mem_owner_,
-                                              col_buffer,
-                                              fragment.getNumTuples(),
-                                              chunk_meta_it->second->type(),
-                                              thread_idx));
+        chunk_iter_holder.push_back(chunk->begin_iterator(chunk_meta_it->second));
+        auto& chunk_iter = chunk_iter_holder.back();
+        col_buffer = reinterpret_cast<int8_t*>(&chunk_iter);
+      } else {
+        auto ab = chunk->getBuffer();
+        CHECK(ab->getMemoryPtr());
+        col_buffer = ab->getMemoryPtr();  // @TODO(alex) change to use ChunkIter
       }
-      auto merged_results =
-          ColumnarResults::mergeResults(executor_->row_set_mem_owner_, column_frags);
-      table_column = merged_results.get();
-      columnarized_scan_table_cache_.emplace(std::make_pair(table_id, col_id),
-                                             std::move(merged_results));
-    } else {
-      table_column = column_it->second.get();
+      memcpy(write_ptrs[frag_id].first, col_buffer, write_ptrs[frag_id].second);
     }
+
+    std::vector<int8_t*> raw_write_ptrs;
+    raw_write_ptrs.reserve(frag_count);
+    for (uint i = 0; i < frag_count; i++) {
+      raw_write_ptrs.emplace_back(write_ptrs[i].first);
+    }
+
+    std::unique_ptr<ColumnarResults> merged_results(
+        new ColumnarResults(raw_write_ptrs, total_row_count, col_info->type, thread_idx));
+
+    table_column = merged_results.get();
+    columnarized_scan_table_cache_.emplace(std::make_pair(table_id, col_id),
+                                           std::move(merged_results));
   }
   return ColumnFetcher::transferColumnIfNeeded(
       table_column, 0, memory_level, device_id, device_allocator);
diff --git a/omniscidb/ResultSetRegistry/ColumnarResults.cpp b/omniscidb/ResultSetRegistry/ColumnarResults.cpp
@@ -132,6 +132,26 @@ ColumnarResults::ColumnarResults(std::shared_ptr<RowSetMemoryOwner> row_set_mem_
   memcpy(((void*)column_buffers_[0]), one_col_buffer, buf_size);
 }
 
+ColumnarResults::ColumnarResults(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(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);
diff --git a/omniscidb/ResultSetRegistry/ResultSetRegistry.cpp b/omniscidb/ResultSetRegistry/ResultSetRegistry.cpp
diff --git a/omniscidb/ResultSetRegistry/ResultSetRegistry.h b/omniscidb/ResultSetRegistry/ResultSetRegistry.h
