pax: suport index cluster

The pax table supports the cluster syntax based on the btree index
and behaves the same as the aocs table.

Author: gongxun0928

contrib/pax_storage/.gitignore

@@ -14,11 +14,13 @@ Thumbs.db
 # Temp files dir
 .tmp/**
 build*/**
+results/**
 clang-tidy.result
 # pb
 **/*.pb.h
 **/*.pb.cc
-
+tools/**/*.data
+tools/**/*.txt
 # Executables
 /*.out
 !src/data/expected/*.out

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

@@ -5,6 +5,7 @@
 #include "access/pax_dml_state.h"
 #include "access/pax_partition.h"
 #include "access/pax_scanner.h"
+#include "access/pax_table_cluster.h"
 #include "access/pax_updater.h"
 #include "access/paxc_rel_options.h"
 #include "access/paxc_scanner.h"
@@ -242,24 +243,55 @@ void CCPaxAccessMethod::RelationCopyData(Relation rel,
 }
 
 /*
- * Used by rebuild_relation, like CLUSTER, VACUUM FULL, etc.
+ * Copy data from `OldTable` into `NewTable`, as part of a CLUSTER or VACUUM
+ * FULL.
  *
  * PAX does not have dead tuples, but the core framework requires
  * to implement this callback to do CLUSTER/VACUUM FULL/etc.
  * PAX may have re-organize semantics for this function.
  *
- * TODO: how to split the set of micro-partitions to several QE handlers.
+ *
+ * Additional Input parameters:
+ * - use_sort - if true, the table contents are sorted appropriate for
+ *   `OldIndex`; if false and OldIndex is not InvalidOid, the data is copied
+ *   in that index's order; if false and OldIndex is InvalidOid, no sorting is
+ *   performed
+ * - OldIndex - see use_sort
+ * - OldestXmin - computed by vacuum_set_xid_limits(), even when
+ *   not needed for the relation's AM
+ * - *xid_cutoff - ditto
+ * - *multi_cutoff - ditto
+ *
+ * Output parameters:
+ * - *xid_cutoff - rel's new relfrozenxid value, may be invalid
+ * - *multi_cutoff - rel's new relminmxid value, may be invalid
+ * - *tups_vacuumed - stats, for logging, if appropriate for AM
+ * - *tups_recently_dead - stats, for logging, if appropriate for AM
  */
 void CCPaxAccessMethod::RelationCopyForCluster(
-    Relation old_heap, Relation new_heap, Relation /*old_index*/,
-    bool /*use_sort*/, TransactionId /*oldest_xmin*/,
-    TransactionId * /*xid_cutoff*/, MultiXactId * /*multi_cutoff*/,
-    double * /*num_tuples*/, double * /*tups_vacuumed*/,
-    double * /*tups_recently_dead*/) {
-  Assert(RELATION_IS_PAX(old_heap));
-  Assert(RELATION_IS_PAX(new_heap));
+    Relation old_rel, Relation new_rel, Relation old_index, bool use_sort,
+    TransactionId /*oldest_xmin*/, TransactionId * /*xid_cutoff*/,
+    MultiXactId * /*multi_cutoff*/, double * /*num_tuples*/,
+    double * /*tups_vacuumed*/, double * /* tups_recently_dead*/) {
+  Assert(RELATION_IS_PAX(old_rel));
+  Assert(RELATION_IS_PAX(new_rel));
   CBDB_TRY();
-  { pax::CCPaxAuxTable::PaxAuxRelationCopyDataForCluster(old_heap, new_heap); }
+  {
+    //  if false and OldIndex is InvalidOid, no sorting is performed, just copy
+    if (!use_sort && old_index == NULL) {
+      pax::CCPaxAuxTable::PaxAuxRelationCopyDataForCluster(old_rel, new_rel);
+      return;
+    }
+
+    // TODO(gongxun): should we support index's order to cluster table? ao/aocs
+    // does not support.
+
+    // like aocs, pax tables can only be clustered against a btree index
+    CBDB_CHECK(old_index && (old_index->rd_rel->relam == BTREE_AM_OID),
+               cbdb::CException::kExTypeInvalidIndexType,
+               "PAX tables can only be clustered against a btree index");
+    pax::IndexCluster(old_rel, new_rel, old_index, GetActiveSnapshot());
+  }
   CBDB_CATCH_DEFAULT();
   CBDB_FINALLY({});
   CBDB_END_TRY();

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

@@ -0,0 +1,38 @@
+#include "access/pax_table_cluster.h"
+
+#include "clustering/clustering.h"
+#include "clustering/index_clustering.h"
+#include "clustering/pax_clustering_reader.h"
+#include "clustering/pax_clustering_writer.h"
+#include "storage/micro_partition_iterator.h"
+#include "storage/micro_partition_metadata.h"
+
+#define CLUSTER_SORT_MEMORY 128000
+
+namespace pax {
+void IndexCluster(Relation old_rel, Relation new_rel, Relation index,
+                  Snapshot snapshot) {
+  auto iter = MicroPartitionInfoIterator::New(old_rel, snapshot);
+
+  auto reader =
+      PAX_NEW<clustering::PaxClusteringReader>(old_rel, std::move(iter));
+
+  auto writer = PAX_NEW<clustering::PaxClusteringWriter>(new_rel);
+
+  auto cluster = clustering::DataClustering::CreateDataClustering(
+      clustering::DataClustering::kClusterTypeIndex);
+
+  clustering::IndexClustering::IndexClusteringOptions options;
+
+  options.tup_desc = old_rel->rd_att;
+  options.index_rel = index;
+  options.work_mem = CLUSTER_SORT_MEMORY;
+  cluster->Clustering(reader, writer, &options);
+
+  writer->Close();
+  PAX_DELETE(writer);
+  reader->Close();
+  PAX_DELETE(reader);
+}
+
+}  // namespace pax

contrib/pax_storage/src/cpp/access/pax_table_cluster.h

@@ -0,0 +1,8 @@
+#pragma once
+
+#include "comm/cbdb_api.h"
+
+namespace pax {
+void IndexCluster(Relation old_rel, Relation new_rel, Relation index,
+                  Snapshot snapshot);
+}  // namespace pax

contrib/pax_storage/src/cpp/clustering/clustering.cc

@@ -0,0 +1,21 @@
+#include "clustering/clustering.h"
+
+#include "clustering/index_clustering.h"
+#include "comm/pax_memory.h"
+
+namespace pax {
+namespace clustering {
+
+DataClustering *DataClustering::CreateDataClustering(
+    const DataClustering::ClusterType type) {
+  switch (type) {
+    case DataClustering::kClusterTypeIndex:
+      static IndexClustering index_clustering;
+      return &index_clustering;
+    default:
+      return nullptr;
+  }
+}
+
+}  // namespace clustering
+}  // namespace pax

contrib/pax_storage/src/cpp/clustering/clustering.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "clustering/clustering_reader.h"
+#include "clustering/clustering_writer.h"
+#include "clustering/sorter.h"
+
+namespace pax {
+
+namespace clustering {
+
+class DataClustering {
+ public:
+  enum ClusterType {
+    kClusterTypeIndex,
+  };
+
+  struct DataClusteringOptions {
+    ClusterType type;
+  };
+
+ public:
+  DataClustering() = default;
+  virtual ~DataClustering() = default;
+  virtual void Clustering(ClusteringDataReader *reader, ClusteringDataWriter *writer,
+                          const DataClusteringOptions *options) = 0;
+  static DataClustering *CreateDataClustering(const ClusterType type);
+};
+
+}  // namespace clustering
+}  // namespace pax

contrib/pax_storage/src/cpp/clustering/clustering_reader.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "comm/cbdb_api.h"
+namespace pax {
+
+namespace clustering {
+class ClusteringDataReader {
+ public:
+  ClusteringDataReader() = default;
+  virtual ~ClusteringDataReader() = default;
+  // TODO(gongxun): support record batch
+  // return false if no more tuples
+  virtual bool GetNextTuple(TupleTableSlot *) = 0;
+  virtual void Close() = 0;
+};
+
+}  // namespace clustering
+
+}  // namespace pax

contrib/pax_storage/src/cpp/clustering/clustering_writer.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "comm/cbdb_api.h"
+namespace pax {
+
+namespace clustering {
+
+class ClusteringDataWriter {
+ public:
+  ClusteringDataWriter() = default;
+  virtual ~ClusteringDataWriter() = default;
+  /**
+   * slot->tts_value[tuple_desc->natts] is zorder_value
+   */
+  virtual void WriteTuple(TupleTableSlot *tuple) = 0;
+  virtual void Close() = 0;
+};
+
+}  // namespace clustering
+
+}  // namespace pax

contrib/pax_storage/src/cpp/clustering/index_clustering.cc

@@ -0,0 +1,50 @@
+
+#include "clustering/index_clustering.h"
+
+#include "clustering/sorter_index.h"
+#include "comm/cbdb_wrappers.h"
+
+namespace pax {
+namespace clustering {
+
+IndexClustering::IndexClustering() {}
+IndexClustering::~IndexClustering() {}
+void IndexClustering::Clustering(ClusteringDataReader *reader, ClusteringDataWriter *writer,
+                                const DataClusteringOptions *options) {
+  const IndexClusteringOptions *index_options =
+      reinterpret_cast<const IndexClusteringOptions *>(options);
+
+  CBDB_CHECK(index_options != nullptr && index_options->index_rel != nullptr,
+             cbdb::CException::kExTypeInvalid, "index options is invalid");
+
+  TupleTableSlot *origin_slot;
+  TupleTableSlot *sorted_slot;
+  IndexSorter::IndexTupleSorterOptions sorter_options;
+
+  sorter_options.tup_desc = index_options->tup_desc;
+  sorter_options.work_mem = index_options->work_mem;
+  sorter_options.index_rel = index_options->index_rel;
+
+  IndexSorter sorter(sorter_options);
+
+  origin_slot =
+      MakeSingleTupleTableSlot(index_options->tup_desc, &TTSOpsVirtual);
+  while (reader->GetNextTuple(origin_slot)) {
+    sorter.AppendSortData(origin_slot);
+  }
+
+  ExecDropSingleTupleTableSlot(origin_slot);
+
+  sorter.Sort();
+
+  sorted_slot =
+      MakeSingleTupleTableSlot(sorter_options.tup_desc, &TTSOpsMinimalTuple);
+  while (sorter.GetSortedData(sorted_slot)) {
+    slot_getallattrs(sorted_slot);
+    writer->WriteTuple(sorted_slot);
+  }
+  ExecDropSingleTupleTableSlot(sorted_slot);
+}
+
+}  // namespace clustering
+}  // namespace pax

contrib/pax_storage/src/cpp/clustering/index_clustering.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "clustering/clustering.h"
+namespace pax {
+namespace clustering {
+class IndexClustering final : public DataClustering {
+ public:
+  struct IndexClusteringOptions : public DataClusteringOptions {
+    IndexClusteringOptions() { type = kClusterTypeIndex; }
+    TupleDesc tup_desc;
+    Relation index_rel;
+    int work_mem;
+  };
+
+ public:
+  IndexClustering();
+  virtual ~IndexClustering();
+  void Clustering(ClusteringDataReader *reader, ClusteringDataWriter *writer,
+                  const DataClusteringOptions *options) override;
+};
+
+}  // namespace clustering
+
+}  // namespace pax