From 090ea66357b246dd2f3e0b60127ba972ecf30334 Mon Sep 17 00:00:00 2001
From: Marius van Niekerk <marius.v.niekerk@gmail.com>
Date: Wed, 23 Oct 2024 22:35:05 -0400
Subject: [PATCH] Add missing partitioning typing case (#132)

This should now support the examples in the docstring for partitioning.
---
 pyarrow-stubs/_dataset.pyi | 6 ++++++
 pyarrow-stubs/dataset.pyi  | 4 ++++
 2 files changed, 10 insertions(+)

diff --git a/pyarrow-stubs/_dataset.pyi b/pyarrow-stubs/_dataset.pyi
index f7e801d..21684e6 100644
--- a/pyarrow-stubs/_dataset.pyi
+++ b/pyarrow-stubs/_dataset.pyi
@@ -351,6 +351,12 @@ class DirectoryPartitioning(KeyValuePartitioning):
         schema: lib.Schema | None = None,
         segment_encoding: Literal["uri", "none"] = "uri",
     ) -> PartitioningFactory: ...
+    def __init__(
+        self,
+        schema: lib.Schema,
+        dictionaries: dict[str, lib.Array] | None = None,
+        segment_encoding: Literal["uri", "none"] = "uri",
+    ) -> None: ...
 
 class HivePartitioning(KeyValuePartitioning):
     def __init__(
diff --git a/pyarrow-stubs/dataset.pyi b/pyarrow-stubs/dataset.pyi
index 3a01978..e9da4ec 100644
--- a/pyarrow-stubs/dataset.pyi
+++ b/pyarrow-stubs/dataset.pyi
@@ -111,6 +111,10 @@ __all__ = [
 
 _DatasetFormat: TypeAlias = Literal["parquet", "ipc", "arrow", "feather", "csv"]
 
+@overload
+def partitioning(
+    schema: Schema,
+) -> Partitioning: ...
 @overload
 def partitioning(
     schema: Schema,