Feature/sample go

src/graph/context/Iterator.cpp

@@ -259,6 +259,25 @@ void GetNeighborsIter::erase() {
   next();
 }
 
+void GetNeighborsIter::sample(const int64_t count) {
+  algorithm::ReservoirSampling<std::tuple<List*, List>> sampler_(count);
+  doReset(0);
+  for (; valid(); next()) {
+    // <current List of Edge, value of Edge>
+    std::tuple<List*, List> t =
+        std::make_tuple(const_cast<List*>(currentCol_), std::move(*currentEdge_));
+    sampler_.sampling(std::move(t));
+  }
+  doReset(0);
+  clearEdges();
+  auto samples = std::move(sampler_).samples();
+  for (auto& sample : samples) {
+    auto* col = std::get<0>(sample);
+    col->emplace_back(std::move(std::get<1>(sample)));
+  }
+  doReset(0);
+}
+
 const Value& GetNeighborsIter::getColumn(const std::string& col) const {
   if (!valid()) {
     return Value::kNullValue;
@@ -454,6 +473,51 @@ List GetNeighborsIter::getEdges() {
   return edges;
 }
 
+void GetNeighborsIter::nextCol() {
+  if (!valid()) {
+    return;
+  }
+  DCHECK(!noEdge_);
+
+  // go to next column
+  while (++colIdx_) {
+    if (colIdx_ < currentDs_->colUpperBound) {
+      const auto& currentCol = currentRow_->operator[](colIdx_);
+      if (!currentCol.isList() || currentCol.getList().empty()) {
+        continue;
+      }
+
+      currentCol_ = &currentCol.getList();
+      // edgeIdxUpperBound_ = currentCol_->size();
+      // edgeIdx_ = -1;
+      break;
+    }
+    // go to next row
+    if (++currentRow_ < rowsUpperBound_) {
+      colIdx_ = currentDs_->colLowerBound;
+      continue;
+    }
+
+    // go to next dataset
+    if (++currentDs_ < dsIndices_.end()) {
+      colIdx_ = currentDs_->colLowerBound;
+      currentRow_ = currentDs_->ds->begin();
+      rowsUpperBound_ = currentDs_->ds->end();
+      continue;
+    }
+    break;
+  }
+  if (currentDs_ == dsIndices_.end()) {
+    return;
+  }
+}
+
+void GetNeighborsIter::clearEdges() {
+  for (; colValid(); nextCol()) {
+    const_cast<List*>(currentCol_)->clear();
+  }
+}
+
 SequentialIter::SequentialIter(std::shared_ptr<Value> value) : Iterator(value, Kind::kSequential) {
   DCHECK(value->isDataSet());
   auto& ds = value->mutableDataSet();

src/graph/context/Iterator.h

@@ -11,6 +11,7 @@
 
 #include <boost/dynamic_bitset.hpp>
 
+#include "common/algorithm/ReservoirSampling.h"
 #include "common/datatypes/DataSet.h"
 #include "common/datatypes/List.h"
 #include "common/datatypes/Value.h"
@@ -60,6 +61,12 @@ class Iterator {
   // Warning this will break the origin order of elements!
   virtual void unstableErase() = 0;
 
+  // remain the select data in range
+  virtual void select(std::size_t offset, std::size_t count) = 0;
+
+  // Sample the elements
+  virtual void sample(int64_t count) = 0;
+
   virtual const Row* row() const = 0;
 
   // erase range, no include last position, if last > size(), erase to the end
@@ -142,6 +149,12 @@ class DefaultIter final : public Iterator {
 
   void eraseRange(size_t, size_t) override { return; }
 
+  void select(std::size_t, std::size_t) override {
+    DLOG(FATAL) << "Unimplemented method for default iterator.";
+  }
+
+  void sample(int64_t) override { DLOG(FATAL) << "Unimplemented default iterator."; }
+
   void clear() override { reset(); }
 
   size_t size() const override { return 1; }
@@ -193,12 +206,27 @@ class GetNeighborsIter final : public Iterator {
 
   void unstableErase() override { erase(); }
 
+  // erase [first, last)
   void eraseRange(size_t first, size_t last) override {
-    UNUSED(first);
-    UNUSED(last);
-    DCHECK(false);
+    for (std::size_t i = 0; valid() && i < last; next(), ++i) {
+      if (i >= first || i < last) {
+        erase();
+      }
+    }
+    doReset(0);
   }
 
+  void select(std::size_t offset, std::size_t count) override {
+    for (std::size_t i = 0; valid(); next(), ++i) {
+      if (i < offset || i > (offset + count - 1)) {
+        erase();
+      }
+    }
+    doReset(0);
+  }
+
+  void sample(int64_t count) override;
+
   size_t size() const override { return 0; }
 
   const Value& getColumn(const std::string& col) const override;
@@ -236,6 +264,13 @@ class GetNeighborsIter final : public Iterator {
     return currentDs_->tagEdgeNameIndices.find(colIdx_)->second;
   }
 
+  bool colValid() { return !noEdge_ && valid(); }
+
+  // move to next List of Edge data
+  void nextCol();
+
+  void clearEdges();
+
   struct PropIndex {
     size_t colIdx;
     std::vector<std::string> propList;
@@ -320,6 +355,28 @@ class SequentialIter : public Iterator {
 
   void eraseRange(size_t first, size_t last) override;
 
+  void select(std::size_t offset, std::size_t count) override {
+    auto size = this->size();
+    if (size <= static_cast<size_t>(offset)) {
+      clear();
+    } else if (size > static_cast<size_t>(offset + count)) {
+      eraseRange(0, offset);
+      eraseRange(count, size - offset);
+    } else if (size > static_cast<size_t>(offset) && size <= static_cast<size_t>(offset + count)) {
+      eraseRange(0, offset);
+    }
+  }
+
+  void sample(int64_t count) override {
+    DCHECK_GE(count, 0);
+    algorithm::ReservoirSampling<Row> sampler(count);
+    for (auto& row : *rows_) {
+      sampler.sampling(std::move(row));
+    }
+    *rows_ = std::move(sampler).samples();
+    iter_ = rows_->begin();
+  }
+
   void clear() override {
     rows_->clear();
     reset();

src/graph/context/ast/QueryAstContext.h

@@ -84,6 +84,8 @@ struct GoContext final : AstContext {
   bool distinct{false};
   // true: sample, false: limit
   bool random{false};
+  // step limit value
+  std::vector<int64_t> limits;
   std::vector<std::string> colNames;
 
   std::string vidsVar;

src/graph/executor/CMakeLists.txt

@@ -19,6 +19,7 @@ nebula_add_library(
     query/GetVerticesExecutor.cpp
     query/IntersectExecutor.cpp
     query/LimitExecutor.cpp
+    query/SampleExecutor.cpp
     query/MinusExecutor.cpp
     query/ProjectExecutor.cpp
     query/UnwindExecutor.cpp

src/graph/executor/Executor.cpp

@@ -84,6 +84,7 @@
 #include "graph/executor/query/LimitExecutor.h"
 #include "graph/executor/query/MinusExecutor.h"
 #include "graph/executor/query/ProjectExecutor.h"
+#include "graph/executor/query/SampleExecutor.h"
 #include "graph/executor/query/SortExecutor.h"
 #include "graph/executor/query/TopNExecutor.h"
 #include "graph/executor/query/UnionAllVersionVarExecutor.h"
@@ -178,6 +179,9 @@ Executor *Executor::makeExecutor(QueryContext *qctx, const PlanNode *node) {
     case PlanNode::Kind::kLimit: {
       return pool->add(new LimitExecutor(node, qctx));
     }
+    case PlanNode::Kind::kSample: {
+      return pool->add(new SampleExecutor(node, qctx));
+    }
     case PlanNode::Kind::kProject: {
       return pool->add(new ProjectExecutor(node, qctx));
     }

src/graph/executor/query/GetNeighborsExecutor.cpp

@@ -42,6 +42,7 @@ folly::Future<Status> GetNeighborsExecutor::execute() {
 
   time::Duration getNbrTime;
   GraphStorageClient* storageClient = qctx_->getStorageClient();
+  QueryExpressionContext qec(qctx()->ectx());
   return storageClient
       ->getNeighbors(gn_->space(),
                      qctx()->rctx()->session()->id(),
@@ -57,7 +58,7 @@ folly::Future<Status> GetNeighborsExecutor::execute() {
                      gn_->dedup(),
                      gn_->random(),
                      gn_->orderBy(),
-                     gn_->limit(),
+                     gn_->limit(qec),
                      gn_->filter())
       .via(runner())
       .ensure([this, getNbrTime]() {

src/graph/executor/query/LimitExecutor.cpp

@@ -21,16 +21,9 @@ folly::Future<Status> LimitExecutor::execute() {
   ResultBuilder builder;
   builder.value(result.valuePtr());
   auto offset = limit->offset();
-  auto count = limit->count();
-  auto size = iter->size();
-  if (size <= static_cast<size_t>(offset)) {
-    iter->clear();
-  } else if (size > static_cast<size_t>(offset + count)) {
-    iter->eraseRange(0, offset);
-    iter->eraseRange(count, size - offset);
-  } else if (size > static_cast<size_t>(offset) && size <= static_cast<size_t>(offset + count)) {
-    iter->eraseRange(0, offset);
-  }
+  QueryExpressionContext qec(ectx_);
+  auto count = limit->count(qec);
+  iter->select(offset, count);
   builder.iter(std::move(result).iter());
   return finish(builder.build());
 }

src/graph/executor/query/SampleExecutor.cpp

@@ -0,0 +1,36 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License,
+ * attached with Common Clause Condition 1.0, found in the LICENSES directory.
+ */
+
+#include "graph/executor/query/SampleExecutor.h"
+
+#include "graph/planner/plan/Query.h"
+#include "graph/util/ScopedTimer.h"
+
+namespace nebula {
+namespace graph {
+
+folly::Future<Status> SampleExecutor::execute() {
+  SCOPED_TIMER(&execTime_);
+
+  auto* sample = asNode<Sample>(node());
+  Result result = ectx_->getResult(sample->inputVar());
+  auto* iter = result.iterRef();
+  DCHECK_NE(iter->kind(), Iterator::Kind::kDefault);
+  ResultBuilder builder;
+  builder.value(result.valuePtr());
+  QueryExpressionContext qec(ectx_);
+  auto count = sample->count(qec);
+  if (iter->kind() == Iterator::Kind::kGetNeighbors ||
+      iter->size() > static_cast<std::size_t>(count)) {
+    // Sampling
+    iter->sample(count);
+  }
+  builder.iter(std::move(result).iter());
+  return finish(builder.build());
+}
+
+}  // namespace graph
+}  // namespace nebula

src/graph/executor/query/SampleExecutor.h

@@ -0,0 +1,23 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License,
+ * attached with Common Clause Condition 1.0, found in the LICENSES directory.
+ */
+
+#pragma once
+
+#include "graph/executor/Executor.h"
+
+namespace nebula {
+namespace graph {
+
+class SampleExecutor final : public Executor {
+ public:
+  SampleExecutor(const PlanNode *node, QueryContext *qctx)
+      : Executor("SampleExecutor", node, qctx) {}
+
+  folly::Future<Status> execute() override;
+};
+
+}  // namespace graph
+}  // namespace nebula

src/graph/executor/test/CMakeLists.txt

@@ -71,6 +71,7 @@ nebula_add_test(
         FilterTest.cpp
         DedupTest.cpp
         LimitTest.cpp
+        SampleTest.cpp
         SortTest.cpp
         TopNTest.cpp
         AggregateTest.cpp