Skip to content

Commit

Permalink
feat(python): New PyVelox bindings for Types (facebookincubator#12040)
Browse files Browse the repository at this point in the history
Summary:

Adding new PyVelox bindings for Types. The purpose is not to provide
the full functionality of the underlying C++ structures, but to provide the
minimal API required to enable Python users to execute query plans. The
subsequent PRs/diffs will contain the remaining binding libraries.

Differential Revision: D67957151
  • Loading branch information
pedroerp authored and facebook-github-bot committed Jan 9, 2025
1 parent ff943fd commit b5aa0d9
Show file tree
Hide file tree
Showing 5 changed files with 291 additions and 0 deletions.
70 changes: 70 additions & 0 deletions velox/py/tests/test_type.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# 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.

import unittest
from velox.py.type import (
Type,
BIGINT,
INTEGER,
SMALLINT,
TINYINT,
BOOLEAN,
REAL,
DOUBLE,
VARCHAR,
VARBINARY,
ARRAY,
MAP,
ROW,
)


class TestPyVeloxTypes(unittest.TestCase):
def test_simple_types(self):
self.assertTrue(isinstance(BIGINT(), Type))
self.assertTrue(isinstance(INTEGER(), Type))
self.assertTrue(isinstance(SMALLINT(), Type))
self.assertTrue(isinstance(TINYINT(), Type))
self.assertTrue(isinstance(BOOLEAN(), Type))
self.assertTrue(isinstance(REAL(), Type))
self.assertTrue(isinstance(DOUBLE(), Type))
self.assertTrue(isinstance(VARCHAR(), Type))
self.assertTrue(isinstance(VARBINARY(), Type))

def test_complex_types(self):
self.assertTrue(isinstance(ARRAY(VARCHAR()), Type))
self.assertTrue(isinstance(MAP(BIGINT(), VARCHAR()), Type))
self.assertTrue(isinstance(ROW(["c0", "c1"], [INTEGER(), ARRAY(VARCHAR())]), Type))
self.assertTrue(isinstance(ROW(), Type))

# Invalid complex types.
self.assertRaises(TypeError, ARRAY)
self.assertRaises(TypeError, MAP)
self.assertRaises(TypeError, MAP, BIGINT())
self.assertRaises(RuntimeError, ROW, ["col1"])

def test_to_str(self):
self.assertEqual("BOOLEAN", str(BOOLEAN()))
self.assertEqual("VARBINARY", str(VARBINARY()))

self.assertEqual("ARRAY<DOUBLE>", str(ARRAY(DOUBLE())))
self.assertEqual("MAP<VARCHAR,DOUBLE>", str(MAP(VARCHAR(), DOUBLE())))
self.assertEqual("ROW<c0:INTEGER,c1:TINYINT>", str(ROW(["c0", "c1"], [INTEGER(), TINYINT()])))

def test_equality(self):
self.assertEqual(INTEGER(), INTEGER())
self.assertEqual(ROW(["a"], [INTEGER()]), ROW(["a"], [INTEGER()]))

self.assertNotEqual(BIGINT(), INTEGER())
self.assertNotEqual(ARRAY(BIGINT()), REAL())
24 changes: 24 additions & 0 deletions velox/py/type/PyType.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
/*
* Copyright (c) Facebook, Inc. and its affiliates.
*
* 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.
*/

#include "velox/py/type/PyType.h"
#include <pybind11/stl.h>

namespace facebook::velox::py {

namespace py = pybind11;

} // namespace facebook::velox::py
107 changes: 107 additions & 0 deletions velox/py/type/PyType.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
/*
* Copyright (c) Facebook, Inc. and its affiliates.
*
* 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 <pybind11/embed.h>
#include "velox/type/Type.h"

namespace facebook::velox::py {

class PyType {
public:
explicit PyType(const TypePtr& type = nullptr) : type_(type) {}

std::string toString() const {
if (!type_) {
return "[nullptr]";
}
return type_->toString();
}

TypePtr type() const {
return type_;
}

bool equivalent(const PyType& other) const {
return type_->equivalent(*other.type());
}

// Factory functions:

static PyType createRowType(
const std::vector<std::string>& names,
const std::vector<PyType>& pyTypes) {
std::vector<TypePtr> types;
for (const auto& pyType : pyTypes) {
types.emplace_back(pyType.type());
}

if (names.empty()) {
return PyType{ROW(std::move(types))};
}
return PyType{ROW(folly::copy(names), std::move(types))};
}

static PyType createMapType(const PyType& keyType, const PyType& valueType) {
return PyType{MAP(keyType.type(), valueType.type())};
}

static PyType createArrayType(const PyType& elementsType) {
return PyType{ARRAY(elementsType.type())};
}

static PyType createBigint() {
return PyType{BIGINT()};
}

static PyType createInteger() {
return PyType{INTEGER()};
}

static PyType createSmallint() {
return PyType{SMALLINT()};
}

static PyType createTinyint() {
return PyType{TINYINT()};
}

static PyType createBoolean() {
return PyType{BOOLEAN()};
}

static PyType createReal() {
return PyType{REAL()};
}

static PyType createDouble() {
return PyType{DOUBLE()};
}

static PyType createVarchar() {
return PyType{VARCHAR()};
}

static PyType createVarbinary() {
return PyType{VARBINARY()};
}

private:
TypePtr type_;
};

} // namespace facebook::velox::py
56 changes: 56 additions & 0 deletions velox/py/type/type.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
/*
* Copyright (c) Facebook, Inc. and its affiliates.
*
* 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.
*/

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "velox/py/type/PyType.h"

namespace py = pybind11;

PYBIND11_MODULE(type, m) {
using namespace facebook;

py::class_<velox::py::PyType>(m, "Type")
.def("__str__", &velox::py::PyType::toString)
.def("__eq__", [](velox::py::PyType& u, velox::py::PyType& v) {
return u.equivalent(v);
});

m.def("BIGINT", &velox::py::PyType::createBigint);
m.def("INTEGER", &velox::py::PyType::createInteger);
m.def("SMALLINT", &velox::py::PyType::createSmallint);
m.def("TINYINT", &velox::py::PyType::createTinyint);
m.def("BOOLEAN", &velox::py::PyType::createBoolean);

m.def("REAL", &velox::py::PyType::createReal);
m.def("DOUBLE", &velox::py::PyType::createDouble);

m.def("VARCHAR", &velox::py::PyType::createVarchar);
m.def("VARBINARY", &velox::py::PyType::createVarbinary);

m.def("ARRAY", &velox::py::PyType::createArrayType, py::arg("elements_type"));
m.def(
"MAP",
&velox::py::PyType::createMapType,
py::arg("key_type"),
py::arg("value_type"));
m.def(
"ROW",
&velox::py::PyType::createRowType,
py::arg("names") = std::vector<std::string>{},
py::arg("types") = std::vector<velox::py::PyType>{});
}
34 changes: 34 additions & 0 deletions velox/py/type/type.pyi
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
#
# 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.

# pyre-unsafe

from typing import List


class Type: ...

def BIGINT() -> Type: ...
def INTEGER() -> Type: ...
def SMALLINT() -> Type: ...
def TINYINT() -> Type: ...
def BOOLEAN() -> Type: ...
def REAL() -> Type: ...
def DOUBLE() -> Type: ...
def VARCHAR() -> Type: ...
def ARRAY(elements_type: Type) -> Type: ...
def MAP(key_type: Type, value_type: Type) -> Type: ...
def ROW(names: List[str] = [], types: List[Type] = []) -> Type: ...

0 comments on commit b5aa0d9

Please sign in to comment.