Refactor velox type dependencies to pyvelox.

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(

csrc/velox/lib.cpp

@@ -12,7 +12,6 @@
 #include <velox/common/memory/Memory.h>
 #include <velox/type/Type.h>
 #include <velox/vector/TypeAliases.h>
-#include <iostream>
 #include <memory>
 
 #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 <torch/csrc/utils/pybind.h> // @manual
@@ -136,12 +134,6 @@ py::class_<SimpleColumn<T>, BaseColumn> declareSimpleType(
       });
 
   using I = typename velox::TypeTraits<kind>::ImplType;
-  py::class_<I, velox::Type, std::shared_ptr<I>>(
-      m,
-      (std::string("VeloxType_") + velox::TypeTraits<kind>::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<I> type) {
@@ -681,20 +673,8 @@ void declareArrayType(py::module& m) {
       .def("withElements", &ArrayColumn::withElements);
 
   using I = typename velox::TypeTraits<velox::TypeKind::ARRAY>::ImplType;
-  py::class_<I, velox::Type, std::shared_ptr<I>>(
-      m,
-      "VeloxArrayType",
-      // TODO: Move the Koksi binding of Velox type to OSS
-      py::module_local())
-      .def(py::init<velox::TypePtr>())
-      .def("element_type", &velox::ArrayType::elementType);
-
-  using J = typename velox::FixedSizeArrayType;
-  py::class_<J, velox::Type, std::shared_ptr<J>>(
-      m, "VeloxFixedArrayType", py::module_local())
-      .def(py::init<int, velox::TypePtr>())
-      .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<I> type) {
@@ -737,14 +717,6 @@ void declareMapType(py::module& m) {
       .def("slice", &MapColumn::slice);
 
   using I = typename velox::TypeTraits<velox::TypeKind::MAP>::ImplType;
-  py::class_<I, velox::Type, std::shared_ptr<I>>(
-      m,
-      "VeloxMapType",
-      // TODO: Move the Koksi binding of Velox type to OSS
-      py::module_local())
-      .def(py::init<velox::TypePtr, velox::TypePtr>())
-      .def("key_type", &velox::MapType::keyType)
-      .def("value_type", &velox::MapType::valueType);
 
   m.def("Column", [](std::shared_ptr<I> type) {
     return std::make_unique<MapColumn>(type);
@@ -772,19 +744,6 @@ void declareRowType(py::module& m) {
       });
 
   using I = typename velox::TypeTraits<velox::TypeKind::ROW>::ImplType;
-  py::class_<I, velox::Type, std::shared_ptr<I>>(
-      m,
-      "VeloxRowType",
-      // TODO: Move the Koksi binding of Velox type to OSS
-      py::module_local())
-      .def(py::init<
-           std::vector<std::string>&&,
-           std::vector<std::shared_ptr<const velox::Type>>&&>())
-      .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<I> type) {
     return std::make_unique<RowColumn>(type);
   });
@@ -833,33 +792,8 @@ PYBIND11_MODULE(_torcharrow, m) {
       .def_property_readonly("length", &BaseColumn::getLength)
       .def("__len__", &BaseColumn::getLength);
 
-  py::enum_<velox::TypeKind>(
-      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_<velox::Type, std::shared_ptr<velox::Type>>(
-      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<velox::TypeKind::BIGINT>(m);
   declareIntegralType<velox::TypeKind::INTEGER>(m);
   declareIntegralType<velox::TypeKind::SMALLINT>(m);

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()
+

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 <velox::TypeKind kind>
+    void declareType(py::module& m) {
+        using I = typename velox::TypeTraits<kind>::ImplType;
+        py::class_<I, velox::Type, std::shared_ptr<I>>(
+                m,
+                (std::string("VeloxType_") + velox::TypeTraits<kind>::name).c_str())
+                .def(py::init());
+    }
+
+    void addVeloxBindings(py::module& m) {
+        py::enum_<velox::TypeKind>(
+                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_<velox::Type, std::shared_ptr<facebook::velox::Type>>(
+                m,
+                "VeloxType")
+                .def("kind", &velox::Type::kind)
+                .def("kind_name", &velox::Type::kindName);
+
+        declareType<velox::TypeKind::BIGINT>(m);
+        declareType<velox::TypeKind::BOOLEAN>(m);
+        declareType<velox::TypeKind::TINYINT>(m);
+        declareType<velox::TypeKind::SMALLINT>(m);
+        declareType<velox::TypeKind::INTEGER>(m);
+        declareType<velox::TypeKind::REAL>(m);
+        declareType<velox::TypeKind::DOUBLE>(m);
+        declareType<velox::TypeKind::VARCHAR>(m);
+        declareType<velox::TypeKind::VARBINARY>(m);
+        declareType<velox::TypeKind::TIMESTAMP>(m);
+
+        using I = typename velox::TypeTraits<velox::TypeKind::ARRAY>::ImplType;
+        py::class_<I, velox::Type, std::shared_ptr<I>>(
+                m,
+                "VeloxArrayType")
+                .def(py::init<velox::TypePtr>())
+                .def("element_type", &velox::ArrayType::elementType);
+
+        using J = typename velox::FixedSizeArrayType;
+        py::class_<J, velox::Type, std::shared_ptr<J>>(
+                m, "VeloxFixedArrayType")
+                .def(py::init<int, velox::TypePtr>())
+                .def("element_type", &velox::FixedSizeArrayType::elementType)
+                .def("fixed_width", &velox::FixedSizeArrayType::fixedElementsWidth);
+
+        using M = typename velox::TypeTraits<velox::TypeKind::MAP>::ImplType;
+        py::class_<M, velox::Type, std::shared_ptr<M>>(
+                m,
+                "VeloxMapType")
+                .def(py::init<velox::TypePtr, velox::TypePtr>())
+                .def("key_type", &velox::MapType::keyType)
+                .def("value_type", &velox::MapType::valueType);
+
+        using R = typename velox::TypeTraits<velox::TypeKind::ROW>::ImplType;
+
+        py::class_<R, velox::Type, std::shared_ptr<R>>(
+                m,
+                "VeloxRowType")
+                .def(py::init<
+                        std::vector<std::string>&&,
+                        std::vector<std::shared_ptr<const velox::Type>>&&>())
+                .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
+}
+

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 <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
+#include <velox/type/Type.h>
+
+namespace facebook::pyvelox {
+
+    void addVeloxBindings(pybind11::module& m);
+}