Op: Add a new vectorization MicroPartitionReader to replace ReadVecTuple

The vectorized condition in `MicroPartitionReader` will cause complicate the logic.

- Remove function `ReadVecTuple` which is specialized logic
- Added a new MicroPartitionReader named `PaxVecReader` to adapting read from vectorization version.

Author: jiaqizho

contrib/pax_storage/src/cpp/CMakeLists.txt

@@ -170,7 +170,8 @@ set(pax_catalog_src
     catalog/pax_aux_table.cc)
 
 set(pax_vec_src
-  storage/vec/pax_vec_adapter.cc)
+  storage/vec/pax_vec_adapter.cc
+  storage/vec/pax_vec_reader.cc)
 
 link_directories($ENV{GPHOME}/lib)
 

contrib/pax_storage/src/cpp/access/pax_scanner.cc

@@ -36,6 +36,13 @@ TableScanDesc PaxScanDesc::BeginScan(Relation relation, Snapshot snapshot,
   desc->key_ = key;
   desc->reused_buffer_ = new DataBuffer<char>(32 * 1024 * 1024);  // 32mb
   desc->filter_ = filter;
+#ifdef VEC_BUILD
+  if (flags & (1 << 12)) {
+    desc->vec_adapter_ = new VecAdapter(cbdb::RelationGetTupleDesc(relation));
+    reader_options.is_vec = true;
+    reader_options.adapter = desc->vec_adapter_;
+  }
+#endif
 
   // init shared memory
   cbdb::InitCommandResource();
@@ -60,12 +67,6 @@ TableScanDesc PaxScanDesc::BeginScan(Relation relation, Snapshot snapshot,
   desc->reader_ = new TableReader(std::move(iter), reader_options);
   desc->reader_->Open();
 
-#ifdef VEC_BUILD
-  if (flags & (1 << 12)) {
-    desc->vec_adapter_ = new VecAdapter(cbdb::RelationGetTupleDesc(relation));
-  }
-#endif
-
   MemoryContextSwitchTo(old_ctx);
   return &desc->rs_base_;
 }
@@ -147,15 +148,6 @@ bool PaxScanDesc::ScanGetNextSlot(TableScanDesc scan, TupleTableSlot *slot) {
   CTupleSlot cslot(slot);
   old_ctx = MemoryContextSwitchTo(desc->memory_context_);
 
-#ifdef VEC_BUILD
-  // TODO(vec): replace (1 << 12), and check other place
-  if (desc->rs_base_.rs_flags & (1 << 12)) {
-    ok = desc->reader_->ReadVecTuple(&cslot, desc->vec_adapter_);
-    MemoryContextSwitchTo(old_ctx);
-    return ok;
-  }
-#endif
-
   ok = desc->reader_->ReadTuple(&cslot);
 
   MemoryContextSwitchTo(old_ctx);

contrib/pax_storage/src/cpp/storage/micro_partition.h

@@ -8,6 +8,7 @@
 #include <stdexcept>
 #include <string>
 
+#include "storage/columns/pax_columns.h"
 #include "storage/micro_partition_metadata.h"
 
 namespace pax {
@@ -133,6 +134,22 @@ class MicroPartitionReader {
   // is also created during this stage, no matter the map relation is needed or
   // not. We may optimize to avoid creating the map relation later.
   virtual bool ReadTuple(CTupleSlot *slot) = 0;
+
+ protected:
+  // Allow different MicroPartitionReader shared columns
+  // but should not let export columns out of micro partition
+  //
+  // In MicroPartition writer/reader implementation, all in-memory data should
+  // be accessed by pax column This is because most of the common logic of
+  // column operation is done in pax column, such as type mapping, bitwise
+  // fetch, compression/encoding. At the same time, pax column can also be used
+  // as a general interface for internal using, because it's zero copy from
+  // buffer. more details in `storage/columns`
+  virtual PaxColumns *GetAllColumns() = 0;
+#ifdef VEC_BUILD
+ private:
+  friend class PaxVecReader;
+#endif
 };
 
 }  // namespace pax

contrib/pax_storage/src/cpp/storage/orc/orc.h

@@ -128,12 +128,12 @@ class OrcReader : public MicroPartitionReader {
 
   bool ReadTuple(CTupleSlot *cslot) override;
 
-  PaxColumns *GetAllColumns();
-
 #ifndef RUN_GTEST
  protected:  // NOLINT
 #endif
 
+  PaxColumns *GetAllColumns() override;
+
   orc::proto::StripeFooter ReadStripeWithProjection(
       DataBuffer<char> *data_buffer, OrcReader::StripeInformation *stripe_info,
       const bool *proj_map, size_t proj_len);

contrib/pax_storage/src/cpp/storage/pax.cc

@@ -10,6 +10,10 @@
 #include "storage/micro_partition_file_factory.h"
 #include "storage/micro_partition_metadata.h"
 
+#ifdef VEC_BUILD
+#include "storage/vec/pax_vec_reader.h"
+#endif
+
 namespace pax {
 
 static std::string GenRandomBlockId() {
@@ -182,65 +186,11 @@ bool TableReader::ReadTuple(CTupleSlot *slot) {
     }
     OpenFile();
   }
-  num_tuples_++;
-  slot->SetTableNo(table_no_);
-  slot->SetBlockNumber(current_block_number_);
-  slot->StoreVirtualTuple();
-  return true;
-}
-
-#ifdef VEC_BUILD
-// TODO(jiaqizho): should remove this method but provider a
-// new mirco partition reader which include adapter
-bool TableReader::ReadVecTuple(CTupleSlot *slot, VecAdapter *adapter) {
-  size_t flush_nums_of_rows = 0;
-
-  if (is_empty_) {
-    return false;
-  }
-
-  // Three conditions indicate that there is no data left
-  // - `VecAdapter` must be set
-  // - `VecAdapter` has consumed all buffer data
-  // - current have not next reader
-  if (adapter->IsInitialized() && adapter->IsEnd() && !iterator_->HasNext()) {
-    return false;
-  }
-
-  Assert(reader_);
-
-  // `No set` means that the first reader not setup in apapter
-  if (!adapter->IsInitialized()) {
-    PaxColumns *pax_columns =
-        reinterpret_cast<OrcReader *>(reader_)->GetAllColumns();
-    adapter->SetDataSource(pax_columns);
-  }
-
-  // `VecAdapter` has consumed all buffer data
-  if (adapter->IsEnd()) {
-    if (iterator_->HasNext()) {
-      reader_->Close();
-      OpenFile();
-
-      PaxColumns *pax_columns =
-          reinterpret_cast<OrcReader *>(reader_)->GetAllColumns();
-      adapter->SetDataSource(pax_columns);
-    }
-  }
-
-  auto ok = adapter->AppendToVecBuffer();
-  CBDB_CHECK(ok, cbdb::CException::ExType::kExTypeLogicError);
-
-  flush_nums_of_rows = adapter->FlushVecBuffer(slot);
-  Assert(flush_nums_of_rows);
-
-  num_tuples_ += flush_nums_of_rows;
   slot->SetTableNo(table_no_);
   slot->SetBlockNumber(current_block_number_);
   slot->StoreVirtualTuple();
   return true;
 }
-#endif  // VEC_BUILD
 
 void TableReader::OpenFile() {
   Assert(iterator_->HasNext());
@@ -266,6 +216,14 @@ void TableReader::OpenFile() {
 
   reader_ = new OrcReader(
       Singleton<LocalFileSystem>::GetInstance()->Open(options.file_name));
+
+#ifdef VEC_BUILD
+  if (reader_options_.is_vec) {
+    Assert(reader_options_.adapter);
+    reader_ = new PaxVecReader(reader_, reader_options_.adapter);
+  }
+#endif
+
   reader_->Open(options);
 }
 

contrib/pax_storage/src/cpp/storage/pax.h

@@ -73,6 +73,10 @@ class TableReader final {
     // Will not used in TableReader
     // But pass into micro partition reader
     PaxFilter *filter = nullptr;
+#ifdef VEC_BUILD
+    bool is_vec = false;
+    VecAdapter *adapter = nullptr;
+#endif
   };
 
   TableReader(std::unique_ptr<IteratorBase<MicroPartitionMetadata>> &&iterator,
@@ -85,11 +89,6 @@ class TableReader final {
 
   void Close();
 
-#ifdef VEC_BUILD
-  // std::unique_ptr<VecAdapter> here ?
-  bool ReadVecTuple(CTupleSlot *slot, VecAdapter *adapter);
-#endif
-
   bool ReadTuple(CTupleSlot *slot);
 
   // deprecate:
@@ -104,7 +103,6 @@ class TableReader final {
  private:
   const std::unique_ptr<IteratorBase<MicroPartitionMetadata>> iterator_;
   MicroPartitionReader *reader_ = nullptr;
-  size_t num_tuples_ = 0;
   bool is_empty_ = false;
   const ReaderOptions reader_options_;
   int current_block_number_ = 0;

contrib/pax_storage/src/cpp/storage/vec/pax_vec_reader.cc

@@ -0,0 +1,36 @@
+#include "storage/vec/pax_vec_reader.h"
+
+#include "storage/vec/pax_vec_adapter.h"
+#ifdef VEC_BUILD
+
+namespace pax {
+
+PaxVecReader::PaxVecReader(MicroPartitionReader *reader, VecAdapter *adapter)
+    : reader_(reader), adapter_(adapter) {}
+
+PaxVecReader::~PaxVecReader() { delete reader_; }
+
+void PaxVecReader::Open(const ReaderOptions &options) {
+  reader_->Open(options);
+  PaxColumns *pax_columns = reader_->GetAllColumns();
+  adapter_->SetDataSource(pax_columns);
+}
+
+void PaxVecReader::Close() { reader_->Close(); }
+
+bool PaxVecReader::ReadTuple(CTupleSlot *cslot) {
+  if (!adapter_->AppendToVecBuffer()) {
+    return false;
+  }
+
+  auto flush_nums_of_rows = adapter_->FlushVecBuffer(cslot);
+  Assert(flush_nums_of_rows);
+  return true;
+}
+
+PaxColumns *PaxVecReader::GetAllColumns() {
+  CBDB_RAISE(cbdb::CException::ExType::kExTypeLogicError);
+}
+};  // namespace pax
+
+#endif  // VEC_BUILD

contrib/pax_storage/src/cpp/storage/vec/pax_vec_reader.h

@@ -0,0 +1,35 @@
+#pragma once
+#include "storage/micro_partition.h"
+
+#ifdef VEC_BUILD
+
+namespace pax {
+
+class VecAdapter;
+
+class PaxVecReader : public MicroPartitionReader {
+ public:
+  // If enable read tuple from vec reader,
+  // then OrcReader will be hold by PaxVecReader,
+  // current MicroPartitionReader lifecycle will be bound to the PaxVecReader)
+  PaxVecReader(MicroPartitionReader *reader, VecAdapter *adapter);
+
+  ~PaxVecReader() override;
+
+  void Open(const ReaderOptions &options) override;
+
+  void Close() override;
+
+  bool ReadTuple(CTupleSlot *cslot) override;
+
+ protected:
+  PaxColumns *GetAllColumns() override;
+
+ private:
+  MicroPartitionReader *reader_;
+  VecAdapter *adapter_;
+};
+
+}  // namespace pax
+
+#endif  // VEC_BUILD