Enable SQL queries on result sets.

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

Author: ienkovich

omniscidb/Calcite/SchemaJson.cpp

@@ -188,29 +188,28 @@ std::string schema_to_json(SchemaProvider* schema_provider) {
   if (dbs.empty()) {
     return "{}";
   }
-  // Current JSON format supports a single database only. So, we exclude
-  // ResultSetRegistry from the schema for now (which makes it impossible
-  // to run SQL queries on result sets).
-  int db_id;
-  if (dbs.size() == (size_t)1) {
-    db_id = dbs.front();
-  } else {
+  // Current JSON format supports a single database only. To support result
+  // sets in SQL queries, we add tables from the ResultSetRegistry using
+  // negative table ids.
+  auto tables = schema_provider->listTables(dbs[0]);
+  if (dbs.size() != (size_t)1) {
     CHECK_EQ(dbs.size(), (size_t)2);
     CHECK(dbs[0] == hdk::ResultSetRegistry::DB_ID ||
           dbs[1] == hdk::ResultSetRegistry::DB_ID);
-    db_id = dbs[0] == hdk::ResultSetRegistry::DB_ID ? dbs[1] : dbs[0];
+    auto more_tables = schema_provider->listTables(dbs[1]);
+    tables.insert(tables.end(), more_tables.begin(), more_tables.end());
   }
 
-  auto tables = schema_provider->listTables(db_id);
-
   rapidjson::Document doc(rapidjson::kObjectType);
 
   for (auto tinfo : tables) {
     rapidjson::Value table(rapidjson::kObjectType);
     table.AddMember("name",
                     rapidjson::Value().SetString(rapidjson::StringRef(tinfo->name)),
                     doc.GetAllocator());
-    table.AddMember("id", rapidjson::Value().SetInt(tinfo->table_id), doc.GetAllocator());
+    int table_id = tinfo->db_id == hdk::ResultSetRegistry::DB_ID ? -tinfo->table_id
+                                                                 : tinfo->table_id;
+    table.AddMember("id", rapidjson::Value().SetInt(table_id), doc.GetAllocator());
     table.AddMember(
         "columns", rapidjson::Value(rapidjson::kArrayType), doc.GetAllocator());
 

omniscidb/QueryEngine/RelAlgDagBuilder.cpp

@@ -27,6 +27,7 @@
 #include "JsonAccessors.h"
 #include "RelAlgDagBuilder.h"
 #include "RelAlgOptimizer.h"
+#include "ResultSetRegistry/ResultSetRegistry.h"
 #include "ScalarExprVisitor.h"
 #include "Shared/sqldefs.h"
 
@@ -2019,7 +2020,13 @@ TableInfoPtr getTableFromScanNode(int db_id,
   const auto& table_json = field(scan_ra, "table");
   CHECK(table_json.IsArray());
   CHECK_EQ(unsigned(2), table_json.Size());
-  const auto info = schema_provider->getTableInfo(db_id, table_json[1].GetString());
+  auto info = schema_provider->getTableInfo(db_id, table_json[1].GetString());
+  // If table wasn't found in the default database, then try search in the
+  // result set registry.
+  if (!info) {
+    info = schema_provider->getTableInfo(hdk::ResultSetRegistry::DB_ID,
+                                         table_json[1].GetString());
+  }
   CHECK(info);
   return info;
 }

omniscidb/Tests/ArrowSQLRunner/ArrowSQLRunner.cpp

@@ -88,7 +88,7 @@ class ArrowSQLRunnerImpl {
     std::string query_ra = getSqlQueryRelAlg(sql);
 
     auto dag =
-        std::make_unique<RelAlgDagBuilder>(query_ra, TEST_DB_ID, storage_, config_);
+        std::make_unique<RelAlgDagBuilder>(query_ra, TEST_DB_ID, schema_mgr_, config_);
 
     return std::make_unique<RelAlgExecutor>(executor_.get(), schema_mgr_, std::move(dag));
   }

omniscidb/Tests/QueryBuilderTest.cpp

@@ -4748,6 +4748,29 @@ TEST_F(QueryBuilderTest, RowidOnResult) {
   }
 }
 
+TEST_F(QueryBuilderTest, SqlOnResult) {
+  QueryBuilder builder(ctx(), schema_mgr_, configPtr());
+
+  auto res1 =
+      runSqlQuery("SELECT col_bi, col_i FROM test1;", ExecutorDeviceType::CPU, false);
+  compare_res_data(res1,
+                   std::vector<int64_t>({1, 2, 3, 4, 5}),
+                   std::vector<int32_t>({11, 22, 33, 44, 55}));
+
+  auto dag = builder.scan(res1.tableName()).proj({1, 0}).finalize();
+  auto res2 = runQuery(std::move(dag));
+  compare_res_data(res2,
+                   std::vector<int32_t>({11, 22, 33, 44, 55}),
+                   std::vector<int64_t>({1, 2, 3, 4, 5}));
+
+  auto res3 = runSqlQuery("SELECT col_bi + 1, col_i - 1 FROM " + res2.tableName() + ";",
+                          ExecutorDeviceType::CPU,
+                          false);
+  compare_res_data(res3,
+                   std::vector<int64_t>({2, 3, 4, 5, 6}),
+                   std::vector<int32_t>({10, 21, 32, 43, 54}));
+}
+
 class Taxi : public TestSuite {
  protected:
   static void SetUpTestSuite() {

python/pyhdk/_sql.pyx

@@ -107,10 +107,13 @@ cdef class RelAlgExecutor:
     cdef unique_ptr[CQueryDag] c_dag
     cdef int db_id = 0
 
+    # Choose the default database ID. Ignore ResultSetRegistry.
     db_ids = schema_provider.listDatabases()
-    assert len(db_ids) <= 1
+    assert len(db_ids) <= 2
     if len(db_ids) == 1:
       db_id = db_ids[0]
+    elif len(db_ids) == 2:
+      db_id = db_ids[1] if db_ids[0] == ((100 << 24) + 1) else db_ids[0]
 
     if ra_json is not None:
       c_dag.reset(new CRelAlgDagBuilder(ra_json, db_id, c_schema_provider, c_executor.getConfigPtr()))

python/pyhdk/hdk.py

@@ -11,7 +11,7 @@
     DataMgr,
     SchemaMgr,
 )
-from pyhdk._sql import Calcite, RelAlgExecutor
+from pyhdk._sql import Calcite, RelAlgExecutor, ExecutionResult
 from pyhdk._execute import Executor, ResultSetRegistry
 from pyhdk._builder import QueryBuilder, QueryExpr, QueryNode
 
@@ -2035,7 +2035,10 @@ def sql(self, sql_query, query_opts=None, **kwargs):
 
         parts = []
         for name, orig_table in kwargs.items():
-            if isinstance(orig_table, QueryNode) and orig_table.is_scan:
+            if (
+                isinstance(orig_table, (QueryNode, ExecutionResult))
+                and orig_table.is_scan
+            ):
                 orig_table = orig_table.table_name
             if not isinstance(orig_table, str):
                 raise TypeError(
@@ -2050,8 +2053,12 @@ def sql(self, sql_query, query_opts=None, **kwargs):
 
         sql_query = "".join(parts) + sql_query
         ra = self._calcite.process(sql_query)
-        ra_executor = RelAlgExecutor(self._executor, self._storage, self._data_mgr, ra)
-        return ra_executor.execute(**query_opts)
+        ra_executor = RelAlgExecutor(
+            self._executor, self._schema_mgr, self._data_mgr, ra
+        )
+        res = ra_executor.execute(**query_opts)
+        res.scan = self.scan(res.table_name)
+        return res
 
     def scan(self, table_name):
         """

python/tests/test_pyhdk_api.py

@@ -925,6 +925,21 @@ def test_alias(self):
             },
         )
 
+    def test_run_on_res(self):
+        hdk = pyhdk.init()
+        ht1 = hdk.import_pydict(
+            {"a": [1, 2, 3, 4, 5], "b": [5, 4, 3, 2, 1], "x": [1.1, 2.2, 3.3, 4.4, 5.5]}
+        )
+
+        res1 = hdk.sql("SELECT a, b FROM t1;", t1=ht1)
+        self.check_res(res1, {"a": [1, 2, 3, 4, 5], "b": [5, 4, 3, 2, 1]})
+
+        res2 = hdk.sql("SELECT b + 1 as b, a - 1 as a FROM t1;", t1=res1)
+        self.check_res(res2, {"b": [6, 5, 4, 3, 2], "a": [0, 1, 2, 3, 4]})
+
+        res3 = hdk.sql(f"SELECT b - 1 as b, a + 1 as a FROM {res1.table_name};")
+        self.check_res(res3, {"b": [4, 3, 2, 1, 0], "a": [2, 3, 4, 5, 6]})
+
 
 class BaseTaxiTest:
     @staticmethod