@@ -9,7 +9,8 @@ COPY protos protos
 COPY README.md README.md
 
 # Install Feast for AWS with Lambda dependencies
-RUN pip3 install -e 'sdk/python[aws,redis]'
+# TODO(achals): The additional spark deps should be removed. Details at https://github.com/feast-dev/feast/pull/2480.
+RUN pip3 install -e 'sdk/python[aws,redis,spark]'
 RUN pip3 install -r sdk/python/feast/infra/feature_servers/aws_lambda/requirements.txt --target "${LAMBDA_TASK_ROOT}"
 
 # Set the CMD to your handler (could also be done as a parameter override outside of the Dockerfile)
 
@@ -338,12 +338,14 @@ def to_bigquery(
 
     def _to_arrow_internal(self) -> pyarrow.Table:
         with self._query_generator() as query:
-            return self._execute_query(query).to_arrow()
+            q = self._execute_query(query=query)
+            assert q
+            return q.to_arrow()
 
     @log_exceptions_and_usage
     def _execute_query(
         self, query, job_config=None, timeout: int = 1800
-    ) -> bigquery.job.query.QueryJob:
+    ) -> Optional[bigquery.job.query.QueryJob]:
         bq_job = self.client.query(query, job_config=job_config)
 
         if job_config and job_config.dry_run:
@@ -426,7 +428,7 @@ def _get_table_reference_for_new_entity(
     dataset.location = dataset_location if dataset_location else "US"
 
     try:
-        client.get_dataset(dataset)
+        client.get_dataset(dataset.reference)
     except NotFound:
         # Only create the dataset if it does not exist
         client.create_dataset(dataset, exists_ok=True)
 
@@ -15,6 +15,8 @@
 JsonObject = Any
 
 
+# TODO: These methods need to be updated when bumping the version of protobuf.
+# https://github.com/feast-dev/feast/issues/2484
 def _patch_proto_json_encoding(
     proto_type: Type[ProtoMessage],
     to_json_object: Callable[[_Printer, ProtoMessage], JsonObject],
@@ -68,7 +70,7 @@ def to_json_object(printer: _Printer, message: ProtoMessage) -> JsonObject:
         return value
 
     def from_json_object(
-        parser: _Parser, value: JsonObject, message: ProtoMessage
+        parser: _Parser, value: JsonObject, message: ProtoMessage,
     ) -> None:
         if value is None:
             message.null_val = 0
@@ -140,7 +142,7 @@ def to_json_object(printer: _Printer, message: ProtoMessage) -> JsonObject:
         return [printer._MessageToJsonObject(item) for item in message.val]
 
     def from_json_object(
-        parser: _Parser, value: JsonObject, message: ProtoMessage
+        parser: _Parser, value: JsonObject, message: ProtoMessage,
     ) -> None:
         array = value if isinstance(value, list) else value["val"]
         for item in array:
@@ -181,7 +183,7 @@ def to_json_object(printer: _Printer, message: ProtoMessage) -> JsonObject:
         return list(message.val)
 
     def from_json_object(
-        parser: _Parser, value: JsonObject, message: ProtoMessage
+        parser: _Parser, value: JsonObject, message: ProtoMessage,
     ) -> None:
         array = value if isinstance(value, list) else value["val"]
         message.val.extend(array)
 
@@ -400,7 +400,7 @@ mypy==0.931
     # via feast (setup.py)
 mypy-extensions==0.4.3
     # via mypy
-mypy-protobuf==3.1.0
+mypy-protobuf==3.1
     # via feast (setup.py)
 nbclient==0.5.11
     # via nbconvert
@@ -850,4 +850,4 @@ zipp==3.7.0
 
 # The following packages are considered to be unsafe in a requirements file:
 # pip
-# setuptools
+# setuptools
@@ -53,7 +53,7 @@
     "mmh3",
     "pandas>=1.0.0",
     "pandavro==1.5.*",
-    "protobuf>=3.10",
+    "protobuf>=3.10,<3.20",
     "proto-plus<1.19.7",
     "pyarrow>=4.0.0",
     "pydantic>=1.0.0",
@@ -112,7 +112,7 @@
         "mock==2.0.0",
         "moto",
         "mypy==0.931",
-        "mypy-protobuf==3.1.0",
+        "mypy-protobuf==3.1",
         "avro==1.10.0",
         "gcsfs",
         "urllib3>=1.25.4",
@@ -149,7 +149,7 @@
         + GE_REQUIRED
 )
 
-DEV_REQUIRED = ["mypy-protobuf>=3.1.0", "grpcio-testing==1.*"] + CI_REQUIRED
+DEV_REQUIRED = ["mypy-protobuf==3.1", "grpcio-testing==1.*"] + CI_REQUIRED
 
 # Get git repo root directory
 repo_root = str(pathlib.Path(__file__).resolve().parent.parent.parent)
@@ -264,7 +264,13 @@ def run(self):
     ],
     entry_points={"console_scripts": ["feast=feast.cli:cli"]},
     use_scm_version=use_scm_version,
-    setup_requires=["setuptools_scm", "grpcio", "grpcio-tools==1.34.0", "mypy-protobuf==3.1.0", "sphinx!=4.0.0"],
+    setup_requires=[
+        "setuptools_scm",
+        "grpcio",
+        "grpcio-tools==1.34.0",
+        "mypy-protobuf==3.1",
+        "sphinx!=4.0.0",
+    ],
     package_data={
         "": [
             "protos/feast/**/*.proto",