Skip to content
This repository has been archived by the owner on Nov 1, 2024. It is now read-only.

Commit

Permalink
Create PyVelox package and move type dependencies over.
Browse files Browse the repository at this point in the history
  • Loading branch information
kgpai committed Oct 4, 2022
1 parent 6d6af07 commit 351bdba
Show file tree
Hide file tree
Showing 6 changed files with 171 additions and 70 deletions.
3 changes: 3 additions & 0 deletions csrc/velox/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -103,6 +105,7 @@ target_link_libraries(_torcharrow PRIVATE
velox_function_registry
velox_arrow_bridge
torcharrow_udfs
pyvelox
)

target_compile_definitions(
Expand Down
74 changes: 4 additions & 70 deletions csrc/velox/lib.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
});
Expand Down Expand Up @@ -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);
Expand Down
34 changes: 34 additions & 0 deletions csrc/velox/pyvelox/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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()

Empty file added csrc/velox/pyvelox/__init__.py
Empty file.
111 changes: 111 additions & 0 deletions csrc/velox/pyvelox/pyvelox.cpp
Original file line number Diff line number Diff line change
@@ -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
}

19 changes: 19 additions & 0 deletions csrc/velox/pyvelox/pyvelox.h
Original file line number Diff line number Diff line change
@@ -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);
}

0 comments on commit 351bdba

Please sign in to comment.