fix: Add spark to lambda dockerfile (feast-dev#2480)

achals · felixwang9817 · web-flow · commit 514666f31cf6 · 2022-04-04T13:01:57.000-07:00
* add spark to lambda dockerfile

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* add *args

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Add *args

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Pin protobuf==3.19.4

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Remove *args

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Add a range

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Add a todo

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* cleanup prints

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* lock deps

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* lock deps correctly

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* fix lint

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* fix lint take 2

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Undo general updates

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

Co-authored-by: Felix Wang &lt;wangfelix98@gmail.com&gt;
diff --git a/sdk/python/feast/infra/feature_servers/aws_lambda/Dockerfile b/sdk/python/feast/infra/feature_servers/aws_lambda/Dockerfile
@@ -10,7 +10,8 @@ COPY go go
 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)
diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py
@@ -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)
diff --git a/sdk/python/feast/proto_json.py b/sdk/python/feast/proto_json.py
@@ -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)
diff --git a/sdk/python/requirements/py3.7-ci-requirements.txt b/sdk/python/requirements/py3.7-ci-requirements.txt
@@ -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
@@ -854,4 +854,4 @@ zipp==3.7.0
 
 # The following packages are considered to be unsafe in a requirements file:
 # pip
-# setuptools
+# setuptools
diff --git a/sdk/python/setup.py b/sdk/python/setup.py
@@ -52,7 +52,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",
@@ -111,7 +111,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",
@@ -150,7 +150,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)
@@ -422,7 +422,7 @@ def run(self):
         "setuptools_scm",
         "grpcio",
         "grpcio-tools==1.34.0",
-        "mypy-protobuf==3.1.0",
+        "mypy-protobuf==3.1",
         "sphinx!=4.0.0",
     ],
     package_data={
