Enable execution results usage in the query builder.

ienkovich · ienkovich · commit 145637419203 · 2023-04-25T10:11:30.000-05:00
Signed-off-by: ienkovich &lt;ilya.enkovich@intel.com&gt;
diff --git a/omniscidb/QueryEngine/Descriptors/RelAlgExecutionDescriptor.h b/omniscidb/QueryEngine/Descriptors/RelAlgExecutionDescriptor.h
@@ -51,6 +51,11 @@ class ExecutionResult {
 
   bool empty() const { return !result_token_; }
 
+  const std::string& tableName() const {
+    CHECK(!empty());
+    return result_token_->tableName();
+  }
+
   const std::vector<TargetMetaInfo>& getTargetsMeta() const { return targets_meta_; }
 
   const std::vector<PushedDownFilterInfo>& getPushedDownFilterInfo() const;
diff --git a/omniscidb/QueryEngine/RelAlgExecutor.cpp b/omniscidb/QueryEngine/RelAlgExecutor.cpp
@@ -108,7 +108,7 @@ RelAlgExecutor::RelAlgExecutor(Executor* executor,
                                std::unique_ptr<hdk::ir::QueryDag> query_dag)
     : executor_(executor)
     , query_dag_(std::move(query_dag))
-    , schema_provider_(std::make_shared<RelAlgSchemaProvider>(*query_dag_->getRootNode()))
+    , schema_provider_(schema_provider)
     , data_provider_(executor->getDataMgr()->getDataProvider())
     , config_(executor_->getConfig())
     , now_(0)
@@ -118,18 +118,14 @@ RelAlgExecutor::RelAlgExecutor(Executor* executor,
 
   // Add ResultSetRegistry to the schema provider by wrapping the current provider
   // and the registry in SchemaMgr.
-  std::set<int> used_schemas;
-  for (int db_id : schema_provider_->listDatabases()) {
-    used_schemas.insert(getSchemaId(db_id));
-  }
-  auto schema_mgr = std::make_shared<SchemaMgr>();
-  for (int schema_id : used_schemas) {
-    schema_mgr->registerProvider(schema_id, schema_provider_);
-  }
-  schema_mgr->registerProvider(
-      hdk::ResultSetRegistry::SCHEMA_ID,
-      std::dynamic_pointer_cast<hdk::ResultSetRegistry>(rs_registry_));
-  schema_provider_ = schema_mgr;
+  // TODO: In the future we expect pre-initialized registry and passed schema provider
+  // to cover it.
+  auto db_ids = schema_provider->listDatabases();
+  if (std::find(db_ids.begin(), db_ids.end(), hdk::ResultSetRegistry::DB_ID) ==
+      db_ids.end()) {
+    schema_provider_ =
+        mergeProviders(std::vector<SchemaProviderPtr>({schema_provider, rs_registry_}));
+  }
 }
 
 RelAlgExecutor::~RelAlgExecutor() {
diff --git a/omniscidb/ResultSetRegistry/ResultSetRegistry.cpp b/omniscidb/ResultSetRegistry/ResultSetRegistry.cpp
@@ -48,6 +48,9 @@ TableFragmentsInfo getEmptyTableMetadata(int table_id) {
 
 }  // namespace
 
+ResultSetRegistry::ResultSetRegistry(ConfigPtr config)
+    : ResultSetRegistry(config, "rs_registry") {}
+
 ResultSetRegistry::ResultSetRegistry(ConfigPtr config,
                                      const std::string& schema_name,
                                      int db_id)
@@ -89,10 +92,11 @@ ResultSetTableTokenPtr ResultSetRegistry::put(ResultSetTable table) {
   mapd_unique_lock<mapd_shared_mutex> data_lock(data_mutex_);
 
   auto table_id = next_table_id_++;
+  auto table_name = std::string("__result_set_") + std::to_string(table_id);
   // Add schema information for the ResultSet.
   auto tinfo = addTableInfo(db_id_,
                             table_id,
-                            ResultSetTableToken::tableName(table_id),
+                            table_name,
                             false,
                             Data_Namespace::MemoryLevel::CPU_LEVEL,
                             table.size());
@@ -105,6 +109,7 @@ ResultSetTableTokenPtr ResultSetRegistry::put(ResultSetTable table) {
                   first_rs->colType(col_idx),
                   false);
   }
+  addRowidColumn(db_id_, table_id);
 
   // TODO: lazily compute row count and try to avoid global write
   // locks for that
diff --git a/omniscidb/ResultSetRegistry/ResultSetRegistry.h b/omniscidb/ResultSetRegistry/ResultSetRegistry.h
@@ -29,9 +29,8 @@ class ResultSetRegistry : public SimpleSchemaProvider,
   constexpr static int SCHEMA_ID = 100;
   constexpr static int DB_ID = (SCHEMA_ID << 24) + 1;
 
-  ResultSetRegistry(ConfigPtr config,
-                    const std::string& schema_name = "rs_registry",
-                    int db_id = DB_ID);
+  ResultSetRegistry(ConfigPtr config);
+  ResultSetRegistry(ConfigPtr config, const std::string& schema_name, int db_id = DB_ID);
 
   static std::shared_ptr<ResultSetRegistry> getOrCreate(Data_Namespace::DataMgr* data_mgr,
                                                         ConfigPtr config);
diff --git a/omniscidb/ResultSetRegistry/ResultSetTableToken.h b/omniscidb/ResultSetRegistry/ResultSetTableToken.h
@@ -42,11 +42,7 @@ class ResultSetTableToken {
   size_t resultSetCount() const { return tinfo_->fragments; }
   ResultSetPtr resultSet(size_t idx) const;
 
-  static std::string tableName(int table_id) {
-    return std::string("__result_set_") + std::to_string(table_id);
-  }
-
-  std::string tableName() const { return tableName(tableId()); }
+  const std::string& tableName() const { return tinfo_->name; }
 
   std::string toString() const {
     return "ResultSetTableToken(" + std::to_string(dbId()) + ":" +
diff --git a/omniscidb/SchemaMgr/SchemaMgr.cpp b/omniscidb/SchemaMgr/SchemaMgr.cpp
@@ -80,3 +80,18 @@ const SchemaProvider* SchemaMgr::getMgr(int db_id) const {
   }
   return nullptr;
 }
+
+SchemaProviderPtr mergeProviders(const std::vector<SchemaProviderPtr>& providers) {
+  auto res = std::make_shared<SchemaMgr>();
+  for (auto& provider : providers) {
+    std::set<int> provider_schemas;
+    for (int db_id : provider->listDatabases()) {
+      int schema_id = getSchemaId(db_id);
+      provider_schemas.insert(schema_id);
+    }
+    for (int schema_id : provider_schemas) {
+      res->registerProvider(schema_id, provider);
+    }
+  }
+  return res;
+}
diff --git a/omniscidb/SchemaMgr/SchemaMgr.h b/omniscidb/SchemaMgr/SchemaMgr.h
@@ -38,3 +38,5 @@ class SchemaMgr : public SchemaProvider {
 };
 
 using SchemaMgrPtr = std::shared_ptr<SchemaMgr>;
+
+SchemaProviderPtr mergeProviders(const std::vector<SchemaProviderPtr>& providers);
diff --git a/omniscidb/Tests/ArrowSQLRunner/ArrowSQLRunner.cpp b/omniscidb/Tests/ArrowSQLRunner/ArrowSQLRunner.cpp
@@ -90,7 +90,7 @@ class ArrowSQLRunnerImpl {
     auto dag =
         std::make_unique<RelAlgDagBuilder>(query_ra, TEST_DB_ID, storage_, config_);
 
-    return std::make_unique<RelAlgExecutor>(executor_.get(), storage_, std::move(dag));
+    return std::make_unique<RelAlgExecutor>(executor_.get(), schema_mgr_, std::move(dag));
   }
 
   ExecutionResult runSqlQuery(const std::string& sql,
@@ -290,6 +290,10 @@ class ArrowSQLRunnerImpl {
 
   std::shared_ptr<ArrowStorage> getStorage() { return storage_; }
 
+  SchemaProviderPtr getSchemaProvider() { return schema_mgr_; }
+
+  std::shared_ptr<hdk::ResultSetRegistry> getResultSetRegistry() { return rs_registry_; }
+
   DataMgr* getDataMgr() { return data_mgr_.get(); }
 
   Executor* getExecutor() { return executor_.get(); }
@@ -299,6 +303,8 @@ class ArrowSQLRunnerImpl {
   ~ArrowSQLRunnerImpl() {
     executor_.reset();
     storage_.reset();
+    rs_registry_.reset();
+    schema_mgr_.reset();
     calcite_.reset();
     rel_alg_cache_.reset();
 
@@ -314,16 +320,21 @@ class ArrowSQLRunnerImpl {
     }
 
     storage_ = std::make_shared<ArrowStorage>(TEST_SCHEMA_ID, "test", TEST_DB_ID);
+    rs_registry_ = std::make_shared<hdk::ResultSetRegistry>(config_);
+    schema_mgr_ = std::make_shared<SchemaMgr>();
+    schema_mgr_->registerProvider(TEST_SCHEMA_ID, storage_);
+    schema_mgr_->registerProvider(hdk::ResultSetRegistry::SCHEMA_ID, rs_registry_);
 
     data_mgr_ = std::make_unique<DataMgr>(*config_);
     auto* ps_mgr = data_mgr_->getPersistentStorageMgr();
     ps_mgr->registerDataProvider(TEST_SCHEMA_ID, storage_);
+    ps_mgr->registerDataProvider(hdk::ResultSetRegistry::SCHEMA_ID, rs_registry_);
 
     executor_ = Executor::getExecutor(data_mgr_.get(), config_, "", "");
-    executor_->setSchemaProvider(storage_);
+    executor_->setSchemaProvider(schema_mgr_);
 
     if (config_->debug.use_ra_cache.empty() || !config_->debug.build_ra_cache.empty()) {
-      calcite_ = std::make_shared<CalciteJNI>(storage_, config_, udf_filename, 1024);
+      calcite_ = std::make_shared<CalciteJNI>(schema_mgr_, config_, udf_filename, 1024);
 
       if (config_->debug.use_ra_cache.empty()) {
         ExtensionFunctionsWhitelist::add(calcite_->getExtensionFunctionWhitelist());
@@ -336,13 +347,15 @@ class ArrowSQLRunnerImpl {
       calcite_->setRuntimeExtensionFunctions({}, /*is_runtime=*/false);
     }
 
-    rel_alg_cache_ = std::make_shared<RelAlgCache>(calcite_, storage_, config_);
+    rel_alg_cache_ = std::make_shared<RelAlgCache>(calcite_, schema_mgr_, config_);
   }
 
   ConfigPtr config_;
   std::unique_ptr<DataMgr> data_mgr_;
   std::shared_ptr<Executor> executor_;
   std::shared_ptr<ArrowStorage> storage_;
+  std::shared_ptr<hdk::ResultSetRegistry> rs_registry_;
+  std::shared_ptr<SchemaMgr> schema_mgr_;
   std::shared_ptr<CalciteJNI> calcite_;
   std::shared_ptr<RelAlgCache> rel_alg_cache_;
 
@@ -491,6 +504,14 @@ std::shared_ptr<ArrowStorage> getStorage() {
   return ArrowSQLRunnerImpl::get()->getStorage();
 }
 
+SchemaProviderPtr getSchemaProvider() {
+  return ArrowSQLRunnerImpl::get()->getSchemaProvider();
+}
+
+std::shared_ptr<hdk::ResultSetRegistry> getResultSetRegistry() {
+  return ArrowSQLRunnerImpl::get()->getResultSetRegistry();
+}
+
 DataMgr* getDataMgr() {
   return ArrowSQLRunnerImpl::get()->getDataMgr();
 }
diff --git a/omniscidb/Tests/ArrowSQLRunner/ArrowSQLRunner.h b/omniscidb/Tests/ArrowSQLRunner/ArrowSQLRunner.h
@@ -21,6 +21,7 @@
 #include "QueryEngine/ArrowResultSet.h"
 #include "QueryEngine/CompilationOptions.h"
 #include "QueryEngine/Descriptors/RelAlgExecutionDescriptor.h"
+#include "ResultSetRegistry/ResultSetRegistry.h"
 #include "Shared/Config.h"
 
 #include "BufferPoolStats.h"
@@ -115,6 +116,10 @@ BufferPoolStats getBufferPoolStats(const Data_Namespace::MemoryLevel memory_leve
 
 std::shared_ptr<ArrowStorage> getStorage();
 
+SchemaProviderPtr getSchemaProvider();
+
+std::shared_ptr<hdk::ResultSetRegistry> getResultSetRegistry();
+
 DataMgr* getDataMgr();
 
 Executor* getExecutor();
diff --git a/omniscidb/Tests/QueryBuilderTest.cpp b/omniscidb/Tests/QueryBuilderTest.cpp
@@ -269,6 +269,8 @@ class QueryBuilderTest : public TestSuite {
     schema_mgr_ = std::make_shared<SchemaMgr>();
     schema_mgr_->registerProvider(TEST_SCHEMA_ID, getStorage());
     schema_mgr_->registerProvider(TEST_SCHEMA_ID2, storage2_);
+    schema_mgr_->registerProvider(hdk::ResultSetRegistry::SCHEMA_ID,
+                                  getResultSetRegistry());
 
     createTable("test1",
                 {{"col_bi", ctx().int64()},
@@ -4677,6 +4679,75 @@ TEST_F(Issue355, Reproducer) {
   compare_res_data(res, std::vector<int64_t>({1}), std::vector<int64_t>({12}));
 }
 
+TEST_F(QueryBuilderTest, RunOnResult) {
+  QueryBuilder builder(ctx(), schema_mgr_, configPtr());
+
+  {
+    auto dag1 = builder.scan("test1").proj({0, 1}).finalize();
+    auto res1 = runQuery(std::move(dag1));
+    compare_res_data(res1,
+                     std::vector<int64_t>({1, 2, 3, 4, 5}),
+                     std::vector<int32_t>({11, 22, 33, 44, 55}));
+
+    auto dag2 = builder.scan(res1.tableName()).proj({1, 0}).finalize();
+    auto res2 = runQuery(std::move(dag2));
+    compare_res_data(res2,
+                     std::vector<int32_t>({11, 22, 33, 44, 55}),
+                     std::vector<int64_t>({1, 2, 3, 4, 5}));
+
+    auto scan = builder.scan(res2.tableName());
+    auto dag3 = scan.proj({scan["col_i"] + 1, scan["col_bi"] + 2}).finalize();
+    auto res3 = runQuery(std::move(dag3));
+    compare_res_data(res3,
+                     std::vector<int32_t>({12, 23, 34, 45, 56}),
+                     std::vector<int64_t>({3, 4, 5, 6, 7}));
+  }
+
+  {
+    auto scan1 = builder.scan("test1");
+    auto dag1 = scan1.proj({scan1["col_bi"] + 1, scan1["col_f"]}).finalize();
+    auto res1 = runQuery(std::move(dag1));
+    compare_res_data(res1,
+                     std::vector<int64_t>({2, 3, 4, 5, 6}),
+                     std::vector<float>({1.1, 2.2, 3.3, 4.4, 5.5}));
+
+    auto scan2 = builder.scan("test1");
+    auto dag2 = scan2.proj({scan2["col_bi"] + 2, scan2["col_d"]}).finalize();
+    auto res2 = runQuery(std::move(dag2));
+    compare_res_data(res2,
+                     std::vector<int64_t>({3, 4, 5, 6, 7}),
+                     std::vector<double>({11.11, 22.22, 33.33, 44.44, 55.55}));
+
+    auto dag3 =
+        builder.scan(res1.tableName()).join(builder.scan(res2.tableName())).finalize();
+    auto res3 = runQuery(std::move(dag3));
+    compare_res_data(res3,
+                     std::vector<int64_t>({3, 4, 5, 6}),
+                     std::vector<float>({2.2, 3.3, 4.4, 5.5}),
+                     std::vector<double>({11.11, 22.22, 33.33, 44.44}));
+  }
+}
+
+TEST_F(QueryBuilderTest, RowidOnResult) {
+  QueryBuilder builder(ctx(), schema_mgr_, configPtr());
+
+  {
+    auto dag1 = builder.scan("test1").proj({0, 1}).finalize();
+    auto res1 = runQuery(std::move(dag1));
+    compare_res_data(res1,
+                     std::vector<int64_t>({1, 2, 3, 4, 5}),
+                     std::vector<int32_t>({11, 22, 33, 44, 55}));
+
+    auto dag2 =
+        builder.scan(res1.tableName()).proj({"col_i", "col_bi", "rowid"}).finalize();
+    auto res2 = runQuery(std::move(dag2));
+    compare_res_data(res2,
+                     std::vector<int32_t>({11, 22, 33, 44, 55}),
+                     std::vector<int64_t>({1, 2, 3, 4, 5}),
+                     std::vector<int64_t>({0, 1, 2, 3, 4}));
+  }
+}
+
 class Taxi : public TestSuite {
  protected:
   static void SetUpTestSuite() {
diff --git a/python/pyhdk/_builder.pyx b/python/pyhdk/_builder.pyx
@@ -450,6 +450,8 @@ cdef class QueryNode:
     return res
 
   def join(self, rhs_node, lhs_cols=None, rhs_cols=None, cond=None, how="inner"):
+    if isinstance(rhs_node, ExecutionResult):
+      rhs_node = rhs_node.scan
     if not isinstance(rhs_node, QueryNode):
       raise TypeError(f"Expected QueryNode for 'rhs_node' arg. Provided: {type(rhs_node)}.")
     if not isinstance(how, str):
@@ -604,7 +606,9 @@ cdef class QueryNode:
     dag = QueryDag()
     dag.c_dag.reset(c_dag)
     rel_alg_executor = RelAlgExecutor(self._hdk._executor, self._hdk._storage, self._hdk._data_mgr, dag=dag)
-    return rel_alg_executor.execute(**kwargs)
+    res = rel_alg_executor.execute(**kwargs)
+    res.scan = self._hdk.scan(res.table_name)
+    return res
 
   def finalize(self):
     cdef CQueryDag* c_dag = self.c_node.finalize().release()
diff --git a/python/pyhdk/_execute.pxd b/python/pyhdk/_execute.pxd
@@ -12,7 +12,7 @@ from libcpp.utility cimport move
 from pyarrow.lib cimport CTable as CArrowTable
 
 from pyhdk._common cimport CType, CConfig
-from pyhdk._storage cimport CDataMgr, CBufferProvider
+from pyhdk._storage cimport CDataMgr, CBufferProvider, CSchemaProvider, CAbstractDataProvider
 
 cdef extern from "omniscidb/QueryEngine/Compiler/CodegenTraitsDescriptor.h" namespace "compiler":
   enum CCallingConvDesc "CallingConvDesc":
@@ -101,6 +101,10 @@ cdef extern from "omniscidb/ResultSet/ResultSet.h":
 
 ctypedef shared_ptr[CResultSet] CResultSetPtr
 
+cdef extern from "omniscidb/ResultSetRegistry/ResultSetRegistry.h":
+  cdef cppclass CResultSetRegistry "hdk::ResultSetRegistry"(CSchemaProvider, CAbstractDataProvider):
+    CResultSetRegistry(shared_ptr[CConfig]) except +;
+
 cdef extern from "omniscidb/QueryEngine/TargetMetaInfo.h":
   cdef cppclass CTargetMetaInfo "TargetMetaInfo":
     const string& get_resname()
diff --git a/python/pyhdk/_execute.pyx b/python/pyhdk/_execute.pyx
@@ -3,11 +3,21 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from libcpp.memory cimport shared_ptr, make_shared, static_pointer_cast
+
 from pyhdk._common cimport Config
-from pyhdk._storage cimport DataMgr
+from pyhdk._storage cimport DataMgr, Storage, CAbstractBufferMgr, CSchemaProvider
 
 cdef class Executor:
   def __cinit__(self, DataMgr data_mgr, Config config):
     cdef string debug_dir = "".encode('UTF-8')
     cdef string debug_file = "".encode('UTF-8')
     self.c_executor = CExecutor.getExecutor(data_mgr.c_data_mgr.get(), config.c_config, debug_dir, debug_file)
+
+cdef class ResultSetRegistry(Storage):
+  cdef shared_ptr[CResultSetRegistry] c_registry
+
+  def __cinit__(self, Config config):
+    self.c_registry = make_shared[CResultSetRegistry](config.c_config)
+    self.c_schema_provider = static_pointer_cast[CSchemaProvider, CResultSetRegistry](self.c_registry)
+    self.c_abstract_buffer_mgr = static_pointer_cast[CAbstractBufferMgr, CResultSetRegistry](self.c_registry)
diff --git a/python/pyhdk/_sql.pxd b/python/pyhdk/_sql.pxd
@@ -57,6 +57,7 @@ cdef extern from "omniscidb/QueryEngine/Descriptors/RelAlgExecutionDescriptor.h"
     const CResultSetPtr& getRows()
     const vector[CTargetMetaInfo]& getTargetsMeta()
     string getExplanation()
+    const string& tableName()
 
 cdef class ExecutionResult:
   cdef CExecutionResult c_result
@@ -65,6 +66,10 @@ cdef class ExecutionResult:
   # obects lifetime control. In Python we achieve it by holding DataMgr in
   # each ExecutionResult object.
   cdef shared_ptr[CDataMgr] c_data_mgr
+  # Scan object is used to forward query builder related calls. It is used
+  # to provide an ability to work with an execution result as with a regular
+  # table
+  cdef object _scan
 
 cdef extern from "omniscidb/QueryEngine/RelAlgExecutor.h":
   cdef cppclass CRelAlgExecutor "RelAlgExecutor":
diff --git a/python/pyhdk/_sql.pyx b/python/pyhdk/_sql.pyx
diff --git a/python/pyhdk/_storage.pxd b/python/pyhdk/_storage.pxd
diff --git a/python/pyhdk/_storage.pyx b/python/pyhdk/_storage.pyx
diff --git a/python/pyhdk/hdk.py b/python/pyhdk/hdk.py
diff --git a/python/tests/test_pyhdk_api.py b/python/tests/test_pyhdk_api.py
