expressions.hpp

/*
 * Copyright (c) 2020-2023, 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
 *
 *     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 <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>
 
namespace 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; }
 
  std::vector<std::reference_wrapper<expression 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 const op;
  std::vector<std::reference_wrapper<expression const>> 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;
};
 
}  // namespace ast
 
}  // namespace cudf