diff --git a/python/setup.py b/python/setup.py
index 2282b2535e73f..6642ec5042fa1 100755
--- a/python/setup.py
+++ b/python/setup.py
@@ -238,6 +238,7 @@ def run(self):
             "pyspark.mllib.linalg",
             "pyspark.mllib.stat",
             "pyspark.ml",
+            "pyspark.ml.connect",
             "pyspark.ml.linalg",
             "pyspark.ml.param",
             "pyspark.ml.torch",
@@ -245,13 +246,16 @@ def run(self):
             "pyspark.sql",
             "pyspark.sql.avro",
             "pyspark.sql.connect",
+            "pyspark.sql.connect.avro",
+            "pyspark.sql.connect.client",
             "pyspark.sql.connect.proto",
+            "pyspark.sql.connect.streaming",
+            "pyspark.sql.connect.streaming.worker",
             "pyspark.sql.pandas",
             "pyspark.sql.protobuf",
             "pyspark.sql.streaming",
             "pyspark.sql.worker",
             "pyspark.streaming",
-            "pyspark.sql.connect.streaming.worker",
             "pyspark.bin",
             "pyspark.sbin",
             "pyspark.jars",