expressions.hpp

转到此文件的文档。

/*
 * Copyright (c) 2020-2025, NVIDIA CORPORATION.
 *
 * 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
 *
 * https://apache.ac.cn/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 <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 <cstdint>
#include <memory>
#include <vector>

namespace CUDF_EXPORT cudf {
namespace ast {
// Forward declaration.
namespace detail {
class expression_parser;
class expression_transformer;
} // namespace detail

struct 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]] bool may_evaluate_null(table_view const& left, rmm::cuda_stream_view stream) const
 {
 return may_evaluate_null(left, left, stream);
 }

 [[nodiscard]] virtual bool may_evaluate_null(table_view const& left,
 table_view const& right,
 rmm::cuda_stream_view stream) const = 0;

 virtual ~expression() {}
};

enum class ast_operator : int32_t {
 // Binary operators
 ADD,
 SUB,
 MUL,
 DIV,
 TRUE_DIV,
 FLOOR_DIV,
 MOD,
 PYMOD,
 POW,
 EQUAL,
 NULL_EQUAL,
 NOT_EQUAL,
 LESS,
 GREATER,
 LESS_EQUAL,
 GREATER_EQUAL,
 BITWISE_AND,
 BITWISE_OR,
 BITWISE_XOR,
 LOGICAL_AND,
 NULL_LOGICAL_AND,
 LOGICAL_OR,
 NULL_LOGICAL_OR,
 // Unary operators
 IDENTITY,
 IS_NULL,
 SIN,
 COS,
 TAN,
 ARCSIN,
 ARCCOS,
 ARCTAN,
 SINH,
 COSH,
 TANH,
 ARCSINH,
 ARCCOSH,
 ARCTANH,
 EXP,
 LOG,
 SQRT,
 CBRT,
 CEIL,
 FLOOR,
 ABS,
 RINT,
 BIT_INVERT,
 NOT,
 CAST_TO_INT64,
 CAST_TO_UINT64,
 CAST_TO_FLOAT64
};

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);
 }
 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())
 {
 }

 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(value), value(value)
 {
 }

 template <typename T>
 literal(cudf::timestamp_scalar<T>& value) : scalar(value), value(value)
 {
 }

 template <typename T>
 literal(cudf::duration_scalar<T>& value) : scalar(value), value(value)
 {
 }

 literal(cudf::string_scalar& value) : scalar(value), value(value) {}

 [[nodiscard]] cudf::data_type get_data_type() const { return get_value().type(); }

 [[nodiscard]] generic_scalar_device_view get_value() const { return value; }

 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,
 rmm::cuda_stream_view stream) const override
 {
 return !is_valid(stream);
 }

 [[nodiscard]] bool is_valid(rmm::cuda_stream_view stream) const
 {
 return scalar.is_valid(stream);
 }

 private
 cudf::scalar const& scalar;
 generic_scalar_device_view const value;
};

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
 {
 return table.column(get_column_index()).type();
 }

 [[nodiscard]] cudf::data_type get_data_type(table_view const& left_table,
 table_view const& right_table) const
 {
 auto const table = [&] {
 if (get_table_source() == table_reference::LEFT) {
 return left_table;
 } else if (get_table_source() == table_reference::RIGHT) {
 return right_table;
 } else {
 CUDF_FAIL("Column reference data type cannot be determined from unknown table.");
 }
 }();
 return table.column(get_column_index()).type();
 }

 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,
 rmm::cuda_stream_view stream) const override
 {
 return (table_source == table_reference::LEFT ? left : right).column(column_index).has_nulls();
 }

 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,
 rmm::cuda_stream_view stream) const override
 {
 return std::any_of(operands.cbegin(),
 operands.cend(),
 [&left, &right, &stream](std::reference_wrapper<expression const> subexpr) {
 return subexpr.get().may_evaluate_null(left, right, stream);
 });
 };

 private
 ast_operator op;
 std::vector<std::reference_wrapper<expression const>> operands;
};

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,
 rmm::cuda_stream_view stream) const override
 {
 return true;
 }

 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;
};
 // end of group
} // namespace ast

} // namespace CUDF_EXPORT cudf
cudf::ast::column_name_reference