feat(python): New PyVelox bindings for Types

velox/py/tests/test_type.py

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

velox/py/type/PyType.cpp

@@ -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

velox/py/type/PyType.h

@@ -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

velox/py/type/type.cpp

@@ -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>{});
+}

velox/py/type/type.pyi

@@ -0,0 +1,35 @@
+#!/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 VARBINARY() -> 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: ...