[Join] Inline and parallelize tbb in getAllTableColumnFragments.

Devjiu · Devjiu · commit e342fae5a982 · 2023-09-27T14:01:22.000Z
This commit refactors and simplifies method `getAllTableColumnFragments`. Also some parallelization added. Partially resolves: #574 Signed-off-by: Dmitrii Makarenko <dmitrii.makarenko@intel.com>
diff --git a/omniscidb/DataMgr/BufferMgr/Buffer.h b/omniscidb/DataMgr/BufferMgr/Buffer.h
@@ -135,9 +135,10 @@ class Buffer : public AbstractBuffer {
   }
 
   inline int unPin() override {
-    std::lock_guard<std::mutex> pin_lock(pin_mutex_);
+    std::unique_lock<std::mutex> pin_lock(pin_mutex_);
     int res = (--pin_count_);
     if (!res && delete_on_unpin_) {
+      pin_lock.unlock();
       delete this;
     }
     return res;
diff --git a/omniscidb/DataMgr/BufferMgr/BufferMgr.cpp b/omniscidb/DataMgr/BufferMgr/BufferMgr.cpp
@@ -102,9 +102,12 @@ void BufferMgr::clear() {
   // pinned. Here we delete unpinned buffers and mark pinned buffers
   // for removal to have them deleted when unpinned.
   for (auto& buf : chunk_index_) {
-    if (buf.second->buffer) {
-      buf.second->buffer->deleteWhenUnpinned();
-      buf.second->buffer = nullptr;
+    // Iterator can be empty, check Buffer c-tors
+    if (buf.second == BufferList::iterator()) {
+      if (buf.second->buffer) {
+        buf.second->buffer->deleteWhenUnpinned();
+        buf.second->buffer = nullptr;
+      }
     }
   }
 
@@ -623,6 +626,7 @@ void BufferMgr::deleteBuffer(const ChunkKey& key, const bool) {
   chunk_index_lock.unlock();
   std::lock_guard<std::mutex> sized_segs_lock(sized_segs_mutex_);
   if (seg_it->buffer) {
+    CHECK_EQ(seg_it->buffer->getPinCount(), 0);
     delete seg_it->buffer;  // Delete Buffer for segment
     seg_it->buffer = 0;
   }
@@ -831,6 +835,10 @@ void BufferMgr::free(AbstractBuffer* buffer) {
   if (casted_buffer == 0) {
     LOG(FATAL) << "Wrong buffer type - expects base class pointer to Buffer type.";
   }
+  if (casted_buffer->seg_it_ == BufferList::iterator()) {
+    delete buffer;
+  }
+
   deleteBuffer(casted_buffer->seg_it_->chunk_key);
 }
 
diff --git a/omniscidb/QueryEngine/ColumnFetcher.cpp b/omniscidb/QueryEngine/ColumnFetcher.cpp
@@ -16,15 +16,16 @@
 
 #include "QueryEngine/ColumnFetcher.h"
 
-#include <memory>
-
 #include "DataMgr/ArrayNoneEncoder.h"
 #include "QueryEngine/ErrorHandling.h"
 #include "QueryEngine/Execute.h"
 #include "Shared/Intervals.h"
 #include "Shared/likely.h"
 #include "Shared/sqltypes.h"
 
+#include <tbb/parallel_for.h>
+#include <memory>
+
 namespace {
 
 std::string getMemoryLevelString(Data_Namespace::MemoryLevel memoryLevel) {
@@ -239,6 +240,11 @@ const int8_t* ColumnFetcher::getAllTableColumnFragments(
   int db_id = col_info->db_id;
   int table_id = col_info->table_id;
   int col_id = col_info->column_id;
+  if (col_info->type->isString() || col_info->type->isArray()) {
+    throw std::runtime_error(
+        "Array type passed to getAllTableColumnFragments. Should be handled in "
+        "linearization.");
+  }
   const auto fragments_it = all_tables_fragments.find({db_id, table_id});
   CHECK(fragments_it != all_tables_fragments.end());
   const auto fragments = fragments_it->second;
@@ -248,7 +254,6 @@ 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();
@@ -262,44 +267,95 @@ const int8_t* ColumnFetcher::getAllTableColumnFragments(
     }
 
     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;
-        }
-        auto chunk_meta_it = fragment.getChunkMetadataMap().find(col_id);
-        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));
+    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);
       }
-      auto merged_results =
-          ColumnarResults::mergeResults(executor_->row_set_mem_owner_, column_frags);
+      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);
+    std::vector<std::pair<int8_t*, size_t>> write_ptrs;
+    std::vector<size_t> valid_fragments;
+    for (size_t frag_id = 0; frag_id < frag_count; ++frag_id) {
+      const auto& fragment = (*fragments)[frag_id];
+      if (fragment.isEmptyPhysicalFragment()) {
+        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;
+      valid_fragments.push_back(frag_id);
+    }
+
+    if (write_ptrs.empty()) {
+      std::unique_ptr<ColumnarResults> merged_results(nullptr);
+
       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();
+
+      return ColumnFetcher::transferColumnIfNeeded(
+          table_column, 0, memory_level, device_id, device_allocator);
     }
+
+    size_t valid_frag_count = valid_fragments.size();
+    tbb::parallel_for(
+        tbb::blocked_range<size_t>(0, valid_frag_count),
+        [&](const tbb::blocked_range<size_t>& frag_ids) {
+          for (size_t v_frag_id = frag_ids.begin(); v_frag_id < frag_ids.end();
+               ++v_frag_id) {
+            std::list<std::shared_ptr<Chunk_NS::Chunk>> chunk_holder;
+            std::list<ChunkIter> chunk_iter_holder;
+            size_t frag_id = valid_fragments[v_frag_id];
+            const auto& fragment = (*fragments)[frag_id];
+            auto chunk_meta_it = fragment.getChunkMetadataMap().find(col_id);
+            CHECK(chunk_meta_it != fragment.getChunkMetadataMap().end());
+            std::shared_ptr<Chunk_NS::Chunk> chunk;
+            {
+              ChunkKey chunk_key{
+                  db_id, fragment.physicalTableId, col_id, fragment.fragmentId};
+              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);
+            }
+            auto ab = chunk->getBuffer();
+            CHECK(ab->getMemoryPtr());
+            int8_t* col_buffer =
+                ab->getMemoryPtr();  // @TODO(alex) change to use ChunkIter
+            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 (size_t i = 0; i < frag_count; i++) {
+      raw_write_ptrs.emplace_back(write_ptrs[i].first);
+    }
+
+    std::unique_ptr<ColumnarResults> merged_results(new ColumnarResults(
+        std::move(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);
@@ -469,8 +525,10 @@ const int8_t* ColumnFetcher::linearizeColumnFragments(
     }
   }
   CHECK(res.first);  // check merged data buffer
+  res.first->deleteWhenUnpinned();
   if (!type->isFixedLenArray()) {
     CHECK(res.second);  // check merged index buffer
+    res.second->deleteWhenUnpinned();
   }
   auto merged_data_buffer = res.first;
   auto merged_index_buffer = res.second;
@@ -1009,34 +1067,3 @@ ChunkIter ColumnFetcher::prepareChunkIter(AbstractBuffer* merged_data_buf,
   merged_chunk_iter.elem_type_size = chunk_iter.elem_type_size;
   return merged_chunk_iter;
 }
-
-void ColumnFetcher::freeLinearizedBuf() {
-  std::lock_guard<std::mutex> linearized_col_cache_guard(linearized_col_cache_mutex_);
-  auto buffer_provider = executor_->getBufferProvider();
-
-  if (!linearized_data_buf_cache_.empty()) {
-    for (auto& kv : linearized_data_buf_cache_) {
-      for (auto& kv2 : kv.second) {
-        buffer_provider->free(kv2.second);
-      }
-    }
-  }
-
-  if (!linearized_idx_buf_cache_.empty()) {
-    for (auto& kv : linearized_idx_buf_cache_) {
-      for (auto& kv2 : kv.second) {
-        buffer_provider->free(kv2.second);
-      }
-    }
-  }
-}
-
-void ColumnFetcher::freeTemporaryCpuLinearizedIdxBuf() {
-  std::lock_guard<std::mutex> linearized_col_cache_guard(linearized_col_cache_mutex_);
-  auto buffer_provider = executor_->getBufferProvider();
-  if (!linearlized_temporary_cpu_index_buf_cache_.empty()) {
-    for (auto& kv : linearlized_temporary_cpu_index_buf_cache_) {
-      buffer_provider->free(kv.second);
-    }
-  }
-}
diff --git a/omniscidb/QueryEngine/ColumnFetcher.h b/omniscidb/QueryEngine/ColumnFetcher.h
@@ -93,9 +93,6 @@ class ColumnFetcher {
       DeviceAllocator* device_allocator,
       const size_t thread_idx) const;
 
-  void freeTemporaryCpuLinearizedIdxBuf();
-  void freeLinearizedBuf();
-
   DataProvider* getDataProvider() const { return data_provider_; }
 
  private:
diff --git a/omniscidb/QueryEngine/Execute.cpp b/omniscidb/QueryEngine/Execute.cpp
@@ -2186,10 +2186,6 @@ hdk::ResultSetTable Executor::executeWorkUnitImpl(
   do {
     SharedKernelContext shared_context(query_infos);
     ColumnFetcher column_fetcher(this, data_provider, column_cache);
-    ScopeGuard scope_guard = [&column_fetcher] {
-      column_fetcher.freeLinearizedBuf();
-      column_fetcher.freeTemporaryCpuLinearizedIdxBuf();
-    };
 
     if (ra_exe_unit.isShuffle()) {
       allocateShuffleBuffers(query_infos, ra_exe_unit, row_set_mem_owner, shared_context);
diff --git a/omniscidb/QueryEngine/JoinHashTable/Runtime/JoinColumnIterator.h b/omniscidb/QueryEngine/JoinHashTable/Runtime/JoinColumnIterator.h
@@ -73,6 +73,7 @@ struct JoinColumnIterator {
   DEVICE FORCE_INLINE JoinColumnIterator& operator++() {
     index += step;
     index_inside_chunk += step;
+    // this loop is made to find index_of_chunk by total index of element
     while (chunk_data &&
            index_inside_chunk >= join_chunk_array[index_of_chunk].num_elems) {
       index_inside_chunk -= join_chunk_array[index_of_chunk].num_elems;
diff --git a/omniscidb/ResultSetRegistry/ColumnarResults.cpp b/omniscidb/ResultSetRegistry/ColumnarResults.cpp
@@ -122,28 +122,6 @@ ColumnarResults::ColumnarResults(std::shared_ptr<RowSetMemoryOwner> row_set_mem_
   }
 }
 
-ColumnarResults::ColumnarResults(std::shared_ptr<RowSetMemoryOwner> row_set_mem_owner,
-                                 const 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__);
-
-  if (target_type->isVarLen()) {
-    throw ColumnarConversionNotSupported();
-  }
-  const auto buf_size = num_rows * target_type->size();
-  column_buffers_[0] =
-      reinterpret_cast<int8_t*>(row_set_mem_owner->allocate(buf_size, thread_idx_));
-  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,
diff --git a/omniscidb/ResultSetRegistry/ColumnarResults.h b/omniscidb/ResultSetRegistry/ColumnarResults.h
@@ -67,12 +67,6 @@ class ColumnarResults {
                   const Config& config,
                   const bool is_parallel_execution_enforced = false);
 
-  ColumnarResults(const std::shared_ptr<RowSetMemoryOwner> row_set_mem_owner,
-                  const int8_t* one_col_buffer,
-                  const size_t num_rows,
-                  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,
diff --git a/omniscidb/Tests/ArrowBasedExecuteTest.cpp b/omniscidb/Tests/ArrowBasedExecuteTest.cpp
@@ -9268,6 +9268,15 @@ TEST_F(Select, Joins_Fixed_Size_Array_Multi_Frag) {
         v<int64_t>(run_simple_agg(
             "SELECT COUNT(1) FROM mf_t_arr r1, mf_t_arr r2 WHERE r1.t2[1] = r2.t2[1]",
             dt)));
+    // {
+    //   std::ostringstream oss_l;
+    //   oss_l << "SELECT * FROM  mf_d_arr t1, mf_d_arr t2 WHERE ";
+    //   auto part = oss_l.str();
+    //   auto q_{part + "t1.c2[1] = t2.c2[1];"};
+    //   const auto rows = run_multiple_agg(q_, dt);
+    //   LOG(ERROR) << rows->contentToString(true);
+    //   LOG(ERROR) << "should be done.";
+    // }
   }
 
   dropTable("mf_f_arr");
diff --git a/omniscidb/Tests/JoinHashTableTest.cpp b/omniscidb/Tests/JoinHashTableTest.cpp
diff --git a/python/pyhdk/hdk.py b/python/pyhdk/hdk.py

Original file line number	Diff line number	Diff line change
`@@ -135,9 +135,10 @@ class Buffer : public AbstractBuffer {`
`135`	`135`	`}`
`136`	`136`
`137`	`137`	`inline int unPin() override {`
`138`		`- std::lock_guard<std::mutex> pin_lock(pin_mutex_);`
	`138`	`+ std::unique_lock<std::mutex> pin_lock(pin_mutex_);`
`139`	`139`	`int res = (--pin_count_);`
`140`	`140`	`if (!res && delete_on_unpin_) {`
	`141`	`+ pin_lock.unlock();`
`141`	`142`	`delete this;`
`142`	`143`	`}`
`143`	`144`	`return res;`