expressions.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
 
#include <cudf/ast/ast_operator.hpp>
#include <cudf/column/scalar_column_view.hpp>
#include <cudf/fixed_point/fixed_point.hpp>
#include <cudf/scalar/scalar.hpp>
#include <cudf/scalar/scalar_device_view.cuh>
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/error.hpp>
 
#include <cuda/stream>
 
#include <cstdint>
#include <functional>
#include <initializer_list>
#include <memory>
#include <optional>
#include <vector>
 
namespace CUDF_EXPORT cudf {
 
namespace detail {
namespace row_ir {
 
struct node;
 
struct ast_converter;
 
}  // namespace row_ir
}  // namespace detail
 
namespace ast {
// Forward declaration.
namespace detail {
class expression_parser;
class expression_transformer;
}  // namespace detail
 
struct [[nodiscard]] expression {
  virtual cudf::size_type accept(detail::expression_parser& visitor) const = 0;
 
  virtual std::reference_wrapper<expression const> accept(
    detail::expression_transformer& visitor) const = 0;
 
  [[nodiscard]] virtual std::unique_ptr<cudf::detail::row_ir::node> accept(
    cudf::detail::row_ir::ast_converter& visitor) const = 0;
 
  [[nodiscard]] bool may_evaluate_null(table_view const& left, cuda::stream_ref stream) const
  {
    return may_evaluate_null(left, left, stream);
  }
 
  [[nodiscard]] virtual bool may_evaluate_null(table_view const& left,
                                               table_view const& right,
                                               cuda::stream_ref stream) const = 0;
 
  virtual ~expression() {}
};
 
enum class table_reference {
  LEFT,   
  RIGHT,  
  OUTPUT  
};
 
class generic_scalar_device_view : public cudf::detail::scalar_device_view_base {
 public:
  template <typename T>
  __device__ T const value() const noexcept
  {
    if constexpr (std::is_same_v<T, cudf::string_view>) {
      return string_view(static_cast<char const*>(_data), _size);
    }
    if constexpr (cudf::is_fixed_point<T>()) {
      using rep_type   = typename T::rep;
      auto const rep   = *static_cast<rep_type const*>(_data);
      auto const scale = numeric::scale_type{type().scale()};
      return T{numeric::scaled_integer<rep_type>{rep, scale}};
    }
    return *static_cast<T const*>(_data);
  }
 
  template <typename T>
  generic_scalar_device_view(numeric_scalar<T>& s)
    : generic_scalar_device_view(s.type(), s.data(), s.validity_data())
  {
  }
 
  template <typename T>
  generic_scalar_device_view(timestamp_scalar<T>& s)
    : generic_scalar_device_view(s.type(), s.data(), s.validity_data())
  {
  }
 
  template <typename T>
  generic_scalar_device_view(duration_scalar<T>& s)
    : generic_scalar_device_view(s.type(), s.data(), s.validity_data())
  {
  }
 
  generic_scalar_device_view(string_scalar& s)
    : generic_scalar_device_view(s.type(), s.data(), s.validity_data(), s.size())
  {
  }
 
  template <typename T>
  generic_scalar_device_view(cudf::fixed_point_scalar<T>& s)
    : generic_scalar_device_view{s.type(), s.data(), s.validity_data()}
  {
  }
 
 protected:
  void const* _data{};      
  size_type const _size{};  
 
  generic_scalar_device_view(data_type type, void const* data, bool* is_valid)
    : cudf::detail::scalar_device_view_base(type, is_valid), _data(data)
  {
  }
 
  generic_scalar_device_view(data_type type, void const* data, bool* is_valid, size_type size)
    : cudf::detail::scalar_device_view_base(type, is_valid), _data(data), _size(size)
  {
  }
};
 
class literal : public expression {
 public:
  template <typename T>
  literal(cudf::numeric_scalar<T>& value)
    : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}}
  {
  }
 
  template <typename T>
  literal(cudf::timestamp_scalar<T>& value)
    : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}}
  {
  }
 
  template <typename T>
  literal(cudf::duration_scalar<T>& value)
    : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}}
  {
  }
 
  literal(cudf::string_scalar& value)
    : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}}
  {
  }
 
  template <typename T>
  literal(cudf::fixed_point_scalar<T>& value)
    : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}}
  {
  }
 
  literal(scalar_column_view value) : scalar{std::move(value)} {}
 
  [[nodiscard]] cudf::data_type get_data_type() const
  {
    return std::visit(
      [](auto const& value) {
        if constexpr (std::is_same_v<std::decay_t<decltype(value)>, ast_scalar>) {
          return value.value.type();
        } else {
          return value.type();
        }
      },
      scalar);
  }
 
  [[nodiscard]] bool is_scalar_column_view() const noexcept
  {
    return std::holds_alternative<scalar_column_view>(scalar);
  }
 
  [[nodiscard]] generic_scalar_device_view get_value() const
  {
    return std::get<ast_scalar>(scalar).value;
  }
 
  [[nodiscard]] cudf::scalar const& get_scalar() const
  {
    return std::get<ast_scalar>(scalar).scalar.get();
  }
 
  [[nodiscard]] scalar_column_view const& get_scalar_column_view() const
  {
    return std::get<scalar_column_view>(scalar);
  }
 
  cudf::size_type accept(detail::expression_parser& visitor) const override;
 
  std::reference_wrapper<expression const> accept(
    detail::expression_transformer& visitor) const override;
 
  [[nodiscard]] std::unique_ptr<cudf::detail::row_ir::node> accept(
    cudf::detail::row_ir::ast_converter& visitor) const override;
 
  [[nodiscard]] bool may_evaluate_null(table_view const& left,
                                       table_view const& right,
                                       cuda::stream_ref stream) const override
  {
    return !is_valid(stream);
  }
 
  [[nodiscard]] bool is_valid(cuda::stream_ref stream) const
  {
    if (auto* s = std::get_if<ast_scalar>(&scalar)) {
      return s->scalar.get().is_valid(stream);
    } else {
      auto& c = std::get<scalar_column_view>(scalar);
      return c.null_count() == 0;
    }
  }
 
 private:
  struct ast_scalar {
    std::reference_wrapper<cudf::scalar const> scalar;
    generic_scalar_device_view value;
  };
 
  std::variant<ast_scalar, scalar_column_view> scalar;
};
 
class column_reference : public expression {
 public:
  column_reference(cudf::size_type column_index,
                   table_reference table_source = table_reference::LEFT)
    : column_index(column_index), table_source(table_source)
  {
  }
 
  [[nodiscard]] cudf::size_type get_column_index() const { return column_index; }
 
  [[nodiscard]] table_reference get_table_source() const { return table_source; }
 
  [[nodiscard]] cudf::data_type get_data_type(table_view const& table) const;
 
  [[nodiscard]] cudf::data_type get_data_type(table_view const& left_table,
                                              table_view const& right_table) const;
 
  cudf::size_type accept(detail::expression_parser& visitor) const override;
 
  std::reference_wrapper<expression const> accept(
    detail::expression_transformer& visitor) const override;
 
  [[nodiscard]] bool may_evaluate_null(table_view const& left,
                                       table_view const& right,
                                       cuda::stream_ref stream) const override
  {
    return (table_source == table_reference::LEFT ? left : right).column(column_index).has_nulls();
  }
 
  [[nodiscard]] std::unique_ptr<cudf::detail::row_ir::node> accept(
    cudf::detail::row_ir::ast_converter& visitor) const override;
 
 private:
  cudf::size_type column_index;
  table_reference table_source;
};
 
class operation : public expression {
 public:
  operation(ast_operator op, expression const& input);
 
  operation(ast_operator op, expression const& left, expression const& right);
 
  // operation only stores references to expressions, so it does not accept r-value
  // references: the calling code must own the expressions.
  operation(ast_operator op, expression&& input)                         = delete;
  operation(ast_operator op, expression&& left, expression&& right)      = delete;
  operation(ast_operator op, expression&& left, expression const& right) = delete;
  operation(ast_operator op, expression const& left, expression&& right) = delete;
 
  [[nodiscard]] ast_operator get_operator() const { return op; }
 
  [[nodiscard]] std::vector<std::reference_wrapper<expression const>> const& get_operands() const
  {
    return operands;
  }
 
  cudf::size_type accept(detail::expression_parser& visitor) const override;
 
  std::reference_wrapper<expression const> accept(
    detail::expression_transformer& visitor) const override;
 
  [[nodiscard]] bool may_evaluate_null(table_view const& left,
                                       table_view const& right,
                                       cuda::stream_ref stream) const override;
 
  [[nodiscard]] std::unique_ptr<cudf::detail::row_ir::node> accept(
    cudf::detail::row_ir::ast_converter& visitor) const override;
 
 private:
  ast_operator op;
  std::vector<std::reference_wrapper<expression const>> operands;
};
 
namespace detail {
 
class predicate : public expression {
 public:
  predicate(expression const& source) : source_{source} {}
 
  cudf::size_type accept(detail::expression_parser& visitor) const override;
 
  std::reference_wrapper<expression const> accept(
    detail::expression_transformer& visitor) const override;
 
  [[nodiscard]] bool may_evaluate_null(table_view const& left,
                                       table_view const& right,
                                       cuda::stream_ref stream) const override;
 
  [[nodiscard]] std::unique_ptr<cudf::detail::row_ir::node> accept(
    cudf::detail::row_ir::ast_converter& visitor) const override;
 
  [[nodiscard]] expression const& get_operand() const { return source_; }
 
 private:
  std::reference_wrapper<expression const> source_;
};
 
}  // namespace detail
 
class column_name_reference : public expression {
 public:
  column_name_reference(std::string column_name) : column_name(std::move(column_name)) {}
 
  [[nodiscard]] std::string get_column_name() const { return column_name; }
 
  cudf::size_type accept(detail::expression_parser& visitor) const override;
 
  std::reference_wrapper<expression const> accept(
    detail::expression_transformer& visitor) const override;
 
  [[nodiscard]] bool may_evaluate_null(table_view const& left,
                                       table_view const& right,
                                       cuda::stream_ref stream) const override
  {
    return true;
  }
 
  [[nodiscard]] std::unique_ptr<cudf::detail::row_ir::node> accept(
    cudf::detail::row_ir::ast_converter& visitor) const override;
 
 private:
  std::string column_name;
};
 
class tree {
 public:
  tree() = default;
 
  tree(tree&&) = default;
 
  tree& operator=(tree&&) = default;
 
  ~tree() = default;
 
  // the tree is not copyable
  tree(tree const&)            = delete;
  tree& operator=(tree const&) = delete;
 
  template <typename Expr, typename... Args>
  std::enable_if_t<std::is_base_of_v<expression, Expr>, Expr const&> emplace(Args&&... args)
  {
    auto expr            = std::make_unique<Expr>(std::forward<Args>(args)...);
    Expr const& expr_ref = *expr;
    expressions.emplace_back(std::move(expr));
    return expr_ref;
  }
 
  template <typename Expr>
  decltype(auto) push(Expr expr)
  {
    return emplace<Expr>(std::move(expr));
  }
 
  [[nodiscard]] expression const& front() const { return *expressions.front(); }
 
  [[nodiscard]] expression const& back() const { return *expressions.back(); }
 
  [[nodiscard]] size_t size() const { return expressions.size(); }
 
  expression const& at(size_t index) { return *expressions.at(index); }
 
  expression const& operator[](size_t index) const { return *expressions[index]; }
 
 private:
  // TODO: use better ownership semantics, the unique_ptr here is redundant. Consider using a bump
  // allocator with type-erased deleters.
  std::vector<std::unique_ptr<expression>> expressions;
};
 
namespace jit {
 
enum class op : uint8_t {
  // Identity operators
  IDENTITY,
 
  // Null handling operators
  IS_NULL,
 
  COALESCE,
  PREDICATE,
 
  ADD,
  SUB,
  MUL,
  DIV,
  NEG,
  ABS,
  MOD,
  PYMOD,
  TRUE_DIV,
  FLOOR_DIV,
 
  ADD_OVERFLOW,
  SUB_OVERFLOW,
  MUL_OVERFLOW,
  DIV_OVERFLOW,
  NEG_OVERFLOW,
  ABS_OVERFLOW,
  MOD_OVERFLOW,
  CHECK_PRECISION,
 
  BITWISE_AND,
  BITWISE_INVERT,
  BITWISE_OR,
  BITWISE_XOR,
  BITWISE_SHIFT_LEFT,
  BITWISE_SHIFT_RIGHT,
 
  CAST_TO_BOOL8,
  CAST_TO_INT8,
  CAST_TO_INT16,
  CAST_TO_INT32,
  CAST_TO_INT64,
  CAST_TO_UINT8,
  CAST_TO_UINT16,
  CAST_TO_UINT32,
  CAST_TO_UINT64,
  CAST_TO_FLOAT32,
  CAST_TO_FLOAT64,
  CAST_TO_DECIMAL32,
  CAST_TO_DECIMAL64,
  CAST_TO_DECIMAL128,
  RESCALE,
 
  EQUAL,
  NOT_EQUAL,
  GREATER,
  GREATER_EQUAL,
  LESS,
  LESS_EQUAL,
  NULL_EQUAL,
  NULL_LOGICAL_AND,
  NULL_LOGICAL_OR,
  LOGICAL_AND,
  LOGICAL_OR,
  LOGICAL_NOT,
  IF_ELSE,
 
  CBRT,
  CEIL,
  FLOOR,
  RINT,
  SQRT,
  POW,
  EXP,
  LOG,
 
  ARCCOS,
  ARCCOSH,
  ARCSIN,
  ARCSINH,
  ARCTAN,
  ARCTANH,
  COS,
  COSH,
  SIN,
  SINH,
  TAN,
  TANH,
};
 
expression const& operation(ast::tree& tree,
                            op operator_id,
                            std::vector<std::reference_wrapper<expression const>> const& args,
                            cudf::error_policy error_policy     = cudf::error_policy::PROPAGATE,
                            std::optional<int32_t> target_scale = std::nullopt);
 
expression const& operation(ast::tree& tree,
                            op operator_id,
                            std::initializer_list<std::reference_wrapper<expression const>> args,
                            cudf::error_policy error_policy     = cudf::error_policy::PROPAGATE,
                            std::optional<int32_t> target_scale = std::nullopt);
 
}  // namespace jit
  // end of group
}  // namespace ast
}  // namespace CUDF_EXPORT cudf