Add expand op in velox

velox/core/PlanNode.cpp

@@ -292,6 +292,69 @@ PlanNodePtr AggregationNode::create(const folly::dynamic& obj, void* context) {
       deserializeSingleSource(obj, context));
 }
 
+namespace {
+RowTypePtr getSparkExpandOutputType(
+    const std::vector<std::vector<TypedExprPtr>>& projectSets,
+    const std::vector<std::string>& names) {
+  std::vector<std::string> outputs;
+  outputs.reserve(names.size());
+  std::vector<TypePtr> types;
+  types.reserve(names.size());
+  for (int32_t i = 0; i < names.size(); ++i) {
+    outputs.push_back(names[i]);
+    auto expr = projectSets[0][i];
+    types.push_back(expr->type());
+  }
+
+  return ROW(std::move(outputs), std::move(types));
+}
+} // namespace
+
+ExpandNode::ExpandNode(
+    PlanNodeId id,
+    std::vector<std::vector<TypedExprPtr>> projectSets,
+    std::vector<std::string> names,
+    PlanNodePtr source)
+    : PlanNode(std::move(id)),
+      sources_{source},
+      outputType_(getSparkExpandOutputType(
+          projectSets,
+          names)),
+      projectSets_(std::move(projectSets)),
+      names_(std::move(names)) {}
+
+void ExpandNode::addDetails(std::stringstream& stream) const {
+  for (auto i = 0; i < projectSets_.size(); ++i) {
+    if (i > 0) {
+      stream << ", ";
+    }
+    stream << "[";
+    addKeys(stream, projectSets_[i]);
+    stream << "]";
+  }
+}
+
+folly::dynamic ExpandNode::serialize() const {
+  auto obj = PlanNode::serialize();
+  obj["projectSets"] = ISerializable::serialize(projectSets_);
+  obj["names"] = ISerializable::serialize(names_);
+
+  return obj;
+}
+
+// static
+PlanNodePtr ExpandNode::create(const folly::dynamic& obj, void* context) {
+  auto source = deserializeSingleSource(obj, context);
+  auto names = deserializeStrings(obj["names"]);
+    auto projectSets = ISerializable::deserialize<
+      std::vector<std::vector<ITypedExpr>>>(obj["projectSets"], context);
+  return std::make_shared<ExpandNode>(
+      deserializePlanNodeId(obj),
+      std::move(projectSets),
+      std::move(names),
+      std::move(source));
+}
+
 namespace {
 RowTypePtr getGroupIdOutputType(
     const std::vector<GroupIdNode::GroupingKeyInfo>& groupingKeyInfos,

velox/core/PlanNode.h

@@ -668,6 +668,58 @@ inline std::string mapAggregationStepToName(const AggregationNode::Step& step) {
   return ss.str();
 }
 
+/// Plan node used to apply all of the projections expressions to every input
+/// row, hence we will get mulitple output row for an input rows. This has
+/// similar behavior to spark ExpandExec.
+class ExpandNode : public PlanNode {
+ public:
+
+  /// @param id Plan node ID.
+  /// @param projectSets A list of project sets. The output conatins one cloumn
+  /// for each project expr. The project expr may be cloumn reference, null or
+  /// int constant.
+  /// @param names The names and order of the projects in the output.
+  /// @param source Input plan node.
+  ExpandNode(
+      PlanNodeId id,
+      std::vector<std::vector<TypedExprPtr>> projectSets,
+      std::vector<std::string> names,
+      PlanNodePtr source);
+
+  const RowTypePtr& outputType() const override {
+    return outputType_;
+  }
+
+  const std::vector<PlanNodePtr>& sources() const override {
+    return sources_;
+  }
+
+  const std::vector<std::vector<TypedExprPtr>>& projectSets()
+      const {
+    return projectSets_;
+  }
+
+  const std::vector<std::string>& names() const {
+    return names_;
+  }
+
+  std::string_view name() const override {
+    return "Expand";
+  }
+
+  folly::dynamic serialize() const override;
+
+  static PlanNodePtr create(const folly::dynamic& obj, void* context);
+
+ private:
+  void addDetails(std::stringstream& stream) const override;
+
+  const std::vector<PlanNodePtr> sources_;
+  const RowTypePtr outputType_;
+  const std::vector<std::vector<TypedExprPtr>> projectSets_;
+  const std::vector<std::string> names_;
+};
+
 /// Plan node used to implement aggregations over grouping sets. Duplicates the
 /// aggregation input for each set of grouping keys. The output contains one
 /// column for each grouping key, followed by aggregation inputs, followed by a

velox/exec/CMakeLists.txt

@@ -25,6 +25,7 @@ add_library(
   Driver.cpp
   EnforceSingleRow.cpp
   Exchange.cpp
+  Expand.cpp
   FilterProject.cpp
   GroupId.cpp
   GroupingSet.cpp

velox/exec/Expand.cpp

@@ -0,0 +1,117 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "velox/exec/Expand.h"
+
+namespace facebook::velox::exec {
+
+Expand::Expand(
+    int32_t operatorId,
+    DriverCtx* driverCtx,
+    const std::shared_ptr<const core::ExpandNode>& expandNode)
+    : Operator(
+          driverCtx,
+          expandNode->outputType(),
+          operatorId,
+          expandNode->id(),
+          "Expand") {
+  const auto& inputType = expandNode->sources()[0]->outputType();
+  auto numProjectSets = expandNode->projectSets().size();
+  projectMappings_.reserve(numProjectSets);
+  constantMappings_.reserve(numProjectSets);
+  auto numProjects = expandNode->names().size();
+  for (const auto& projectSet : expandNode->projectSets()) {
+    std::vector<column_index_t> projectMapping;
+    projectMapping.reserve(numProjects);
+    std::vector<ConstantTypedExprPtr> constantMapping;
+    constantMapping.reserve(numProjects);
+    for (const auto& project : projectSet) {
+      if (auto field =
+        std::dynamic_pointer_cast<const core::FieldAccessTypedExpr>(project)) {
+        projectMapping.push_back(inputType->getChildIdx(field->name()));
+        constantMapping.push_back(nullptr);
+      } else if (
+        auto constant =
+          std::dynamic_pointer_cast<const core::ConstantTypedExpr>(project)) {
+        projectMapping.push_back(kUnMapedProject);
+        constantMapping.push_back(constant);
+      } else {
+        VELOX_FAIL("Unexpted expression for Expand");
+      }
+    }
+
+    projectMappings_.emplace_back(std::move(projectMapping));
+    constantMappings_.emplace_back(std::move(constantMapping));
+  }
+}
+
+bool Expand::needsInput() const {
+  return !noMoreInput_ && input_ == nullptr;
+}
+
+void Expand::addInput(RowVectorPtr input) {
+  // Load Lazy vectors.
+  for (auto& child : input->children()) {
+    child->loadedVector();
+  }
+
+  input_ = std::move(input);
+}
+
+RowVectorPtr Expand::getOutput() {
+  if (!input_) {
+    return nullptr;
+  }
+
+  // Make a copy of input for the grouping set at 'projectSetIndex_'.
+  auto numInput = input_->size();
+
+  std::vector<VectorPtr> outputColumns(outputType_->size());
+
+  const auto& projectMapping = projectMappings_[projectSetIndex_];
+  const auto& constantMapping = constantMappings_[projectSetIndex_];
+  auto numGroupingKeys = projectMapping.size();
+
+  for (auto i = 0; i < numGroupingKeys; ++i) {
+    if (projectMapping[i] == kUnMapedProject) {
+      auto constantExpr = constantMapping[i];
+      if (constantExpr->value().isNull()) {
+        // Add null column.
+        outputColumns[i] = BaseVector::createNullConstant(
+          outputType_->childAt(i), numInput, pool());
+      } else {
+        // Add constant column: gid, gpos, etc.
+        outputColumns[i] = BaseVector::createConstant(
+          constantExpr->type(),
+          constantExpr->value(),
+          numInput,
+          pool());
+      }
+    } else {
+      outputColumns[i] = input_->childAt(projectMapping[i]);
+    }
+  }
+
+  ++projectSetIndex_;
+  if (projectSetIndex_ == projectMappings_.size()) {
+    projectSetIndex_ = 0;
+    input_ = nullptr;
+  }
+
+  return std::make_shared<RowVector>(
+      pool(), outputType_, nullptr, numInput, std::move(outputColumns));
+}
+
+} // namespace facebook::velox::exec

velox/exec/Expand.h

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+#include "velox/core/Expressions.h"
+#include "velox/exec/Operator.h"
+
+namespace facebook::velox::exec {
+
+using ConstantTypedExprPtr = std::shared_ptr<const core::ConstantTypedExpr>;
+
+class Expand : public Operator {
+ public:
+  Expand(
+      int32_t operatorId,
+      DriverCtx* driverCtx,
+      const std::shared_ptr<const core::ExpandNode>& expandNode);
+
+  bool needsInput() const override;
+
+  void addInput(RowVectorPtr input) override;
+
+  RowVectorPtr getOutput() override;
+
+  BlockingReason isBlocked(ContinueFuture* /*future*/) override {
+    return BlockingReason::kNotBlocked;
+  }
+
+  bool isFinished() override {
+    return finished_ || (noMoreInput_ && input_ == nullptr);
+  }
+
+ private:
+  static constexpr column_index_t kUnMapedProject =
+      std::numeric_limits<column_index_t>::max();
+
+  bool finished_{false};
+
+  std::vector<std::vector<column_index_t>> projectMappings_;
+
+  std::vector<std::vector<ConstantTypedExprPtr>> constantMappings_;
+
+  /// 'getOutput()' returns 'input_' for one grouping set at a time.
+  /// 'groupingSetIndex_' contains the index of the grouping set to output in
+  /// the next 'getOutput' call. This index is used to generate groupId column
+  /// and lookup the input-to-output column mappings in the
+  /// projectMappings_.
+  int32_t projectSetIndex_{0};
+};
+} // namespace facebook::velox::exec

velox/exec/LocalPlanner.cpp

@@ -32,6 +32,7 @@
 #include "velox/exec/MergeJoin.h"
 #include "velox/exec/OrderBy.h"
 #include "velox/exec/PartitionedOutput.h"
+#include "velox/exec/Expand.h"
 #include "velox/exec/StreamingAggregation.h"
 #include "velox/exec/TableScan.h"
 #include "velox/exec/TableWriter.h"
@@ -402,6 +403,11 @@ std::shared_ptr<Driver> DriverFactory::createDriver(
         operators.push_back(
             std::make_unique<HashAggregation>(id, ctx.get(), aggregationNode));
       }
+    } else if (
+        auto expandNode =
+            std::dynamic_pointer_cast<const core::ExpandNode>(planNode)) {
+      operators.push_back(
+          std::make_unique<Expand>(id, ctx.get(), expandNode));
     } else if (
         auto groupIdNode =
             std::dynamic_pointer_cast<const core::GroupIdNode>(planNode)) {