From 82661afed9baf0cf16d49e8101a805ade5058456 Mon Sep 17 00:00:00 2001
From: Igor Dvorzhak <idv@google.com>
Date: Fri, 5 Dec 2025 00:28:35 -0800
Subject: [PATCH] fix: Switch to pyspark package as a dependency

Fixes compatibility with Colab environment
---
 requirements-dev.txt | 2 +-
 setup.py             | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/requirements-dev.txt b/requirements-dev.txt
index 72b14e49..5cf7026e 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -4,7 +4,7 @@ ipython~=9.1
 ipywidgets>=8.0.0
 packaging>=20.0
 pyink~=24.0
-pyspark-client~=4.0.0
+pyspark[connect]~=4.0.0
 setuptools>=72.0
 sparksql-magic>=0.0.3
 tqdm>=4.67
diff --git a/setup.py b/setup.py
index 655e63db..2e5c016c 100644
--- a/setup.py
+++ b/setup.py
@@ -31,7 +31,7 @@
         "google-api-core>=2.19",
         "google-cloud-dataproc>=5.18",
         "packaging>=20.0",
-        "pyspark-client~=4.0.0",
+        "pyspark[connect]~=4.0.0",
         "tqdm>=4.67",
         "websockets>=14.0",
     ],