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

Commit a4aa08f

Browse files
kgpaikgpai
committed
Create PyVelox package and move type dependencies over.
1 parent 608258f commit a4aa08f

File tree

6 files changed

+171
-70
lines changed

6 files changed

+171
-70
lines changed

‎csrc/velox/CMakeLists.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ cmake_minimum_required(VERSION 3.15)
99

1010
# _torcharrow is a shared library as it's a Python extension
1111
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
12+
set(CREATE_PYVELOX_MODULE OFF)
1213

1314
# To make the right CPython is built with on GitHub Actions,
1415
# see https://github.com/actions/setup-python/issues/121#issuecomment-1014500503
@@ -91,6 +92,7 @@ endif()
9192
# set_target_properties(_torcharrow PROPERTIES CXX_VISIBILITY_PRESET default)
9293
add_subdirectory(velox)
9394
add_subdirectory(functions)
95+
add_subdirectory(pyvelox)
9496

9597

9698
# Link with Velox:
@@ -103,6 +105,7 @@ target_link_libraries(_torcharrow PRIVATE
103105
velox_function_registry
104106
velox_arrow_bridge
105107
torcharrow_udfs
108+
pyvelox
106109
)
107110

108111
target_compile_definitions(

‎csrc/velox/lib.cpp

Lines changed: 4 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
#include <velox/common/memory/Memory.h>
1313
#include <velox/type/Type.h>
1414
#include <velox/vector/TypeAliases.h>
15-
#include <iostream>
1615
#include <memory>
1716

1817
#include "bindings.h"
@@ -23,10 +22,9 @@
2322
#include "velox/buffer/StringViewBufferHolder.h"
2423
#include "velox/common/base/Exceptions.h"
2524
#include "velox/functions/prestosql/registration/RegistrationFunctions.h"
26-
#include "velox/type/Type.h"
27-
#include "velox/vector/TypeAliases.h"
2825
#include "velox/vector/arrow/Abi.h"
2926
#include "velox/vector/arrow/Bridge.h"
27+
#include "pyvelox/pyvelox.h"
3028

3129
#ifdef USE_TORCH
3230
#include <torch/csrc/utils/pybind.h> // @manual
@@ -136,12 +134,6 @@ py::class_<SimpleColumn<T>, BaseColumn> declareSimpleType(
136134
});
137135

138136
using I = typename velox::TypeTraits<kind>::ImplType;
139-
py::class_<I, velox::Type, std::shared_ptr<I>>(
140-
m,
141-
(std::string("VeloxType_") + velox::TypeTraits<kind>::name).c_str(),
142-
// TODO: Move the Koksi binding of Velox type to OSS
143-
py::module_local())
144-
.def(py::init());
145137

146138
// Empty Column
147139
m.def("Column", [](std::shared_ptr<I> type) {
@@ -681,20 +673,8 @@ void declareArrayType(py::module& m) {
681673
.def("withElements", &ArrayColumn::withElements);
682674

683675
using I = typename velox::TypeTraits<velox::TypeKind::ARRAY>::ImplType;
684-
py::class_<I, velox::Type, std::shared_ptr<I>>(
685-
m,
686-
"VeloxArrayType",
687-
// TODO: Move the Koksi binding of Velox type to OSS
688-
py::module_local())
689-
.def(py::init<velox::TypePtr>())
690-
.def("element_type", &velox::ArrayType::elementType);
691-
692-
using J = typename velox::FixedSizeArrayType;
693-
py::class_<J, velox::Type, std::shared_ptr<J>>(
694-
m, "VeloxFixedArrayType", py::module_local())
695-
.def(py::init<int, velox::TypePtr>())
696-
.def("element_type", &velox::FixedSizeArrayType::elementType)
697-
.def("fixed_width", &velox::FixedSizeArrayType::fixedElementsWidth);
676+
677+
using J = typename velox::FixedSizeArrayType;
698678

699679
// Empty Column
700680
m.def("Column", [](std::shared_ptr<I> type) {
@@ -737,14 +717,6 @@ void declareMapType(py::module& m) {
737717
.def("slice", &MapColumn::slice);
738718

739719
using I = typename velox::TypeTraits<velox::TypeKind::MAP>::ImplType;
740-
py::class_<I, velox::Type, std::shared_ptr<I>>(
741-
m,
742-
"VeloxMapType",
743-
// TODO: Move the Koksi binding of Velox type to OSS
744-
py::module_local())
745-
.def(py::init<velox::TypePtr, velox::TypePtr>())
746-
.def("key_type", &velox::MapType::keyType)
747-
.def("value_type", &velox::MapType::valueType);
748720

749721
m.def("Column", [](std::shared_ptr<I> type) {
750722
return std::make_unique<MapColumn>(type);
@@ -772,19 +744,6 @@ void declareRowType(py::module& m) {
772744
});
773745

774746
using I = typename velox::TypeTraits<velox::TypeKind::ROW>::ImplType;
775-
py::class_<I, velox::Type, std::shared_ptr<I>>(
776-
m,
777-
"VeloxRowType",
778-
// TODO: Move the Koksi binding of Velox type to OSS
779-
py::module_local())
780-
.def(py::init<
781-
std::vector<std::string>&&,
782-
std::vector<std::shared_ptr<const velox::Type>>&&>())
783-
.def("size", &I::size)
784-
.def("get_child_idx", &I::getChildIdx)
785-
.def("contains_child", &I::containsChild)
786-
.def("name_of", &I::nameOf)
787-
.def("child_at", &I::childAt);
788747
m.def("Column", [](std::shared_ptr<I> type) {
789748
return std::make_unique<RowColumn>(type);
790749
});
@@ -833,33 +792,8 @@ PYBIND11_MODULE(_torcharrow, m) {
833792
.def_property_readonly("length", &BaseColumn::getLength)
834793
.def("__len__", &BaseColumn::getLength);
835794

836-
py::enum_<velox::TypeKind>(
837-
m,
838-
"TypeKind", // TODO: Move the Koksi binding of Velox type to OSS
839-
py::module_local())
840-
.value("BOOLEAN", velox::TypeKind::BOOLEAN)
841-
.value("TINYINT", velox::TypeKind::TINYINT)
842-
.value("SMALLINT", velox::TypeKind::SMALLINT)
843-
.value("INTEGER", velox::TypeKind::INTEGER)
844-
.value("BIGINT", velox::TypeKind::BIGINT)
845-
.value("REAL", velox::TypeKind::REAL)
846-
.value("DOUBLE", velox::TypeKind::DOUBLE)
847-
.value("VARCHAR", velox::TypeKind::VARCHAR)
848-
.value("VARBINARY", velox::TypeKind::VARBINARY)
849-
.value("TIMESTAMP", velox::TypeKind::TIMESTAMP)
850-
.value("ARRAY", velox::TypeKind::ARRAY)
851-
.value("MAP", velox::TypeKind::MAP)
852-
.value("ROW", velox::TypeKind::ROW)
853-
.export_values();
854-
855-
py::class_<velox::Type, std::shared_ptr<velox::Type>>(
856-
m,
857-
"VeloxType",
858-
// TODO: Move the Koksi binding of Velox type to OSS
859-
py::module_local())
860-
.def("kind", &velox::Type::kind)
861-
.def("kind_name", &velox::Type::kindName);
862795

796+
pyvelox::addVeloxBindings(m);
863797
declareIntegralType<velox::TypeKind::BIGINT>(m);
864798
declareIntegralType<velox::TypeKind::INTEGER>(m);
865799
declareIntegralType<velox::TypeKind::SMALLINT>(m);

‎csrc/velox/pyvelox/CMakeLists.txt

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
if(${CREATE_PYVELOX_MODULE})
8+
9+
# Define our Python module:
10+
pybind11_add_module(
11+
_pyvelox
12+
MODULE
13+
NO_EXTRAS # TODO: LTO crashes GCC9.2. File issues to pybind11
14+
pyvelox.cpp
15+
pyvelox.h
16+
)
17+
18+
# Link with Velox:
19+
target_link_libraries(_pyvelox PRIVATE
20+
velox_type
21+
)
22+
23+
install(
24+
TARGETS _pyvelox
25+
LIBRARY DESTINATION .
26+
)
27+
else()
28+
add_library(pyvelox pyvelox.cpp pyvelox.h)
29+
target_link_libraries(
30+
pyvelox
31+
velox_type
32+
pybind11::module)
33+
endif()
34+

‎csrc/velox/pyvelox/__init__.py

Whitespace-only changes.

‎csrc/velox/pyvelox/pyvelox.cpp

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include "pyvelox.h"
10+
11+
12+
namespace facebook::pyvelox {
13+
namespace py = pybind11;
14+
15+
template <velox::TypeKind kind>
16+
void declareType(py::module& m) {
17+
using I = typename velox::TypeTraits<kind>::ImplType;
18+
py::class_<I, velox::Type, std::shared_ptr<I>>(
19+
m,
20+
(std::string("VeloxType_") + velox::TypeTraits<kind>::name).c_str())
21+
.def(py::init());
22+
}
23+
24+
void addVeloxBindings(py::module& m) {
25+
py::enum_<velox::TypeKind>(
26+
m,
27+
"TypeKind",
28+
py::module_local())
29+
.value("BOOLEAN", velox::TypeKind::BOOLEAN)
30+
.value("TINYINT", velox::TypeKind::TINYINT)
31+
.value("SMALLINT", velox::TypeKind::SMALLINT)
32+
.value("INTEGER", velox::TypeKind::INTEGER)
33+
.value("BIGINT", velox::TypeKind::BIGINT)
34+
.value("REAL", velox::TypeKind::REAL)
35+
.value("DOUBLE", velox::TypeKind::DOUBLE)
36+
.value("VARCHAR", velox::TypeKind::VARCHAR)
37+
.value("VARBINARY", velox::TypeKind::VARBINARY)
38+
.value("TIMESTAMP", velox::TypeKind::TIMESTAMP)
39+
.value("ARRAY", velox::TypeKind::ARRAY)
40+
.value("MAP", velox::TypeKind::MAP)
41+
.value("ROW", velox::TypeKind::ROW)
42+
.export_values();
43+
44+
py::class_<velox::Type, std::shared_ptr<facebook::velox::Type>>(
45+
m,
46+
"VeloxType")
47+
.def("kind", &velox::Type::kind)
48+
.def("kind_name", &velox::Type::kindName);
49+
50+
declareType<velox::TypeKind::BIGINT>(m);
51+
declareType<velox::TypeKind::BOOLEAN>(m);
52+
declareType<velox::TypeKind::TINYINT>(m);
53+
declareType<velox::TypeKind::SMALLINT>(m);
54+
declareType<velox::TypeKind::INTEGER>(m);
55+
declareType<velox::TypeKind::REAL>(m);
56+
declareType<velox::TypeKind::DOUBLE>(m);
57+
declareType<velox::TypeKind::VARCHAR>(m);
58+
declareType<velox::TypeKind::VARBINARY>(m);
59+
declareType<velox::TypeKind::TIMESTAMP>(m);
60+
61+
using I = typename velox::TypeTraits<velox::TypeKind::ARRAY>::ImplType;
62+
py::class_<I, velox::Type, std::shared_ptr<I>>(
63+
m,
64+
"VeloxArrayType")
65+
.def(py::init<velox::TypePtr>())
66+
.def("element_type", &velox::ArrayType::elementType);
67+
68+
using J = typename velox::FixedSizeArrayType;
69+
py::class_<J, velox::Type, std::shared_ptr<J>>(
70+
m, "VeloxFixedArrayType")
71+
.def(py::init<int, velox::TypePtr>())
72+
.def("element_type", &velox::FixedSizeArrayType::elementType)
73+
.def("fixed_width", &velox::FixedSizeArrayType::fixedElementsWidth);
74+
75+
using M = typename velox::TypeTraits<velox::TypeKind::MAP>::ImplType;
76+
py::class_<M, velox::Type, std::shared_ptr<M>>(
77+
m,
78+
"VeloxMapType")
79+
.def(py::init<velox::TypePtr, velox::TypePtr>())
80+
.def("key_type", &velox::MapType::keyType)
81+
.def("value_type", &velox::MapType::valueType);
82+
83+
using R = typename velox::TypeTraits<velox::TypeKind::ROW>::ImplType;
84+
85+
py::class_<R, velox::Type, std::shared_ptr<R>>(
86+
m,
87+
"VeloxRowType")
88+
.def(py::init<
89+
std::vector<std::string>&&,
90+
std::vector<std::shared_ptr<const velox::Type>>&&>())
91+
.def("size", &R::size)
92+
.def("get_child_idx", &R::getChildIdx)
93+
.def("contains_child", &R::containsChild)
94+
.def("name_of", &R::nameOf)
95+
.def("child_at", &R::childAt);
96+
}
97+
98+
#ifdef CREATE_PYVELOX_MODULE
99+
PYBIND11_MODULE(_pyvelox, m) {
100+
m.doc() = R"pbdoc(
101+
PyVelox native code module
102+
-----------------------
103+
)pbdoc";
104+
105+
addVeloxBindings(m);
106+
107+
m.attr("__version__") = "dev";
108+
}
109+
#endif
110+
}
111+

‎csrc/velox/pyvelox/pyvelox.h

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <pybind11/pybind11.h>
12+
#include <pybind11/stl.h>
13+
#include <pybind11/stl_bind.h>
14+
#include <velox/type/Type.h>
15+
16+
namespace facebook::pyvelox {
17+
18+
void addVeloxBindings(pybind11::module& m);
19+
}

0 commit comments

Comments
 (0)