From 589d657fbfcd82555e3349f4dada7dbbb0611b67 Mon Sep 17 00:00:00 2001
From: Pedro Eugenio Rocha Pedreira <pedroerp@meta.com>
Date: Wed, 5 Feb 2025 16:14:30 -0800
Subject: [PATCH] feat(python): Add missing .pyi symbols (#12269)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/12269

Adding missing symbols required for Python static type checks

Reviewed By: kgpai

Differential Revision: D69196070

fbshipit-source-id: 5981c5cb63353a81371e59471f2f0cd2d339ffa7
---
 velox/py/plan_builder/plan_builder.pyi | 15 ++++++++++++++-
 velox/py/runner/runner.pyi             |  5 +++--
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/velox/py/plan_builder/plan_builder.pyi b/velox/py/plan_builder/plan_builder.pyi
index 41e44d36da18..be9c06bc1f30 100644
--- a/velox/py/plan_builder/plan_builder.pyi
+++ b/velox/py/plan_builder/plan_builder.pyi
@@ -16,11 +16,21 @@
 
 # pyre-unsafe
 
+from enum import Enum
 from typing import List, Dict, Type
 
 # pyre-fixme[21]: Could not find `velox.py.type`.
 from velox.py.type import Type
 
+# pyre-fixme[21]: Could not find `velox.py.type`.
+from velox.py.file import File
+
+
+class JoinType(Enum):
+    INNER = 1
+    LEFT = 2
+    RIGHT = 3
+    FULL = 4
 
 class PlanNode: ...
 class PlanBuilder:
@@ -28,10 +38,13 @@ class PlanBuilder:
     def table_scan(
         self,
         # pyre-fixme[11]: Annotation `Type` is not defined as a type.
-        output: Type,
+        output_schema: Type,
         aliases: Dict[str, str] = {},
         subfields: Dict[str, List[int]] = {},
         row_index: str = "",
+        connector_id: str = "prism",
+        # pyre-fixme[11]: Annotation `File` is not defined as a type.
+        input_files: List[File] = [],
     ) -> PlanBuilder: ...
     def get_plan_node(self) -> PlanBuilder: ...
     def new_builder(self) -> PlanBuilder: ...
diff --git a/velox/py/runner/runner.pyi b/velox/py/runner/runner.pyi
index f1e289a235b7..12d700f2a7b8 100644
--- a/velox/py/runner/runner.pyi
+++ b/velox/py/runner/runner.pyi
@@ -18,11 +18,12 @@
 
 from typing import Iterator
 
+# pyre-fixme[21]: Could not find `velox.py.type`.
+from velox.py.cector import Vector
 
-class Vector:
-    def size(self) -> int: ...
 
 class LocalRunner:
     def __init__(self, PlanNode) -> None: ...
+    # pyre-fixme[11]: Annotation `Vector` is not defined as a type.
     def execute(self) -> Iterator[Vector]: ...
     def add_file_split(self, plan_id: str, file_path: str) -> None: ...