Support dictionary proxies in ResultSetRegistry.

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

Author: ienkovich

omniscidb/QueryEngine/Execute.cpp

@@ -1170,7 +1170,7 @@ ResultSetPtr Executor::reduceMultiDeviceResults(
     return std::make_shared<ResultSet>(targets,
                                        ExecutorDeviceType::CPU,
                                        QueryMemoryDescriptor(),
-                                       nullptr,
+                                       row_set_mem_owner,
                                        data_mgr_,
                                        blockSize(),
                                        gridSize());

omniscidb/ResultSet/ResultSet.cpp

@@ -493,6 +493,11 @@ StringDictionary* ResultSet::getStringDictionaryProxy(int const dict_id) const {
   return row_set_mem_owner_->getOrAddStringDictProxy(dict_id);
 }
 
+std::shared_ptr<StringDictionary> ResultSet::getStringDictionaryProxyOwned(
+    int const dict_id) const {
+  return row_set_mem_owner_->getStringDictProxyOwned(dict_id);
+}
+
 class ResultSet::CellCallback {
   std::vector<int32_t> const id_map_;
   int64_t const null_int_;

omniscidb/ResultSet/ResultSet.h

@@ -524,6 +524,8 @@ class ResultSet {
   getUniqueStringsForDictEncodedTargetCol(const size_t col_idx) const;
 
   StringDictionary* getStringDictionaryProxy(int const dict_id) const;
+  std::shared_ptr<StringDictionary> getStringDictionaryProxyOwned(
+      int const dict_id) const;
 
   template <typename ENTRY_TYPE, QueryDescriptionType QUERY_TYPE, bool COLUMNAR_FORMAT>
   ENTRY_TYPE getEntryAt(const size_t row_idx,

omniscidb/ResultSet/RowSetMemoryOwner.cpp

@@ -48,6 +48,14 @@ StringDictionary* RowSetMemoryOwner::getOrAddStringDictProxy(const int dict_id_i
     CHECK_LE(dd->dictNBits, 32);
     return addStringDict(dd->stringDict, dict_id, generation);
   }
+  // It's possible the original dictionary has been removed from its storage
+  // but we still have it in a proxy.
+  if (dict_id != DictRef::literalsDictId) {
+    auto res = getStringDictProxyOwned(dict_id_in);
+    CHECK(res) << "Cannot find dict or proxy " << dict_id_in;
+    CHECK(generation < 0 || res->getBaseGeneration() == generation);
+    return res.get();
+  }
   CHECK_EQ(dict_id, DictRef::literalsDictId);
   if (!lit_str_dict_proxy_) {
     DictRef literal_dict_ref(DictRef::invalidDbId, DictRef::literalsDictId);
@@ -58,6 +66,19 @@ StringDictionary* RowSetMemoryOwner::getOrAddStringDictProxy(const int dict_id_i
   return lit_str_dict_proxy_.get();
 }
 
+std::shared_ptr<StringDictionary> RowSetMemoryOwner::getStringDictProxyOwned(
+    const int dict_id) {
+  std::lock_guard<std::mutex> lock(state_mutex_);
+  if (dict_id == DictRef::literalsDictId) {
+    return lit_str_dict_proxy_;
+  }
+  auto it = str_dict_proxy_owned_.find(dict_id);
+  if (it != str_dict_proxy_owned_.end()) {
+    return it->second;
+  }
+  return nullptr;
+}
+
 quantile::TDigest* RowSetMemoryOwner::nullTDigest(double const q) {
   std::lock_guard<std::mutex> lock(state_mutex_);
   return t_digests_

omniscidb/ResultSet/RowSetMemoryOwner.h

@@ -232,6 +232,7 @@ class RowSetMemoryOwner final : public SimpleAllocator, boost::noncopyable {
 
   StringDictionary* getOrAddStringDictProxy(const int dict_id_in,
                                             const int64_t generation = -1);
+  std::shared_ptr<StringDictionary> getStringDictProxyOwned(const int dict_id);
 
   void addLiteralStringDictProxy(std::shared_ptr<StringDictionary> lit_str_dict_proxy) {
     std::lock_guard<std::mutex> lock(state_mutex_);

omniscidb/ResultSetRegistry/ResultSetRegistry.cpp

@@ -113,14 +113,55 @@ ResultSetTableTokenPtr ResultSetRegistry::put(ResultSetTable table) {
   bool has_varlen = false;
   bool has_array = false;
   for (size_t col_idx = 0; col_idx < first_rs->colCount(); ++col_idx) {
-    addColumnInfo(db_id_,
-                  table_id,
-                  columnId(col_idx),
-                  first_rs->colName(col_idx),
-                  first_rs->colType(col_idx),
-                  false);
-    has_varlen = has_varlen || first_rs->colType(col_idx)->isVarLen();
-    has_array = has_array || first_rs->colType(col_idx)->isArray();
+    auto col_type = first_rs->colType(col_idx);
+    if (col_type->isExtDictionary()) {
+      auto dict_type = col_type->as<hdk::ir::ExtDictionaryType>();
+      auto dict_proxy = first_rs->getStringDictionaryProxyOwned(dict_type->dictId());
+      // If there is no proxy and dictionary is owned by the registry then simply
+      // re-use it.
+      if (!dict_proxy && getSchemaId(dict_type->dictId()) == schema_id_) {
+        CHECK(dicts_.count(dict_type->dictId()));
+        dict_proxy = dicts_.at(dict_type->dictId())->dict_descriptor->stringDict;
+      }
+      // If we have a dictionary with no proxy then it's not safe to use because it can
+      // simply be removed from its storage. To avoid this, we create a proxy dictionary.
+      // We require result set to have a row set memory owner for that.
+      if (!dict_proxy) {
+        CHECK(dict_type->dictId() > 0) << dict_type->toString();
+        CHECK(first_rs->getRowSetMemOwner());
+        CHECK(first_rs->getStringDictionaryProxy(dict_type->dictId()));
+        dict_proxy = first_rs->getStringDictionaryProxyOwned(dict_type->dictId());
+        CHECK(dict_proxy);
+      }
+      // Add proxy as a new dictionary if it exists and is not registered yet.
+      if (dict_proxy->getDictId() <= 0) {
+        auto dict_id = addSchemaIdChecked(next_dict_id_++, schema_id_);
+        dict_proxy->setDictId(dict_id);
+        auto dict_desc = std::make_unique<DictDescriptor>(db_id_,
+                                                          dict_id,
+                                                          first_rs->colName(col_idx),
+                                                          /*nbits=*/dict_type->size() * 8,
+                                                          /*is_shared=*/true,
+                                                          /*refcount=*/1,
+                                                          table_name,
+                                                          /*temp=*/true);
+        dict_desc->stringDict = dict_proxy;
+        auto dict_data = std::make_unique<DictionaryData>();
+        dict_data->dict_descriptor = std::move(dict_desc);
+        CHECK(!dicts_.count(dict_id));
+        dicts_.emplace(dict_id, std::move(dict_data));
+      }
+      // Register table as a dictionary user and fix-up dictionary type.
+      auto dict_id = dict_proxy->getDictId();
+      CHECK_EQ(getSchemaId(dict_id), schema_id_);
+      CHECK(dicts_.count(dict_id));
+      dicts_.at(dict_id)->table_ids.insert(table_id);
+      col_type = hdk::ir::Context::defaultCtx().extDict(dict_type->elemType(), dict_id);
+    }
+    addColumnInfo(
+        db_id_, table_id, columnId(col_idx), first_rs->colName(col_idx), col_type, false);
+    has_varlen = has_varlen || col_type->isVarLen();
+    has_array = has_array || col_type->isArray();
   }
   addRowidColumn(db_id_, table_id, columnId(first_rs->colCount()));
 
@@ -169,11 +210,31 @@ void ResultSetRegistry::drop(const ResultSetTableToken& token) {
   mapd_unique_lock<mapd_shared_mutex> schema_lock(schema_mutex_);
   mapd_unique_lock<mapd_shared_mutex> data_lock(data_mutex_);
 
+  // Drop data.
   CHECK(tables_.count(token.tableId()));
   std::unique_ptr<TableData> table = std::move(tables_.at(token.tableId()));
   mapd_unique_lock<mapd_shared_mutex> table_lock(table->mutex);
   tables_.erase(token.tableId());
 
+  // Drop dicts.
+  auto cols = listColumnsNoLock(db_id_, token.tableId());
+  std::unordered_set<int> used_dicts;
+  for (auto& col : cols) {
+    if (auto dict_type = col->type->as<hdk::ir::ExtDictionaryType>()) {
+      used_dicts.insert(dict_type->dictId());
+    }
+  }
+  for (int dict_id : used_dicts) {
+    CHECK_EQ(getSchemaId(dict_id), schema_id_);
+    CHECK(dicts_.count(dict_id));
+    auto& desc = *dicts_.at(dict_id);
+    desc.table_ids.erase(token.tableId());
+    if (desc.table_ids.empty()) {
+      dicts_.erase(dict_id);
+    }
+  }
+
+  // Drop schema.
   SimpleSchemaProvider::dropTable(token.dbId(), token.tableId());
 }
 
@@ -514,8 +575,11 @@ TableFragmentsInfo ResultSetRegistry::getTableMetadata(int db_id, int table_id)
 }
 
 const DictDescriptor* ResultSetRegistry::getDictMetadata(int dict_id, bool load_dict) {
-  // Currently, we don't hold any dictionaries in the registry.
-  UNREACHABLE();
+  mapd_shared_lock<mapd_shared_mutex> data_lock(data_mutex_);
+  CHECK_EQ(getSchemaId(dict_id), schema_id_);
+  if (dicts_.count(dict_id)) {
+    return dicts_.at(dict_id)->dict_descriptor.get();
+  }
   return nullptr;
 }
 

omniscidb/ResultSetRegistry/ResultSetRegistry.h

@@ -81,10 +81,17 @@ class ResultSetRegistry : public SimpleSchemaProvider,
     TableStats table_stats;
   };
 
+  struct DictionaryData {
+    std::unique_ptr<DictDescriptor> dict_descriptor;
+    std::set<int> table_ids;
+  };
+
   const int db_id_;
   const int schema_id_;
   int next_table_id_ = 1;
+  int next_dict_id_ = 1;
   std::unordered_map<int, std::unique_ptr<TableData>> tables_;
+  std::unordered_map<int, std::unique_ptr<DictionaryData>> dicts_;
   const ConfigPtr config_;
   mutable mapd_shared_mutex data_mutex_;
 };

omniscidb/StringDictionary/StringDictionary.h

@@ -89,6 +89,11 @@ class StringDictionary {
 
   int32_t getDbId() const noexcept;
   int32_t getDictId() const noexcept;
+  void setDictId(int32_t new_id) {
+    CHECK_EQ(dict_ref_.dictId, -1);
+    CHECK_GT(new_id, 0);
+    dict_ref_.dictId = new_id;
+  }
 
   StringDictionary* getBaseDictionary() const noexcept { return base_dict_.get(); }
   int64_t getBaseGeneration() const noexcept { return base_generation_; }
@@ -249,7 +254,7 @@ class StringDictionary {
   uint32_t hashById(int string_id) const { return hashByIndex(idToIndex(string_id)); }
   uint32_t hashByIndex(int string_idx) const { return hash_cache_[string_idx]; }
 
-  const DictRef dict_ref_;
+  DictRef dict_ref_;
   const std::shared_ptr<StringDictionary> base_dict_;
   const int64_t base_generation_;
   size_t str_count_;

omniscidb/Tests/QueryBuilderTest.cpp

@@ -7020,13 +7020,19 @@ class Issue588 : public TestSuite {
     insertCsvValues("test_588_3", "1,1\n2,2\n3,3");
     createTable("test_588_4", {{"id", ctx().int64()}, {"D", ctx().int64()}});
     insertCsvValues("test_588_4", "1,2\n2,3\n3,3");
+    createTable("test_588_5",
+                {{"id", ctx().int32()},
+                 {"s1", ctx().extDict(ctx().text(), 0)},
+                 {"s2", ctx().extDict(ctx().text(), 0)}});
+    insertCsvValues("test_588_5", "1,str1,str22\n2,str2,str22\n3,str3,str33");
   }
 
   static void TearDownTestSuite() {
     dropTable("test_588_1");
     dropTable("test_588_2");
     dropTable("test_588_3");
     dropTable("test_588_4");
+    dropTable("test_588_5");
   }
 };
 
@@ -7049,6 +7055,89 @@ TEST_F(Issue588, Reproducer1) {
   auto res3 = runQuery(std::move(dag3));
 }
 
+TEST_F(Issue588, Reproducer2) {
+  QueryBuilder builder(ctx(), getSchemaProvider(), configPtr());
+  auto scan1 = builder.scan("test_588_1");
+  auto dag1 = scan1.proj({scan1.ref("A"), builder.cst("str")}).finalize();
+  auto res1 = runQuery(std::move(dag1));
+
+  auto scan2 = builder.scan(res1.tableName());
+  auto dag2 = scan2.proj({0, 1}).finalize();
+  auto res2 = runQuery(std::move(dag2));
+}
+
+TEST_F(Issue588, Reproducer3) {
+  QueryBuilder builder(ctx(), getSchemaProvider(), configPtr());
+  auto scan1 = builder.scan("test_588_5");
+  auto dict1_type = scan1.ref("s1").type();
+  auto dag1 = scan1
+                  .proj({scan1.ref("id"),
+                         builder.ifThenElse(scan1.ref("id").gt(2),
+                                            scan1.ref("s1"),
+                                            builder.cst("str").cast(dict1_type))})
+                  .finalize();
+  auto res1 = runQuery(std::move(dag1));
+
+  auto scan2 = builder.scan(res1.tableName());
+  auto dag2 = scan2.proj({0, 1}).finalize();
+  auto res2 = runQuery(std::move(dag2));
+}
+
+TEST_F(Issue588, Reproducer4) {
+  QueryBuilder builder(ctx(), getSchemaProvider(), configPtr());
+  auto scan1 = builder.scan("test_588_5");
+  auto dict1_type = scan1.ref("s1").type();
+  auto dag1 = scan1
+                  .proj({scan1.ref("id"),
+                         builder.ifThenElse(scan1.ref("id").gt(2),
+                                            scan1.ref("s1"),
+                                            scan1.ref("s2").cast(dict1_type))})
+                  .finalize();
+  auto res1 = runQuery(std::move(dag1));
+
+  auto scan2 = builder.scan(res1.tableName());
+  auto dag2 = scan2.proj({0, 1}).finalize();
+  auto res2 = runQuery(std::move(dag2));
+}
+
+TEST_F(Issue588, Reproducer5) {
+  QueryBuilder builder(ctx(), getSchemaProvider(), configPtr());
+  auto scan1 = builder.scan("test_588_5");
+  auto dict1_type = scan1.ref("s1").type();
+  auto dag1 = scan1
+                  .proj({scan1.ref("id"),
+                         builder.ifThenElse(
+                             scan1.ref("id").gt(2),
+                             scan1.ref("s1"),
+                             builder.ifThenElse(scan1.ref("id").gt(1),
+                                                scan1.ref("s2").cast(dict1_type),
+                                                builder.cst("str").cast(dict1_type)))})
+                  .finalize();
+  auto res1 = runQuery(std::move(dag1));
+
+  auto scan2 = builder.scan(res1.tableName());
+  auto dag2 = scan2.proj({0, 1}).finalize();
+  auto res2 = runQuery(std::move(dag2));
+}
+
+class Issue667 : public TestSuite {
+ protected:
+  static void SetUpTestSuite() {
+    createTable("test_667", {{"str", ctx().extDict(ctx().text(), 0)}});
+    insertCsvValues("test_667", "1\n2\n3");
+  }
+
+  static void TearDownTestSuite() {}
+};
+
+TEST_F(Issue667, DropInputThenConvertResult) {
+  QueryBuilder builder(ctx(), getSchemaProvider(), configPtr());
+  auto dag = builder.scan("test_667").proj("str").finalize();
+  auto res = runQuery(std::move(dag));
+  dropTable("test_667");
+  compare_res_data(res, std::vector<std::string>({"1"s, "2"s, "3"s}));
+}
+
 TEST_F(QueryBuilderTest, RunOnResult) {
   QueryBuilder builder(ctx(), schema_mgr_, configPtr());