diff --git a/csrc/velox/CMakeLists.txt b/csrc/velox/CMakeLists.txt index 9aa97ce5e..7788b148f 100644 --- a/csrc/velox/CMakeLists.txt +++ b/csrc/velox/CMakeLists.txt @@ -9,6 +9,7 @@ cmake_minimum_required(VERSION 3.15) # _torcharrow is a shared library as it's a Python extension set(CMAKE_POSITION_INDEPENDENT_CODE ON) +set(CREATE_PYVELOX_MODULE OFF) # To make the right CPython is built with on GitHub Actions, # see https://github.com/actions/setup-python/issues/121#issuecomment-1014500503 @@ -91,6 +92,7 @@ endif() # set_target_properties(_torcharrow PROPERTIES CXX_VISIBILITY_PRESET default) add_subdirectory(velox) add_subdirectory(functions) +add_subdirectory(pyvelox) # Link with Velox: @@ -103,6 +105,7 @@ target_link_libraries(_torcharrow PRIVATE velox_function_registry velox_arrow_bridge torcharrow_udfs + pyvelox ) target_compile_definitions( diff --git a/csrc/velox/lib.cpp b/csrc/velox/lib.cpp index 10952daec..272a92a01 100644 --- a/csrc/velox/lib.cpp +++ b/csrc/velox/lib.cpp @@ -12,7 +12,6 @@ #include #include #include -#include #include #include "bindings.h" @@ -23,10 +22,9 @@ #include "velox/buffer/StringViewBufferHolder.h" #include "velox/common/base/Exceptions.h" #include "velox/functions/prestosql/registration/RegistrationFunctions.h" -#include "velox/type/Type.h" -#include "velox/vector/TypeAliases.h" #include "velox/vector/arrow/Abi.h" #include "velox/vector/arrow/Bridge.h" +#include "pyvelox/pyvelox.h" #ifdef USE_TORCH #include // @manual @@ -136,12 +134,6 @@ py::class_, BaseColumn> declareSimpleType( }); using I = typename velox::TypeTraits::ImplType; - py::class_>( - m, - (std::string("VeloxType_") + velox::TypeTraits::name).c_str(), - // TODO: Move the Koksi binding of Velox type to OSS - py::module_local()) - .def(py::init()); // Empty Column m.def("Column", [](std::shared_ptr type) { @@ -681,20 +673,8 @@ void declareArrayType(py::module& m) { .def("withElements", &ArrayColumn::withElements); using I = typename velox::TypeTraits::ImplType; - py::class_>( - m, - "VeloxArrayType", - // TODO: Move the Koksi binding of Velox type to OSS - py::module_local()) - .def(py::init()) - .def("element_type", &velox::ArrayType::elementType); - - using J = typename velox::FixedSizeArrayType; - py::class_>( - m, "VeloxFixedArrayType", py::module_local()) - .def(py::init()) - .def("element_type", &velox::FixedSizeArrayType::elementType) - .def("fixed_width", &velox::FixedSizeArrayType::fixedElementsWidth); + + using J = typename velox::FixedSizeArrayType; // Empty Column m.def("Column", [](std::shared_ptr type) { @@ -737,14 +717,6 @@ void declareMapType(py::module& m) { .def("slice", &MapColumn::slice); using I = typename velox::TypeTraits::ImplType; - py::class_>( - m, - "VeloxMapType", - // TODO: Move the Koksi binding of Velox type to OSS - py::module_local()) - .def(py::init()) - .def("key_type", &velox::MapType::keyType) - .def("value_type", &velox::MapType::valueType); m.def("Column", [](std::shared_ptr type) { return std::make_unique(type); @@ -772,19 +744,6 @@ void declareRowType(py::module& m) { }); using I = typename velox::TypeTraits::ImplType; - py::class_>( - m, - "VeloxRowType", - // TODO: Move the Koksi binding of Velox type to OSS - py::module_local()) - .def(py::init< - std::vector&&, - std::vector>&&>()) - .def("size", &I::size) - .def("get_child_idx", &I::getChildIdx) - .def("contains_child", &I::containsChild) - .def("name_of", &I::nameOf) - .def("child_at", &I::childAt); m.def("Column", [](std::shared_ptr type) { return std::make_unique(type); }); @@ -833,33 +792,8 @@ PYBIND11_MODULE(_torcharrow, m) { .def_property_readonly("length", &BaseColumn::getLength) .def("__len__", &BaseColumn::getLength); - py::enum_( - m, - "TypeKind", // TODO: Move the Koksi binding of Velox type to OSS - py::module_local()) - .value("BOOLEAN", velox::TypeKind::BOOLEAN) - .value("TINYINT", velox::TypeKind::TINYINT) - .value("SMALLINT", velox::TypeKind::SMALLINT) - .value("INTEGER", velox::TypeKind::INTEGER) - .value("BIGINT", velox::TypeKind::BIGINT) - .value("REAL", velox::TypeKind::REAL) - .value("DOUBLE", velox::TypeKind::DOUBLE) - .value("VARCHAR", velox::TypeKind::VARCHAR) - .value("VARBINARY", velox::TypeKind::VARBINARY) - .value("TIMESTAMP", velox::TypeKind::TIMESTAMP) - .value("ARRAY", velox::TypeKind::ARRAY) - .value("MAP", velox::TypeKind::MAP) - .value("ROW", velox::TypeKind::ROW) - .export_values(); - - py::class_>( - m, - "VeloxType", - // TODO: Move the Koksi binding of Velox type to OSS - py::module_local()) - .def("kind", &velox::Type::kind) - .def("kind_name", &velox::Type::kindName); + pyvelox::addVeloxBindings(m); declareIntegralType(m); declareIntegralType(m); declareIntegralType(m); diff --git a/csrc/velox/pyvelox/CMakeLists.txt b/csrc/velox/pyvelox/CMakeLists.txt new file mode 100644 index 000000000..33feb3eb3 --- /dev/null +++ b/csrc/velox/pyvelox/CMakeLists.txt @@ -0,0 +1,34 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +if(${CREATE_PYVELOX_MODULE}) + + # Define our Python module: + pybind11_add_module( + _pyvelox + MODULE + NO_EXTRAS # TODO: LTO crashes GCC9.2. File issues to pybind11 + pyvelox.cpp + pyvelox.h + ) + + # Link with Velox: + target_link_libraries(_pyvelox PRIVATE + velox_type + ) + + install( + TARGETS _pyvelox + LIBRARY DESTINATION . + ) +else() + add_library(pyvelox pyvelox.cpp pyvelox.h) + target_link_libraries( + pyvelox + velox_type + pybind11::module) +endif() + diff --git a/csrc/velox/pyvelox/__init__.py b/csrc/velox/pyvelox/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/csrc/velox/pyvelox/pyvelox.cpp b/csrc/velox/pyvelox/pyvelox.cpp new file mode 100644 index 000000000..33ddedb44 --- /dev/null +++ b/csrc/velox/pyvelox/pyvelox.cpp @@ -0,0 +1,111 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "pyvelox.h" + + +namespace facebook::pyvelox { + namespace py = pybind11; + + template + void declareType(py::module& m) { + using I = typename velox::TypeTraits::ImplType; + py::class_>( + m, + (std::string("VeloxType_") + velox::TypeTraits::name).c_str()) + .def(py::init()); + } + + void addVeloxBindings(py::module& m) { + py::enum_( + m, + "TypeKind", + py::module_local()) + .value("BOOLEAN", velox::TypeKind::BOOLEAN) + .value("TINYINT", velox::TypeKind::TINYINT) + .value("SMALLINT", velox::TypeKind::SMALLINT) + .value("INTEGER", velox::TypeKind::INTEGER) + .value("BIGINT", velox::TypeKind::BIGINT) + .value("REAL", velox::TypeKind::REAL) + .value("DOUBLE", velox::TypeKind::DOUBLE) + .value("VARCHAR", velox::TypeKind::VARCHAR) + .value("VARBINARY", velox::TypeKind::VARBINARY) + .value("TIMESTAMP", velox::TypeKind::TIMESTAMP) + .value("ARRAY", velox::TypeKind::ARRAY) + .value("MAP", velox::TypeKind::MAP) + .value("ROW", velox::TypeKind::ROW) + .export_values(); + + py::class_>( + m, + "VeloxType") + .def("kind", &velox::Type::kind) + .def("kind_name", &velox::Type::kindName); + + declareType(m); + declareType(m); + declareType(m); + declareType(m); + declareType(m); + declareType(m); + declareType(m); + declareType(m); + declareType(m); + declareType(m); + + using I = typename velox::TypeTraits::ImplType; + py::class_>( + m, + "VeloxArrayType") + .def(py::init()) + .def("element_type", &velox::ArrayType::elementType); + + using J = typename velox::FixedSizeArrayType; + py::class_>( + m, "VeloxFixedArrayType") + .def(py::init()) + .def("element_type", &velox::FixedSizeArrayType::elementType) + .def("fixed_width", &velox::FixedSizeArrayType::fixedElementsWidth); + + using M = typename velox::TypeTraits::ImplType; + py::class_>( + m, + "VeloxMapType") + .def(py::init()) + .def("key_type", &velox::MapType::keyType) + .def("value_type", &velox::MapType::valueType); + + using R = typename velox::TypeTraits::ImplType; + + py::class_>( + m, + "VeloxRowType") + .def(py::init< + std::vector&&, + std::vector>&&>()) + .def("size", &R::size) + .def("get_child_idx", &R::getChildIdx) + .def("contains_child", &R::containsChild) + .def("name_of", &R::nameOf) + .def("child_at", &R::childAt); + } + +#ifdef CREATE_PYVELOX_MODULE + PYBIND11_MODULE(_pyvelox, m) { + m.doc() = R"pbdoc( + PyVelox native code module + ----------------------- + )pbdoc"; + + addVeloxBindings(m); + + m.attr("__version__") = "dev"; + } +#endif +} + diff --git a/csrc/velox/pyvelox/pyvelox.h b/csrc/velox/pyvelox/pyvelox.h new file mode 100644 index 000000000..0b6b7d334 --- /dev/null +++ b/csrc/velox/pyvelox/pyvelox.h @@ -0,0 +1,19 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +namespace facebook::pyvelox { + + void addVeloxBindings(pybind11::module& m); +} \ No newline at end of file