Support UNNEST in projections.

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

Author: ienkovich

omniscidb/IR/Node.cpp

@@ -6,7 +6,9 @@
  */
 
 #include "Node.h"
+#include "ExprCollector.h"
 #include "InputRewriter.h"
+#include "UnnestDetector.h"
 
 namespace hdk::ir {
 
@@ -128,6 +130,15 @@ bool Project::hasWindowFunctionExpr() const {
   return false;
 }
 
+bool Project::hasUnnestExpr() const {
+  for (auto& expr : exprs_) {
+    if (UnnestDetector::collect(expr.get())) {
+      return true;
+    }
+  }
+  return false;
+}
+
 bool Project::isIdentity() const {
   if (!isSimple()) {
     return false;

omniscidb/IR/Node.h

@@ -337,6 +337,7 @@ class Project : public Node {
   void appendInput(std::string new_field_name, ExprPtr expr);
 
   bool hasWindowFunctionExpr() const;
+  bool hasUnnestExpr() const;
 
   std::string toString() const override {
     return cat(::typeName(this),

omniscidb/IR/UnnestDetector.h

@@ -0,0 +1,27 @@
+/**
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include "ExprCollector.h"
+
+namespace hdk::ir {
+
+class UnnestDetector : public ExprCollector<bool, UnnestDetector> {
+ public:
+  UnnestDetector() { result_ = false; }
+
+ protected:
+  void visitUOper(const hdk::ir::UOper* uoper) override {
+    if (uoper->opType() == OpType::kUnnest) {
+      result_ = true;
+      return;
+    }
+    BaseClass::visitUOper(uoper);
+  }
+};
+
+}  // namespace hdk::ir

omniscidb/QueryEngine/Execute.cpp

@@ -69,6 +69,7 @@
 #include "QueryEngine/RuntimeFunctions.h"
 #include "QueryEngine/SpeculativeTopN.h"
 #include "QueryEngine/StringDictionaryGenerations.h"
+#include "QueryEngine/UnnestedVarsCollector.h"
 #include "QueryEngine/Visitors/TransientStringLiteralsVisitor.h"
 #include "ResultSet/ColRangeInfo.h"
 #include "Shared/checked_alloc.h"
@@ -1759,10 +1760,12 @@ hdk::ResultSetTable Executor::executeWorkUnit(
     }
   };
 
+  bool has_proj_unnest =
+      !UnnestedVarsCollector::collect(ra_exe_unit_in.target_exprs).empty();
   try {
     auto result = executeWorkUnitImpl(max_groups_buffer_entry_guess,
                                       is_agg,
-                                      true,
+                                      !has_proj_unnest,
                                       query_infos,
                                       ra_exe_unit_in,
                                       co,

omniscidb/QueryEngine/Execute.h

@@ -825,7 +825,8 @@ class Executor : public StringDictionaryProxyProvider {
   llvm::Value* arrayLoopCodegen(const hdk::ir::Expr* array_expr,
                                 std::stack<llvm::BasicBlock*>& array_loops,
                                 DiamondCodegen& diamond_codegen,
-                                const CompilationOptions& co);
+                                const CompilationOptions& co,
+                                llvm::Value* array_size = nullptr);
 
   llvm::Value* castToFP(llvm::Value*,
                         const hdk::ir::Type* from_type,

omniscidb/QueryEngine/IRCodegen.cpp

@@ -1226,7 +1226,8 @@ Executor::GroupColLLVMValue Executor::groupByColumnCodegen(
 llvm::Value* Executor::arrayLoopCodegen(const hdk::ir::Expr* array_expr,
                                         std::stack<llvm::BasicBlock*>& array_loops,
                                         DiamondCodegen& diamond_codegen,
-                                        const CompilationOptions& co) {
+                                        const CompilationOptions& co,
+                                        llvm::Value* array_size) {
   AUTOMATIC_IR_METADATA(cgen_state_.get());
   CodeGenerator code_generator(this, co.codegen_traits_desc);
   auto array_lv = code_generator.codegen(array_expr, true, co).front();
@@ -1257,7 +1258,8 @@ llvm::Value* Executor::arrayLoopCodegen(const hdk::ir::Expr* array_expr,
   CHECK(array_type->isArray());
   auto elem_type = array_type->as<hdk::ir::ArrayBaseType>()->elemType();
   auto array_len =
-      (array_type->size() > 0)
+      array_size ? array_size
+      : (array_type->size() > 0)
           ? cgen_state_->llInt(array_type->size() / elem_type->size())
           : cgen_state_->emitExternalCall(
                 "array_size",

omniscidb/QueryEngine/QueryExecutionSequence.cpp

@@ -7,6 +7,8 @@
 #include "QueryExecutionSequence.h"
 #include "ScalarExprVisitor.h"
 
+#include "IR/UnnestDetector.h"
+
 #include <boost/graph/adjacency_list.hpp>
 #include <boost/graph/topological_sort.hpp>
 
@@ -115,6 +117,40 @@ class QueryExecutionSequenceImpl {
       execution_points_.insert(node->getInput(0));
     }
 
+    // Similar limitation applies to the UNNEST operation. We support
+    // only column references as an unnest operand in projections and
+    // therefore have to execution its input. Also, we don't have any
+    // special unnesting expression in the execution unit and rely on
+    // target expression. So, subsequent nodes that drop unnested arrays
+    // can completely remove unnesting from execution units, so all
+    // projections with UNNEST operation should be executed directly.
+    // We also support UNNEST operation in aggregation keys. So, allow
+    // to merge a projection with appropriate aggregate node.
+    // TODO: add unnest_exprs to execution unit to enable merge with
+    // other nodes.
+    if (node->is<ir::Project>() && node->as<ir::Project>()->hasUnnestExpr()) {
+      execution_points_.insert(node->getInput(0));
+      bool can_be_merged = false;
+      auto [start, end] = boost::in_edges(node_to_vertex_[node], graph_);
+      if (start != end) {
+        auto parent_node = graph_[start->m_source];
+        if (parent_node->is<ir::Aggregate>()) {
+          size_t keys_count = parent_node->as<ir::Aggregate>()->getGroupByCount();
+          auto proj = node->as<ir::Project>();
+          can_be_merged = true;
+          for (size_t i = keys_count; i < proj->size(); ++i) {
+            if (ir::UnnestDetector::collect(proj->getExpr(i))) {
+              can_be_merged = false;
+              break;
+            }
+          }
+        }
+      }
+      if (!can_be_merged) {
+        execution_points_.insert(node);
+      }
+    }
+
     // Currently, we cannot merge union code into any other execution
     // module. Therefore, mark it and all its inputs as execution points.
     // TODO: UNION ALL should be able to be merged into other execution
@@ -175,7 +211,8 @@ class QueryExecutionSequenceImpl {
       CHECK(start != end);
       auto node = graph_[start->m_source];
 
-      if (node->is<ir::Project>() && !node->as<ir::Project>()->hasWindowFunctionExpr()) {
+      if (node->is<ir::Project>() && !node->as<ir::Project>()->hasWindowFunctionExpr() &&
+          !node->as<ir::Project>()->hasUnnestExpr()) {
         // In case of aggregation we allow only 'simple' projections which
         // don't have complex expressions referencing aggregate exprs.
         bool is_simple = true;

omniscidb/QueryEngine/RelAlgExecutor.cpp

@@ -17,6 +17,7 @@
 #include "RelAlgExecutor.h"
 #include "DataMgr/DataMgr.h"
 #include "IR/TypeUtils.h"
+#include "QueryBuilder/QueryBuilder.h"
 #include "QueryEngine/CalciteDeserializerUtils.h"
 #include "QueryEngine/CardinalityEstimator.h"
 #include "QueryEngine/ColumnFetcher.h"
@@ -35,6 +36,7 @@
 #include "QueryEngine/RelAlgVisitor.h"
 #include "QueryEngine/ResultSetBuilder.h"
 #include "QueryEngine/ResultSetSort.h"
+#include "QueryEngine/UnnestedVarsCollector.h"
 #include "QueryEngine/WindowContext.h"
 #include "QueryEngine/WorkUnitBuilder.h"
 #include "QueryOptimizer/CanonicalizeQuery.h"
@@ -1505,14 +1507,27 @@ std::optional<size_t> RelAlgExecutor::getFilteredCountAll(const WorkUnit& work_u
                                                           const bool is_agg,
                                                           const CompilationOptions& co,
                                                           const ExecutionOptions& eo) {
-  const auto count = hdk::ir::makeExpr<hdk::ir::AggExpr>(
-      hdk::ir::Context::defaultCtx().integer(config_.exec.group_by.bigint_count ? 8 : 4),
-      hdk::ir::AggType::kCount,
-      nullptr,
-      false,
-      nullptr);
+  auto unnested_vars = UnnestedVarsCollector::collect(work_unit.exe_unit.target_exprs);
+  hdk::ir::QueryBuilder builder(
+      hdk::ir::Context::defaultCtx(), schema_provider_, executor_->getConfigPtr());
+  hdk::ir::BuilderExpr count_all_agg;
+  if (!unnested_vars.empty()) {
+    hdk::ir::BuilderExpr total_count;
+    for (auto var : unnested_vars) {
+      hdk::ir::BuilderExpr var_expr(&builder, var->shared());
+      if (!total_count.expr()) {
+        total_count = var_expr.cardinality();
+      } else {
+        total_count = total_count.mul(var_expr.cardinality());
+      }
+    }
+    count_all_agg = total_count.sum();
+  } else {
+    count_all_agg = builder.count();
+  }
+
   const auto count_all_exe_unit =
-      create_count_all_execution_unit(work_unit.exe_unit, count);
+      create_count_all_execution_unit(work_unit.exe_unit, count_all_agg.expr());
   size_t one{1};
   hdk::ResultSetTable count_all_result;
   try {