feat(cuDF): Add CudfExpression interface to allow nesting of cuDF AST and cuDF standalone functions (facebookincubator#15145)

Summary:
This PR adds the following features in cuDF based expression evaluation.
1. `CudfExpression` class which is extended to implement cuDF based AST expression evaluation and standalone cuDF functions based evaluations. This allows for arbitrary nesting of AST and cudf functions. Previously only one level of nesting was supported and only AST could nest functions based evaluation. This class will be used in the future to add in JIT based evaluation without changing ast or function expressions.
2. Function registration: Now `CudfFunction`s can be registered without adding a corresponding entry in AST's pushExprToTree.
3. Expressions now return `ColumnOrView`. This puts the onus on the caller for materialization if the expression only refers to a field. This makes accessing nested fields zero-copy.
4. Expression evaluator registration: Now ASTExpression can be optionally not registered to force standalone function based evaluation. Registration also takes a priority value to prioritize one type of expression translation over the other. In case of duplicate functionality e.g when replacing a velox Expr arithmetic binary op, one can prioritize replacing with cudf AST over cudf::binary_operation

This PR also refactors SubfieldFilterToAST converter and ASTExpression into separate source files.

Contributes to facebookincubator#14149

Pull Request resolved: facebookincubator#15145

Reviewed By: kgpai

Differential Revision: D85887084

Pulled By: kevinwilfong

fbshipit-source-id: 4e0785e60b453a796cd1aed2e3c6eb08cac875ac

Author: devavret

velox/docs/configs.rst

@@ -1143,6 +1143,14 @@ Note: These configurations are experimental and subject to change.
      - string
      - ""
      - The prefix to use for the function names in cuDF.
+   * - cudf.ast_expression_enabled
+     - bool
+     - true
+     - If true, enable using cuDF AST-based expression evaluation when supported.
+   * - cudf.ast_expression_priority
+     - integer
+     - 100
+     - Priority of cuDF AST expressions. Higher value wins when multiple cuDF execution options are available for the same Velox expression. Standalone cuDF functions have priority 50. If enabled, with a default priority of 100, AST will be chosen as replacement for cudf execution.
    * - cudf.allow_cpu_fallback
      - bool
      - true

velox/experimental/cudf/CMakeLists.txt

@@ -15,6 +15,7 @@
 add_subdirectory(exec)
 add_subdirectory(connectors)
 add_subdirectory(vector)
+add_subdirectory(expression)
 
 if(VELOX_BUILD_TESTING)
   add_subdirectory(tests)

velox/experimental/cudf/CudfConfig.h

@@ -29,6 +29,10 @@ struct CudfConfig {
   static constexpr const char* kCudfMemoryPercent{"cudf.memory_percent"};
   static constexpr const char* kCudfFunctionNamePrefix{
       "cudf.function_name_prefix"};
+  static constexpr const char* kCudfAstExpressionEnabled{
+      "cudf.ast_expression_enabled"};
+  static constexpr const char* kCudfAstExpressionPriority{
+      "cudf.ast_expression_priority"};
   static constexpr const char* kCudfAllowCpuFallback{"cudf.allow_cpu_fallback"};
   static constexpr const char* kCudfLogFallback{"cudf.log_fallback"};
 
@@ -60,6 +64,18 @@ struct CudfConfig {
   /// Register all the functions with the functionNamePrefix.
   std::string functionNamePrefix;
 
+  /// Enable AST in expression evaluation
+  bool astExpressionEnabled{true};
+
+  /// Priority of AST expression. Expression with higher priority is chosen for
+  /// a given root expression.
+  /// Example:
+  /// Priority of expression that uses individual cuDF functions is 50.
+  /// If AST priority is 100 then for a velox expression node that is supported
+  /// by both, AST will be chosen as replacement for cudf execution, if AST
+  /// priority is 25 then standalone cudf function is chosen.
+  int astExpressionPriority{100};
+
   /// Whether to log a reason for falling back to Velox CPU execution.
   bool logFallback{true};
 };

velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp

@@ -18,10 +18,11 @@
 #include "velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.h"
 #include "velox/experimental/cudf/connectors/hive/CudfHiveDataSource.h"
 #include "velox/experimental/cudf/connectors/hive/CudfHiveTableHandle.h"
-#include "velox/experimental/cudf/exec/ExpressionEvaluator.h"
 #include "velox/experimental/cudf/exec/ToCudf.h"
 #include "velox/experimental/cudf/exec/Utilities.h"
 #include "velox/experimental/cudf/exec/VeloxCudfInterop.h"
+#include "velox/experimental/cudf/expression/ExpressionEvaluator.h"
+#include "velox/experimental/cudf/expression/SubfieldFiltersToAst.h"
 #include "velox/experimental/cudf/vector/CudfVector.h"
 
 #include "velox/common/time/Timer.h"
@@ -130,8 +131,8 @@ CudfHiveDataSource::CudfHiveDataSource(
       }
     }();
 
-    cudfExpressionEvaluator_ = velox::cudf_velox::ExpressionEvaluator(
-        remainingFilterExprSet_->exprs(), remainingFilterType_);
+    cudfExpressionEvaluator_ = velox::cudf_velox::createCudfExpression(
+        remainingFilterExprSet_->exprs()[0], remainingFilterType_);
     // TODO(kn): Get column names and subfields from remaining filter and add to
     // readColumnNames_
   }
@@ -180,7 +181,7 @@ std::optional<RowVectorPtr> CudfHiveDataSource::next(
     const auto originalNumColumns = cudfTableColumns.size();
     // Filter may need addtional computed columns which are added to
     // cudfTableColumns
-    auto filterResult = cudfExpressionEvaluator_.compute(
+    auto filterResult = cudfExpressionEvaluator_->eval(
         cudfTableColumns, stream_, cudf::get_current_device_resource_ref());
     // discard computed columns
     std::vector<std::unique_ptr<cudf::column>> originalColumns;
@@ -194,7 +195,7 @@ std::optional<RowVectorPtr> CudfHiveDataSource::next(
     // Keep only rows where the filter is true
     cudfTable = cudf::apply_boolean_mask(
         *originalTable,
-        *filterResult[0],
+        asView(filterResult),
         stream_,
         cudf::get_current_device_resource_ref());
   }

velox/experimental/cudf/connectors/hive/CudfHiveDataSource.h

@@ -18,8 +18,8 @@
 
 #include "velox/experimental/cudf/connectors/hive/CudfHiveConfig.h"
 #include "velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.h"
-#include "velox/experimental/cudf/exec/ExpressionEvaluator.h"
 #include "velox/experimental/cudf/exec/NvtxHelper.h"
+#include "velox/experimental/cudf/expression/ExpressionEvaluator.h"
 
 #include "velox/common/base/RandomUtil.h"
 #include "velox/common/io/IoStatistics.h"
@@ -128,7 +128,7 @@ class CudfHiveDataSource : public DataSource, public NvtxHelper {
   // Expression evaluator for remaining filter.
   core::ExpressionEvaluator* const expressionEvaluator_;
   std::unique_ptr<exec::ExprSet> remainingFilterExprSet_;
-  velox::cudf_velox::ExpressionEvaluator cudfExpressionEvaluator_;
+  std::shared_ptr<velox::cudf_velox::CudfExpression> cudfExpressionEvaluator_;
 
   // Expression evaluator for subfield filter.
   std::vector<std::unique_ptr<cudf::scalar>> subfieldScalars_;

velox/experimental/cudf/exec/CMakeLists.txt

@@ -23,7 +23,6 @@ add_library(
   CudfLocalPartition.cpp
   CudfOrderBy.cpp
   DebugUtil.cpp
-  ExpressionEvaluator.cpp
   ToCudf.cpp
   Utilities.cpp
   VeloxCudfInterop.cpp
@@ -40,6 +39,7 @@ target_link_libraries(
     velox_cudf_vector
     velox_exec
     velox_cudf_hive_connector
+    velox_cudf_expression
 )
 
 target_compile_options(velox_cudf_exec PRIVATE -Wno-missing-field-initializers)

velox/experimental/cudf/exec/CudfFilterProject.cpp

@@ -173,15 +173,26 @@ void CudfFilterProject::initialize() {
       ++i;
     }
   }
-  std::vector<std::shared_ptr<velox::exec::Expr>> projectExprs;
   if (hasFilter_) {
-    filterEvaluator_ = ExpressionEvaluator({expr->exprs()[0]}, inputType);
-    projectExprs = {expr->exprs().begin() + 1, expr->exprs().end()};
+    // First expr is Filter, rest are Project
+    filterEvaluator_ = createCudfExpression(expr->exprs()[0], inputType);
+    std::transform(
+        expr->exprs().begin() + 1,
+        expr->exprs().end(),
+        std::back_inserter(projectEvaluators_),
+        [inputType](const auto& expr) {
+          return createCudfExpression(expr, inputType);
+        });
+  } else {
+    std::transform(
+        expr->exprs().begin(),
+        expr->exprs().end(),
+        std::back_inserter(projectEvaluators_),
+        [inputType](const auto& expr) {
+          return createCudfExpression(expr, inputType);
+        });
   }
 
-  projectEvaluator_ =
-      ExpressionEvaluator(hasFilter_ ? projectExprs : expr->exprs(), inputType);
-
   filter_.reset();
   project_.reset();
 }
@@ -233,16 +244,16 @@ void CudfFilterProject::filter(
     std::vector<std::unique_ptr<cudf::column>>& inputTableColumns,
     rmm::cuda_stream_view stream) {
   // Evaluate the Filter
-  auto filterColumns = filterEvaluator_.compute(
-      inputTableColumns, stream, cudf::get_current_device_resource_ref());
-  auto filterColumn = filterColumns[0]->view();
+  auto filterColumn = filterEvaluator_->eval(
+      inputTableColumns, stream, cudf::get_current_device_resource_ref(), true);
+  auto filterColumnView = asView(filterColumn);
   bool shouldApplyFilter = [&]() {
-    if (filterColumn.has_nulls()) {
+    if (filterColumnView.has_nulls()) {
       return true;
     }
-    // check if all values in filterColumn are true
+    // check if all values in filterColumnView are true
     auto isAllTrue = cudf::reduce(
-        filterColumn,
+        filterColumnView,
         *cudf::make_all_aggregation<cudf::reduce_aggregation>(),
         cudf::data_type(cudf::type_id::BOOL8),
         stream,
@@ -256,23 +267,39 @@ void CudfFilterProject::filter(
     auto filterTable =
         std::make_unique<cudf::table>(std::move(inputTableColumns));
     auto filteredTable =
-        cudf::apply_boolean_mask(*filterTable, filterColumn, stream);
+        cudf::apply_boolean_mask(*filterTable, filterColumnView, stream);
     inputTableColumns = filteredTable->release();
   }
 }
 
 std::vector<std::unique_ptr<cudf::column>> CudfFilterProject::project(
     std::vector<std::unique_ptr<cudf::column>>& inputTableColumns,
     rmm::cuda_stream_view stream) {
-  auto columns = projectEvaluator_.compute(
-      inputTableColumns, stream, cudf::get_current_device_resource_ref());
+  std::vector<ColumnOrView> columns;
+  for (auto& projectEvaluator : projectEvaluators_) {
+    columns.push_back(projectEvaluator->eval(
+        inputTableColumns,
+        stream,
+        cudf::get_current_device_resource_ref(),
+        true));
+  }
 
   // Rearrange columns to match outputType_
   std::vector<std::unique_ptr<cudf::column>> outputColumns(outputType_->size());
   // computed resultProjections
   for (int i = 0; i < resultProjections_.size(); i++) {
-    VELOX_CHECK_NOT_NULL(columns[i]);
-    outputColumns[resultProjections_[i].outputChannel] = std::move(columns[i]);
+    auto& columnOrView = columns[i];
+    if (std::holds_alternative<std::unique_ptr<cudf::column>>(columnOrView)) {
+      // Move the owned column
+      outputColumns[resultProjections_[i].outputChannel] =
+          std::move(std::get<std::unique_ptr<cudf::column>>(columnOrView));
+    } else {
+      // Materialize the column_view into an owned column
+      auto view = std::get<cudf::column_view>(columnOrView);
+      outputColumns[resultProjections_[i].outputChannel] =
+          std::make_unique<cudf::column>(
+              view, stream, cudf::get_current_device_resource_ref());
+    }
   }
 
   // Count occurrences of each inputChannel, and move columns if they occur only
@@ -300,16 +327,6 @@ std::vector<std::unique_ptr<cudf::column>> CudfFilterProject::project(
     inputChannelCount[identity.inputChannel]--;
   }
 
-  for (auto i = 0; i < outputColumns.size(); ++i) {
-    const auto cudfOutputType =
-        cudf::data_type(cudf_velox::veloxToCudfTypeId(outputType_->childAt(i)));
-
-    if (outputColumns[i]->type() != cudfOutputType) {
-      outputColumns[i] =
-          cudf::cast(*(outputColumns[i]), cudfOutputType, stream);
-    }
-  }
-
   return outputColumns;
 }
 

velox/experimental/cudf/exec/CudfFilterProject.h

@@ -16,8 +16,8 @@
 
 #pragma once
 
-#include "velox/experimental/cudf/exec/ExpressionEvaluator.h"
 #include "velox/experimental/cudf/exec/NvtxHelper.h"
+#include "velox/experimental/cudf/expression/ExpressionEvaluator.h"
 
 #include "velox/core/Expressions.h"
 #include "velox/core/PlanNode.h"
@@ -63,8 +63,8 @@ class CudfFilterProject : public exec::Operator, public NvtxHelper {
 
   void close() override {
     Operator::close();
-    projectEvaluator_.close();
-    filterEvaluator_.close();
+    projectEvaluators_.clear();
+    filterEvaluator_.reset();
   }
 
  private:
@@ -78,8 +78,8 @@ class CudfFilterProject : public exec::Operator, public NvtxHelper {
   std::shared_ptr<const core::ProjectNode> project_;
   std::shared_ptr<const core::FilterNode> filter_;
 
-  ExpressionEvaluator projectEvaluator_;
-  ExpressionEvaluator filterEvaluator_;
+  std::vector<CudfExpressionPtr> projectEvaluators_;
+  CudfExpressionPtr filterEvaluator_;
 
   std::vector<velox::exec::IdentityProjection> resultProjections_;
   std::vector<velox::exec::IdentityProjection> identityProjections_;

velox/experimental/cudf/exec/CudfHashJoin.cpp

@@ -16,10 +16,11 @@
 
 #include "velox/experimental/cudf/CudfConfig.h"
 #include "velox/experimental/cudf/exec/CudfHashJoin.h"
-#include "velox/experimental/cudf/exec/ExpressionEvaluator.h"
 #include "velox/experimental/cudf/exec/ToCudf.h"
 #include "velox/experimental/cudf/exec/Utilities.h"
 #include "velox/experimental/cudf/exec/VeloxCudfInterop.h"
+#include "velox/experimental/cudf/expression/AstExpression.h"
+#include "velox/experimental/cudf/expression/ExpressionEvaluator.h"
 
 #include "velox/exec/Task.h"