Copy table stats for shuffled data.

Signed-off-by: ienkovich <[email protected]>

Author: ienkovich

omniscidb/QueryEngine/RelAlgExecutor.cpp

@@ -790,6 +790,10 @@ void RelAlgExecutor::executeStepWithPartitionedAggregation(const hdk::ir::Node*
     proj.reset();
   }
 
+  // Currently, we merge shuffle node with simple projections only. Therefore, we can
+  // assign original table stats to shuffling results to avoid metadata computation.
+  maybeCopyTableStatsFromInput(shuffle_node.get());
+
   // Create new aggregation node and execute it.
   auto part_agg = std::make_shared<hdk::ir::Aggregate>(
       agg->getGroupByCount(), agg->getAggs(), agg->getFields(), agg_input_shared);
@@ -826,6 +830,74 @@ void RelAlgExecutor::executeStepWithPartitionedAggregation(const hdk::ir::Node*
   temporary_tables_.erase(-new_root->getId());
 }
 
+void RelAlgExecutor::maybeCopyTableStatsFromInput(const hdk::ir::Node* node) {
+  std::vector<int> col_mapping;
+  col_mapping.reserve(node->size());
+  // Stats copy is supported for shuffle and simple projections only.
+  if (node->is<hdk::ir::Project>()) {
+    auto proj = node->as<hdk::ir::Project>();
+    if (!proj->isSimple()) {
+      VLOG(1) << "Cannot copy table stats for non-simple projection.";
+      return;
+    }
+    for (auto& expr : proj->getExprs()) {
+      col_mapping.push_back(expr->as<hdk::ir::ColumnRef>()->index());
+    }
+  } else if (node->is<hdk::ir::Shuffle>()) {
+    for (auto& expr : node->as<hdk::ir::Shuffle>()->exprs()) {
+      CHECK(expr->is<hdk::ir::ColumnRef>());
+      col_mapping.push_back(expr->as<hdk::ir::ColumnRef>()->index());
+    }
+  } else {
+    VLOG(1) << "Cannot copy table stats for node " << node->toString();
+    return;
+  }
+
+  // We can traverse through a chain of simple projections to the original data source.
+  auto data_source = node->getInput(0);
+  while (!data_source->getResult() && !data_source->is<hdk::ir::Scan>()) {
+    auto proj = data_source->as<hdk::ir::Project>();
+    if (!proj || !proj->isSimple()) {
+      VLOG(1) << "Cannot copy table stats due to non-simple projection. "
+              << node->toString();
+      return;
+    }
+    for (size_t i = 0; i < col_mapping.size(); ++i) {
+      auto idx = static_cast<size_t>(col_mapping[i]);
+      CHECK_LT(idx, proj->size());
+      col_mapping[i] = proj->getExpr(idx)->as<hdk::ir::ColumnRef>()->index();
+    }
+    data_source = data_source->getInput(0);
+  }
+
+  auto input_token =
+      data_source->getResult() ? data_source->getResult()->getToken().get() : nullptr;
+  auto input_scan = data_source->getResult() ? nullptr : data_source->as<hdk::ir::Scan>();
+  int input_db_id = input_token ? input_token->dbId() : input_scan->getDatabaseId();
+  int input_table_id = input_token ? input_token->tableId() : input_scan->getTableId();
+  auto input_meta = data_provider_->getTableMetadata(input_db_id, input_table_id);
+  if (input_meta.hasComputedTableStats()) {
+    auto& orig_stats = input_meta.getTableStats();
+    auto target_token = node->getResult()->getToken();
+    TableStats stats;
+    for (size_t i = 0; i < col_mapping.size(); ++i) {
+      auto target_col_id = target_token->columnId(i);
+      auto input_col_id = input_token
+                              ? input_token->columnId(col_mapping[i])
+                              : input_scan->getColumnInfo(col_mapping[i])->column_id;
+      CHECK(orig_stats.count(input_col_id))
+          << "Cannot find stats for column " << input_col_id
+          << ". data_source=" << data_source->toString();
+      stats.emplace(target_col_id, orig_stats.at(input_col_id));
+    }
+    target_token->setTableStats(std::move(stats));
+    VLOG(1) << "Copy table stats from " << input_db_id << ":" << input_table_id << " to "
+            << target_token->dbId() << ":" << target_token->tableId();
+  } else {
+    VLOG(1) << "Cannot copy table stats because original table stats are unavailable.";
+  }
+}
+
 void RelAlgExecutor::executeStep(const hdk::ir::Node* step_root,
                                  const CompilationOptions& co,
                                  const ExecutionOptions& eo,

omniscidb/QueryEngine/RelAlgExecutor.h

@@ -131,6 +131,7 @@ class RelAlgExecutor {
                                              const ExecutionOptions& eo,
                                              size_t estimated_buffer_size,
                                              const int64_t queue_time_ms);
+  void maybeCopyTableStatsFromInput(const hdk::ir::Node* node);
   ExecutionResult executeStep(const hdk::ir::Node* step_root,
                               const CompilationOptions& co,
                               const ExecutionOptions& eo,

omniscidb/ResultSetRegistry/ResultSetRegistry.cpp

@@ -176,6 +176,17 @@ void ResultSetRegistry::drop(const ResultSetTableToken& token) {
   SimpleSchemaProvider::dropTable(token.dbId(), token.tableId());
 }
 
+void ResultSetRegistry::setTableStats(const ResultSetTableToken& token,
+                                      TableStats stats) {
+  mapd_shared_lock<mapd_shared_mutex> data_lock(data_mutex_);
+  CHECK(tables_.count(token.tableId()));
+  auto* table = tables_.at(token.tableId()).get();
+  data_lock.unlock();
+
+  mapd_unique_lock<mapd_shared_mutex> table_lock(table->mutex);
+  table->table_stats = std::move(stats);
+}
+
 ResultSetTableTokenPtr ResultSetRegistry::head(const ResultSetTableToken& token,
                                                size_t n) {
   mapd_shared_lock<mapd_shared_mutex> data_lock(data_mutex_);

omniscidb/ResultSetRegistry/ResultSetRegistry.h

@@ -39,6 +39,8 @@ class ResultSetRegistry : public SimpleSchemaProvider,
   ResultSetPtr get(const ResultSetTableToken& token, size_t frag_id) const;
   void drop(const ResultSetTableToken& token);
 
+  void setTableStats(const ResultSetTableToken& token, TableStats stats);
+
   ResultSetTableTokenPtr head(const ResultSetTableToken& token, size_t n);
   ResultSetTableTokenPtr tail(const ResultSetTableToken& token, size_t n);
 

omniscidb/ResultSetRegistry/ResultSetTableToken.cpp

@@ -38,6 +38,10 @@ void ResultSetTableToken::reset() {
   }
 }
 
+void ResultSetTableToken::setTableStats(TableStats stats) const {
+  registry_->setTableStats(*this, std::move(stats));
+}
+
 ResultSetTableTokenPtr ResultSetTableToken::head(size_t n) const {
   return registry_->head(*this, n);
 }

omniscidb/ResultSetRegistry/ResultSetTableToken.h

@@ -9,6 +9,7 @@
 #include "ResultSetTable.h"
 
 #include "DataMgr/ChunkMetadata.h"
+#include "DataProvider/TableFragmentsInfo.h"
 #include "SchemaMgr/TableInfo.h"
 
 #include "arrow/api.h"
@@ -54,6 +55,8 @@ class ResultSetTableToken : public std::enable_shared_from_this<ResultSetTableTo
 
   const std::string& tableName() const { return tinfo_->name; }
 
+  void setTableStats(TableStats stats) const;
+
   ResultSetTableTokenPtr head(size_t n) const;
   ResultSetTableTokenPtr tail(size_t n) const;