body

Author: Devjiu

omniscidb/ArrowStorage/ArrowStorage.cpp

@@ -187,6 +187,50 @@ std::unique_ptr<AbstractDataToken> ArrowStorage::getZeroCopyBufferMemory(
   return nullptr;
 }
 
+std::unique_ptr<AbstractDataToken> ArrowStorage::getZeroCopyColumnData(
+    const ColumnRef& col_ref) {
+  mapd_shared_lock<mapd_shared_mutex> data_lock(data_mutex_);
+  CHECK_EQ(col_ref.db_id, db_id_);
+  CHECK_EQ(tables_.count(col_ref.table_id), (size_t)1);
+  auto& table = *tables_.at(col_ref.table_id);
+  mapd_shared_lock<mapd_shared_mutex> table_lock(table.mutex);
+  data_lock.unlock();
+
+  auto col_type = getColumnInfo(col_ref.db_id, col_ref.table_id, col_ref.column_id)->type;
+
+  if (col_type->isExtDictionary()) {
+    auto dict_id = col_type->as<hdk::ir::ExtDictionaryType>()->dictId();
+    auto dict_descriptor = getDictMetadata(
+        dict_id);  // this will force materialize the dictionary. it is thread safe
+    CHECK(dict_descriptor);
+  }
+
+  if (!col_type->isVarLen()) {
+    size_t col_idx = columnIndex(col_ref.column_id);
+    size_t elem_size = col_type->size();
+    const auto* fixed_type =
+        dynamic_cast<const arrow::FixedWidthType*>(table.col_data[col_idx]->type().get());
+    CHECK(fixed_type) << table.col_data[col_idx]->type()->ToString() << " (table "
+                      << col_ref.table_id << ", column " << col_idx << ")";
+    size_t arrow_elem_size = fixed_type->bit_width() / 8;
+    // For fixed size arrays we simply use elem type in arrow and therefore have to scale
+    // to get a proper slice.
+    size_t elems = elem_size / arrow_elem_size;
+    CHECK_GT(elems, (size_t)0);
+    auto data_to_fetch = table.col_data[col_idx];
+    if (data_to_fetch->num_chunks() == 1) {
+      auto chunk = data_to_fetch->chunk(0);
+      const int8_t* ptr =
+          chunk->data()->GetValues<int8_t>(1, chunk->data()->offset * arrow_elem_size);
+      size_t chunk_size = chunk->length() * arrow_elem_size;
+      return std::make_unique<ArrowChunkDataToken>(
+          std::move(chunk), col_type, ptr, chunk_size);
+    }
+  }
+
+  return nullptr;
+}
+
 void ArrowStorage::fetchFixedLenData(const TableData& table,
                                      size_t frag_idx,
                                      size_t col_idx,

omniscidb/ArrowStorage/ArrowStorage.h

@@ -70,6 +70,9 @@ class ArrowStorage : public SimpleSchemaProvider, public AbstractDataProvider {
       const ChunkKey& key,
       size_t num_bytes) override;
 
+  std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override;
+
   TableFragmentsInfo getTableMetadata(int db_id, int table_id) const override;
 
   const DictDescriptor* getDictMetadata(int dict_id, bool load_dict = true) override;

omniscidb/DataMgr/AbstractBufferMgr.h

@@ -84,8 +84,11 @@ class AbstractBufferMgr {
   virtual void deleteBuffersWithPrefix(const ChunkKey& keyPrefix,
                                        const bool purge = true) = 0;
   virtual AbstractBuffer* getBuffer(const ChunkKey& key, const size_t numBytes = 0) = 0;
+  // virtual AbstractBuffer* getZeroCopyColumnBuffer(const ColumnRef &col_ref) = 0;
   virtual std::unique_ptr<AbstractDataToken> getZeroCopyBufferMemory(const ChunkKey& key,
                                                                      size_t numBytes) = 0;
+  virtual std::unique_ptr<AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) = 0;
   virtual void fetchBuffer(const ChunkKey& key,
                            AbstractBuffer* destBuffer,
                            const size_t numBytes = 0) = 0;

omniscidb/DataMgr/AbstractDataProvider.h

@@ -42,6 +42,13 @@ class AbstractDataProvider : public Data_Namespace::AbstractBufferMgr {
   std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyBufferMemory(
       const ChunkKey& key,
       size_t numBytes) override {
+    UNREACHABLE();
+    return nullptr;
+  }
+
+  std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override {
+    UNREACHABLE();
     return nullptr;
   }
 

omniscidb/DataMgr/BufferMgr/BufferMgr.cpp

@@ -865,6 +865,20 @@ std::unique_ptr<AbstractDataToken> BufferMgr::getZeroCopyBufferMemory(const Chun
   return parent_mgr_->getZeroCopyBufferMemory(key, numBytes);
 }
 
+std::unique_ptr<AbstractDataToken> BufferMgr::getZeroCopyColumnData(
+    const ColumnRef& col_ref) {
+  return parent_mgr_->getZeroCopyColumnData(col_ref);
+}
+
+// AbstractBuffer* BufferMgr::getZeroCopyColumnBuffer(const ColumnRef& col_ref) {
+//   AbstractBuffer* res = nullptr;
+
+//   if (auto token = getZeroCopyColumnData(col_ref)) {
+//     res = createZeroCopyBuffer({}, std::move(token));
+//   }
+//   return res;
+// }
+
 MemoryInfo BufferMgr::getMemoryInfo() {
   std::unique_lock<std::mutex> sized_segs_lock(sized_segs_mutex_);
   MemoryInfo mi;

omniscidb/DataMgr/BufferMgr/BufferMgr.h

@@ -159,8 +159,12 @@ class BufferMgr : public AbstractBufferMgr {  // implements
   /// Returns the a pointer to the chunk with the specified key.
   AbstractBuffer* getBuffer(const ChunkKey& key, const size_t num_bytes = 0) override;
 
+  // AbstractBuffer* BufferMgr::getZeroCopyColumnBuffer(const ColumnRef &col_ref);
+
   std::unique_ptr<AbstractDataToken> getZeroCopyBufferMemory(const ChunkKey& key,
                                                              size_t numBytes) override;
+  std::unique_ptr<AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override;
 
   /**
    * @brief Puts the contents of d into the Buffer with ChunkKey key.

omniscidb/DataMgr/DataMgr.cpp

@@ -458,6 +458,13 @@ AbstractBuffer* DataMgr::getChunkBuffer(const ChunkKey& key,
   return bufferMgrs_[level][deviceId]->getBuffer(key, numBytes);
 }
 
+std::unique_ptr<AbstractDataToken> DataMgr::getZeroCopyColumnData(
+    const ColumnRef& col_ref) {
+  const auto level = static_cast<size_t>(Data_Namespace::CPU_LEVEL);
+  CHECK_LT(level, levelSizes_.size());  // make sure we have a legit buffermgr
+  return bufferMgrs_[level][0]->getZeroCopyColumnData(col_ref);
+}
+
 void DataMgr::deleteChunksWithPrefix(const ChunkKey& keyPrefix) {
   int numLevels = bufferMgrs_.size();
   for (int level = numLevels - 1; level >= 0; --level) {

omniscidb/DataMgr/DataMgr.h

@@ -177,6 +177,7 @@ class DataMgr {
                                  const MemoryLevel memoryLevel,
                                  const int deviceId = 0,
                                  const size_t numBytes = 0);
+  std::unique_ptr<AbstractDataToken> getZeroCopyColumnData(const ColumnRef& col_ref);
   void deleteChunksWithPrefix(const ChunkKey& keyPrefix);
   void deleteChunksWithPrefix(const ChunkKey& keyPrefix, const MemoryLevel memLevel);
   AbstractBuffer* alloc(const MemoryLevel memoryLevel,

omniscidb/DataMgr/DataMgrDataProvider.cpp

@@ -29,9 +29,16 @@ std::shared_ptr<Chunk_NS::Chunk> DataMgrDataProvider::getChunk(
   return Chunk_NS::Chunk::getChunk(
       col_info, data_mgr_, key, memory_level, device_id, num_bytes, num_elems);
 }
+
+std::unique_ptr<Data_Namespace::AbstractDataToken>
+DataMgrDataProvider::getZeroCopyColumnData(const ColumnRef& col_ref) {
+  return data_mgr_->getZeroCopyColumnData(col_ref);
+}
+
 TableFragmentsInfo DataMgrDataProvider::getTableMetadata(int db_id, int table_id) const {
   return data_mgr_->getTableMetadata(db_id, table_id);
 }
+
 const DictDescriptor* DataMgrDataProvider::getDictMetadata(int dict_id,
                                                            bool load_dict) const {
   return data_mgr_->getDictMetadata(dict_id, load_dict);

omniscidb/DataMgr/DataMgrDataProvider.h

@@ -35,6 +35,9 @@ class DataMgrDataProvider : public DataProvider {
       const size_t num_bytes,
       const size_t num_elems) override;
 
+  std::unique_ptr<Data_Namespace::AbstractDataToken> getZeroCopyColumnData(
+      const ColumnRef& col_ref) override;
+
   TableFragmentsInfo getTableMetadata(int db_id, int table_id) const override;
 
   const DictDescriptor* getDictMetadata(int dict_id,