Reuse column_cache in OOM handling logic (#7646)

Signed-off-by: Misiu Godfrey <misiu.godfrey@kraken.mapd.com>
SINHASantos · Aug 26, 2024 · c1eb564 · c1eb564
1 parent ad61921
commit c1eb564
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 5 deletions.
diff --git a/QueryEngine/ColumnFetcher.cpp b/QueryEngine/ColumnFetcher.cpp
@@ -62,7 +62,7 @@ std::string getMemoryLevelString(Data_Namespace::MemoryLevel memoryLevel) {
 }
 }  // namespace
 
-ColumnFetcher::ColumnFetcher(Executor* executor, const ColumnCacheMap& column_cache)
+ColumnFetcher::ColumnFetcher(Executor* executor, ColumnCacheMap& column_cache)
     : executor_(executor), columnarized_table_cache_(column_cache) {}
 
 //! Gets a column fragment chunk on CPU or on GPU depending on the effective

diff --git a/QueryEngine/ColumnFetcher.h b/QueryEngine/ColumnFetcher.h
@@ -48,7 +48,7 @@ using MergedChunk = std::pair<AbstractBuffer*, AbstractBuffer*>;
 
 class ColumnFetcher {
  public:
-  ColumnFetcher(Executor* executor, const ColumnCacheMap& column_cache);
+  ColumnFetcher(Executor* executor, ColumnCacheMap& column_cache);
 
   //! Gets one chunk's pointer and element count on either CPU or GPU.
   static std::pair<const int8_t*, size_t> getOneColumnFragment(
@@ -183,7 +183,7 @@ class ColumnFetcher {
   mutable std::mutex linearization_mutex_;
   mutable std::mutex chunk_list_mutex_;
   mutable std::mutex linearized_col_cache_mutex_;
-  mutable ColumnCacheMap columnarized_table_cache_;
+  ColumnCacheMap& columnarized_table_cache_;
   mutable std::unordered_map<InputColDescriptor, std::unique_ptr<const ColumnarResults>>
       columnarized_scan_table_cache_;
   using DeviceMergedChunkIterMap = std::unordered_map<int, int8_t*>;

diff --git a/QueryEngine/RelAlgExecutor.cpp b/QueryEngine/RelAlgExecutor.cpp
@@ -3877,6 +3877,7 @@ ExecutionResult RelAlgExecutor::executeWorkUnit(
       handlePersistentError(e.getErrorCode());
       return handleOutOfMemoryRetry(
           {ra_exe_unit, work_unit.body, local_groups_buffer_entry_guess},
+          column_cache,
           targets_meta,
           is_agg,
           co,
@@ -4155,6 +4156,7 @@ bool RelAlgExecutor::isRowidLookup(const WorkUnit& work_unit) {
 
 ExecutionResult RelAlgExecutor::handleOutOfMemoryRetry(
     const RelAlgExecutor::WorkUnit& work_unit,
+    ColumnCacheMap& column_cache,
     const std::vector<TargetMetaInfo>& targets_meta,
     const bool is_agg,
     const CompilationOptions& co,
@@ -4191,7 +4193,6 @@ ExecutionResult RelAlgExecutor::handleOutOfMemoryRetry(
                       "kernels disabled.";
       const auto ra_exe_unit = decide_approx_count_distinct_implementation(
           ra_exe_unit_in, table_infos, executor_, co.device_type, target_exprs_owned_);
-      ColumnCacheMap column_cache;
       result = {executor_->executeWorkUnit(max_groups_buffer_entry_guess,
                                            is_agg,
                                            table_infos,
@@ -4227,7 +4228,6 @@ ExecutionResult RelAlgExecutor::handleOutOfMemoryRetry(
     iteration_ctr++;
     auto ra_exe_unit = decide_approx_count_distinct_implementation(
         ra_exe_unit_in, table_infos, executor_, co_cpu.device_type, target_exprs_owned_);
-    ColumnCacheMap column_cache;
     try {
       result = {executor_->executeWorkUnit(max_groups_buffer_entry_guess,
                                            is_agg,

diff --git a/QueryEngine/RelAlgExecutor.h b/QueryEngine/RelAlgExecutor.h
@@ -357,6 +357,7 @@ class RelAlgExecutor : private StorageIOFacility {
   bool isRowidLookup(const WorkUnit& work_unit);
 
   ExecutionResult handleOutOfMemoryRetry(const RelAlgExecutor::WorkUnit& work_unit,
+                                         ColumnCacheMap& column_cache,
                                          const std::vector<TargetMetaInfo>& targets_meta,
                                          const bool is_agg,
                                          const CompilationOptions& co,