introduce col fetch callback

Author: akroviakov

omniscidb/QueryEngine/Execute.cpp

@@ -2987,173 +2987,147 @@ FetchResult Executor::fetchChunks(
   std::vector<std::vector<int64_t>> all_num_rows;
   std::vector<std::vector<uint64_t>> all_frag_offsets;
 
-  // in MT case we want to preserve "the order of insertion" into all_frag_col_buffers
+  auto fetch_column_callback = [&](std::shared_ptr<const InputColDescriptor> col_id,
+                                   const std::vector<size_t>& selected_frag_ids,
+                                   std::vector<const int8_t*>& frag_col_buffers,
+                                   const bool parallelized =
+                                       false) -> bool /*empty_frag*/ {
+    if (interrupted_.load()) {
+      throw QueryExecutionError(ERR_INTERRUPTED);
+    }
+    const auto fragments_it = all_tables_fragments.find(col_id->getTableRef());
+    CHECK(fragments_it != all_tables_fragments.end());
+    const auto fragments = fragments_it->second;
+    auto it = plan_state_->global_to_local_col_ids_.find(*col_id);
+    CHECK(it != plan_state_->global_to_local_col_ids_.end());
+    CHECK_LT(static_cast<size_t>(it->second),
+             plan_state_->global_to_local_col_ids_.size());
+    const size_t frag_id = selected_frag_ids[local_col_to_frag_pos[it->second]];
+    if (!fragments->size()) {
+      return true;
+    }
+    auto memory_level_for_column = memory_level;
+    if (plan_state_->columns_to_fetch_.find(*col_id) ==
+        plan_state_->columns_to_fetch_.end()) {
+      memory_level_for_column = Data_Namespace::CPU_LEVEL;
+    }
+    if (needFetchAllFragments(*col_id, ra_exe_unit, selected_fragments)) {
+      // determine if we need special treatment to linearlize multi-frag table
+      // i.e., a column that is classified as varlen type, i.e., array
+      // for now, we can support more types in this way
+      CHECK(!parallelized);  // otherwise recursive tbb parallel for with deadlocks
+      if (needLinearizeAllFragments(
+              *col_id, ra_exe_unit, selected_fragments, memory_level)) {
+        bool for_lazy_fetch = false;
+        if (plan_state_->columns_to_not_fetch_.find(*col_id) !=
+            plan_state_->columns_to_not_fetch_.end()) {
+          for_lazy_fetch = true;
+          VLOG(2) << "Try to linearize lazy fetch column (col_id: " << col_id->getColId()
+                  << ")";
+        }
+        frag_col_buffers[it->second] = column_fetcher.linearizeColumnFragments(
+            col_id->getColInfo(),
+            all_tables_fragments,
+            chunks,
+            chunk_iterators,
+            for_lazy_fetch ? Data_Namespace::CPU_LEVEL : memory_level,
+            for_lazy_fetch ? 0 : device_id,
+            device_allocator);
+      } else {
+        frag_col_buffers[it->second] =
+            column_fetcher.getAllTableColumnFragments(col_id->getColInfo(),
+                                                      all_tables_fragments,
+                                                      memory_level_for_column,
+                                                      device_id,
+                                                      device_allocator,
+                                                      thread_idx);
+      }
+    } else {
+      frag_col_buffers[it->second] =
+          column_fetcher.getOneTableColumnFragment(col_id->getColInfo(),
+                                                   frag_id,
+                                                   all_tables_fragments,
+                                                   chunks,
+                                                   chunk_iterators,
+                                                   memory_level_for_column,
+                                                   device_id,
+                                                   device_allocator);
+    }
+    return false;
+  };
+
+  // in MT fetching for GPU, we want to preserve "the order of insertion" into
+  // all_frag_col_buffers
   std::vector<std::vector<size_t>> selected_frag_ids_vec;
-  if(memory_level == Data_Namespace::MemoryLevel::GPU_LEVEL){
-    std::mutex all_frag;
+  if (memory_level == Data_Namespace::MemoryLevel::GPU_LEVEL) {
     std::atomic<bool> empty_frags{false};
-    tbb::task_arena limitedArena(1);
+    tbb::task_arena limitedArena(8);
+    std::vector<size_t> idx_frags_to_inearize;
     for (const auto& selected_frag_ids : frag_ids_crossjoin) {
       selected_frag_ids_vec.push_back(selected_frag_ids);
+      for (const auto& col_id : col_global_ids) {
+        CHECK(col_id);
+        if (!col_id->isVirtual() &&
+            needFetchAllFragments(*col_id, ra_exe_unit, selected_fragments)) {
+          idx_frags_to_inearize.push_back(selected_frag_ids_vec.size() - 1);
+        }
+      }
     }
     all_frag_col_buffers.resize(selected_frag_ids_vec.size());
 
+    // Try MT fetching for frags that do not need linearization
     limitedArena.execute([&]() {
-      tbb::parallel_for(
-          0ul, selected_frag_ids_vec.size(), [&](const size_t idx) {
-            const auto& selected_frag_ids = selected_frag_ids_vec[idx];
-            std::vector<const int8_t*> frag_col_buffers(
-                plan_state_->global_to_local_col_ids_.size());
-            for (const auto& col_id : col_global_ids) {
-              if (interrupted_.load()) {
-                throw QueryExecutionError(ERR_INTERRUPTED);
-              }
-              CHECK(col_id);
-              if (col_id->isVirtual()) {
-                continue;
-              }
-              const auto fragments_it = all_tables_fragments.find(col_id->getTableRef());
-              CHECK(fragments_it != all_tables_fragments.end());
-              const auto fragments = fragments_it->second;
-              auto it = plan_state_->global_to_local_col_ids_.find(*col_id);
-              CHECK(it != plan_state_->global_to_local_col_ids_.end());
-              CHECK_LT(static_cast<size_t>(it->second),
-                      plan_state_->global_to_local_col_ids_.size());
-              const size_t frag_id = selected_frag_ids[local_col_to_frag_pos[it->second]];
-              if (!fragments->size()) {
-                empty_frags = true;
-                tbb::task::current_context()->cancel_group_execution();
-              }
-              auto memory_level_for_column = memory_level;
-              if (plan_state_->columns_to_fetch_.find(*col_id) ==
-                  plan_state_->columns_to_fetch_.end()) {
-                memory_level_for_column = Data_Namespace::CPU_LEVEL;
-              }
-              if (needFetchAllFragments(*col_id, ra_exe_unit, selected_fragments)) {
-                // determine if we need special treatment to linearlize multi-frag table
-                // i.e., a column that is classified as varlen type, i.e., array
-                // for now, we can support more types in this way
-                all_frag.lock();
-                if (needLinearizeAllFragments(
-                        *col_id, ra_exe_unit, selected_fragments, memory_level)) {
-                  bool for_lazy_fetch = false;
-                  if (plan_state_->columns_to_not_fetch_.find(*col_id) !=
-                      plan_state_->columns_to_not_fetch_.end()) {
-                    for_lazy_fetch = true;
-                    VLOG(2) << "Try to linearize lazy fetch column (col_id: "
-                            << col_id->getColId() << ")";
-                  }
-                  frag_col_buffers[it->second] = column_fetcher.linearizeColumnFragments(
-                      col_id->getColInfo(),
-                      all_tables_fragments,
-                      chunks,
-                      chunk_iterators,
-                      for_lazy_fetch ? Data_Namespace::CPU_LEVEL : memory_level,
-                      for_lazy_fetch ? 0 : device_id,
-                      device_allocator);
-                } else {
-                  frag_col_buffers[it->second] =
-                      column_fetcher.getAllTableColumnFragments(col_id->getColInfo(),
-                                                                all_tables_fragments,
-                                                                memory_level_for_column,
-                                                                device_id,
-                                                                device_allocator,
-                                                                /*thread_idx=*/0);
-                }
-                all_frag.unlock();
-              } else {
-                LOG(INFO) << "Pushing to idx " << idx;
-                frag_col_buffers[it->second] =
-                    column_fetcher.getOneTableColumnFragment(col_id->getColInfo(),
-                                                            frag_id,
-                                                            all_tables_fragments,
-                                                            chunks,
-                                                            chunk_iterators,
-                                                            memory_level_for_column,
-                                                            device_id,
-                                                            device_allocator);
-              }
+      tbb::parallel_for(0ul, selected_frag_ids_vec.size(), [&](const size_t idx) {
+        if (std::find(idx_frags_to_inearize.begin(), idx_frags_to_inearize.end(), idx) ==
+            idx_frags_to_inearize.end()) {
+          const auto& selected_frag_ids = selected_frag_ids_vec[idx];
+          std::vector<const int8_t*> frag_col_buffers(
+              plan_state_->global_to_local_col_ids_.size());
+          for (const auto& col_id : col_global_ids) {
+            CHECK(col_id);
+            if (!col_id->isVirtual() &&
+                fetch_column_callback(
+                    col_id, selected_frag_ids, frag_col_buffers, true)) {
+              empty_frags = true;  // not virtual, but empty frags
+              tbb::task::current_context()->cancel_group_execution();
             }
-            LOG(INFO) << "frag_col_buffers size to push: " << frag_col_buffers.size();
-            all_frag_col_buffers[idx] = frag_col_buffers;
-          });
+          }
+          all_frag_col_buffers[idx] = frag_col_buffers;
+        }
+      });
     });
     if (empty_frags) {
       return {};
     }
+    for (const size_t idx :
+         idx_frags_to_inearize) {  // linear frags materialization is already
+                                   // parallelized, avoid nested tbb
+      const auto& selected_frag_ids = selected_frag_ids_vec[idx];
+      std::vector<const int8_t*> frag_col_buffers(
+          plan_state_->global_to_local_col_ids_.size());
+      for (const auto& col_id : col_global_ids) {
+        CHECK(col_id);
+        if (!col_id->isVirtual() &&
+            fetch_column_callback(col_id, selected_frag_ids, frag_col_buffers)) {
+          return {};  // not virtual, but empty frags
+        }
+      }
+      all_frag_col_buffers[idx] = frag_col_buffers;
+    }
   } else {
     for (const auto& selected_frag_ids : frag_ids_crossjoin) {
-            std::vector<const int8_t*> frag_col_buffers(
-                plan_state_->global_to_local_col_ids_.size());
-            for (const auto& col_id : col_global_ids) {
-              if (interrupted_.load()) {
-                throw QueryExecutionError(ERR_INTERRUPTED);
-              }
-              CHECK(col_id);
-              if (col_id->isVirtual()) {
-                continue;
-              }
-              const auto fragments_it = all_tables_fragments.find(col_id->getTableRef());
-              CHECK(fragments_it != all_tables_fragments.end());
-              const auto fragments = fragments_it->second;
-              auto it = plan_state_->global_to_local_col_ids_.find(*col_id);
-              CHECK(it != plan_state_->global_to_local_col_ids_.end());
-              CHECK_LT(static_cast<size_t>(it->second),
-                      plan_state_->global_to_local_col_ids_.size());
-              const size_t frag_id = selected_frag_ids[local_col_to_frag_pos[it->second]];
-              if (!fragments->size()) {
-                return{};
-              }
-              auto memory_level_for_column = memory_level;
-              if (plan_state_->columns_to_fetch_.find(*col_id) ==
-                  plan_state_->columns_to_fetch_.end()) {
-                memory_level_for_column = Data_Namespace::CPU_LEVEL;
-              }
-              if (needFetchAllFragments(*col_id, ra_exe_unit, selected_fragments)) {
-                // determine if we need special treatment to linearlize multi-frag table
-                // i.e., a column that is classified as varlen type, i.e., array
-                // for now, we can support more types in this way
-                if (needLinearizeAllFragments(
-                        *col_id, ra_exe_unit, selected_fragments, memory_level)) {
-                  bool for_lazy_fetch = false;
-                  if (plan_state_->columns_to_not_fetch_.find(*col_id) !=
-                      plan_state_->columns_to_not_fetch_.end()) {
-                    for_lazy_fetch = true;
-                    VLOG(2) << "Try to linearize lazy fetch column (col_id: "
-                            << col_id->getColId() << ")";
-                  }
-                  frag_col_buffers[it->second] = column_fetcher.linearizeColumnFragments(
-                      col_id->getColInfo(),
-                      all_tables_fragments,
-                      chunks,
-                      chunk_iterators,
-                      for_lazy_fetch ? Data_Namespace::CPU_LEVEL : memory_level,
-                      for_lazy_fetch ? 0 : device_id,
-                      device_allocator);
-                } else {
-                  frag_col_buffers[it->second] =
-                      column_fetcher.getAllTableColumnFragments(col_id->getColInfo(),
-                                                                all_tables_fragments,
-                                                                memory_level_for_column,
-                                                                device_id,
-                                                                device_allocator,
-                                                                thread_idx);
-                }
-              } else {
-                frag_col_buffers[it->second] =
-                    column_fetcher.getOneTableColumnFragment(col_id->getColInfo(),
-                                                             frag_id,
-                                                             all_tables_fragments,
-                                                             chunks,
-                                                             chunk_iterators,
-                                                             memory_level_for_column,
-                                                             device_id,
-                                                             device_allocator);
-              }
-            }
-            selected_frag_ids_vec.push_back(selected_frag_ids);
-            all_frag_col_buffers.push_back(frag_col_buffers);
-          }
+      std::vector<const int8_t*> frag_col_buffers(
+          plan_state_->global_to_local_col_ids_.size());
+      for (const auto& col_id : col_global_ids) {
+        CHECK(col_id);
+        if (!col_id->isVirtual() &&
+            fetch_column_callback(col_id, selected_frag_ids, frag_col_buffers)) {
+          return {};  // not virtual, but empty frags
+        }
+      }
+      selected_frag_ids_vec.push_back(selected_frag_ids);
+      all_frag_col_buffers.push_back(frag_col_buffers);
+    }
   }
   std::tie(all_num_rows, all_frag_offsets) = getRowCountAndOffsetForAllFrags(
       ra_exe_unit, selected_frag_ids_vec, ra_exe_unit.input_descs, all_tables_fragments);