Enable users to pass in Pandas Dataframe when calling import_data() and batch_predict() from AutoML Tables client (googleapis#9116)

Author: TrucHLe

automl/google/cloud/automl_v1beta1/__init__.py

@@ -21,6 +21,7 @@
 from google.cloud.automl_v1beta1.gapic import auto_ml_client
 from google.cloud.automl_v1beta1.gapic import enums
 from google.cloud.automl_v1beta1.gapic import prediction_service_client
+from google.cloud.automl_v1beta1.tables import gcs_client
 from google.cloud.automl_v1beta1.tables import tables_client
 
 
@@ -38,4 +39,15 @@ class PredictionServiceClient(prediction_service_client.PredictionServiceClient)
  enums = enums
 
 
-__all__ = ("enums", "types", "AutoMlClient", "PredictionServiceClient", "TablesClient")
+class GcsClient(gcs_client.GcsClient):
+ __doc__ = gcs_client.GcsClient.__doc__
+
+
+__all__ = (
+ "enums",
+ "types",
+ "AutoMlClient",
+ "PredictionServiceClient",
+ "TablesClient",
+ "GcsClient",
+)

automl/google/cloud/automl_v1beta1/tables/gcs_client.py

@@ -0,0 +1,112 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Wraps the Google Cloud Storage client library for use in tables helper."""
+
+import time
+
+from google.api_core import exceptions
+
+try:
+ import pandas
+except ImportError:  # pragma: NO COVER
+ pandas = None
+
+try:
+ from google.cloud import storage
+except ImportError:  # pragma: NO COVER
+ storage = None
+
+_PANDAS_REQUIRED = "pandas is required to verify type DataFrame."
+_STORAGE_REQUIRED = (
+ "google-cloud-storage is required to create Google Cloud Storage client."
+)
+
+
+class GcsClient(object):
+ """Uploads Pandas DataFrame to a bucket in Google Cloud Storage."""
+
+ def __init__(self, client=None, credentials=None):
+ """Constructor.
+
+        Args:
+            client (Optional[storage.Client]): A Google Cloud Storage Client
+                instance.
+            credentials (Optional[google.auth.credentials.Credentials]): The
+                authorization credentials to attach to requests. These
+                credentials identify this application to the service. If none
+                are specified, the client will attempt to ascertain the
+                credentials from the environment.
+        """
+ if storage is None:
+ raise ImportError(_STORAGE_REQUIRED)
+
+ if client is not None:
+ self.client = client
+ elif credentials is not None:
+ self.client = storage.Client(credentials=credentials)
+ else:
+ self.client = storage.Client()
+
+ def ensure_bucket_exists(self, project, region):
+ """Checks if a bucket named '{project}-automl-tables-staging' exists.
+
+        Creates this bucket if it doesn't exist.
+
+        Args:
+            project (str): The project that stores the bucket.
+            region (str): The region of the bucket.
+
+        Returns:
+            A string representing the created bucket name.
+        """
+ bucket_name = "{}-automl-tables-staging".format(project)
+
+ try:
+ self.client.get_bucket(bucket_name)
+ except exceptions.NotFound:
+ bucket = self.client.bucket(bucket_name)
+ bucket.create(project=project, location=region)
+ return bucket_name
+
+ def upload_pandas_dataframe(self, bucket_name, dataframe, uploaded_csv_name=None):
+ """Uploads a Pandas DataFrame as CSV to the bucket.
+
+        Args:
+            bucket_name (str): The bucket name to upload the CSV to.
+            dataframe (pandas.DataFrame): The Pandas Dataframe to be uploaded.
+            uploaded_csv_name (Optional[str]): The name for the uploaded CSV.
+
+        Returns:
+            A string representing the GCS URI of the uploaded CSV.
+        """
+ if pandas is None:
+ raise ImportError(_PANDAS_REQUIRED)
+
+ if not isinstance(dataframe, pandas.DataFrame):
+ raise ValueError("'dataframe' must be a pandas.DataFrame instance.")
+
+ if uploaded_csv_name is None:
+ uploaded_csv_name = "automl-tables-dataframe-{}.csv".format(
+ int(time.time())
+            )
+ csv_string = dataframe.to_csv()
+
+ bucket = self.client.get_bucket(bucket_name)
+ blob = bucket.blob(uploaded_csv_name)
+ blob.upload_from_string(csv_string)
+
+ return "gs://{}/{}".format(bucket_name, uploaded_csv_name)

automl/google/cloud/automl_v1beta1/tables/tables_client.py

@@ -18,11 +18,13 @@
 
 import pkg_resources
 import logging
+import google.auth
 
 from google.api_core.gapic_v1 import client_info
 from google.api_core import exceptions
 from google.cloud.automl_v1beta1 import gapic
 from google.cloud.automl_v1beta1.proto import data_types_pb2
+from google.cloud.automl_v1beta1.tables import gcs_client
 
 _GAPIC_LIBRARY_VERSION = pkg_resources.get_distribution("google-cloud-automl").version
 _LOGGER = logging.getLogger(__name__)
@@ -43,6 +45,7 @@ def __init__(
  region="us-central1",
  client=None,
  prediction_client=None,
+ gcs_client=None,
  **kwargs
     ):
  """Constructor.
@@ -118,6 +121,7 @@ def __init__(
 
  self.project = project
  self.region = region
+ self.gcs_client = gcs_client
 
  def __lookup_by_display_name(self, object_type, items, display_name):
  relevant_items = [i for i in items if i.display_name == display_name]
@@ -403,6 +407,21 @@ def __type_code_to_value_type(self, type_code, value):
  else:
  raise ValueError("Unknown type_code: {}".format(type_code))
 
+ def __ensure_gcs_client_is_initialized(self, credentials=None):
+ """Checks if GCS client is initialized. Initializes it if not.
+
+        Args:
+            credentials (google.auth.credentials.Credentials): The
+                authorization credentials to attach to requests. These
+                credentials identify this application to the service. If none
+                are specified, the client will attempt to ascertain the
+                credentials from the environment.
+        """
+ if self.gcs_client is None:
+ if credentials is None:
+ credentials, _ = google.auth.default()
+ self.gcs_client = gcs_client.GcsClient(credentials=credentials)
+
  def list_datasets(self, project=None, region=None, **kwargs):
  """List all datasets in a particular project and region.
 
@@ -642,10 +661,12 @@ def import_data(
  dataset=None,
  dataset_display_name=None,
  dataset_name=None,
+ pandas_dataframe=None,
  gcs_input_uris=None,
  bigquery_input_uri=None,
  project=None,
  region=None,
+ credentials=None,
  **kwargs
     ):
  """Imports data into a dataset.
@@ -679,6 +700,11 @@ def import_data(
             region (Optional[string]):
                 If you have initialized the client with a value for `region` it
                 will be used if this parameter is not supplied.
+            credentials (Optional[google.auth.credentials.Credentials]): The
+                authorization credentials to attach to requests. These
+                credentials identify this application to the service. If none
+                are specified, the client will attempt to ascertain the
+                credentials from the environment.
             dataset_display_name (Optional[string]):
                 The human-readable name given to the dataset you want to import
                 data into. This must be supplied if `dataset` or `dataset_name`
@@ -691,13 +717,21 @@ def import_data(
                 The `Dataset` instance you want to import data into. This must
                 be supplied if `dataset_display_name` or `dataset_name` are not
                 supplied.
+            pandas_dataframe (Optional[pandas.DataFrame]):
+                A Pandas Dataframe object containing the data to import. The data
+                will be converted to CSV, and this CSV will be staged to GCS in
+                `gs://{project}-automl-tables-staging/{uploaded_csv_name}`
+                This parameter must be supplied if neither `gcs_input_uris` nor
+                `bigquery_input_uri` is supplied.
             gcs_input_uris (Optional[Union[string, Sequence[string]]]):
                 Either a single `gs://..` prefixed URI, or a list of URIs
                 referring to GCS-hosted CSV files containing the data to
-                import. This must be supplied if `bigquery_input_uri` is not.
+                import. This must be supplied if neither `bigquery_input_uri`
+                nor `pandas_dataframe` is supplied.
             bigquery_input_uri (Optional[string]):
                 A URI pointing to the BigQuery table containing the data to
-                import. This must be supplied if `gcs_input_uris` is not.
+                import. This must be supplied if neither `gcs_input_uris` nor
+                `pandas_dataframe` is supplied.
 
         Returns:
             A :class:`~google.cloud.automl_v1beta1.types._OperationFuture`
@@ -720,15 +754,23 @@ def import_data(
         )
 
  request = {}
- if gcs_input_uris is not None:
+
+ if pandas_dataframe is not None:
+ self.__ensure_gcs_client_is_initialized(credentials)
+ bucket_name = self.gcs_client.ensure_bucket_exists(project, region)
+ gcs_input_uri = self.gcs_client.upload_pandas_dataframe(
+ bucket_name, pandas_dataframe
+            )
+ request = {"gcs_source": {"input_uris": [gcs_input_uri]}}
+ elif gcs_input_uris is not None:
  if type(gcs_input_uris) != list:
  gcs_input_uris = [gcs_input_uris]
  request = {"gcs_source": {"input_uris": gcs_input_uris}}
  elif bigquery_input_uri is not None:
  request = {"bigquery_source": {"input_uri": bigquery_input_uri}}
  else:
  raise ValueError(
- "One of 'gcs_input_uris', or " "'bigquery_input_uri' must be set."
+ "One of 'gcs_input_uris', or 'bigquery_input_uri', or 'pandas_dataframe' must be set."
             )
 
  op = self.auto_ml_client.import_data(dataset_name, request, **kwargs)
@@ -2605,6 +2647,7 @@ def predict(
 
  def batch_predict(
  self,
+ pandas_dataframe=None,
  bigquery_input_uri=None,
  bigquery_output_uri=None,
  gcs_input_uris=None,
@@ -2614,6 +2657,7 @@ def batch_predict(
  model_display_name=None,
  project=None,
  region=None,
+ credentials=None,
  inputs=None,
  **kwargs
     ):
@@ -2645,15 +2689,30 @@ def batch_predict(
             region (Optional[string]):
                 If you have initialized the client with a value for `region` it
                 will be used if this parameter is not supplied.
+            credentials (Optional[google.auth.credentials.Credentials]): The
+                authorization credentials to attach to requests. These
+                credentials identify this application to the service. If none
+                are specified, the client will attempt to ascertain the
+                credentials from the environment.
+            pandas_dataframe (Optional[pandas.DataFrame]):
+                A Pandas Dataframe object containing the data you want to predict
+                off of. The data will be converted to CSV, and this CSV will be
+                staged to GCS in `gs://{project}-automl-tables-staging/{uploaded_csv_name}`
+                This must be supplied if neither `gcs_input_uris` nor
+                `bigquery_input_uri` is supplied.
             gcs_input_uris (Optional(Union[List[string], string]))
                 Either a list of or a single GCS URI containing the data you
-                want to predict off of.
+                want to predict off of. This must be supplied if neither
+                `pandas_dataframe` nor `bigquery_input_uri` is supplied.
             gcs_output_uri_prefix (Optional[string])
-                The folder in GCS you want to write output to.
+                The folder in GCS you want to write output to. This must be
+                supplied if `bigquery_output_uri` is not.
             bigquery_input_uri (Optional[string])
-                The BigQuery table to input data from.
+                The BigQuery table to input data from. This must be supplied if
+                neither `pandas_dataframe` nor `gcs_input_uris` is supplied.
             bigquery_output_uri (Optional[string])
-                The BigQuery table to output data to.
+                The BigQuery table to output data to. This must be supplied if
+                `gcs_output_uri_prefix` is not.
             model_display_name (Optional[string]):
                 The human-readable name given to the model you want to predict
                 with.  This must be supplied if `model` or `model_name` are not
@@ -2688,7 +2747,15 @@ def batch_predict(
         )
 
  input_request = None
- if gcs_input_uris is not None:
+
+ if pandas_dataframe is not None:
+ self.__ensure_gcs_client_is_initialized(credentials)
+ bucket_name = self.gcs_client.ensure_bucket_exists(project, region)
+ gcs_input_uri = self.gcs_client.upload_pandas_dataframe(
+ bucket_name, pandas_dataframe
+            )
+ input_request = {"gcs_source": {"input_uris": [gcs_input_uri]}}
+ elif gcs_input_uris is not None:
  if type(gcs_input_uris) != list:
  gcs_input_uris = [gcs_input_uris]
  input_request = {"gcs_source": {"input_uris": gcs_input_uris}}

automl/noxfile.py

@@ -70,6 +70,7 @@ def default(session):
  for local_dep in LOCAL_DEPS:
  session.install("-e", local_dep)
  session.install("-e", ".")
+ session.install("-e", ".[pandas,storage]")
 
  # Run py.test against the unit tests.
  session.run(
@@ -117,6 +118,7 @@ def system(session):
  session.install("-e", local_dep)
  session.install("-e", "../test_utils/")
  session.install("-e", ".")
+ session.install("-e", ".[pandas,storage]")
 
  # Run py.test against the system tests.
  if system_test_exists:

automl/setup.py

@@ -25,6 +25,10 @@
  "google-api-core[grpc] >= 1.14.0, < 2.0.0dev",
  'enum34; python_version < "3.4"',
 ]
+extras = {
+ "pandas": ["pandas>=0.24.0"],
+ "storage": ["google-cloud-storage >= 1.18.0, < 2.0.0dev"],
+}
 
 package_root = os.path.abspath(os.path.dirname(__file__))
 
@@ -67,6 +71,7 @@
  packages=packages,
  namespace_packages=namespaces,
  install_requires=dependencies,
+ extras_require=extras,
  python_requires=">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*",
  include_package_data=True,
  zip_safe=False,