pax_global_header00006660000000000000000000000064127234025100014506gustar00rootroot0000000000000052 comment=4d230cdff0fedc121d5d5fcd97a863d3bfcfb72e DataflowPythonSDK-0.2.5/000077500000000000000000000000001272340251000150175ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/.gitignore000066400000000000000000000001161272340251000170050ustar00rootroot00000000000000*~ *.py[cod] *.egg-info/ .eggs/ build/ dist/ distribute-* # Test files .tox/ DataflowPythonSDK-0.2.5/.travis.yml000066400000000000000000000001501272340251000171240ustar00rootroot00000000000000language: python sudo: false env: - TOX_ENV=py27 install: - pip install tox script: tox -e $TOX_ENV DataflowPythonSDK-0.2.5/CONTRIBUTING.md000066400000000000000000000025701272340251000172540ustar00rootroot00000000000000Want to contribute? Great! First, read this page (including the small print at the end). ### Before you contribute Before we can use your code, you must sign the [Google Individual Contributor License Agreement](https://developers.google.com/open-source/cla/individual?csw=1) (CLA), which you can do online. The CLA is necessary mainly because you own the copyright to your changes, even after your contribution becomes part of our codebase, so we need your permission to use and distribute your code. We also need to be sure of various other things. For instance that you'll tell us if you know that your code infringes on other people's patents. You don't have to sign the CLA until after you've submitted your code for review and a member has approved it, but you must do it before we can put your code into our codebase. Before you start working on a larger contribution, we recommend to get in touch with us first through the issue tracker with your idea so that we can help out and possibly guide you. Coordinating up front makes it much easier to avoid frustration later on. ### Code reviews All submissions, including submissions by project members, require review. We use Github pull requests for this purpose. ### The small print Contributions made by corporations are covered by a different agreement from the one above, the Software Grant and Corporate Contributor License Agreement. DataflowPythonSDK-0.2.5/LICENSE000066400000000000000000000261361272340251000160340ustar00rootroot00000000000000 Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] 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 http://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. DataflowPythonSDK-0.2.5/README.md000066400000000000000000000414261272340251000163050ustar00rootroot00000000000000# Cloud Dataflow SDK for Python [Google Cloud Dataflow](https://cloud.google.com/dataflow/) provides a simple, powerful programming model for building both batch and streaming parallel data processing pipelines. The Dataflow SDK for Python provides access to Dataflow capabilities from the Python programming language. ## Table of Contents * [Status of this Release](#status-of-this-release) * [Signing up for Alpha Batch Cloud Execution](#signing-up-for-alpha-batch-cloud-execution) * [Overview of Dataflow Programming](#overview-of-dataflow-programming) * [Getting Started](#getting-started) * [Setting up an environment](#setting-up-an-environment) * [Install ``pip``](#install-pip) * [Install ``virtualenv``](#install-virtualenv) * [Install ``setuptools``](#install-setuptools) * [Getting the Dataflow software](#getting-the-dataflow-software) * [Create and activate virtual environment](#create-and-activate-virtual-environment) * [Download and install](#download-and-install) * [Notes on installing with ``setup.py install``](#notes-on-installing-with-setuppy-install) * [Local execution of a pipeline](#local-execution-of-a-pipeline) * [A Quick Tour of the Source Code](#a-quick-tour-of-the-source-code) * [Some Simple Examples](#some-simple-examples) * [Hello world](#hello-world) * [Hello world (with Map)](#hello-world-with-map) * [Hello world (with FlatMap)](#hello-world-with-flatmap) * [Hello world (with FlatMap and yield)](#hello-world-with-flatmap-and-yield) * [Counting words](#counting-words) * [Counting words with GroupByKey](#counting-words-with-groupbykey) * [Type hints](#type-hints) * [BigQuery](#bigquery) * [Combiner Examples](#combiner-examples) * [More Examples](#more-examples) * [Organizing Your Code](#organizing-your-code) * [Contact Us](#contact-us) ## Status of this Release This is a version of Google Cloud Dataflow SDK for Python that is still early in its development, and significant changes should be expected before the first stable version. Google recently [announced its intention](http://googlecloudplatform.blogspot.com/2016/01/Dataflow-and-open-source-proposal-to-join-the-Apache-Incubator.html) to donate the Google Cloud Dataflow SDKs and Programming Model to the Apache Software Foundation (ASF), after which they will be called the Apache Beam SDKs. The SDK for Java is actively transitioning to [Apache Beam](http://beam.incubator.apache.org/), an ASF incubator project. The SDK for Python will be added to Apache Beam soon after. Expect many renames. ## Signing up for Alpha Batch Cloud Execution Google Cloud Dataflow now provides Alpha support for Batch pipelines written with the SDK for Python. This Alpha program is designed to give customers access to the service for early testing. Customers are advised not to use this feature in production systems. If you are interested in being considered to participate in the Alpha program, please submit this [form](http://goo.gl/forms/o4w14whz9x). Note that filling the form does not guarantee entry to the Alpha program. ## Overview of Dataflow Programming For an introduction to the programming model, please read [Dataflow Programming Model](https://cloud.google.com/dataflow/model/programming-model) but note that some examples on that site use only Java. The key concepts of the programming model are * [`PCollection`](https://cloud.google.com/dataflow/model/pcollection): represents a collection of data, which could be bounded or unbounded in size. * [`PTransform`](https://cloud.google.com/dataflow/model/transforms): represents a computation that transforms input PCollections into output PCollections. * [`Pipeline`](https://cloud.google.com/dataflow/model/pipelines): manages a directed acyclic graph of PTransforms and PCollections that is ready for execution. * `Runner`: specifies where and how the Pipeline should execute. This release has some significant limitations: * We provide only one PipelineRunner, the `DirectPipelineRunner`. * The Google Cloud Dataflow service does not yet accept jobs from this SDK. * Triggers are not supported. * The SDK works only on Python 2.7. ## Getting Started ### Setting up an environment If this is the first time you are installing the Dataflow SDK, you may need to set up your machine's Python development environment. #### Install ``pip`` `pip` is Python's package manager. If you already have `pip` installed (type `pip -V` to check), skip this step. There are several ways to install `pip`; use whichever works for you. Preferred option: install using your system's package manager, which may be *one* of the following commands, depending on your Linux distribution: ```sh sudo yum install python-pip sudo apt-get install python-pip sudo zypper install python-pip ``` Otherwise, if you have `easy_install` (likely if you are on MacOS): sudo easy_install pip Or you may have to install the bootstrapper. Download the following script to your system: https://bootstrap.pypa.io/get-pip.py You can fetch it with your browser or use a command-line program, such as *one* of the following: ```sh curl -O https://bootstrap.pypa.io/get-pip.py wget https://bootstrap.pypa.io/get-pip.py ``` After downloading `get-pip.py`, run it to install `pip`: ```sh python ./get-pip.py ``` #### Install ``virtualenv`` We recommend installing in a [Python virtual environment](http://docs.python-guide.org/en/latest/dev/virtualenvs/) for initial experiments. If you do not have `virtualenv` version 13.1.0 or later (type `virtualenv --version` to check), it will install a too-old version of `setuptools` in the virtual environment. To install (or upgrade) your `virtualenv`: pip install --upgrade virtualenv #### Install ``setuptools`` If you are not going to use a Python virtual environment (but we recommend you do; see the previous section), ensure `setuptools` version 17.1 or newer is installed (type `easy_install --version` to check). If you do not have that installed: pip install --upgrade setuptools ### Getting the Dataflow software #### Create and activate virtual environment A virtual environment is a directory tree containing its own Python distribution. To create a virtual environment: virtualenv /path/to/directory A virtual environment needs to be activated for each shell that is to use it; activating sets some environment variables that point to the virtual environment's directories. To activate a virtual environment in Bash: . /path/to/directory/bin/activate That is, source the script `bin/activate` under the virtual environment directory you created. #### Download and install Install the latest tarball from GitHub by browsing to and copying one of the "Source code" links. The `.tar.gz` file is smaller; we'll assume you use that one. With a virtual environment active, paste the URL into a ``pip install`` shell command, executing something like this: ```sh pip install https://github.com/GoogleCloudPlatform/DataflowPythonSDK/vX.Y.Z.tar.gz ``` #### Notes on installing with ``setup.py install`` We recommend installing using ``pip install``, as described above. However, you also may install from an unpacked source code tree. You can get such a tree by un-tarring the ``.tar.gz`` file or by using ``git clone``. From a source tree, you can install by running cd DataflowPythonSDK* python setup.py install --root / python setup.py test The ``--root /`` prevents Dataflow from being installed as an ``egg`` package. This workaround prevents failures if Dataflow is installed in the same virtual environment as another package under the ``google`` top-level package. If you get import errors during or after installing with ``setup.py``, uninstall the package: pip uninstall python-dataflow and use the ``pip install`` method described above to re-install it. ## Local execution of a pipeline The `$VIRTUAL_ENV/lib/python2.7/site-packages/google/cloud/dataflow/examples` subdirectory (the `google/cloud/dataflow/examples` subdirectory in the source distribution) has many examples large and small. All examples can be run locally by passing the arguments required by the example script. For instance, to run `wordcount.py`, try: python -m google.cloud.dataflow.examples.wordcount --output OUTPUT_FILE ## A Quick Tour of the Source Code You can follow along this tour by, with your virtual environment active, running a `pydoc` server on a local port of your choosing (this example uses port 8888). pydoc -p 8888 Now open your browser and go to http://localhost:8888/google.cloud.dataflow.html Some interesting classes to navigate to: * `PCollection`, in file [`google/cloud/dataflow/pvalue.py`](http://localhost:8888/google.cloud.dataflow.pvalue.html) * `PTransform`, in file [`google/cloud/dataflow/transforms/ptransform.py`](http://localhost:8888/google.cloud.dataflow.transforms.ptransform.html) * `FlatMap`, `GroupByKey`, and `Map`, in file [`google/cloud/dataflow/transforms/core.py`](http://localhost:8888/google.cloud.dataflow.transforms.core.html) * combiners, in file [`google/cloud/dataflow/transforms/combiners.py`](http://localhost:8888/google.cloud.dataflow.transforms.combiners.html) ## Some Simple Examples ### Hello world Create a transform from an iterable and use the pipe operator to chain transforms: ```python # Standard imports import google.cloud.dataflow as df # Create a pipeline executing on a direct runner (local, non-cloud). p = df.Pipeline('DirectPipelineRunner') # Create a PCollection with names and write it to a file. (p | df.Create('add names', ['Ann', 'Joe']) | df.Write('save', df.io.TextFileSink('./names'))) # Execute the pipeline. p.run() ``` ### Hello world (with Map) The `Map` transform takes a callable, which will be applied to each element of the input `PCollection` and must return an element to go into the output `PCollection`. ```python import google.cloud.dataflow as df p = df.Pipeline('DirectPipelineRunner') # Read file with names, add a greeting for each, and write results. (p | df.Read('load messages', df.io.TextFileSource('./names')) | df.Map('add greeting', lambda name, msg: '%s %s!' % (msg, name), 'Hello') | df.Write('save', df.io.TextFileSink('./greetings'))) p.run() ``` ### Hello world (with FlatMap) A `FlatMap` is like a `Map` except its callable returns a (possibly empty) iterable of elements for the output `PCollection`. ```python import google.cloud.dataflow as df p = df.Pipeline('DirectPipelineRunner') # Read previous file, add a name to each greeting and write results. (p | df.Read('load messages', df.io.TextFileSource('./names')) | df.FlatMap('add greetings', lambda name, msgs: ['%s %s!' % (m, name) for m in msgs], ['Hello', 'Hola']) | df.Write('save', df.io.TextFileSink('./greetings'))) p.run() ``` ### Hello world (with FlatMap and yield) The callable of a `FlatMap` can be a generator, that is, a function using `yield`. ```python import google.cloud.dataflow as df p = df.Pipeline('DirectPipelineRunner') # Add greetings using a FlatMap function using yield. def add_greetings(name, messages): for m in messages: yield '%s %s!' % (m, name) (p | df.Read('load names', df.io.TextFileSource('./names')) | df.FlatMap('greet', add_greetings, ['Hello', 'Hola']) | df.Write('save', df.io.TextFileSink('./greetings'))) p.run() ``` ### Counting words This example counts the words in a text and also shows how to read a text file from [Google Cloud Storage](https://cloud.google.com/storage/). ```python import re import google.cloud.dataflow as df p = df.Pipeline('DirectPipelineRunner') (p | df.Read('read', df.io.TextFileSource( 'gs://dataflow-samples/shakespeare/kinglear.txt')) | df.FlatMap('split', lambda x: re.findall(r'\w+', x)) | df.combiners.Count.PerElement('count words') | df.Write('write', df.io.TextFileSink('./results'))) p.run() ``` ### Counting words with GroupByKey Here we use `GroupByKey` to count the words. This is a somewhat forced example of `GroupByKey`; normally one would use the transform `df.combiners.Count.PerElement`, as in the previous example. The example also shows the use of a wild-card in specifying the text file source. ```python import re import google.cloud.dataflow as df p = df.Pipeline('DirectPipelineRunner') class MyCountTransform(df.PTransform): def apply(self, pcoll): return (pcoll | df.Map('one word', lambda w: (w, 1)) # GroupByKey accepts a PCollection of (w, 1) and # outputs a PCollection of (w, (1, 1, ...)) | df.GroupByKey('group words') | df.Map('count words', lambda (word, counts): (word, len(counts)))) (p | df.Read('read', df.io.TextFileSource('./names*')) | df.FlatMap('split', lambda x: re.findall(r'\w+', x)) | MyCountTransform() | df.Write('write', df.io.TextFileSink('./results'))) p.run() ``` ### Type hints In some cases, you can improve the efficiency of the data encoding by providing type hints. For example: ```python import google.cloud.dataflow as df from google.cloud.dataflow.typehints import typehints p = df.Pipeline('DirectPipelineRunner') (p | df.Read('A', df.io.TextFileSource('./names')) | df.Map('B1', lambda x: (x, 1)).with_output_types(typehints.KV[str, int]) | df.GroupByKey('GBK') | df.Write('C', df.io.TextFileSink('./results'))) p.run() ``` ### BigQuery Here is a pipeline that reads input from a BigQuery table and writes the result to a different table. This example calculates the number of tornadoes per month from weather data. To run it you will need to provide an output table that you can write to. ```python import google.cloud.dataflow as df input_table = 'clouddataflow-readonly:samples.weather_stations' project = 'YOUR-PROJECT' output_table = 'DATASET.TABLENAME' p = df.Pipeline(argv=['--project', project]) (p | df.Read('read', df.io.BigQuerySource(input_table)) | df.FlatMap( 'months with tornadoes', lambda row: [(int(row['month']), 1)] if row['tornado'] else []) | df.CombinePerKey('monthly count', sum) | df.Map('format', lambda (k, v): {'month': k, 'tornado_count': v}) | df.Write('write', df.io.BigQuerySink( output_table, schema='month:INTEGER, tornado_count:INTEGER', create_disposition=df.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=df.io.BigQueryDisposition.WRITE_TRUNCATE))) p.run() ``` Here is a pipeline that achieves the same functionality, i.e., calculates the number of tornadoes per month, but uses a query to filter out input instead of using the whole table. ```python import google.cloud.dataflow as df project = 'YOUR-PROJECT' output_table = 'DATASET.TABLENAME' input_query = 'SELECT month, COUNT(month) AS tornado_count ' \ 'FROM [clouddataflow-readonly:samples.weather_stations] ' \ 'WHERE tornado=true GROUP BY month' p = df.Pipeline(argv=['--project', project]) (p | df.Read('read', df.io.BigQuerySource(query=input_query)) | df.Write('write', df.io.BigQuerySink( output_table, schema='month:INTEGER, tornado_count:INTEGER', create_disposition=df.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=df.io.BigQueryDisposition.WRITE_TRUNCATE))) p.run() ``` ### Combiner Examples A common case for Dataflow combiners is to sum (or max or min) over the values of each key. Such standard Python functions can be used directly as combiner functions. In fact, any function "reducing" an iterable to a single value can be used. ```python import google.cloud.dataflow as df p = df.Pipeline('DirectPipelineRunner') SAMPLE_DATA = [('a', 1), ('b', 10), ('a', 2), ('a', 3), ('b', 20)] (p | df.Create(SAMPLE_DATA) | df.CombinePerKey(sum) | df.Write(df.io.TextFileSink('./results'))) p.run() ``` The `google/cloud/dataflow/examples/cookbook/combiners_test.py` file in the source distribution contains more combiner examples. ### More Examples The `google/cloud/dataflow/examples` subdirectory in the source distribution has some larger examples. ## Organizing Your Code Many projects will grow to multiple source code files. It is beneficial to organize the project so that all the code involved in running a workflow can be built as a Python package so that it can be installed in the VM workers executing a job. Please follow the example in `google/cloud/dataflow/examples/complete/juliaset`. If the code is organized in this fashion then you can use the `--setup_file` command line option to create a source distribution out of the project files, stage the resulting tarball and later install it in the workers executing the job. ## Contact Us We welcome all usage-related questions on [Stack Overflow](https://stackoverflow.com/questions/tagged/google-cloud-dataflow) tagged with `google-cloud-dataflow`. Please use the [issue tracker](https://github.com/GoogleCloudPlatform/DataflowPythonSDK/issues) on GitHub to report any bugs, comments or questions regarding SDK development. DataflowPythonSDK-0.2.5/google/000077500000000000000000000000001272340251000162735ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/__init__.py000066400000000000000000000002261272340251000204040ustar00rootroot00000000000000try: __import__('pkg_resources').declare_namespace(__name__) except ImportError: __path__ = __import__('pkgutil').extend_path(__path__, __name__) DataflowPythonSDK-0.2.5/google/cloud/000077500000000000000000000000001272340251000174015ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/__init__.py000066400000000000000000000002261272340251000215120ustar00rootroot00000000000000try: __import__('pkg_resources').declare_namespace(__name__) except ImportError: __path__ = __import__('pkgutil').extend_path(__path__, __name__) DataflowPythonSDK-0.2.5/google/cloud/dataflow/000077500000000000000000000000001272340251000212025ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/__init__.py000066400000000000000000000046571272340251000233270ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Google Cloud Dataflow SDK for Python. Google Cloud Dataflow provides a simple, powerful programming model for building both batch and streaming parallel data processing pipelines. The Dataflow SDK for Python provides access to Dataflow capabilities from the Python programming language. Status ------ The SDK is still early in its development, and significant changes should be expected before the first stable version. Overview -------- The key concepts in this programming model are * PCollection: represents a collection of data, which could be bounded or unbounded in size. * PTransform: represents a computation that transforms input PCollections into output PCollections. * Pipeline: manages a directed acyclic graph of PTransforms and PCollections that is ready for execution. * Runner: specifies where and how the Pipeline should execute. * Reading and Writing Data: your pipeline can read from an external source and write to an external data sink. Typical usage ------------- At the top of your source file:: import google.cloud.dataflow as df After this import statement * transform classes are available as df.FlatMap, df.GroupByKey, etc. * Pipeline class is available as df.Pipeline * text source/sink classes are available as df.io.TextFileSource, df.io.TextFileSink Examples -------- The examples subdirectory has some examples. """ import sys if sys.version_info.major != 2: raise RuntimeError( 'Dataflow SDK for Python is supported only on Python 2.7. ' 'It is not supported on Python [%s].' % sys.version) import google.cloud.dataflow.internal.pickler from google.cloud.dataflow import coders from google.cloud.dataflow import io from google.cloud.dataflow import typehints from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.transforms import * DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/000077500000000000000000000000001272340251000224615ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/__init__.py000066400000000000000000000013051272340251000245710ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. from google.cloud.dataflow.coders.coders import * from google.cloud.dataflow.coders.typecoders import registry DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/coder_impl.pxd000066400000000000000000000051601272340251000253150ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. # cython: profile=True cimport cython cimport cpython.ref cimport cpython.tuple cimport libc.stdint cimport libc.stdlib cimport libc.string from .stream cimport InputStream, OutputStream cdef object loads, dumps, create_InputStream, create_OutputStream cdef type WindowedValue cdef class CoderImpl(object): cpdef encode_to_stream(self, value, OutputStream stream, bint nested) cpdef decode_from_stream(self, InputStream stream, bint nested) cpdef bytes encode(self, value) cpdef decode(self, bytes encoded) cdef class SimpleCoderImpl(CoderImpl): pass cdef class StreamCoderImpl(CoderImpl): pass cdef class CallbackCoderImpl(CoderImpl): cdef object _encoder cdef object _decoder cdef class DeterministicPickleCoderImpl(CoderImpl): cdef CoderImpl _pickle_coder cdef object _step_label cdef bint _check_safe(self, value) except -1 cdef class BytesCoderImpl(CoderImpl): pass cdef class FloatCoderImpl(StreamCoderImpl): pass cdef class TimestampCoderImpl(StreamCoderImpl): cdef object timestamp_class cdef list small_ints cdef class VarIntCoderImpl(StreamCoderImpl): @cython.locals(ivalue=libc.stdint.int64_t) cpdef bytes encode(self, value) cdef class SingletonCoderImpl(CoderImpl): cdef object _value cdef class AbstractComponentCoderImpl(StreamCoderImpl): cdef tuple _coder_impls cpdef _extract_components(self, value) cpdef _construct_from_components(self, components) @cython.locals(c=CoderImpl) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) @cython.locals(c=CoderImpl) cpdef decode_from_stream(self, InputStream stream, bint nested) cdef class TupleCoderImpl(AbstractComponentCoderImpl): pass cdef class SequenceCoderImpl(StreamCoderImpl): cdef CoderImpl _elem_coder cpdef _construct_from_sequence(self, values) cdef class TupleSequenceCoderImpl(SequenceCoderImpl): pass cdef class WindowedValueCoderImpl(StreamCoderImpl): """A coder for windowed values.""" cdef CoderImpl _value_coder cdef CoderImpl _timestamp_coder cdef CoderImpl _windows_coder DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/coder_impl.py000066400000000000000000000223601272340251000251530ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Coder implementations. The actual encode/decode implementations are split off from coders to allow conditional (compiled/pure) implementations, which can be used to encode many elements with minimal overhead. This module may be optionally compiled with Cython, using the corresponding coder_impl.pxd file for type hints. """ import collections from cPickle import loads, dumps # pylint: disable=g-import-not-at-top try: # Don't depend on the full dataflow sdk to test coders. from google.cloud.dataflow.transforms.window import WindowedValue except ImportError: WindowedValue = collections.namedtuple( 'WindowedValue', ('value', 'timestamp', 'windows')) try: from stream import InputStream as create_InputStream from stream import OutputStream as create_OutputStream except ImportError: from slow_stream import InputStream as create_InputStream from slow_stream import OutputStream as create_OutputStream # pylint: enable=g-import-not-at-top class CoderImpl(object): def encode_to_stream(self, value, stream, nested): """Reads object from potentially-nested encoding in stream.""" raise NotImplementedError def decode_from_stream(self, stream, nested): """Reads object from potentially-nested encoding in stream.""" raise NotImplementedError def encode(self, value): """Encodes an object to an unnested string.""" raise NotImplementedError def decode(self, encoded): """Encodes an object to an unnested string.""" raise NotImplementedError class SimpleCoderImpl(CoderImpl): """Subclass of CoderImpl implementing stream methods using encode/decode.""" def encode_to_stream(self, value, stream, nested): """Reads object from potentially-nested encoding in stream.""" stream.write(self.encode(value), nested) def decode_from_stream(self, stream, nested): """Reads object from potentially-nested encoding in stream.""" return self.decode(stream.read_all(nested)) class StreamCoderImpl(CoderImpl): """Subclass of CoderImpl implementing encode/decode using stream methods.""" def encode(self, value): out = create_OutputStream() self.encode_to_stream(value, out, False) return out.get() def decode(self, encoded): return self.decode_from_stream(create_InputStream(encoded), False) class CallbackCoderImpl(CoderImpl): """A CoderImpl that calls back to the _impl methods on the Coder itself. This is the default implementation used if Coder._get_impl() is not overwritten. """ def __init__(self, encoder, decoder): self._encoder = encoder self._decoder = decoder def encode_to_stream(self, value, stream, nested): return stream.write(self._encoder(value), nested) def decode_from_stream(self, stream, nested): return self._decoder(stream.read_all(nested)) def encode(self, value): return self._encoder(value) def decode(self, encoded): return self._decoder(encoded) class DeterministicPickleCoderImpl(CoderImpl): def __init__(self, pickle_coder, step_label): self._pickle_coder = pickle_coder self._step_label = step_label def _check_safe(self, value): if isinstance(value, (str, unicode, long, int, float)): pass elif value is None: pass elif isinstance(value, (tuple, list)): for x in value: self._check_safe(x) else: raise TypeError( "Unable to deterministically code '%s' of type '%s', " "please provide a type hint for the input of '%s'" % ( value, type(value), self._step_label)) def encode_to_stream(self, value, stream, nested): self._check_safe(value) return self._pickle_coder.encode_to_stream(value, stream, nested) def decode_from_stream(self, stream, nested): return self._pickle_coder.decode_from_stream(stream, nested) def encode(self, value): self._check_safe(value) return self._pickle_coder.encode(value) def decode(self, encoded): return self._pickle_coder.decode(encoded) class BytesCoderImpl(CoderImpl): """A coder for bytes/str objects.""" def encode_to_stream(self, value, out, nested): out.write(value, nested) def decode_from_stream(self, in_stream, nested): return in_stream.read_all(nested) def encode(self, value): assert isinstance(value, bytes), (value, type(value)) return value def decode(self, encoded): return encoded class FloatCoderImpl(StreamCoderImpl): def encode_to_stream(self, value, out, nested): out.write_bigendian_double(value) def decode_from_stream(self, in_stream, nested): return in_stream.read_bigendian_double() class TimestampCoderImpl(StreamCoderImpl): def __init__(self, timestamp_class): self.timestamp_class = timestamp_class def encode_to_stream(self, value, out, nested): out.write_bigendian_int64(value.micros) def decode_from_stream(self, in_stream, nested): return self.timestamp_class(micros=in_stream.read_bigendian_int64()) small_ints = [chr(_) for _ in range(128)] class VarIntCoderImpl(StreamCoderImpl): """A coder for long/int objects.""" def encode_to_stream(self, value, out, nested): out.write_var_int64(value) def decode_from_stream(self, in_stream, nested): return in_stream.read_var_int64() def encode(self, value): ivalue = value # type cast if 0 <= ivalue < len(small_ints): return small_ints[ivalue] else: return StreamCoderImpl.encode(self, value) def decode(self, encoded): if len(encoded) == 1: i = ord(encoded) if 0 <= i < 128: return i return StreamCoderImpl.decode(self, encoded) class SingletonCoderImpl(CoderImpl): """A coder that always encodes exactly one value.""" def __init__(self, value): self._value = value def encode_to_stream(self, value, stream, nested): pass def decode_from_stream(self, stream, nested): return self._value def encode(self, value): b = '' # avoid byte vs str vs unicode error return b def decode(self, encoded): return self._value class AbstractComponentCoderImpl(StreamCoderImpl): def __init__(self, coder_impls): for c in coder_impls: assert isinstance(c, CoderImpl), c self._coder_impls = tuple(coder_impls) def _extract_components(self, value): raise NotImplementedError def _construct_from_components(self, components): raise NotImplementedError def encode_to_stream(self, value, out, nested): values = self._extract_components(value) if len(self._coder_impls) != len(values): raise ValueError( 'Number of components does not match number of coders.') for i in range(0, len(self._coder_impls)): c = self._coder_impls[i] # type cast c.encode_to_stream(values[i], out, True) def decode_from_stream(self, in_stream, nested): return self._construct_from_components( [c.decode_from_stream(in_stream, True) for c in self._coder_impls]) class TupleCoderImpl(AbstractComponentCoderImpl): """A coder for tuple objects.""" def _extract_components(self, value): return value def _construct_from_components(self, components): return tuple(components) class SequenceCoderImpl(StreamCoderImpl): """A coder for sequences of known length.""" def __init__(self, elem_coder): self._elem_coder = elem_coder def _construct_from_sequence(self, values): raise NotImplementedError def encode_to_stream(self, value, out, nested): # Compatible with Java's IterableLikeCoder. out.write_bigendian_int32(len(value)) for elem in value: self._elem_coder.encode_to_stream(elem, out, True) def decode_from_stream(self, in_stream, nested): size = in_stream.read_bigendian_int32() return self._construct_from_sequence( [self._elem_coder.decode_from_stream(in_stream, True) for _ in range(size)]) class TupleSequenceCoderImpl(SequenceCoderImpl): """A coder for homogeneous tuple objects.""" def _construct_from_sequence(self, components): return tuple(components) class WindowedValueCoderImpl(StreamCoderImpl): """A coder for windowed values.""" def __init__(self, value_coder, timestamp_coder, window_coder): self._value_coder = value_coder self._timestamp_coder = timestamp_coder self._windows_coder = TupleSequenceCoderImpl(window_coder) def encode_to_stream(self, value, out, nested): self._value_coder.encode_to_stream(value.value, out, True) self._timestamp_coder.encode_to_stream(value.timestamp, out, True) self._windows_coder.encode_to_stream(value.windows, out, True) def decode_from_stream(self, in_stream, nested): return WindowedValue( self._value_coder.decode_from_stream(in_stream, True), self._timestamp_coder.decode_from_stream(in_stream, True), self._windows_coder.decode_from_stream(in_stream, True)) DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/coders.py000066400000000000000000000351201272340251000243130ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Collection of useful coders.""" import base64 import collections import cPickle as pickle from google.cloud.dataflow.coders import coder_impl # pylint: disable=g-import-not-at-top # Avoid dependencies on the full SDK. try: # Import dill from the pickler module to make sure our monkey-patching of dill # occurs. from google.cloud.dataflow.internal.pickler import dill from google.cloud.dataflow.transforms.timeutil import Timestamp except ImportError: # We fall back to using the stock dill library in tests that don't use the # full Python SDK. import dill Timestamp = collections.namedtuple('Timestamp', 'micros') def serialize_coder(coder): from google.cloud.dataflow.internal import pickler return '%s$%s' % (coder.__class__.__name__, pickler.dumps(coder)) def deserialize_coder(serialized): from google.cloud.dataflow.internal import pickler return pickler.loads(serialized.split('$', 1)[1]) # pylint: enable=g-import-not-at-top class Coder(object): """Base class for coders.""" def encode(self, value): """Encodes the given object into a byte string.""" raise NotImplementedError('Encode not implemented: %s.' % self) def decode(self, encoded): """Decodes the given byte string into the corresponding object.""" raise NotImplementedError('Decode not implemented: %s.' % self) def is_deterministic(self): """Whether this coder is guaranteed to encode values deterministically. A deterministic coder is required for key coders in GroupByKey operations to produce consistent results. For example, note that the default coder, the PickleCoder, is not deterministic: the ordering of picked entries in maps may vary across executions since there is no defined order, and such a coder is not in general suitable for usage as a key coder in GroupByKey operations, since each instance of the same key may be encoded differently. Returns: Whether coder is deterministic. """ return False # =========================================================================== # Methods below are internal SDK details that don't need to be modified for # user-defined coders. # =========================================================================== def _create_impl(self): """Creates a CoderImpl to do the actual encoding and decoding. """ return coder_impl.CallbackCoderImpl(self.encode, self.decode) def get_impl(self): if not hasattr(self, '_impl'): self._impl = self._create_impl() assert isinstance(self._impl, coder_impl.CoderImpl) return self._impl def __getstate__(self): return self._dict_without_impl() def _dict_without_impl(self): if hasattr(self, '_impl'): d = dict(self.__dict__) del d['_impl'] return d else: return self.__dict__ @classmethod def from_type_hint(cls, unused_typehint, unused_registry): # If not overridden, just construct the coder without arguments. return cls() def is_kv_coder(self): return False def key_coder(self): if self.is_kv_coder(): raise NotImplementedError('key_coder: %s' % self) else: raise ValueError('Not a KV coder: %s.' % self) def value_coder(self): if self.is_kv_coder(): raise NotImplementedError('value_coder: %s' % self) else: raise ValueError('Not a KV coder: %s.' % self) def _get_component_coders(self): """Returns the internal component coders of this coder.""" # This is an internal detail of the Coder API and does not need to be # refined in user-defined Coders. return [] def as_cloud_object(self): """Returns Google Cloud Dataflow API description of this coder.""" # This is an internal detail of the Coder API and does not need to be # refined in user-defined Coders. value = { # We pass coders in the form "$" to make the # job description JSON more readable. Data before the $ is ignored by # the worker. '@type': serialize_coder(self), 'component_encodings': list( component.as_cloud_object() for component in self._get_component_coders()) } return value def __repr__(self): return self.__class__.__name__ def __eq__(self, other): # pylint: disable=protected-access return (self.__class__ == other.__class__ and self._dict_without_impl() == other._dict_without_impl()) # pylint: enable=protected-access class StrUtf8Coder(Coder): """A coder used for reading and writing strings as UTF-8.""" def encode(self, value): return value.encode('utf-8') def decode(self, value): return value.decode('utf-8') def is_deterministic(self): return True class ToStringCoder(Coder): """A default string coder used if no sink coder is specified.""" def encode(self, value): if isinstance(value, unicode): return value.encode('utf-8') elif isinstance(value, str): return value else: return str(value) def decode(self, _): raise NotImplementedError('ToStringCoder cannot be used for decoding.') def is_deterministic(self): return True class FastCoder(Coder): """Coder subclass used when a (faster) CoderImpl is supplied directly. The Coder class defines _create_impl in terms of encode() and decode(); this class inverts that defining encode() and decode() in terms of _create_impl(). """ def encode(self, value): """Encodes the given object into a byte string.""" return self.get_impl().encode(value) def decode(self, encoded): """Decodes the given byte string into the corresponding object.""" return self.get_impl().decode(encoded) def _create_impl(self): raise NotImplementedError class BytesCoder(FastCoder): """Byte string coder.""" def _create_impl(self): return coder_impl.BytesCoderImpl() def is_deterministic(self): return True class VarIntCoder(FastCoder): """Variable-length integer coder.""" def _create_impl(self): return coder_impl.VarIntCoderImpl() def is_deterministic(self): return True class FloatCoder(FastCoder): """A coder used for floating-point values.""" def _create_impl(self): return coder_impl.FloatCoderImpl() def is_deterministic(self): return True class TimestampCoder(FastCoder): """A coder used for timeutil.Timestamp values.""" def _create_impl(self): return coder_impl.TimestampCoderImpl(Timestamp) def is_deterministic(self): return True class SingletonCoder(FastCoder): """A coder that always encodes exactly one value.""" def __init__(self, value): self._value = value def _create_impl(self): return coder_impl.SingletonCoderImpl(self._value) def is_deterministic(self): return True def maybe_dill_dumps(o): """Pickle using cPickle or the Dill pickler as a fallback.""" # We need to use the dill pickler for objects of certain custom classes, # including, for example, ones that contain lambdas. try: return pickle.dumps(o) except Exception: # pylint: disable=broad-except return dill.dumps(o) def maybe_dill_loads(o): """Unpickle using cPickle or the Dill pickler as a fallback.""" try: return pickle.loads(o) except Exception: # pylint: disable=broad-except return dill.loads(o) class _PickleCoderBase(FastCoder): """Base class for pickling coders.""" def is_deterministic(self): # Note that the default coder, the PickleCoder, is not deterministic (for # example, the ordering of picked entries in maps may vary across # executions), and so is not in general suitable for usage as a key coder in # GroupByKey operations. return False def as_cloud_object(self, is_pair_like=True): value = super(_PickleCoderBase, self).as_cloud_object() # We currently use this coder in places where we cannot infer the coder to # use for the value type in a more granular way. In places where the # service expects a pair, it checks for the "is_pair_like" key, in which # case we would fail without the hack below. if is_pair_like: value['is_pair_like'] = True value['component_encodings'] = [ self.as_cloud_object(is_pair_like=False), self.as_cloud_object(is_pair_like=False) ] return value # We allow .key_coder() and .value_coder() to be called on PickleCoder since # we can't always infer the return values of lambdas in ParDo operations, the # result of which may be used in a GroupBykey. def is_kv_coder(self): return True def key_coder(self): return self def value_coder(self): return self class PickleCoder(_PickleCoderBase): """Coder using Python's pickle functionality.""" def _create_impl(self): return coder_impl.CallbackCoderImpl(pickle.dumps, pickle.loads) class DillCoder(_PickleCoderBase): """Coder using dill's pickle functionality.""" def _create_impl(self): return coder_impl.CallbackCoderImpl(maybe_dill_dumps, maybe_dill_loads) class DeterministicPickleCoder(FastCoder): """Throws runtime errors when pickling non-deterministic values.""" def __init__(self, pickle_coder, step_label): self._pickle_coder = pickle_coder self._step_label = step_label def _create_impl(self): return coder_impl.DeterministicPickleCoderImpl( self._pickle_coder.get_impl(), self._step_label) def is_deterministic(self): return True def is_kv_coder(self): return True def key_coder(self): return self def value_coder(self): return self class Base64PickleCoder(Coder): """Coder of objects by Python pickle, then base64 encoding.""" # TODO(robertwb): Do base64 encoding where it's needed (e.g. in json) rather # than via a special Coder. def encode(self, value): return base64.b64encode(pickle.dumps(value)) def decode(self, encoded): return pickle.loads(base64.b64decode(encoded)) def is_deterministic(self): # Note that the Base64PickleCoder is not deterministic. See the # corresponding comments for PickleCoder above. return False # We allow .key_coder() and .value_coder() to be called on Base64PickleCoder # since we can't always infer the return values of lambdas in ParDo # operations, the result of which may be used in a GroupBykey. # # TODO(ccy): this is currently only used for KV values from Create transforms. # Investigate a way to unify this with PickleCoder. def is_kv_coder(self): return True def key_coder(self): return self def value_coder(self): return self class TupleCoder(FastCoder): """Coder of tuple objects.""" def __init__(self, components): self._coders = tuple(components) def _create_impl(self): return coder_impl.TupleCoderImpl([c.get_impl() for c in self._coders]) def is_deterministic(self): return all(c.is_deterministic() for c in self._coders) @staticmethod def from_type_hint(typehint, registry): return TupleCoder([registry.get_coder(t) for t in typehint.tuple_types]) def as_cloud_object(self): value = super(TupleCoder, self).as_cloud_object() value['is_pair_like'] = True return value def _get_component_coders(self): return self.coders() def coders(self): return self._coders def is_kv_coder(self): return len(self._coders) == 2 def key_coder(self): if len(self._coders) != 2: raise ValueError('TupleCoder does not have exactly 2 components.') return self._coders[0] def value_coder(self): if len(self._coders) != 2: raise ValueError('TupleCoder does not have exactly 2 components.') return self._coders[1] def __repr__(self): return 'TupleCoder[%s]' % ', '.join(str(c) for c in self._coders) class TupleSequenceCoder(FastCoder): """Coder of homogeneous tuple objects.""" def __init__(self, elem_coder): self._elem_coder = elem_coder def _create_impl(self): return coder_impl.TupleSequenceCoderImpl(self._elem_coder.get_impl()) def is_deterministic(self): return self._elem_coder.is_deterministic() @staticmethod def from_type_hint(typehint, registry): return TupleSequenceCoder(registry.get_coder(typehint.inner_type)) def _get_component_coders(self): return (self._elem_coder,) def __repr__(self): return 'TupleSequenceCoder[%r]' % self._elem_coder class WindowCoder(PickleCoder): """Coder for windows in windowed values.""" def _create_impl(self): return coder_impl.CallbackCoderImpl(pickle.dumps, pickle.loads) def is_deterministic(self): # Note that WindowCoder as implemented is not deterministic because the # implementation simply pickles windows. See the corresponding comments # on PickleCoder for more details. return False def as_cloud_object(self): return super(WindowCoder, self).as_cloud_object(is_pair_like=False) class WindowedValueCoder(FastCoder): """Coder for windowed values.""" def __init__(self, wrapped_value_coder, timestamp_coder=None, window_coder=None): if not timestamp_coder: timestamp_coder = TimestampCoder() if not window_coder: window_coder = PickleCoder() self.wrapped_value_coder = wrapped_value_coder self.timestamp_coder = timestamp_coder self.window_coder = window_coder def _create_impl(self): return coder_impl.WindowedValueCoderImpl( self.wrapped_value_coder.get_impl(), self.timestamp_coder.get_impl(), self.window_coder.get_impl()) def is_deterministic(self): return all(c.is_deterministic() for c in [self.wrapped_value_coder, self.timestamp_coder, self.window_coder]) def as_cloud_object(self): value = super(WindowedValueCoder, self).as_cloud_object() value['is_wrapper'] = True return value def _get_component_coders(self): return [self.wrapped_value_coder, self.timestamp_coder, self.window_coder] def is_kv_coder(self): return self.wrapped_value_coder.is_kv_coder() def key_coder(self): return self.wrapped_value_coder.key_coder() def value_coder(self): return self.wrapped_value_coder.value_coder() def __repr__(self): return 'WindowedValueCoder[%s]' % self.wrapped_value_coder DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/coders_test.py000066400000000000000000000040421272340251000253510ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. import base64 import logging import unittest from google.cloud.dataflow import coders class PickleCoderTest(unittest.TestCase): def test_basics(self): v = ('a' * 10, 'b' * 90) pickler = coders.PickleCoder() self.assertEquals(v, pickler.decode(pickler.encode(v))) pickler = coders.Base64PickleCoder() self.assertEquals(v, pickler.decode(pickler.encode(v))) self.assertEquals( coders.Base64PickleCoder().encode(v), base64.b64encode(coders.PickleCoder().encode(v))) def test_equality(self): self.assertEquals(coders.PickleCoder(), coders.PickleCoder()) self.assertEquals(coders.Base64PickleCoder(), coders.Base64PickleCoder()) self.assertNotEquals(coders.Base64PickleCoder(), coders.PickleCoder()) self.assertNotEquals(coders.Base64PickleCoder(), object()) class CodersTest(unittest.TestCase): def test_str_utf8_coder(self): real_coder = coders.registry.get_coder(str) expected_coder = coders.BytesCoder() self.assertEqual( real_coder.encode('abc'), expected_coder.encode('abc')) self.assertEqual('abc', real_coder.decode(real_coder.encode('abc'))) real_coder = coders.registry.get_coder(bytes) expected_coder = coders.BytesCoder() self.assertEqual( real_coder.encode('abc'), expected_coder.encode('abc')) self.assertEqual('abc', real_coder.decode(real_coder.encode('abc'))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/coders_test_common.py000066400000000000000000000140531272340251000267240ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests common to all coder implementations.""" import logging import math import sys import unittest import dill import coders # Defined out of line for picklability. class CustomCoder(coders.Coder): def encode(self, x): return str(x+1) def decode(self, encoded): return int(encoded) - 1 class CodersTest(unittest.TestCase): # These class methods ensure that we test each defined coder in both # nested and unnested context. @classmethod def setUpClass(cls): cls.seen = set() cls.seen_nested = set() @classmethod def tearDownClass(cls): standard = set(c for c in coders.__dict__.values() if isinstance(c, type) and issubclass(c, coders.Coder) and 'Base' not in c.__name__) standard -= set([coders.Coder, coders.FastCoder, coders.Base64PickleCoder, coders.FloatCoder, coders.TimestampCoder, coders.ToStringCoder, coders.WindowCoder, coders.WindowedValueCoder]) assert not standard - cls.seen, standard - cls.seen assert not standard - cls.seen_nested, standard - cls.seen_nested @classmethod def _observe(cls, coder): cls.seen.add(type(coder)) cls._observe_nested(coder) @classmethod def _observe_nested(cls, coder): if isinstance(coder, coders.TupleCoder): for c in coder.coders(): cls.seen_nested.add(type(c)) cls._observe_nested(c) def check_coder(self, coder, *values): self._observe(coder) for v in values: self.assertEqual(v, coder.decode(coder.encode(v))) copy1 = dill.loads(dill.dumps(coder)) copy2 = dill.loads(dill.dumps(coder)) for v in values: self.assertEqual(v, copy1.decode(copy2.encode(v))) def test_custom_coder(self): self.check_coder(CustomCoder(), 1, -10, 5) self.check_coder(coders.TupleCoder((CustomCoder(), coders.BytesCoder())), (1, 'a'), (-10, 'b'), (5, 'c')) def test_pickle_coder(self): self.check_coder(coders.PickleCoder(), 'a', 1, 1.5, (1, 2, 3)) def test_deterministic_pickle_coder(self): coder = coders.DeterministicPickleCoder(coders.PickleCoder(), 'step') self.check_coder(coder, 'a', 1, 1.5, (1, 2, 3)) with self.assertRaises(TypeError): self.check_coder(coder, dict()) with self.assertRaises(TypeError): self.check_coder(coder, [1, dict()]) self.check_coder(coders.TupleCoder((coder, coders.PickleCoder())), (1, dict()), ('a', [dict()])) def test_dill_coder(self): cell_value = (lambda x: lambda: x)(0).func_closure[0] self.check_coder(coders.DillCoder(), 'a', 1, cell_value) self.check_coder( coders.TupleCoder((coders.VarIntCoder(), coders.DillCoder())), (1, cell_value)) def test_bytes_coder(self): self.check_coder(coders.BytesCoder(), 'a', '\0', 'z' * 1000) def test_varint_coder(self): # Small ints. self.check_coder(coders.VarIntCoder(), *range(-10, 10)) # Multi-byte encoding starts at 128 self.check_coder(coders.VarIntCoder(), *range(120, 140)) # Large values self.check_coder(coders.VarIntCoder(), *[int(math.pow(-1, k) * math.exp(k)) for k in range(0, int(math.log(sys.maxint)))]) def test_float_coder(self): self.check_coder(coders.FloatCoder(), *[float(0.1 * x) for x in range(-100, 100)]) self.check_coder(coders.FloatCoder(), *[float(2 ** (0.1 * x)) for x in range(-100, 100)]) self.check_coder(coders.FloatCoder(), float('-Inf'), float('Inf')) def test_singleton_coder(self): a = 'anything' b = 'something else' self.check_coder(coders.SingletonCoder(a), a) self.check_coder(coders.SingletonCoder(b), b) self.check_coder(coders.TupleCoder((coders.SingletonCoder(a), coders.SingletonCoder(b))), (a, b)) def test_timestamp_coder(self): self.check_coder(coders.TimestampCoder(), *[coders.Timestamp(micros=x) for x in range(-100, 100)]) self.check_coder(coders.TimestampCoder(), coders.Timestamp(micros=-1234567890), coders.Timestamp(micros=1234567890)) self.check_coder(coders.TimestampCoder(), coders.Timestamp(micros=-1234567890123456789), coders.Timestamp(micros=1234567890123456789)) def test_tuple_coder(self): self.check_coder( coders.TupleCoder((coders.VarIntCoder(), coders.BytesCoder())), (1, 'a'), (-2, 'a' * 100), (300, 'abc\0' * 5)) self.check_coder( coders.TupleCoder( (coders.TupleCoder((coders.PickleCoder(), coders.VarIntCoder())), coders.StrUtf8Coder())), ((1, 2), 'a'), ((-2, 5), u'a\u0101' * 100), ((300, 1), 'abc\0' * 5)) def test_tuple_sequence_coder(self): int_tuple_coder = coders.TupleSequenceCoder(coders.VarIntCoder()) self.check_coder(int_tuple_coder, (1, -1, 0), (), tuple(range(1000))) self.check_coder( coders.TupleCoder((coders.VarIntCoder(), int_tuple_coder)), (1, (1, 2, 3))) def test_base64_pickle_coder(self): self.check_coder(coders.Base64PickleCoder(), 'a', 1, 1.5, (1, 2, 3)) def test_utf8_coder(self): self.check_coder(coders.StrUtf8Coder(), 'a', u'ab\u00FF', u'\u0101\0') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/fast_coders_test.py000066400000000000000000000021001272340251000263570ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for compiled implementation of coder impls.""" import logging import unittest # Run all the standard coder test cases. from google.cloud.dataflow.coders.coders_test_common import * class FastCoders(unittest.TestCase): def test_using_fast_impl(self): # pylint: disable=g-import-not-at-top # pylint: disable=unused-variable import google.cloud.dataflow.coders.stream if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/observable.py000066400000000000000000000017631272340251000251660ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Observable base class for iterables.""" class ObservableMixin(object): """An observable iterable. Subclasses need to call self.notify_observers with any object yielded. """ def __init__(self): self.observers = [] def register_observer(self, callback): self.observers.append(callback) def notify_observers(self, value, **kwargs): for o in self.observers: o(value, **kwargs) DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/observable_test.py000066400000000000000000000030231272340251000262140ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for the Observable mixin class.""" import logging import unittest from google.cloud.dataflow.coders import observable class ObservableMixinTest(unittest.TestCase): observed_count = 0 observed_sum = 0 observed_keys = [] def observer(self, value, key=None): self.observed_count += 1 self.observed_sum += value self.observed_keys.append(key) def test_observable(self): class Watched(observable.ObservableMixin): def __iter__(self): for i in (1, 4, 3): self.notify_observers(i, key='a%d' % i) yield i watched = Watched() watched.register_observer(lambda v, key: self.observer(v, key=key)) for _ in watched: pass self.assertEquals(3, self.observed_count) self.assertEquals(8, self.observed_sum) self.assertEquals(['a1', 'a3', 'a4'], sorted(self.observed_keys)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/slow_coders_test.py000066400000000000000000000022611272340251000264160ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for uncompiled implementation of coder impls.""" import logging import unittest # Run all the standard coder test cases. from google.cloud.dataflow.coders.coders_test_common import * class SlowCoders(unittest.TestCase): def test_using_slow_impl(self): # Assert that we are not using the compiled implementation. with self.assertRaises(ImportError): # pylint: disable=g-import-not-at-top # pylint: disable=unused-variable import google.cloud.dataflow.coders.stream if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/slow_stream.py000066400000000000000000000065351272340251000254030ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A pure Python implementation of stream.pyx.""" import struct class OutputStream(object): """A pure Python implementation of stream.OutputStream.""" def __init__(self): self.data = [] def write(self, b, nested=False): assert isinstance(b, str) if nested: self.write_var_int64(len(b)) self.data.append(b) def write_byte(self, val): self.data.append(chr(val)) def write_var_int64(self, v): if v < 0: v += 1 << 64 if v <= 0: raise ValueError('Value too large (negative).') while True: bits = v & 0x7F v >>= 7 if v: bits |= 0x80 self.write_byte(bits) if not v: break def write_bigendian_int64(self, v): self.write(struct.pack('>q', v)) def write_bigendian_int32(self, v): self.write(struct.pack('>i', v)) def write_bigendian_double(self, v): self.write(struct.pack('>d', v)) def get(self): return ''.join(self.data) class ByteCountingOutputStream(OutputStream): """A pure Python implementation of stream.ByteCountingOutputStream.""" def __init__(self): # Note that we don't actually use any of the data initialized by our super. super(ByteCountingOutputStream, self).__init__() self.count = 0 def write(self, byte_array, nested=False): blen = len(byte_array) if nested: self.write_var_int64(blen) self.count += blen def write_byte(self, _): self.count += 1 def get_count(self): return self.count def get(self): raise NotImplementedError def __str__(self): return '<%s %s>' % (self.__class__.__name__, self.count) class InputStream(object): """A pure Python implementation of stream.InputStream.""" def __init__(self, data): self.data = data self.pos = 0 def size(self): return len(self.data) - self.pos def read(self, size): self.pos += size return self.data[self.pos - size : self.pos] def read_all(self, nested): return self.read(self.read_var_int64() if nested else self.size()) def read_byte(self): self.pos += 1 return ord(self.data[self.pos - 1]) def read_var_int64(self): shift = 0 result = 0 while True: byte = self.read_byte() if byte < 0: raise RuntimeError('VarLong not terminated.') bits = byte & 0x7F if shift >= 64 or (shift >= 63 and bits > 1): raise RuntimeError('VarLong too long.') result |= bits << shift shift += 7 if not byte & 0x80: break if result >= 1 << 63: result -= 1 << 64 return result def read_bigendian_int64(self): return struct.unpack('>q', self.read(8))[0] def read_bigendian_int32(self): return struct.unpack('>i', self.read(4))[0] def read_bigendian_double(self): return struct.unpack('>d', self.read(8))[0] DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/stream.pxd000066400000000000000000000036021272340251000244720ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. cimport libc.stdint cdef class OutputStream(object): cdef char* data cdef size_t size cdef size_t pos cpdef write(self, bytes b, bint nested=*) cpdef write_byte(self, unsigned char val) cpdef write_var_int64(self, libc.stdint.int64_t v) cpdef write_bigendian_int64(self, libc.stdint.int64_t signed_v) cpdef write_bigendian_int32(self, libc.stdint.int32_t signed_v) cpdef write_bigendian_double(self, double d) cpdef bytes get(self) cdef extend(self, size_t missing) cdef class ByteCountingOutputStream(OutputStream): cdef size_t count cpdef write(self, bytes b, bint nested=*) cpdef write_byte(self, unsigned char val) cpdef write_bigendian_int64(self, libc.stdint.int64_t val) cpdef write_bigendian_int32(self, libc.stdint.int32_t val) cpdef size_t get_count(self) cpdef bytes get(self) cdef class InputStream(object): cdef size_t pos cdef bytes all cdef char* allc cpdef size_t size(self) except? -1 cpdef bytes read(self, size_t len) cpdef long read_byte(self) except? -1 cpdef libc.stdint.int64_t read_var_int64(self) except? -1 cpdef libc.stdint.int64_t read_bigendian_int64(self) except? -1 cpdef libc.stdint.int32_t read_bigendian_int32(self) except? -1 cpdef double read_bigendian_double(self) except? -1 cpdef bytes read_all(self, bint nested=*) DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/stream.pyx000066400000000000000000000150511272340251000245200ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. cimport libc.stdlib cimport libc.string cdef class OutputStream(object): """An output string stream implementation supporting write() and get().""" #TODO(robertwb): Consider using raw C++ streams. def __cinit__(self): self.size = 1024 self.pos = 0 self.data = libc.stdlib.malloc(self.size) def __dealloc__(self): if self.data: libc.stdlib.free(self.data) cpdef write(self, bytes b, bint nested=False): cdef size_t blen = len(b) if nested: self.write_var_int64(blen) if blen > self.size - self.pos: self.extend(blen) libc.string.memcpy(self.data + self.pos, b, blen) self.pos += blen cpdef write_byte(self, unsigned char val): if self.size <= self.pos: self.extend(1) self.data[self.pos] = val self.pos += 1 cpdef write_var_int64(self, libc.stdint.int64_t signed_v): """Encode a long using variable-length encoding to a stream.""" cdef libc.stdint.uint64_t v = signed_v cdef long bits while True: bits = v & 0x7F v >>= 7 if v: bits |= 0x80 self.write_byte(bits) if not v: break cpdef write_bigendian_int64(self, libc.stdint.int64_t signed_v): cdef libc.stdint.uint64_t v = signed_v if self.size < self.pos - 8: self.extend(8) self.data[self.pos ] = (v >> 56) self.data[self.pos + 1] = (v >> 48) self.data[self.pos + 2] = (v >> 40) self.data[self.pos + 3] = (v >> 32) self.data[self.pos + 4] = (v >> 24) self.data[self.pos + 5] = (v >> 16) self.data[self.pos + 6] = (v >> 8) self.data[self.pos + 7] = (v ) self.pos += 8 cpdef write_bigendian_int32(self, libc.stdint.int32_t signed_v): cdef libc.stdint.uint32_t v = signed_v if self.size < self.pos - 4: self.extend(4) self.data[self.pos ] = (v >> 24) self.data[self.pos + 1] = (v >> 16) self.data[self.pos + 2] = (v >> 8) self.data[self.pos + 3] = (v ) self.pos += 4 cpdef write_bigendian_double(self, double d): self.write_bigendian_int64((&d)[0]) cpdef bytes get(self): return self.data[:self.pos] cdef extend(self, size_t missing): while missing > self.size - self.pos: self.size *= 2 self.data = libc.stdlib.realloc(self.data, self.size) cdef class ByteCountingOutputStream(OutputStream): """An output string stream implementation that only counts the bytes. This implementation counts the number of bytes it "writes" but doesn't actually write them anyway. Thus it has write() but not get(). get_count() returns how many bytes were written. This is useful for sizing an encoding. """ def __cinit__(self): self.count = 0 cpdef write(self, bytes b, bint nested=False): cdef size_t blen = len(b) if nested: self.write_var_int64(blen) self.count += blen cpdef write_byte(self, unsigned char _): self.count += 1 cpdef write_bigendian_int64(self, libc.stdint.int64_t _): self.count += 8 cpdef write_bigendian_int32(self, libc.stdint.int32_t _): self.count += 4 cpdef size_t get_count(self): return self.count cpdef bytes get(self): raise NotImplementedError def __str__(self): return '<%s %s>' % (self.__class__.__name__, self.count) cdef class InputStream(object): """An input string stream implementation supporting read() and size().""" def __init__(self, all): self.allc = self.all = all cpdef bytes read(self, size_t size): self.pos += size return self.allc[self.pos - size : self.pos] cpdef long read_byte(self) except? -1: self.pos += 1 # Note: the C++ compiler on Dataflow workers treats the char array below as # a signed char. This causes incorrect coder behavior unless explicitly # cast to an unsigned char here. return ( self.allc[self.pos - 1]) cpdef size_t size(self) except? -1: return len(self.all) - self.pos cpdef bytes read_all(self, bint nested=False): return self.read(self.read_var_int64() if nested else self.size()) cpdef libc.stdint.int64_t read_var_int64(self) except? -1: """Decode a variable-length encoded long from a stream.""" cdef long byte cdef long bits cdef long shift = 0 cdef libc.stdint.int64_t result = 0 while True: byte = self.read_byte() if byte < 0: raise RuntimeError('VarInt not terminated.') bits = byte & 0x7F if (shift >= sizeof(long) * 8 or (shift >= (sizeof(long) * 8 - 1) and bits > 1)): raise RuntimeError('VarLong too long.') result |= bits << shift shift += 7 if not (byte & 0x80): break return result cpdef libc.stdint.int64_t read_bigendian_int64(self) except? -1: self.pos += 8 return (self.allc[self.pos - 1] | self.allc[self.pos - 2] << 8 | self.allc[self.pos - 3] << 16 | self.allc[self.pos - 4] << 24 | self.allc[self.pos - 5] << 32 | self.allc[self.pos - 6] << 40 | self.allc[self.pos - 7] << 48 | self.allc[self.pos - 8] << 56) cpdef libc.stdint.int32_t read_bigendian_int32(self) except? -1: self.pos += 4 return (self.allc[self.pos - 1] | self.allc[self.pos - 2] << 8 | self.allc[self.pos - 3] << 16 | self.allc[self.pos - 4] << 24) cpdef double read_bigendian_double(self) except? -1: cdef libc.stdint.int64_t as_long = self.read_bigendian_int64() return (&as_long)[0] DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/stream_test.py000066400000000000000000000120541272340251000253670ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for the stream implementations.""" import logging import math import unittest from google.cloud.dataflow.coders import slow_stream class StreamTest(unittest.TestCase): # pylint: disable=invalid-name InputStream = slow_stream.InputStream OutputStream = slow_stream.OutputStream ByteCountingOutputStream = slow_stream.ByteCountingOutputStream # pylint: enable=invalid-name def test_read_write(self): out_s = self.OutputStream() out_s.write('abc') out_s.write('\0\t\n') out_s.write('xyz', True) out_s.write('', True) in_s = self.InputStream(out_s.get()) self.assertEquals('abc\0\t\n', in_s.read(6)) self.assertEquals('xyz', in_s.read_all(True)) self.assertEquals('', in_s.read_all(True)) def test_read_all(self): out_s = self.OutputStream() out_s.write('abc') in_s = self.InputStream(out_s.get()) self.assertEquals('abc', in_s.read_all(False)) def test_read_write_byte(self): out_s = self.OutputStream() out_s.write_byte(1) out_s.write_byte(0) out_s.write_byte(0xFF) in_s = self.InputStream(out_s.get()) self.assertEquals(1, in_s.read_byte()) self.assertEquals(0, in_s.read_byte()) self.assertEquals(0xFF, in_s.read_byte()) def run_read_write_var_int64(self, values): out_s = self.OutputStream() for v in values: out_s.write_var_int64(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEquals(v, in_s.read_var_int64()) def test_small_var_int64(self): self.run_read_write_var_int64(range(-10, 30)) def test_medium_var_int64(self): base = -1.7 self.run_read_write_var_int64( [int(base**pow) for pow in range(1, int(63 * math.log(2) / math.log(-base)))]) def test_large_var_int64(self): self.run_read_write_var_int64([0, 2**63 - 1, -2**63, 2**63 - 3]) def test_read_write_double(self): values = 0, 1, -1, 1e100, 1.0/3, math.pi, float('inf') out_s = self.OutputStream() for v in values: out_s.write_bigendian_double(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEquals(v, in_s.read_bigendian_double()) def test_read_write_bigendian_int64(self): values = 0, 1, -1, 2**63-1, -2**63, int(2**61 * math.pi) out_s = self.OutputStream() for v in values: out_s.write_bigendian_int64(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEquals(v, in_s.read_bigendian_int64()) def test_read_write_bigendian_int32(self): values = 0, 1, -1, 2**31-1, -2**31, int(2**29 * math.pi) out_s = self.OutputStream() for v in values: out_s.write_bigendian_int32(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEquals(v, in_s.read_bigendian_int32()) def test_byte_counting(self): bc_s = self.ByteCountingOutputStream() self.assertEquals(0, bc_s.get_count()) bc_s.write('def') self.assertEquals(3, bc_s.get_count()) bc_s.write('') self.assertEquals(3, bc_s.get_count()) bc_s.write_byte(10) self.assertEquals(4, bc_s.get_count()) # "nested" also writes the length of the string, which should # cause 1 extra byte to be counted. bc_s.write('2345', nested=True) self.assertEquals(9, bc_s.get_count()) bc_s.write_var_int64(63) self.assertEquals(10, bc_s.get_count()) bc_s.write_bigendian_int64(42) self.assertEquals(18, bc_s.get_count()) bc_s.write_bigendian_int32(36) self.assertEquals(22, bc_s.get_count()) bc_s.write_bigendian_double(6.25) self.assertEquals(30, bc_s.get_count()) try: # pylint: disable=g-import-not-at-top from google.cloud.dataflow.coders import stream class FastStreamTest(StreamTest): """Runs the test with the compiled stream classes.""" InputStream = stream.InputStream OutputStream = stream.OutputStream ByteCountingOutputStream = stream.ByteCountingOutputStream class SlowFastStreamTest(StreamTest): """Runs the test with compiled and uncompiled stream classes.""" InputStream = stream.InputStream OutputStream = slow_stream.OutputStream ByteCountingOutputStream = slow_stream.ByteCountingOutputStream class FastSlowStreamTest(StreamTest): """Runs the test with uncompiled and compiled stream classes.""" InputStream = slow_stream.InputStream OutputStream = stream.OutputStream ByteCountingOutputStream = stream.ByteCountingOutputStream except ImportError: pass if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/typecoders.py000066400000000000000000000153061272340251000252210ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Type coders registration. This module contains functionality to define and use coders for custom classes. Let's say we have a class Xyz and we are processing a PCollection with elements of type Xyz. If we do not register a coder for Xyz, a default pickle-based fallback coder will be used. This can be undesirable for two reasons. First, we may want a faster coder or a more space efficient one. Second, the pickle-based coder is not deterministic in the sense that objects like dictionaries or sets are not guaranteed to be encoded in the same way every time (elements are not really ordered). Two (sometimes three) steps are needed to define and use a custom coder: - define the coder class - associate the code with the class (a.k.a. coder registration) - typehint DoFns or transforms with the new class or composite types using the class. A coder class is defined by subclassing from CoderBase and defining the encode_to_bytes and decode_from_bytes methods. The framework uses duck-typing for coders so it is not strictly required to subclass from CoderBase as long as the encode/decode methods are defined. Registering a coder class is made with a register_coder() call:: from google.cloud.dataflow import coders ... coders.registry.register_coder(Xyz, XyzCoder) Additionally, DoFns and PTransforms may need type hints. This is not always necessary since there is functionality to infer the return types of DoFns by analyzing the code. For instance, for the function below the return type of 'Xyz' will be inferred:: def MakeXyzs(v): return Xyz(v) If Xyz is inferred then its coder will be used whenever the framework needs to serialize data (e.g., writing to the shuffler subsystem responsible for group by key operations). If a typehint is needed it can be specified by decorating the DoFns or using with_input_types/with_output_types methods on PTransforms. For example, the above function can be decorated:: @with_output_types(Xyz) def MakeXyzs(v): return complex_operation_returning_Xyz(v) See google.cloud.dataflow.typehints.decorators module for more details. """ import logging from google.cloud.dataflow.coders import coders from google.cloud.dataflow.typehints import typehints class CoderRegistry(object): """A coder registry for typehint/coder associations.""" def __init__(self, fallback_coder=None): self._coders = {} self.custom_types = [] self.register_standard_coders(fallback_coder) def register_standard_coders(self, fallback_coder): """Register coders for all basic and composite types.""" self._register_coder_internal(int, coders.VarIntCoder) self._register_coder_internal(float, coders.FloatCoder) self._register_coder_internal(str, coders.BytesCoder) self._register_coder_internal(bytes, coders.BytesCoder) self._register_coder_internal(unicode, coders.StrUtf8Coder) self._register_coder_internal(typehints.TupleConstraint, coders.TupleCoder) self._register_coder_internal(typehints.AnyTypeConstraint, coders.PickleCoder) self._fallback_coder = fallback_coder or coders.PickleCoder def _register_coder_internal(self, typehint_type, typehint_coder_class): self._coders[typehint_type] = typehint_coder_class def register_coder(self, typehint_type, typehint_coder_class): if not isinstance(typehint_coder_class, type): raise TypeError('Coder registration requires a coder class object. ' 'Received %r instead.' % typehint_coder_class) if typehint_type not in self.custom_types: self.custom_types.append(typehint_type) self._register_coder_internal(typehint_type, typehint_coder_class) def get_coder(self, typehint): coder = self._coders.get( typehint.__class__ if isinstance(typehint, typehints.TypeConstraint) else typehint, None) if isinstance(typehint, typehints.TypeConstraint) and coder is not None: return coder.from_type_hint(typehint, self) if coder is None: # We use the fallback coder when there is no coder registered for a # typehint. For example a user defined class with no coder specified. if not hasattr(self, '_fallback_coder'): raise RuntimeError( 'Coder registry has no fallback coder. This can happen if the ' 'fast_coders module could not be imported.') if isinstance(typehint, typehints.IterableTypeConstraint): # In this case, we suppress the warning message for using the fallback # coder, since Iterable is hinted as the output of a GroupByKey # operation and that direct output will not be coded. # TODO(ccy): refine this behavior. pass elif typehint is None: # In some old code, None is used for Any. # TODO(robertwb): Clean this up. pass elif isinstance(typehint, typehints.TypeVariable): # TODO(robertwb): Clean this up when type inference is fully enabled. pass else: logging.warning('Using fallback coder for typehint: %r.', typehint) coder = self._fallback_coder return coder.from_type_hint(typehint, self) def get_custom_type_coder_tuples(self, types): """Returns type/coder tuples for all custom types passed in.""" return [(t, self._coders[t]) for t in types if t in self.custom_types] def verify_deterministic(self, key_coder, op_name, silent=True): if not key_coder.is_deterministic(): error_msg = ('The key coder "%s" for %s ' 'is not deterministic. This may result in incorrect ' 'pipeline output. This can be fixed by adding a type ' 'hint to the operation preceding the GroupByKey step, ' 'and for custom key classes, by writing a ' 'deterministic custom Coder. Please see the ' 'documentation for more details.' % (key_coder, op_name)) if isinstance(key_coder, (coders.PickleCoder, self._fallback_coder)): if not silent: logging.warning(error_msg) return coders.DeterministicPickleCoder(key_coder, op_name) else: raise ValueError(error_msg) else: return key_coder registry = CoderRegistry() DataflowPythonSDK-0.2.5/google/cloud/dataflow/coders/typecoders_test.py000066400000000000000000000101021272340251000262450ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the typecoders module.""" import unittest from google.cloud.dataflow.coders import coders from google.cloud.dataflow.coders import typecoders from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.typehints import typehints class CustomClass(object): def __init__(self, n): self.number = n def __eq__(self, other): return self.number == other.number class CustomCoder(coders.Coder): def encode(self, value): return str(value.number) def decode(self, encoded): return CustomClass(int(encoded)) def is_deterministic(self): # This coder is deterministic. Though we don't use need this coder to be # deterministic for this test, we annotate this as such to follow best # practices. return True class TypeCodersTest(unittest.TestCase): def test_register_non_type_coder(self): coder = CustomCoder() with self.assertRaises(TypeError) as e: # When registering a coder the coder class must be specified. typecoders.registry.register_coder(CustomClass, coder) self.assertEqual(e.exception.message, 'Coder registration requires a coder class object. ' 'Received %r instead.' % coder) def test_get_coder_with_custom_coder(self): typecoders.registry.register_coder(CustomClass, CustomCoder) self.assertEqual(CustomCoder, typecoders.registry.get_coder(CustomClass).__class__) def test_get_coder_with_composite_custom_coder(self): typecoders.registry.register_coder(CustomClass, CustomCoder) coder = typecoders.registry.get_coder(typehints.KV[CustomClass, str]) revived_coder = pickler.loads(pickler.dumps(coder)) self.assertEqual( (CustomClass(123), 'abc'), revived_coder.decode(revived_coder.encode((CustomClass(123), 'abc')))) def test_get_coder_with_standard_coder(self): self.assertEqual(coders.BytesCoder, typecoders.registry.get_coder(str).__class__) def test_fallbackcoder(self): coder = typecoders.registry.get_coder(typehints.Any) self.assertEqual(('abc', 123), coder.decode(coder.encode(('abc', 123)))) def test_get_coder_can_be_pickled(self): coder = typecoders.registry.get_coder(typehints.Tuple[str, int]) revived_coder = pickler.loads(pickler.dumps(coder)) self.assertEqual(('abc', 123), revived_coder.decode(revived_coder.encode(('abc', 123)))) def test_standard_int_coder(self): real_coder = typecoders.registry.get_coder(int) expected_coder = coders.VarIntCoder() self.assertEqual( real_coder.encode(0x0404), expected_coder.encode(0x0404)) self.assertEqual(0x0404, real_coder.decode(real_coder.encode(0x0404))) self.assertEqual( real_coder.encode(0x040404040404), expected_coder.encode(0x040404040404)) self.assertEqual(0x040404040404, real_coder.decode(real_coder.encode(0x040404040404))) def test_standard_str_coder(self): real_coder = typecoders.registry.get_coder(str) expected_coder = coders.BytesCoder() self.assertEqual( real_coder.encode('abc'), expected_coder.encode('abc')) self.assertEqual('abc', real_coder.decode(real_coder.encode('abc'))) real_coder = typecoders.registry.get_coder(bytes) expected_coder = coders.BytesCoder() self.assertEqual( real_coder.encode('abc'), expected_coder.encode('abc')) self.assertEqual('abc', real_coder.decode(real_coder.encode('abc'))) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/dataflow_test.py000066400000000000000000000372431272340251000244250ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Integration tests for the dataflow package.""" from __future__ import absolute_import import logging import re import unittest from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.pvalue import AsDict from google.cloud.dataflow.pvalue import AsIter as AllOf from google.cloud.dataflow.pvalue import AsList from google.cloud.dataflow.pvalue import AsSingleton from google.cloud.dataflow.pvalue import EmptySideInput from google.cloud.dataflow.pvalue import SideOutputValue from google.cloud.dataflow.transforms import Create from google.cloud.dataflow.transforms import DoFn from google.cloud.dataflow.transforms import FlatMap from google.cloud.dataflow.transforms import GroupByKey from google.cloud.dataflow.transforms import Map from google.cloud.dataflow.transforms import ParDo from google.cloud.dataflow.transforms import WindowInto from google.cloud.dataflow.transforms.util import assert_that from google.cloud.dataflow.transforms.util import equal_to from google.cloud.dataflow.transforms.window import IntervalWindow from google.cloud.dataflow.transforms.window import WindowFn class DataflowTest(unittest.TestCase): """Dataflow integration tests.""" SAMPLE_DATA = 'aa bb cc aa bb aa \n' * 10 SAMPLE_RESULT = [('cc', 10), ('bb', 20), ('aa', 30)] # TODO(silviuc): Figure out a nice way to specify labels for stages so that # internal steps get prepended with surorunding stage names. @staticmethod def Count(pcoll): # pylint: disable=invalid-name """A Count transform: v, ... => (v, n), ...""" return (pcoll | Map('AddCount', lambda x: (x, 1)) | GroupByKey('GroupCounts') | Map('AddCounts', lambda (x, ones): (x, sum(ones)))) def test_word_count(self): pipeline = Pipeline('DirectPipelineRunner') lines = pipeline | Create('SomeWords', [DataflowTest.SAMPLE_DATA]) result = ( (lines | FlatMap('GetWords', lambda x: re.findall(r'\w+', x))) .apply('CountWords', DataflowTest.Count)) assert_that(result, equal_to(DataflowTest.SAMPLE_RESULT)) pipeline.run() def test_map(self): pipeline = Pipeline('DirectPipelineRunner') lines = pipeline | Create('input', ['a', 'b', 'c']) result = (lines | Map('upper', str.upper) | Map('prefix', lambda x, prefix: prefix + x, 'foo-')) assert_that(result, equal_to(['foo-A', 'foo-B', 'foo-C'])) pipeline.run() def test_word_count_using_get(self): pipeline = Pipeline('DirectPipelineRunner') lines = pipeline | Create('SomeWords', [DataflowTest.SAMPLE_DATA]) result = ( (lines | FlatMap('GetWords', lambda x: re.findall(r'\w+', x))) .apply('CountWords', DataflowTest.Count)) assert_that(result, equal_to(DataflowTest.SAMPLE_RESULT)) pipeline.run() def test_par_do_with_side_input_as_arg(self): pipeline = Pipeline('DirectPipelineRunner') words_list = ['aa', 'bb', 'cc'] words = pipeline | Create('SomeWords', words_list) prefix = pipeline | Create('SomeString', ['xyz']) # side in suffix = 'zyx' result = words | FlatMap( 'DecorateWords', lambda x, pfx, sfx: ['%s-%s-%s' % (pfx, x, sfx)], AsSingleton(prefix), suffix) assert_that(result, equal_to(['xyz-%s-zyx' % x for x in words_list])) pipeline.run() def test_par_do_with_side_input_as_keyword_arg(self): pipeline = Pipeline('DirectPipelineRunner') words_list = ['aa', 'bb', 'cc'] words = pipeline | Create('SomeWords', words_list) prefix = 'zyx' suffix = pipeline | Create('SomeString', ['xyz']) # side in result = words | FlatMap( 'DecorateWords', lambda x, pfx, sfx: ['%s-%s-%s' % (pfx, x, sfx)], prefix, sfx=AsSingleton(suffix)) assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list])) pipeline.run() def test_par_do_with_do_fn_object(self): class SomeDoFn(DoFn): """A custom DoFn for a FlatMap transform.""" def process(self, context, prefix, suffix): return ['%s-%s-%s' % (prefix, context.element, suffix)] pipeline = Pipeline('DirectPipelineRunner') words_list = ['aa', 'bb', 'cc'] words = pipeline | Create('SomeWords', words_list) prefix = 'zyx' suffix = pipeline | Create('SomeString', ['xyz']) # side in result = words | ParDo('DecorateWordsDoFn', SomeDoFn(), prefix, suffix=AsSingleton(suffix)) assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list])) pipeline.run() def test_par_do_with_multiple_outputs_and_using_yield(self): class SomeDoFn(DoFn): """A custom DoFn using yield.""" def process(self, context): yield context.element if context.element % 2 == 0: yield SideOutputValue('even', context.element) else: yield SideOutputValue('odd', context.element) pipeline = Pipeline('DirectPipelineRunner') nums = pipeline | Create('Some Numbers', [1, 2, 3, 4]) results = nums | ParDo( 'ClassifyNumbers', SomeDoFn()).with_outputs('odd', 'even', main='main') assert_that(results.main, equal_to([1, 2, 3, 4])) assert_that(results.odd, equal_to([1, 3]), label='assert:odd') assert_that(results.even, equal_to([2, 4]), label='assert:even') pipeline.run() def test_par_do_with_multiple_outputs_and_using_return(self): def some_fn(v): if v % 2 == 0: return [v, SideOutputValue('even', v)] else: return [v, SideOutputValue('odd', v)] pipeline = Pipeline('DirectPipelineRunner') nums = pipeline | Create('Some Numbers', [1, 2, 3, 4]) results = nums | FlatMap( 'ClassifyNumbers', some_fn).with_outputs('odd', 'even', main='main') assert_that(results.main, equal_to([1, 2, 3, 4])) assert_that(results.odd, equal_to([1, 3]), label='assert:odd') assert_that(results.even, equal_to([2, 4]), label='assert:even') pipeline.run() def test_empty_singleton_side_input(self): pipeline = Pipeline('DirectPipelineRunner') pcol = pipeline | Create('start', [1, 2]) side = pipeline | Create('side', []) # Empty side input. def my_fn(k, s): v = ('empty' if isinstance(s, EmptySideInput) else 'full') return [(k, v)] result = pcol | FlatMap('compute', my_fn, AsSingleton(side)) assert_that(result, equal_to([(1, 'empty'), (2, 'empty')])) pipeline.run() def test_multi_valued_singleton_side_input(self): pipeline = Pipeline('DirectPipelineRunner') pcol = pipeline | Create('start', [1, 2]) side = pipeline | Create('side', [3, 4]) # 2 values in side input. pcol | FlatMap('compute', lambda x, s: [x * s], AsSingleton(side)) with self.assertRaises(ValueError) as e: pipeline.run() def test_default_value_singleton_side_input(self): pipeline = Pipeline('DirectPipelineRunner') pcol = pipeline | Create('start', [1, 2]) side = pipeline | Create('side', []) # 0 values in side input. result = ( pcol | FlatMap('compute', lambda x, s: [x * s], AsSingleton(side, 10))) assert_that(result, equal_to([10, 20])) pipeline.run() def test_iterable_side_input(self): pipeline = Pipeline('DirectPipelineRunner') pcol = pipeline | Create('start', [1, 2]) side = pipeline | Create('side', [3, 4]) # 2 values in side input. result = pcol | FlatMap('compute', lambda x, s: [x * y for y in s], AllOf(side)) assert_that(result, equal_to([3, 4, 6, 8])) pipeline.run() def test_undeclared_side_outputs(self): pipeline = Pipeline('DirectPipelineRunner') nums = pipeline | Create('Some Numbers', [1, 2, 3, 4]) results = nums | FlatMap( 'ClassifyNumbers', lambda x: [x, SideOutputValue('even' if x % 2 == 0 else 'odd', x)] ).with_outputs() # TODO(silviuc): Revisit this test to check for undeclared side outputs. # This should work with .with_outputs() without any tags declared and # the results[None] should work also. assert_that(results[None], equal_to([1, 2, 3, 4])) assert_that(results.odd, equal_to([1, 3]), label='assert:odd') assert_that(results.even, equal_to([2, 4]), label='assert:even') pipeline.run() def test_empty_side_outputs(self): pipeline = Pipeline('DirectPipelineRunner') nums = pipeline | Create('Some Numbers', [1, 3, 5]) results = nums | FlatMap( 'ClassifyNumbers', lambda x: [x, SideOutputValue('even' if x % 2 == 0 else 'odd', x)] ).with_outputs() assert_that(results[None], equal_to([1, 3, 5])) assert_that(results.odd, equal_to([1, 3, 5]), label='assert:odd') assert_that(results.even, equal_to([]), label='assert:even') pipeline.run() def test_as_list_and_as_dict_side_inputs(self): a_list = [5, 1, 3, 2, 9] some_pairs = [('crouton', 17), ('supreme', None)] pipeline = Pipeline('DirectPipelineRunner') main_input = pipeline | Create('main input', [1]) side_list = pipeline | Create('side list', a_list) side_pairs = pipeline | Create('side pairs', some_pairs) results = main_input | FlatMap( 'concatenate', lambda x, the_list, the_dict: [[x, the_list, the_dict]], AsList(side_list), AsDict(side_pairs)) def matcher(expected_elem, expected_list, expected_pairs): def match(actual): [[actual_elem, actual_list, actual_dict]] = actual equal_to([expected_elem])([actual_elem]) equal_to(expected_list)(actual_list) equal_to(expected_pairs)(actual_dict.iteritems()) return match assert_that(results, matcher(1, a_list, some_pairs)) pipeline.run() def test_as_singleton_without_unique_labels(self): # This should succeed as calling AsSingleton on the same PCollection twice # with the same defaults will return the same PCollectionView. a_list = [2] pipeline = Pipeline('DirectPipelineRunner') main_input = pipeline | Create('main input', [1]) side_list = pipeline | Create('side list', a_list) results = main_input | FlatMap( 'test', lambda x, s1, s2: [[x, s1, s2]], AsSingleton(side_list), AsSingleton(side_list)) def matcher(expected_elem, expected_singleton): def match(actual): [[actual_elem, actual_singleton1, actual_singleton2]] = actual equal_to([expected_elem])([actual_elem]) equal_to([expected_singleton])([actual_singleton1]) equal_to([expected_singleton])([actual_singleton2]) return match assert_that(results, matcher(1, 2)) pipeline.run() def test_as_singleton_with_different_defaults_without_unique_labels(self): # This should fail as AsSingleton with distinct default values should create # distinct PCollectionViews with the same full_label. a_list = [2] pipeline = Pipeline('DirectPipelineRunner') main_input = pipeline | Create('main input', [1]) side_list = pipeline | Create('side list', a_list) with self.assertRaises(RuntimeError) as e: _ = main_input | FlatMap( 'test', lambda x, s1, s2: [[x, s1, s2]], AsSingleton(side_list), AsSingleton(side_list, default_value=3)) self.assertTrue( e.exception.message.startswith( 'Transform "ViewAsSingleton(side list.None)" does not have a ' 'stable unique label.')) def test_as_singleton_with_different_defaults_with_unique_labels(self): a_list = [] pipeline = Pipeline('DirectPipelineRunner') main_input = pipeline | Create('main input', [1]) side_list = pipeline | Create('side list', a_list) results = main_input | FlatMap( 'test', lambda x, s1, s2: [[x, s1, s2]], AsSingleton('si1', side_list, default_value=2), AsSingleton('si2', side_list, default_value=3)) def matcher(expected_elem, expected_singleton1, expected_singleton2): def match(actual): [[actual_elem, actual_singleton1, actual_singleton2]] = actual equal_to([expected_elem])([actual_elem]) equal_to([expected_singleton1])([actual_singleton1]) equal_to([expected_singleton2])([actual_singleton2]) return match assert_that(results, matcher(1, 2, 3)) pipeline.run() def test_as_list_without_unique_labels(self): # This should succeed as calling AsList on the same PCollection twice will # return the same PCollectionView. a_list = [1, 2, 3] pipeline = Pipeline('DirectPipelineRunner') main_input = pipeline | Create('main input', [1]) side_list = pipeline | Create('side list', a_list) results = main_input | FlatMap( 'test', lambda x, ls1, ls2: [[x, ls1, ls2]], AsList(side_list), AsList(side_list)) def matcher(expected_elem, expected_list): def match(actual): [[actual_elem, actual_list1, actual_list2]] = actual equal_to([expected_elem])([actual_elem]) equal_to(expected_list)(actual_list1) equal_to(expected_list)(actual_list2) return match assert_that(results, matcher(1, [1, 2, 3])) pipeline.run() def test_as_list_with_unique_labels(self): a_list = [1, 2, 3] pipeline = Pipeline('DirectPipelineRunner') main_input = pipeline | Create('main input', [1]) side_list = pipeline | Create('side list', a_list) results = main_input | FlatMap( 'test', lambda x, ls1, ls2: [[x, ls1, ls2]], AsList(side_list), AsList(side_list, label='label')) def matcher(expected_elem, expected_list): def match(actual): [[actual_elem, actual_list1, actual_list2]] = actual equal_to([expected_elem])([actual_elem]) equal_to(expected_list)(actual_list1) equal_to(expected_list)(actual_list2) return match assert_that(results, matcher(1, [1, 2, 3])) pipeline.run() def test_as_dict_with_unique_labels(self): some_kvs = [('a', 1), ('b', 2)] pipeline = Pipeline('DirectPipelineRunner') main_input = pipeline | Create('main input', [1]) side_kvs = pipeline | Create('side kvs', some_kvs) results = main_input | FlatMap( 'test', lambda x, dct1, dct2: [[x, dct1, dct2]], AsDict(side_kvs), AsDict(side_kvs, label='label')) def matcher(expected_elem, expected_kvs): def match(actual): [[actual_elem, actual_dict1, actual_dict2]] = actual equal_to([expected_elem])([actual_elem]) equal_to(expected_kvs)(actual_dict1.iteritems()) equal_to(expected_kvs)(actual_dict2.iteritems()) return match assert_that(results, matcher(1, some_kvs)) pipeline.run() def test_window_transform(self): class TestWindowFn(WindowFn): """Windowing function adding two disjoint windows to each element.""" def assign(self, assign_context): _ = assign_context return [IntervalWindow(10, 20), IntervalWindow(20, 30)] def merge(self, existing_windows): return existing_windows pipeline = Pipeline('DirectPipelineRunner') numbers = pipeline | Create('KVs', [(1, 10), (2, 20), (3, 30)]) result = (numbers | WindowInto('W', windowfn=TestWindowFn()) | GroupByKey('G')) assert_that( result, equal_to([(1, [10]), (1, [10]), (2, [20]), (2, [20]), (3, [30]), (3, [30])])) pipeline.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/error.py000066400000000000000000000023371272340251000227120ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Python Dataflow error classes.""" class DataflowError(Exception): """Base class for all Dataflow errors.""" class PipelineError(DataflowError): """An error in the pipeline object (e.g. a PValue not linked to it).""" class PValueError(DataflowError): """An error related to a PValue object (e.g. value is not computed).""" class RunnerError(DataflowError): """An error related to a Runner object (e.g. cannot find a runner to run).""" class SideInputError(DataflowError): """An error related to a side input to a parallel Do operation.""" class TransformError(DataflowError): """An error related to a PTransform object.""" DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/000077500000000000000000000000001272340251000230205ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/__init__.py000066400000000000000000000000001272340251000251170ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/000077500000000000000000000000001272340251000246305ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/autocomplete.py000066400000000000000000000045421272340251000277100ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A workflow emitting the top k most common words for each prefix.""" from __future__ import absolute_import import argparse import logging import re import google.cloud.dataflow as df def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file to process.') parser.add_argument('--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) (p # pylint: disable=expression-not-assigned | df.io.Read('read', df.io.TextFileSource(known_args.input)) | df.FlatMap('split', lambda x: re.findall(r'[A-Za-z\']+', x)) | TopPerPrefix('TopPerPrefix', 5) | df.Map('format', lambda (prefix, candidates): '%s: %s' % (prefix, candidates)) | df.io.Write('write', df.io.TextFileSink(known_args.output))) p.run() class TopPerPrefix(df.PTransform): def __init__(self, label, count): super(TopPerPrefix, self).__init__(label) self._count = count def apply(self, words): """Compute the most common words for each possible prefixes. Args: words: a PCollection of strings Returns: A PCollection of most common words with each prefix, in the form (prefix, [(count, word), (count, word), ...]) """ return (words | df.combiners.Count.PerElement() | df.FlatMap(extract_prefixes) | df.combiners.Top.LargestPerKey(self._count)) def extract_prefixes((word, count)): for k in range(1, len(word) + 1): prefix = word[:k] yield prefix, (count, word) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/autocomplete_test.py000066400000000000000000000045521272340251000307500ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the wordcount example.""" import collections import logging import re import tempfile import unittest import google.cloud.dataflow as df from google.cloud.dataflow.examples.complete import autocomplete from google.cloud.dataflow.pvalue import AsIter from google.cloud.dataflow.utils import options # TODO(robertwb): Move to testing utilities. def assert_that(pcoll, matcher): """Asserts that the give PCollection satisfies the constraints of the matcher in a way that is runnable locally or on a remote service. """ singleton = pcoll.pipeline | df.Create('create_singleton', [None]) def check_matcher(_, side_value): assert matcher(side_value) return [] singleton | df.FlatMap(check_matcher, AsIter(pcoll)) def contains_in_any_order(expected): def matcher(value): vs = collections.Counter(value) es = collections.Counter(expected) if vs != es: raise ValueError( 'extra: %s, missing: %s' % (vs - es, es - vs)) return True return matcher class WordCountTest(unittest.TestCase): WORDS = ['this', 'this', 'that', 'to', 'to', 'to'] def test_top_prefixes(self): p = df.Pipeline('DirectPipelineRunner') words = p | df.Create('create', self.WORDS) result = words | autocomplete.TopPerPrefix('test', 5) # values must be hashable for now result = result | df.Map(lambda (k, vs): (k, tuple(vs))) assert_that(result, contains_in_any_order( [ ('t', ((3, 'to'), (2, 'this'), (1, 'that'))), ('to', ((3, 'to'), )), ('th', ((2, 'this'), (1, 'that'))), ('thi', ((2, 'this'), )), ('this', ((2, 'this'), )), ('tha', ((1, 'that'), )), ('that', ((1, 'that'), )), ])) p.run() if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/estimate_pi.py000066400000000000000000000070071272340251000275110ustar00rootroot00000000000000# -*- coding: utf-8 -*- # Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A workflow that uses a simple Monte Carlo method to estimate π. The algorithm computes the fraction of points drawn uniformly within the unit square that also fall in the quadrant of the unit circle that overlaps the square. A simple area calculation shows that this fraction should be π/4, so we multiply our counts ratio by four to estimate π. """ from __future__ import absolute_import import argparse import json import logging import random import google.cloud.dataflow as df from google.cloud.dataflow.typehints import Any from google.cloud.dataflow.typehints import Iterable from google.cloud.dataflow.typehints import Tuple @df.typehints.with_output_types(Tuple[int, int, int]) @df.typehints.with_input_types(int) def run_trials(runs): """Run trials and return a 3-tuple representing the results. Args: runs: Number of trial runs to be executed. Returns: A 3-tuple (total trials, inside trials, 0). The final zero is needed solely to make sure that the combine_results function has same type for inputs and outputs (a requirement for combiner functions). """ inside_runs = 0 for _ in xrange(runs): x = random.uniform(0, 1) y = random.uniform(0, 1) inside_runs += 1 if x * x + y * y <= 1.0 else 0 return runs, inside_runs, 0 @df.typehints.with_output_types(Tuple[int, int, float]) @df.typehints.with_input_types(Iterable[Tuple[int, int, Any]]) def combine_results(results): """Combiner function to sum up trials and compute the estimate. Args: results: An iterable of 3-tuples (total trials, inside trials, ignored). Returns: A 3-tuple containing the sum of total trials, sum of inside trials, and the probability computed from the two numbers. """ # TODO(silviuc): Do we guarantee that argument can be iterated repeatedly? # Should document one way or the other. total, inside = sum(r[0] for r in results), sum(r[1] for r in results) return total, inside, 4 * float(inside) / total class JsonCoder(object): """A JSON coder used to format the final result.""" def encode(self, x): return json.dumps(x) def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument('--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # A thousand work items of a million tries each. (p # pylint: disable=expression-not-assigned | df.Create('Initialize', [100000] * 100).with_output_types(int) | df.Map('Run trials', run_trials) | df.CombineGlobally('Sum', combine_results).without_defaults() | df.io.Write('Write', df.io.TextFileSink(known_args.output, coder=JsonCoder()))) # Actually run the pipeline (all operations above are deferred). p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/estimate_pi_test.py000066400000000000000000000030231272340251000305420ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the estimate_pi example.""" import json import logging import tempfile import unittest from google.cloud.dataflow.examples.complete import estimate_pi class EstimatePiTest(unittest.TestCase): def create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def test_basics(self): temp_path = self.create_temp_file('result') estimate_pi.run([ '--output=%s' % temp_path]) # Parse result file and compare. with open(temp_path + '-00000-of-00001') as result_file: estimated_pi = json.loads(result_file.readline())[2] # Note: Probabilistically speaking this test can fail with a probability # that is very small (VERY) given that we run at least 10 million trials. self.assertTrue(estimated_pi > 3.13 and estimated_pi < 3.15) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/juliaset/000077500000000000000000000000001272340251000264505ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/juliaset/juliaset/000077500000000000000000000000001272340251000302705ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/juliaset/juliaset/__init__.py000066400000000000000000000000001272340251000323670ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/juliaset/juliaset/juliaset.py000066400000000000000000000102251272340251000324620ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A Julia set computing workflow: https://en.wikipedia.org/wiki/Julia_set. We use the quadratic polinomial f(z) = z*z + c, with c = -.62772 +.42193i """ from __future__ import absolute_import import argparse import google.cloud.dataflow as df def from_pixel(x, y, n): """Converts a NxN pixel position to a (-1..1, -1..1) complex number.""" return complex(2.0 * x / n - 1.0, 2.0 * y / n - 1.0) def get_julia_set_point_color((x, y), c, n, max_iterations): """Given an pixel, convert it into a point in our julia set.""" z = from_pixel(x, y, n) for i in xrange(max_iterations): if z.real * z.real + z.imag * z.imag > 2.0: break z = z * z + c return x, y, i # pylint: disable=undefined-loop-variable def generate_julia_set_colors(pipeline, c, n, max_iterations): """Compute julia set coordinates for each point in our set.""" def point_set(n): for x in range(n): for y in range(n): yield (x, y) julia_set_colors = (pipeline | df.Create('add points', point_set(n)) | df.Map(get_julia_set_point_color, c, n, max_iterations)) return julia_set_colors def generate_julia_set_visualization(data, n, max_iterations): """Generate the pixel matrix for rendering the julia set as an image.""" import numpy as np # pylint: disable=g-import-not-at-top colors = [] for r in range(0, 256, 16): for g in range(0, 256, 16): for b in range(0, 256, 16): colors.append((r, g, b)) xy = np.zeros((n, n, 3), dtype=np.uint8) for x, y, iteration in data: xy[x, y] = colors[iteration * len(colors) / max_iterations] return xy def save_julia_set_visualization(out_file, image_array): """Save the fractal image of our julia set as a png.""" from matplotlib import pyplot as plt # pylint: disable=g-import-not-at-top plt.imsave(out_file, image_array, format='png') def run(argv=None): # pylint: disable=missing-docstring parser = argparse.ArgumentParser() parser.add_argument('--grid_size', dest='grid_size', default=1000, help='Size of the NxN matrix') parser.add_argument( '--coordinate_output', dest='coordinate_output', required=True, help='Output file to write the color coordinates of the image to.') parser.add_argument('--image_output', dest='image_output', default=None, help='Output file to write the resulting image to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) n = int(known_args.grid_size) coordinates = generate_julia_set_colors(p, complex(-.62772, .42193), n, 100) # Group each coordinate triplet by its x value, then write the coordinates to # the output file with an x-coordinate grouping per line. # pylint: disable=expression-not-assigned # pylint: disable=g-long-lambda (coordinates | df.Map('x coord key', lambda (x, y, i): (x, (x, y, i))) | df.GroupByKey('x coord') | df.Map( 'format', lambda (k, coords): ' '.join('(%s, %s, %s)' % coord for coord in coords)) | df.io.Write('write', df.io.TextFileSink(known_args.coordinate_output))) # pylint: enable=g-long-lambda # pylint: enable=expression-not-assigned p.run() # Optionally render the image and save it to a file. # TODO(silviuc): Add this functionality. # if p.options.image_output is not None: # julia_set_image = generate_julia_set_visualization( # file_with_coordinates, n, 100) # save_julia_set_visualization(p.options.image_output, julia_set_image) DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/juliaset/juliaset/juliaset_test.py000066400000000000000000000052211272340251000335210ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the juliaset example.""" import logging import os import re import tempfile import unittest from google.cloud.dataflow.examples.complete.juliaset.juliaset import juliaset class JuliaSetTest(unittest.TestCase): def setUp(self): self.test_files = {} self.test_files['output_coord_file_name'] = self.generate_temp_file() self.test_files['output_image_file_name'] = self.generate_temp_file() def tearDown(self): for test_file in self.test_files.values(): if os.path.exists(test_file): os.remove(test_file) def generate_temp_file(self): with tempfile.NamedTemporaryFile(delete=False) as temp_file: return temp_file.name def run_example(self, grid_size, image_file_name=None): args = [ '--coordinate_output=%s' % self.test_files['output_coord_file_name'], '--grid_size=%s' % grid_size, ] if image_file_name is not None: args.append('--image_output=%s' % image_file_name) juliaset.run(args) def test_output_file_format(self): grid_size = 5 self.run_example(grid_size) # Parse the results from the file, and ensure it was written in the proper # format. with open(self.test_files['output_coord_file_name'] + '-00000-of-00001') as result_file: output_lines = result_file.readlines() # Should have a line for each x-coordinate. self.assertEqual(grid_size, len(output_lines)) for line in output_lines: coordinates = re.findall(r'(\(\d+, \d+, \d+\))', line) # Should have 5 coordinates on each line. self.assertTrue(coordinates) self.assertEqual(grid_size, len(coordinates)) def test_generate_fractal_image(self): temp_image_file = self.test_files['output_image_file_name'] self.run_example(10, image_file_name=temp_image_file) # Ensure that the image was saved properly. # TODO(silviuc): Reactivate the test when --image_output is supported. # self.assertTrue(os.stat(temp_image_file).st_size > 0) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/juliaset/juliaset_main.py000066400000000000000000000040131272340251000316440ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A Julia set computing workflow: https://en.wikipedia.org/wiki/Julia_set. This example has in the juliaset/ folder all the code needed to execute the workflow. It is organized in this way so that it can be packaged as a Python package and later installed in the VM workers executing the job. The root directory for the example contains just a "driver" script to launch the job and the setup.py file needed to create a package. The advantages for organizing the code is that large projects will naturally evolve beyond just one module and you will have to make sure the additional modules are present in the worker. In Python Dataflow, using the --setup_file option when submitting a job, will trigger creating a source distribution (as if running python setup.py sdist) and then staging the resulting tarball in the staging area. The workers, upon startup, will install the tarball. Below is a complete command line for running the juliaset workflow remotely as an example: python juliaset_main.py \ --job_name juliaset-$USER \ --project YOUR-PROJECT \ --runner BlockingDataflowPipelineRunner \ --setup_file ./setup.py \ --staging_location gs://YOUR-BUCKET/juliaset/staging \ --temp_location gs://YOUR-BUCKET/juliaset/temp \ --coordinate_output gs://YOUR-BUCKET/juliaset/out \ --grid_size 20 \ """ import logging from juliaset import juliaset if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) juliaset.run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/juliaset/setup.py000066400000000000000000000100471272340251000301640ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Setup.py module for the workflow's worker utilities. All the workflow related code is gathered in a package that will be built as a source distribution, staged in the staging area for the workflow being run and then installed in the workers when they start running. This behavior is triggered by specifying the --setup_file command line option when running the workflow for remote execution. """ import subprocess import setuptools from setuptools.command.bdist_egg import bdist_egg as _bdist_egg class bdist_egg(_bdist_egg): # pylint: disable=invalid-name """A bdist_egg command class that will be invoked during package install. The package built using the current setup.py will be staged and later installed in the worker using `easy_install package'. This class will be instantiated during install for this specific scenario and will trigger running the custom commands specified. """ def run(self): self.run_command('CustomCommands') _bdist_egg.run(self) # Some custom command to run during setup. The command is not essential for this # workflow. It is used here as an example. Each command will spawn a child # process. Typically, these commands will include steps to install non-Python # packages. For instance, to install a C++-based library libjpeg62 the following # two commands will have to be added: # # ['apt-get', 'update'], # ['apt-get', '--assume-yes', install', 'libjpeg62'], # # First, note that there is no need to use the sudo command because the setup # script runs with appropriate access. # Second, if apt-get tool is used then the first command needs to be 'apt-get # update' so the tool refreshes itself and initializes links to download # repositories. Without this initial step the other apt-get install commands # will fail with package not found errors. Note also --assume-yes option which # shortcuts the interactive confirmation. # # The output of custom commands (including failures) will be logged in the # worker-startup log. CUSTOM_COMMANDS = [ ['echo', 'Custom command worked!']] class CustomCommands(setuptools.Command): """A setuptools Command class able to run arbitrary commands.""" def initialize_options(self): pass def finalize_options(self): pass def RunCustomCommand(self, command_list): print 'Running command: %s' % command_list p = subprocess.Popen( command_list, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) # Can use communicate(input='y\n'.encode()) if the command run requires # some confirmation. stdout_data, _ = p.communicate() print 'Command output: %s' % stdout_data if p.returncode != 0: raise RuntimeError( 'Command %s failed: exit code: %s' % (command_list, p.returncode)) def run(self): for command in CUSTOM_COMMANDS: self.RunCustomCommand(command) # Configure the required packages and scripts to install. # Note that the Python Dataflow containers come with numpy already installed # so this dependency will not trigger anything to be installed unless a version # restriction is specified. REQUIRED_PACKAGES = [ 'numpy', ] setuptools.setup( name='juliaset', version='0.0.1', description='Julia set workflow package.', install_requires=REQUIRED_PACKAGES, packages=setuptools.find_packages(), cmdclass={ # Command class instantiated and run during easy_install scenarios. 'bdist_egg': bdist_egg, 'CustomCommands': CustomCommands, } ) DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/tfidf.py000066400000000000000000000167301272340251000263050ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A TF-IDF workflow (term frequency - inverse document frequency). For an explanation of the TF-IDF algorithm see the following link: http://en.wikipedia.org/wiki/Tf-idf """ from __future__ import absolute_import import argparse import glob import math import re import google.cloud.dataflow as df from google.cloud.dataflow.pvalue import AsSingleton def read_documents(pipeline, uris): """Read the documents at the provided uris and returns (uri, line) pairs.""" pcolls = [] for uri in uris: pcolls.append( pipeline | df.io.Read('read: %s' % uri, df.io.TextFileSource(uri)) | df.Map('withkey: %s' % uri, lambda v, uri: (uri, v), uri)) return pcolls | df.Flatten('flatten read pcolls') class TfIdf(df.PTransform): """A transform containing a basic TF-IDF pipeline. The input consists of KV objects where the key is the document's URI and the value is a piece of the document's content. The output is mapping from terms to scores for each document URI. """ def apply(self, uri_to_content): # Compute the total number of documents, and prepare a singleton # PCollection to use as side input. total_documents = ( uri_to_content | df.Keys('get uris') | df.RemoveDuplicates('get unique uris') | df.combiners.Count.Globally(' count uris')) # Create a collection of pairs mapping a URI to each of the words # in the document associated with that that URI. def split_into_words((uri, line)): return [(uri, w.lower()) for w in re.findall(r'[A-Za-z\']+', line)] uri_to_words = ( uri_to_content | df.FlatMap('split words', split_into_words)) # Compute a mapping from each word to the total number of documents # in which it appears. word_to_doc_count = ( uri_to_words | df.RemoveDuplicates('get unique words per doc') | df.Values('get words') | df.combiners.Count.PerElement('count docs per word')) # Compute a mapping from each URI to the total number of words in the # document associated with that URI. uri_to_word_total = ( uri_to_words | df.Keys(' get uris') | df.combiners.Count.PerElement('count words in doc')) # Count, for each (URI, word) pair, the number of occurrences of that word # in the document associated with the URI. uri_and_word_to_count = ( uri_to_words | df.combiners.Count.PerElement('count word-doc pairs')) # Adjust the above collection to a mapping from (URI, word) pairs to counts # into an isomorphic mapping from URI to (word, count) pairs, to prepare # for a join by the URI key. uri_to_word_and_count = ( uri_and_word_to_count | df.Map('shift keys', lambda ((uri, word), count): (uri, (word, count)))) # Perform a CoGroupByKey (a sort of pre-join) on the prepared # uri_to_word_total and uri_to_word_and_count tagged by 'word totals' and # 'word counts' strings. This yields a mapping from URI to a dictionary # that maps the above mentioned tag strings to an iterable containing the # word total for that URI and word and count respectively. # # A diagram (in which '[]' just means 'iterable'): # # URI: {'word totals': [count], # Total words within this URI's document. # 'word counts': [(word, count), # Counts of specific words # (word, count), # within this URI's document. # ... ]} uri_to_word_and_count_and_total = ( {'word totals': uri_to_word_total, 'word counts': uri_to_word_and_count} | df.CoGroupByKey('cogroup by uri')) # Compute a mapping from each word to a (URI, term frequency) pair for each # URI. A word's term frequency for a document is simply the number of times # that word occurs in the document divided by the total number of words in # the document. def compute_term_frequency((uri, count_and_total)): word_and_count = count_and_total['word counts'] # We have an iterable for one element that we want extracted. [word_total] = count_and_total['word totals'] for word, count in word_and_count: yield word, (uri, float(count) / word_total) word_to_uri_and_tf = ( uri_to_word_and_count_and_total | df.FlatMap('compute term frequencies', compute_term_frequency)) # Compute a mapping from each word to its document frequency. # A word's document frequency in a corpus is the number of # documents in which the word appears divided by the total # number of documents in the corpus. # # This calculation uses a side input, a Dataflow-computed auxiliary value # presented to each invocation of our MapFn lambda. The second argument to # the lambda (called total---note that we are unpacking the first argument) # receives the value we listed after the lambda in Map(). Additional side # inputs (and ordinary Python values, too) can be provided to MapFns and # DoFns in this way. word_to_df = ( word_to_doc_count | df.Map('compute doc frequencies', lambda (word, count), total: (word, float(count) / total), AsSingleton(total_documents))) # Join the term frequency and document frequency collections, # each keyed on the word. word_to_uri_and_tf_and_df = ( {'tf': word_to_uri_and_tf, 'df': word_to_df} | df.CoGroupByKey('cogroup words by tf-df')) # Compute a mapping from each word to a (URI, TF-IDF) score for each URI. # There are a variety of definitions of TF-IDF # ("term frequency - inverse document frequency") score; here we use a # basic version that is the term frequency divided by the log of the # document frequency. def compute_tf_idf((word, tf_and_df)): [docf] = tf_and_df['df'] for uri, tf in tf_and_df['tf']: yield word, (uri, tf * math.log(1 / docf)) word_to_uri_and_tfidf = ( word_to_uri_and_tf_and_df | df.FlatMap('compute tf-idf', compute_tf_idf)) return word_to_uri_and_tfidf def run(argv=None): """Main entry point; defines and runs the tfidf pipeline.""" parser = argparse.ArgumentParser() parser.add_argument('--uris', required=True, help='URIs to process.') parser.add_argument('--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Read documents specified by the uris command line option. pcoll = read_documents(p, glob.glob(known_args.uris)) # Compute TF-IDF information for each word. output = pcoll | TfIdf() # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | df.io.Write('write', df.io.TextFileSink(known_args.output)) p.run() if __name__ == '__main__': run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/tfidf_test.py000066400000000000000000000056341272340251000273450ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the TF-IDF example.""" import logging import os import re import tempfile import unittest import google.cloud.dataflow as df from google.cloud.dataflow.examples.complete import tfidf EXPECTED_RESULTS = set([ ('ghi', '1.txt', 0.3662040962227032), ('abc', '1.txt', 0.0), ('abc', '3.txt', 0.0), ('abc', '2.txt', 0.0), ('def', '1.txt', 0.13515503603605478), ('def', '2.txt', 0.2027325540540822)]) EXPECTED_LINE_RE = r'\(u\'([a-z]*)\', \(\'.*([0-9]\.txt)\', (.*)\)\)' class TfIdfTest(unittest.TestCase): def create_file(self, path, contents): logging.info('Creating temp file: %s', path) with open(path, 'w') as f: f.write(contents) def test_tfidf_transform(self): p = df.Pipeline('DirectPipelineRunner') uri_to_line = p | df.Create( 'create sample', [('1.txt', 'abc def ghi'), ('2.txt', 'abc def'), ('3.txt', 'abc')]) result = ( uri_to_line | tfidf.TfIdf() | df.Map('flatten', lambda (word, (uri, tfidf)): (word, uri, tfidf))) df.assert_that(result, df.equal_to(EXPECTED_RESULTS)) # Run the pipeline. Note that the assert_that above adds to the pipeline # a check that the result PCollection contains expected values. To actually # trigger the check the pipeline must be run. p.run() def test_basics(self): # Setup the files with expected content. temp_folder = tempfile.mkdtemp() self.create_file(os.path.join(temp_folder, '1.txt'), 'abc def ghi') self.create_file(os.path.join(temp_folder, '2.txt'), 'abc def') self.create_file(os.path.join(temp_folder, '3.txt'), 'abc') tfidf.run([ '--uris=%s/*' % temp_folder, '--output', os.path.join(temp_folder, 'result')]) # Parse result file and compare. results = [] with open(os.path.join(temp_folder, 'result-00000-of-00001')) as result_file: for line in result_file: match = re.search(EXPECTED_LINE_RE, line) logging.info('Result line: %s', line) if match is not None: results.append( (match.group(1), match.group(2), float(match.group(3)))) logging.info('Computed results: %s', set(results)) self.assertEqual(set(results), EXPECTED_RESULTS) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/top_wikipedia_sessions.py000066400000000000000000000126241272340251000317650ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """An example that reads Wikipedia edit data and computes strings of edits. An example that reads Wikipedia edit data from Cloud Storage and computes the user with the longest string of edits separated by no more than an hour within each 30 day period. To execute this pipeline locally using the DirectPipelineRunner, specify an output prefix on GCS: --output gs://YOUR_OUTPUT_PREFIX To execute this pipeline using the Google Cloud Dataflow service, specify pipeline configuration in addition to the above: --job_name NAME_FOR_YOUR_JOB --project YOUR_PROJECT_ID --staging_location gs://YOUR_STAGING_DIRECTORY --temp_location gs://YOUR_TEMPORARY_DIRECTORY --runner BlockingDataflowPipelineRunner The default input is gs://dataflow-samples/wikipedia_edits/*.json and can be overridden with --input. """ from __future__ import absolute_import import argparse import json import logging import sys import google.cloud.dataflow as df from google.cloud.dataflow import combiners from google.cloud.dataflow import window ONE_HOUR_IN_SECONDS = 3600 THIRTY_DAYS_IN_SECONDS = 30 * 24 * ONE_HOUR_IN_SECONDS class ExtractUserAndTimestampDoFn(df.DoFn): """Extracts user and timestamp representing a Wikipedia edit.""" def process(self, context): table_row = json.loads(context.element) if 'contributor_username' in table_row: user_name = table_row['contributor_username'] timestamp = table_row['timestamp'] yield window.TimestampedValue(user_name, timestamp) class ComputeSessions(df.PTransform): """Computes the number of edits in each user session. A session is defined as a string of edits where each is separated from the next by less than an hour. """ def __init__(self): super(ComputeSessions, self).__init__() def apply(self, pcoll): return (pcoll | df.WindowInto('ComputeSessionsWindow', window.Sessions(gap_size=ONE_HOUR_IN_SECONDS)) | combiners.Count.PerElement()) class TopPerMonth(df.PTransform): """Computes the longest session ending in each month.""" def __init__(self): super(TopPerMonth, self).__init__() def apply(self, pcoll): return (pcoll | df.WindowInto('TopPerMonthWindow', window.FixedWindows( size=THIRTY_DAYS_IN_SECONDS)) | combiners.core.CombineGlobally( 'Top', combiners.TopCombineFn( 10, lambda first, second: first[1] < second[1])) .without_defaults()) class SessionsToStringsDoFn(df.DoFn): """Adds the session information to be part of the key.""" def process(self, context): yield (context.element[0] + ' : ' + ', '.join([str(w) for w in context.windows]), context.element[1]) class FormatOutputDoFn(df.DoFn): """Formats a string containing the user, count, and session.""" def process(self, context): for kv in context.element: session = kv[0] count = kv[1] yield (session + ' : ' + str(count) + ' : ' + ', '.join([str(w) for w in context.windows])) class ComputeTopSessions(df.PTransform): """Computes the top user sessions for each month.""" def __init__(self, sampling_threshold): super(ComputeTopSessions, self).__init__() self.sampling_threshold = sampling_threshold def apply(self, pcoll): return (pcoll | df.ParDo('ExtractUserAndTimestamp', ExtractUserAndTimestampDoFn()) | df.Filter( lambda x: abs(hash(x)) <= sys.maxint * self.sampling_threshold) | ComputeSessions() | df.ParDo('SessionsToStrings', SessionsToStringsDoFn()) | TopPerMonth() | df.ParDo('FormatOutput', FormatOutputDoFn())) def run(argv=None): """Runs the Wikipedia top edits pipeline. Args: argv: Pipeline options as a list of arguments. """ parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/wikipedia_edits/*.json', help='Input specified as a GCS path containing a BigQuery table exported ' 'as json.') parser.add_argument('--output', required=True, help='Output file to write results to.') parser.add_argument('--sampling_threshold', type=float, default=0.1, help='Fraction of entries used for session tracking') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) (p # pylint: disable=expression-not-assigned | df.Read('read', df.io.TextFileSource(known_args.input)) | ComputeTopSessions(known_args.sampling_threshold) | df.io.Write('write', df.io.TextFileSink(known_args.output))) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/complete/top_wikipedia_sessions_test.py000066400000000000000000000040751272340251000330250ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the top wikipedia sessions example.""" import json import unittest import google.cloud.dataflow as df from google.cloud.dataflow.examples.complete import top_wikipedia_sessions class ComputeTopSessionsTest(unittest.TestCase): EDITS = [ json.dumps({'timestamp': 0.0, 'contributor_username': 'user1'}), json.dumps({'timestamp': 0.001, 'contributor_username': 'user1'}), json.dumps({'timestamp': 0.002, 'contributor_username': 'user1'}), json.dumps({'timestamp': 0.0, 'contributor_username': 'user2'}), json.dumps({'timestamp': 0.001, 'contributor_username': 'user2'}), json.dumps({'timestamp': 3.601, 'contributor_username': 'user2'}), json.dumps({'timestamp': 3.602, 'contributor_username': 'user2'}), json.dumps( {'timestamp': 2 * 3600.0, 'contributor_username': 'user2'}), json.dumps( {'timestamp': 35 * 24 * 3.600, 'contributor_username': 'user3'}) ] EXPECTED = [ 'user1 : [0.0, 3600.002) : 3 : [0.0, 2592000.0)', 'user2 : [0.0, 3603.602) : 4 : [0.0, 2592000.0)', 'user2 : [7200.0, 10800.0) : 1 : [0.0, 2592000.0)', 'user3 : [3024.0, 6624.0) : 1 : [0.0, 2592000.0)', ] def test_compute_top_sessions(self): p = df.Pipeline('DirectPipelineRunner') edits = p | df.Create('create', self.EDITS) result = edits | top_wikipedia_sessions.ComputeTopSessions(1.0) df.assert_that(result, df.equal_to(self.EXPECTED)) p.run() if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/000077500000000000000000000000001272340251000246265ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/bigquery_schema.py000066400000000000000000000101131272340251000303430ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A workflow that writes to a BigQuery table with nested and repeated fields. Demonstrates how to build a bigquery.TableSchema object with nested and repeated fields. Also, shows how to generate data to be written to a BigQuery table with nested and repeated fields. """ from __future__ import absolute_import import argparse import logging import google.cloud.dataflow as df def run(argv=None): """Run the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--output', required=True, help= ('Output BigQuery table for results specified as: PROJECT:DATASET.TABLE ' 'or DATASET.TABLE.')) known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) from google.cloud.dataflow.internal.clients import bigquery # pylint: disable=g-import-not-at-top table_schema = bigquery.TableSchema() # Fields that use standard types. kind_schema = bigquery.TableFieldSchema() kind_schema.name = 'kind' kind_schema.type = 'string' kind_schema.mode = 'nullable' table_schema.fields.append(kind_schema) full_name_schema = bigquery.TableFieldSchema() full_name_schema.name = 'fullName' full_name_schema.type = 'string' full_name_schema.mode = 'required' table_schema.fields.append(full_name_schema) age_schema = bigquery.TableFieldSchema() age_schema.name = 'age' age_schema.type = 'integer' age_schema.mode = 'nullable' table_schema.fields.append(age_schema) gender_schema = bigquery.TableFieldSchema() gender_schema.name = 'gender' gender_schema.type = 'string' gender_schema.mode = 'nullable' table_schema.fields.append(gender_schema) # A nested field phone_number_schema = bigquery.TableFieldSchema() phone_number_schema.name = 'phoneNumber' phone_number_schema.type = 'record' phone_number_schema.mode = 'nullable' area_code = bigquery.TableFieldSchema() area_code.name = 'areaCode' area_code.type = 'integer' area_code.mode = 'nullable' phone_number_schema.fields.append(area_code) number = bigquery.TableFieldSchema() number.name = 'number' number.type = 'integer' number.mode = 'nullable' phone_number_schema.fields.append(number) table_schema.fields.append(phone_number_schema) # A repeated field. children_schema = bigquery.TableFieldSchema() children_schema.name = 'children' children_schema.type = 'string' children_schema.mode = 'repeated' table_schema.fields.append(children_schema) def create_random_record(record_id): return {'kind': 'kind' + record_id, 'fullName': 'fullName'+record_id, 'age': int(record_id) * 10, 'gender': 'male', 'phoneNumber': { 'areaCode': int(record_id) * 100, 'number': int(record_id) * 100000}, 'children': ['child' + record_id + '1', 'child' + record_id + '2', 'child' + record_id + '3'] } # pylint: disable=expression-not-assigned record_ids = p | df.Create('CreateIDs', ['1', '2', '3', '4', '5']) records = record_ids | df.Map('CreateRecords', create_random_record) records | df.io.Write( 'write', df.io.BigQuerySink( known_args.output, schema=table_schema, create_disposition=df.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=df.io.BigQueryDisposition.WRITE_TRUNCATE)) # Run the pipeline (all operations are deferred until run() is called). p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/bigquery_side_input.py000066400000000000000000000075001272340251000312540ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A Dataflow job that uses BigQuery sources as a side inputs. Illustrates how to insert side-inputs into transforms in three different forms, as a singleton, as a iterator, and as a list. This workflow generate a set of tuples of the form (groupId, corpus, word) where groupId is a generated identifier for the group and corpus and word are randomly selected from corresponding rows in BQ dataset 'publicdata:samples.shakespeare'. Users should specify the number of groups to form and optionally a corpus and/or a word that should be ignored when forming groups. """ import argparse import logging from random import randrange import google.cloud.dataflow as df from google.cloud.dataflow.pvalue import AsIter from google.cloud.dataflow.pvalue import AsList from google.cloud.dataflow.pvalue import AsSingleton def create_groups(group_ids, corpus, word, ignore_corpus, ignore_word): """Generate groups given the input PCollections.""" def attach_corpus_fn(group, corpus, ignore): selected = None len_corpus = len(corpus) while not selected: c = corpus[randrange(0, len_corpus - 1)].values()[0] if c != ignore: selected = c yield (group, selected) def attach_word_fn(group, words, ignore): selected = None len_words = len(words) while not selected: c = words[randrange(0, len_words - 1)].values()[0] if c != ignore: selected = c yield group + (selected,) return (group_ids | df.FlatMap( 'attach corpus', attach_corpus_fn, AsList(corpus), AsSingleton(ignore_corpus)) | df.FlatMap( 'attach word', attach_word_fn, AsIter(word), AsSingleton(ignore_word))) def run(argv=None): """Run the workflow.""" parser = argparse.ArgumentParser() parser.add_argument('--output') parser.add_argument('--ignore_corpus', default='') parser.add_argument('--ignore_word', default='') parser.add_argument('--num_groups') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) group_ids = [] for i in xrange(0, int(known_args.num_groups)): group_ids.append('id' + str(i)) query_corpus = 'select UNIQUE(corpus) from publicdata:samples.shakespeare' query_word = 'select UNIQUE(word) from publicdata:samples.shakespeare' ignore_corpus = known_args.ignore_corpus ignore_word = known_args.ignore_word pcoll_corpus = p | df.Read('read corpus', df.io.BigQuerySource(query=query_corpus)) pcoll_word = p | df.Read('read words', df.io.BigQuerySource(query=query_word)) pcoll_ignore_corpus = p | df.Create('create_ignore_corpus', [ignore_corpus]) pcoll_ignore_word = p | df.Create('create_ignore_word', [ignore_word]) pcoll_group_ids = p | df.Create('create groups', group_ids) pcoll_groups = create_groups(pcoll_group_ids, pcoll_corpus, pcoll_word, pcoll_ignore_corpus, pcoll_ignore_word) # pylint:disable=expression-not-assigned pcoll_groups | df.io.Write('WriteToText', df.io.TextFileSink(known_args.output)) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/bigquery_side_input_test.py000066400000000000000000000043151272340251000323140ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the BigQuery side input example.""" import logging import unittest import google.cloud.dataflow as df from google.cloud.dataflow.examples.cookbook import bigquery_side_input class BigQuerySideInputTest(unittest.TestCase): def test_create_groups(self): p = df.Pipeline('DirectPipelineRunner') group_ids_pcoll = p | df.Create('create_group_ids', ['A', 'B', 'C']) corpus_pcoll = p | df.Create('create_corpus', [{'f': 'corpus1'}, {'f': 'corpus2'}, {'f': 'corpus3'}]) words_pcoll = p | df.Create('create_words', [{'f': 'word1'}, {'f': 'word2'}, {'f': 'word3'}]) ignore_corpus_pcoll = p | df.Create('create_ignore_corpus', ['corpus1']) ignore_word_pcoll = p | df.Create('create_ignore_word', ['word1']) groups = bigquery_side_input.create_groups(group_ids_pcoll, corpus_pcoll, words_pcoll, ignore_corpus_pcoll, ignore_word_pcoll) def group_matcher(actual): self.assertEqual(len(actual), 3) for group in actual: self.assertEqual(len(group), 3) self.assertTrue(group[1].startswith('corpus')) self.assertNotEqual(group[1], 'corpus1') self.assertTrue(group[2].startswith('word')) self.assertNotEqual(group[2], 'word1') df.assert_that(groups, group_matcher) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/bigquery_tornadoes.py000066400000000000000000000064771272340251000311230ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A workflow using BigQuery sources and sinks. The workflow will read from a table that has the 'month' and 'tornado' fields as part of the table schema (other additional fields are ignored). The 'month' field is a number represented as a string (e.g., '23') and the 'tornado' field is a boolean field. The workflow will compute the number of tornadoes in each month and output the results to a table (created if needed) with the following schema: - month: number - tornado_count: number This example uses the default behavior for BigQuery source and sinks that represents table rows as plain Python dictionaries. """ from __future__ import absolute_import import argparse import logging import google.cloud.dataflow as df def count_tornadoes(input_data): """Workflow computing the number of tornadoes for each month that had one. Args: input_data: a PCollection of dictionaries representing table rows. Each dictionary will have a 'month' and a 'tornado' key as described in the module comment. Returns: A PCollection of dictionaries containing 'month' and 'tornado_count' keys. Months without tornadoes are skipped. """ return (input_data | df.FlatMap( 'months with tornadoes', lambda row: [(int(row['month']), 1)] if row['tornado'] else []) | df.CombinePerKey('monthly count', sum) | df.Map('format', lambda (k, v): {'month': k, 'tornado_count': v})) def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument('--input', default='clouddataflow-readonly:samples.weather_stations', help=('Input BigQuery table to process specified as: ' 'PROJECT:DATASET.TABLE or DATASET.TABLE.')) parser.add_argument( '--output', required=True, help= ('Output BigQuery table for results specified as: PROJECT:DATASET.TABLE ' 'or DATASET.TABLE.')) known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Read the table rows into a PCollection. rows = p | df.io.Read('read', df.io.BigQuerySource(known_args.input)) counts = count_tornadoes(rows) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned counts | df.io.Write( 'write', df.io.BigQuerySink( known_args.output, schema='month:INTEGER, tornado_count:INTEGER', create_disposition=df.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=df.io.BigQueryDisposition.WRITE_TRUNCATE)) # Run the pipeline (all operations are deferred until run() is called). p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/bigquery_tornadoes_test.py000066400000000000000000000026541272340251000321530ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the BigQuery tornadoes example.""" import logging import unittest import google.cloud.dataflow as df from google.cloud.dataflow.examples.cookbook import bigquery_tornadoes class BigQueryTornadoesTest(unittest.TestCase): def test_basics(self): p = df.Pipeline('DirectPipelineRunner') rows = (p | df.Create('create', [ {'month': 1, 'day': 1, 'tornado': False}, {'month': 1, 'day': 2, 'tornado': True}, {'month': 1, 'day': 3, 'tornado': True}, {'month': 2, 'day': 1, 'tornado': True}])) results = bigquery_tornadoes.count_tornadoes(rows) df.assert_that(results, df.equal_to([{'month': 1, 'tornado_count': 2}, {'month': 2, 'tornado_count': 1}])) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/bigshuffle.py000066400000000000000000000054101272340251000273160ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A BigShuffle workflow.""" from __future__ import absolute_import import argparse import binascii import logging import google.cloud.dataflow as df def crc32line(line): return binascii.crc32(line) & 0xffffffff def run(argv=None): # pylint: disable=expression-not-assigned parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file pattern to process.') parser.add_argument('--output', required=True, help='Output file pattern to write results to.') parser.add_argument('--checksum_output', required=True, help='Checksum output file pattern.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Read the text file[pattern] into a PCollection. lines = p | df.io.Read('read', df.io.TextFileSource(known_args.input)) # Count the occurrences of each word. output = (lines | df.Map('split', lambda x: (x[:10], x[10:99])) | df.GroupByKey('group') | df.FlatMap( 'format', lambda (key, vals): ['%s%s' % (key, val) for val in vals])) input_csum = (lines | df.Map('input-csum', crc32line) | df.CombineGlobally('combine-input-csum', sum) | df.Map('hex-format', lambda x: '%x' % x)) input_csum | df.io.Write( 'write-input-csum', df.io.TextFileSink(known_args.checksum_output + '-input')) # Write the output using a "Write" transform that has side effects. output | df.io.Write('write', df.io.TextFileSink(known_args.output)) # Write the output checksum output_csum = (output | df.Map('output-csum', crc32line) | df.CombineGlobally('combine-output-csum', sum) | df.Map('hex-format-output', lambda x: '%x' % x)) output_csum | df.io.Write( 'write-output-csum', df.io.TextFileSink(known_args.checksum_output + '-output')) # Actually run the pipeline (all operations above are deferred). p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/bigshuffle_test.py000066400000000000000000000041231272340251000303550ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the bigshuffle example.""" import logging import tempfile import unittest from google.cloud.dataflow.examples.cookbook import bigshuffle # TODO(dataflow-python): use gensort to generate input files. class BigShuffleTest(unittest.TestCase): SAMPLE_TEXT = 'a b c a b a\naa bb cc aa bb aa' def create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def test_basics(self): temp_path = self.create_temp_file(self.SAMPLE_TEXT) bigshuffle.run([ '--input=%s*' % temp_path, '--output=%s.result' % temp_path, '--checksum_output=%s.checksum' % temp_path]) # Parse result file and compare. results = [] with open(temp_path + '.result-00000-of-00001') as result_file: for line in result_file: results.append(line.strip()) expected = self.SAMPLE_TEXT.split('\n') self.assertEqual(sorted(results), sorted(expected)) # Check the checksums input_csum = '' with open(temp_path + '.checksum-input-00000-of-00001') as input_csum_file: input_csum = input_csum_file.read().strip() output_csum = '' with open(temp_path + '.checksum-output-00000-of-00001') as output_csum_file: output_csum = output_csum_file.read().strip() expected_csum = 'd629c1f6' self.assertEqual(input_csum, expected_csum) self.assertEqual(input_csum, output_csum) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/coders.py000066400000000000000000000053551272340251000264670ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A workflow using custom JSON-based coders for text sources and sinks. The input file contains a JSON string on each line describing a match record using the following schema: {'guest': [TEAM_NAME, GOALS], 'host': [TEAM_NAME, GOALS]} The output file will contain the computed points for each team with one team per line in the following format: [TEAM_NAME, POINTS] """ from __future__ import absolute_import import argparse import json import logging import google.cloud.dataflow as df class JsonCoder(object): """A JSON coder interpreting each line as a JSON string.""" def encode(self, x): return json.dumps(x) def decode(self, x): return json.loads(x) def compute_points(record): """Compute points based on the record containing the match result. The function assigns 3 points for a win, 1 point for a draw, and 0 points for a loss (see http://en.wikipedia.org/wiki/Three_points_for_a_win). """ host_name, host_goals = record['host'] guest_name, guest_goals = record['guest'] if host_goals == guest_goals: yield host_name, 1 yield guest_name, 1 elif host_goals > guest_goals: yield host_name, 3 yield guest_name, 0 else: yield host_name, 0 yield guest_name, 3 def run(argv=None): """Runs the workflow computing total points from a collection of matches.""" parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file to process.') parser.add_argument('--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) (p # pylint: disable=expression-not-assigned | df.io.Read('read', df.io.TextFileSource(known_args.input, coder=JsonCoder())) | df.FlatMap('points', compute_points) | df.CombinePerKey(sum) | df.io.Write( 'write', df.io.TextFileSink(known_args.output, coder=JsonCoder()))) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/coders_test.py000066400000000000000000000033741272340251000275250ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the coders example.""" import json import logging import tempfile import unittest from google.cloud.dataflow.examples.cookbook import coders class CodersTest(unittest.TestCase): SAMPLE_RECORDS = [ {'host': ['Germany', 1], 'guest': ['Italy', 0]}, {'host': ['Germany', 1], 'guest': ['Brasil', 3]}, {'host': ['Brasil', 1], 'guest': ['Italy', 0]}] def create_temp_file(self, records): with tempfile.NamedTemporaryFile(delete=False) as f: for record in records: f.write('%s\n' % json.dumps(record)) return f.name def test_basics(self): temp_path = self.create_temp_file(self.SAMPLE_RECORDS) coders.run([ '--input=%s*' % temp_path, '--output=%s.result' % temp_path]) # Parse result file and compare. results = [] with open(temp_path + '.result-00000-of-00001') as result_file: for line in result_file: results.append(json.loads(line)) logging.info('result: %s', results) self.assertEqual( sorted(results), sorted([['Italy', 0], ['Brasil', 6], ['Germany', 3]])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/combiners_test.py000066400000000000000000000045301272340251000302220ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Simple tests to showcase combiners. The tests are meant to be "copy/paste" code snippets for the topic they address (combiners in this case). Most examples use neither sources nor sinks. The input data is generated simply with a Create transform and the output is checked directly on the last PCollection produced. """ import logging import unittest import google.cloud.dataflow as df class CombinersTest(unittest.TestCase): """Tests showcasing Dataflow combiners.""" SAMPLE_DATA = [ ('a', 1), ('b', 10), ('a', 2), ('a', 3), ('b', 20), ('c', 100)] def test_combine_per_key_with_callable(self): """CombinePerKey using a standard callable reducing iterables. A common case for Dataflow combiners is to sum (or max or min) over the values of each key. Such standard functions can be used directly as combiner functions. In fact, any function "reducing" an iterable to a single value can be used. """ result = ( df.Pipeline(runner=df.runners.DirectPipelineRunner()) | df.Create(CombinersTest.SAMPLE_DATA) | df.CombinePerKey(sum)) df.assert_that(result, df.equal_to([('a', 6), ('b', 30), ('c', 100)])) result.pipeline.run() def test_combine_per_key_with_custom_callable(self): """CombinePerKey using a custom function reducing iterables.""" def multiply(values): result = 1 for v in values: result *= v return result result = ( df.Pipeline(runner=df.runners.DirectPipelineRunner()) | df.Create(CombinersTest.SAMPLE_DATA) | df.CombinePerKey(multiply)) df.assert_that(result, df.equal_to([('a', 6), ('b', 200), ('c', 100)])) result.pipeline.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/custom_ptransform.py000066400000000000000000000073731272340251000307770ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Various implementations of a Count custom PTransform. These example show the different ways you can write custom PTransforms. """ from __future__ import absolute_import import argparse import logging import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions # pylint doesn't understand our pipeline syntax: # pylint:disable=expression-not-assigned def run_count1(known_args, options): """Runs the first example pipeline.""" class Count(df.PTransform): """Count as a subclass of PTransform, with an apply method.""" def apply(self, pcoll): return ( pcoll | df.Map('Init', lambda v: (v, 1)) | df.CombinePerKey(sum)) logging.info('Running first pipeline') p = df.Pipeline(options=options) (p | df.io.Read(df.io.TextFileSource(known_args.input)) | Count() | df.io.Write(df.io.TextFileSink(known_args.output))) p.run() def run_count2(known_args, options): """Runs the second example pipeline.""" @df.ptransform_fn def Count(label, pcoll): # pylint: disable=invalid-name,unused-argument """Count as a decorated function.""" return ( pcoll | df.Map('Init', lambda v: (v, 1)) | df.CombinePerKey(sum)) logging.info('Running second pipeline') p = df.Pipeline(options=options) (p | df.io.Read(df.io.TextFileSource(known_args.input)) | Count() # pylint: disable=no-value-for-parameter | df.io.Write(df.io.TextFileSink(known_args.output))) p.run() def run_count3(known_args, options): """Runs the third example pipeline.""" @df.ptransform_fn # pylint: disable=invalid-name,unused-argument def Count(label, pcoll, factor=1): """Count as a decorated function with a side input. Args: label: optional label for this transform pcoll: the PCollection passed in from the previous transform factor: the amount by which to count Returns: A PCollection counting the number of times each unique element occurs. """ return ( pcoll | df.Map('Init', lambda v: (v, factor)) | df.CombinePerKey(sum)) logging.info('Running third pipeline') p = df.Pipeline(options=options) (p | df.io.Read(df.io.TextFileSource(known_args.input)) | Count(2) # pylint: disable=no-value-for-parameter | df.io.Write(df.io.TextFileSink(known_args.output))) p.run() def get_args(argv): """Determines user specified arguments from the given list of arguments. Args: argv: all arguments. Returns: A pair of argument lists containing known and remaining arguments. """ parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file to process.') parser.add_argument('--output', required=True, help='Output file to write results to.') return parser.parse_known_args(argv) def run(argv=None): known_args, pipeline_args = get_args(argv) options = PipelineOptions(pipeline_args) run_count1(known_args, options) run_count2(known_args, options) run_count3(known_args, options) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/custom_ptransform_test.py000066400000000000000000000041111272340251000320210ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for the various custom Count implementation examples.""" import logging import tempfile import unittest from google.cloud.dataflow.examples.cookbook import custom_ptransform from google.cloud.dataflow.utils.options import PipelineOptions class CustomCountTest(unittest.TestCase): def test_count1(self): self.run_pipeline(custom_ptransform.run_count1) def test_count2(self): self.run_pipeline(custom_ptransform.run_count2) def test_count3(self): self.run_pipeline(custom_ptransform.run_count3, factor=2) def run_pipeline(self, count_implementation, factor=1): input_path = self.create_temp_file('CAT\nDOG\nCAT\nCAT\nDOG\n') output_path = input_path + '.result' known_args, pipeline_args = custom_ptransform.get_args([ '--input=%s*' % input_path, '--output=%s' % output_path ]) count_implementation(known_args, PipelineOptions(pipeline_args)) self.assertEqual(["(u'CAT', %d)" % (3 * factor), "(u'DOG', %d)" % (2 * factor)], self.get_output(output_path + '-00000-of-00001')) def create_temp_file(self, contents=''): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def get_output(self, path): logging.info('Reading output from "%s"', path) lines = [] with open(path) as f: lines = f.readlines() return sorted(s.rstrip('\n') for s in lines) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/filters.py000066400000000000000000000074601272340251000266570ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """An example workflow that demonstrates filters and other features. - Reading and writing data from BigQuery. - Manipulating BigQuery rows (as Python dicts) in memory. - Global aggregates. - Filtering PCollections using both user-specified parameters as well as global aggregates computed during pipeline execution. """ from __future__ import absolute_import import argparse import logging import google.cloud.dataflow as df from google.cloud.dataflow.pvalue import AsSingleton def filter_cold_days(input_data, month_filter): """Workflow computing rows in a specific month with low temperatures. Args: input_data: a PCollection of dictionaries representing table rows. Each dictionary must have the keys ['year', 'month', 'day', and 'mean_temp']. month_filter: an int representing the month for which colder-than-average days should be returned. Returns: A PCollection of dictionaries with the same keys described above. Each row represents a day in the specified month where temperatures were colder than the global mean temperature in the entire dataset. """ # Project to only the desired fields from a complete input row. # E.g., SELECT f1, f2, f3, ... FROM InputTable. projection_fields = ['year', 'month', 'day', 'mean_temp'] fields_of_interest = ( input_data | df.Map('projected', lambda row: {f: row[f] for f in projection_fields})) # Compute the global mean temperature. global_mean = AsSingleton( fields_of_interest | df.Map('extract mean', lambda row: row['mean_temp']) | df.combiners.Mean.Globally('global mean')) # Filter to the rows representing days in the month of interest # in which the mean daily temperature is below the global mean. return ( fields_of_interest | df.Filter('desired month', lambda row: row['month'] == month_filter) | df.Filter('below mean', lambda row, mean: row['mean_temp'] < mean, global_mean)) def run(argv=None): """Constructs and runs the example filtering pipeline.""" parser = argparse.ArgumentParser() parser.add_argument('--input', help='BigQuery table to read from.', default='clouddataflow-readonly:samples.weather_stations') parser.add_argument('--output', required=True, help='BigQuery table to write to.') parser.add_argument('--month_filter', default=7, help='Numeric value of month to filter on.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) input_data = p | df.Read('input', df.io.BigQuerySource(known_args.input)) # pylint: disable=expression-not-assigned (filter_cold_days(input_data, known_args.month_filter) | df.io.Write('save to BQ', df.io.BigQuerySink( known_args.output, schema='year:INTEGER,month:INTEGER,day:INTEGER,mean_temp:FLOAT', create_disposition=df.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=df.io.BigQueryDisposition.WRITE_TRUNCATE))) # Actually run the pipeline (all operations above are deferred). p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/filters_test.py000066400000000000000000000045321272340251000277130ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the filters example.""" import logging import unittest import google.cloud.dataflow as df from google.cloud.dataflow.examples.cookbook import filters class FiltersTest(unittest.TestCase): # Note that 'removed' should be projected away by the pipeline input_data = [ {'year': 2010, 'month': 1, 'day': 1, 'mean_temp': 3, 'removed': 'a'}, {'year': 2012, 'month': 1, 'day': 2, 'mean_temp': 3, 'removed': 'a'}, {'year': 2011, 'month': 1, 'day': 3, 'mean_temp': 5, 'removed': 'a'}, {'year': 2013, 'month': 2, 'day': 1, 'mean_temp': 3, 'removed': 'a'}, {'year': 2011, 'month': 3, 'day': 3, 'mean_temp': 5, 'removed': 'a'}, ] def _get_result_for_month(self, month): p = df.Pipeline('DirectPipelineRunner') rows = (p | df.Create('create', self.input_data)) results = filters.filter_cold_days(rows, month) return results def test_basic(self): """Test that the correct result is returned for a simple dataset.""" results = self._get_result_for_month(1) df.assert_that( results, df.equal_to([{'year': 2010, 'month': 1, 'day': 1, 'mean_temp': 3}, {'year': 2012, 'month': 1, 'day': 2, 'mean_temp': 3}])) results.pipeline.run() def test_basic_empty(self): """Test that the correct empty result is returned for a simple dataset.""" results = self._get_result_for_month(3) df.assert_that(results, df.equal_to([])) results.pipeline.run() def test_basic_empty_missing(self): """Test that the correct empty result is returned for a missing month.""" results = self._get_result_for_month(4) df.assert_that(results, df.equal_to([])) results.pipeline.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/group_with_coder.py000066400000000000000000000077431272340251000305560ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """An example of using custom classes and coder for grouping operations. This workflow demonstrates registration and usage of a custom coder for a user- defined class. A deterministic custom coder is needed to use a class as a key in a combine or group operation. This example assumes an input file with, on each line, a comma-separated name and score. """ from __future__ import absolute_import import argparse import logging import sys import google.cloud.dataflow as df from google.cloud.dataflow import coders from google.cloud.dataflow.typehints import typehints from google.cloud.dataflow.typehints.decorators import with_output_types class Player(object): """A custom class used as a key in combine/group transforms.""" def __init__(self, name): self.name = name class PlayerCoder(coders.Coder): """A custom coder for the Player class.""" def encode(self, o): """Encode to bytes with a trace that coder was used.""" # Our encoding prepends an 'x:' prefix. return 'x:%s' % str(o.name) def decode(self, s): # To decode, we strip off the prepended 'x:' prefix. assert s[0:2] == 'x:' return Player(s[2:]) def is_deterministic(self): # Since coded Player objects are used as keys below with # df.CombinePerKey(sum), we require that this coder is deterministic # (i.e., two equivalent instances of the classes are encoded into the same # byte string) in order to guarantee consistent results. return True # Annotate the get_players function so that the typehint system knows that the # input to the CombinePerKey operation is a key-value pair of a Player object # and an integer. @with_output_types(typehints.KV[Player, int]) def get_players(descriptor): name, points = descriptor.split(',') return Player(name), int(points) def run(argv=sys.argv[1:]): """Runs the workflow computing total points from a collection of matches.""" parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file to process.') parser.add_argument('--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Register the custom coder for the Player class, so that it will be used in # the computation. coders.registry.register_coder(Player, PlayerCoder) (p # pylint: disable=expression-not-assigned | df.io.Read('read', df.io.TextFileSource(known_args.input)) # The get_players function is annotated with a type hint above, so the type # system knows the output type of the following operation is a key-value pair # of a Player and an int. Please see the documentation for details on # types that are inferred automatically as well as other ways to specify # type hints. | df.Map('get players', get_players) # The output type hint of the previous step is used to infer that the key # type of the following operation is the Player type. Since a custom coder # is registered for the Player class above, a PlayerCoder will be used to # encode Player objects as keys for this combine operation. | df.CombinePerKey(sum) | df.Map(lambda (k, v): '%s,%d' % (k.name, v)) | df.io.Write('write', df.io.TextFileSink(known_args.output))) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/group_with_coder_test.py000066400000000000000000000063361272340251000316120ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the custom coders example.""" import logging import tempfile import unittest from google.cloud.dataflow.examples.cookbook import group_with_coder # Patch group_with_coder.PlayerCoder.decode(). To test that the PlayerCoder was # used, we do not strip the prepended 'x:' string when decoding a Player object. group_with_coder.PlayerCoder.decode = lambda self, s: group_with_coder.Player(s) class GroupWithCoderTest(unittest.TestCase): SAMPLE_RECORDS = [ 'joe,10', 'fred,3', 'mary,7', 'joe,20', 'fred,6', 'ann,5', 'joe,30', 'ann,10', 'mary,1'] def create_temp_file(self, records): with tempfile.NamedTemporaryFile(delete=False) as f: for record in records: f.write('%s\n' % record) return f.name def test_basics_with_type_check(self): # Run the workflow with --pipeline_type_check option. This will make sure # the typehints associated with all transforms will have non-default values # and therefore any custom coders will be used. In our case we want to make # sure the coder for the Player class will be used. temp_path = self.create_temp_file(self.SAMPLE_RECORDS) group_with_coder.run([ '--pipeline_type_check', '--input=%s*' % temp_path, '--output=%s.result' % temp_path]) # Parse result file and compare. results = [] with open(temp_path + '.result-00000-of-00001') as result_file: for line in result_file: name, points = line.split(',') results.append((name, int(points))) logging.info('result: %s', results) self.assertEqual( sorted(results), sorted([('x:ann', 15), ('x:fred', 9), ('x:joe', 60), ('x:mary', 8)])) def test_basics_without_type_check(self): # Run the workflow without --pipeline_type_check option. This will make sure # the typehints associated with all transforms will have default values and # therefore any custom coders will not be used. The default coder (pickler) # will be used instead. temp_path = self.create_temp_file(self.SAMPLE_RECORDS) group_with_coder.run([ '--no_pipeline_type_check', '--input=%s*' % temp_path, '--output=%s.result' % temp_path]) # Parse result file and compare. results = [] with open(temp_path + '.result-00000-of-00001') as result_file: for line in result_file: name, points = line.split(',') results.append((name, int(points))) logging.info('result: %s', results) self.assertEqual( sorted(results), sorted([('ann', 15), ('fred', 9), ('joe', 60), ('mary', 8)])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/mergecontacts.py000066400000000000000000000117601272340251000300430ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Merge phone, email, and mailing address information. A Dataflow pipeline that merges phone, email, and address information associated with the same names. Each input "database" is a tab-delimited text file pairing names with one phone number/email address/mailing address; multiple entries associated with the same name are allowed. Outputs are a tab-delimited text file with the merged information and another file containing some simple statistics. See mergecontacts_test.py for example inputs and outputs. A demonstration of: CoGroupByKey Non-linear pipelines (i.e., pipelines with branches) """ from __future__ import absolute_import import argparse import logging import re import google.cloud.dataflow as df def run(argv=None, assert_results=None): parser = argparse.ArgumentParser() parser.add_argument( '--input_email', required=True, help='Email database, with each line formatted as "nameemail".') parser.add_argument( '--input_phone', required=True, help='Phonebook, with each line formatted as "namephone number".') parser.add_argument( '--input_snailmail', required=True, help='Address database, with each line formatted as "nameaddress".') parser.add_argument('--output_tsv', required=True, help='Tab-delimited output file.') parser.add_argument('--output_stats', required=True, help='Output file for statistics about the input.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Helper: read a tab-separated key-value mapping from a text file, escape all # quotes/backslashes, and convert it a PCollection of (key, value) pairs. def read_kv_textfile(label, textfile): return (p | df.io.Read('read_%s' % label, textfile) | df.Map('backslash_%s' % label, lambda x: re.sub(r'\\', r'\\\\', x)) | df.Map('escape_quotes_%s' % label, lambda x: re.sub(r'"', r'\"', x)) | df.Map('split_%s' % label, lambda x: re.split(r'\t+', x, 1))) # Read input databases. email = read_kv_textfile('email', df.io.TextFileSource(known_args.input_email)) phone = read_kv_textfile('phone', df.io.TextFileSource(known_args.input_phone)) snailmail = read_kv_textfile('snailmail', df.io.TextFileSource(known_args.input_snailmail)) # Group together all entries under the same name. grouped = (email, phone, snailmail) | df.CoGroupByKey('group_by_name') # Prepare tab-delimited output; something like this: # "name""email_1,email_2""phone""first_snailmail_only" tsv_lines = grouped | df.Map( lambda (name, (email, phone, snailmail)): '\t'.join( ['"%s"' % name, '"%s"' % ','.join(email), '"%s"' % ','.join(phone), '"%s"' % next(iter(snailmail), '')])) # Compute some stats about our database of people. luddites = grouped | df.Filter( # People without email. lambda (name, (email, phone, snailmail)): not next(iter(email), None)) writers = grouped | df.Filter( # People without phones. lambda (name, (email, phone, snailmail)): not next(iter(phone), None)) nomads = grouped | df.Filter( # People without addresses. lambda (name, (email, phone, snailmail)): not next(iter(snailmail), None)) num_luddites = luddites | df.combiners.Count.Globally('luddites') num_writers = writers | df.combiners.Count.Globally('writers') num_nomads = nomads | df.combiners.Count.Globally('nomads') # Write tab-delimited output. # pylint: disable=expression-not-assigned tsv_lines | df.io.Write('write_tsv', df.io.TextFileSink(known_args.output_tsv)) # TODO(silviuc): Move the assert_results logic to the unit test. if assert_results is not None: expected_luddites, expected_writers, expected_nomads = assert_results df.assert_that(num_luddites, df.equal_to([expected_luddites]), label='assert:luddites') df.assert_that(num_writers, df.equal_to([expected_writers]), label='assert:writers') df.assert_that(num_nomads, df.equal_to([expected_nomads]), label='assert:nomads') # Execute pipeline. p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/mergecontacts_test.py000066400000000000000000000117651272340251000311070ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the mergecontacts example.""" import logging import tempfile import unittest from google.cloud.dataflow.examples.cookbook import mergecontacts class MergeContactsTest(unittest.TestCase): CONTACTS_EMAIL = '\n'.join(['Nathan Nomad\tnathan@example.com', 'Nicky Nomad\tnicky@example.com', 'Noreen Nomad\tnoreen@example.com', 'Noreen Nomad\tnomad@example.com', 'Robert B\trobert@example.com', 'Silviu C\tsilviu@example.com', 'Tom S\ttom@example.com', 'Wally Writer\twally@example.com', '']) CONTACTS_PHONE = '\n'.join(['Larry Luddite\t724-228-3529', 'Lisa Luddite\t304-277-3504', 'Nathan Nomad\t412-466-8968', 'Nicky Nomad\t724-379-5815', 'Noreen Nomad\t412-472-0145', 'Robert B\t814-865-8799', 'Silviu C\t724-537-0671', 'Tom S\t570-368-3420', 'Tom S\t814-793-9655', '']) CONTACTS_SNAILMAIL = '\n'.join( ['Larry Luddite\t1949 Westcott St, Detroit, MI 48222', 'Lisa Luddite\t1949 Westcott St, Detroit, MI 48222', 'Robert B\t601 N 34th St, Seattle, WA 98103', 'Silviu C\t1600 Amphitheatre Pkwy, Mountain View, CA 94043', 'Tom S\t6425 Penn Ave Ste 700, Pittsburgh, PA 15206', 'Wally Writer\t101 Ridge Rd, Chalkyitsik, AK 99788', '']) EXPECTED_TSV = '\n'.join( ['\t'.join(['"Larry Luddite"', '""', '"724-228-3529"', '"1949 Westcott St, Detroit, MI 48222"']), '\t'.join(['"Lisa Luddite"', '""', '"304-277-3504"', '"1949 Westcott St, Detroit, MI 48222"']), '\t'.join(['"Nathan Nomad"', '"nathan@example.com"', '"412-466-8968"', '""']), '\t'.join(['"Nicky Nomad"', '"nicky@example.com"', '"724-379-5815"', '""']), '\t'.join(['"Noreen Nomad"', '"nomad@example.com,noreen@example.com"', '"412-472-0145"', '""']), '\t'.join(['"Robert B"', '"robert@example.com"', '"814-865-8799"', '"601 N 34th St, Seattle, WA 98103"']), '\t'.join(['"Silviu C"', '"silviu@example.com"', '"724-537-0671"', '"1600 Amphitheatre Pkwy, Mountain View, CA 94043"']), '\t'.join(['"Tom S"', '"tom@example.com"', '"570-368-3420,814-793-9655"', '"6425 Penn Ave Ste 700, Pittsburgh, PA 15206"']), '\t'.join(['"Wally Writer"', '"wally@example.com"', '""', '"101 Ridge Rd, Chalkyitsik, AK 99788"']), '']) EXPECTED_STATS = '\n'.join(['2 luddites', '1 writers', '3 nomads', '']) def create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def normalize_tsv_results(self, tsv_data): """Sort .tsv file data so we can compare it with expected output.""" lines_in = tsv_data.strip().split('\n') lines_out = [] for line in lines_in: name, email, phone, snailmail = line.split('\t') lines_out.append('\t'.join( [name, '"%s"' % ','.join(sorted(email.strip('"').split(','))), '"%s"' % ','.join(sorted(phone.strip('"').split(','))), snailmail])) return '\n'.join(sorted(lines_out)) + '\n' def test_mergecontacts(self): path_email = self.create_temp_file(self.CONTACTS_EMAIL) path_phone = self.create_temp_file(self.CONTACTS_PHONE) path_snailmail = self.create_temp_file(self.CONTACTS_SNAILMAIL) result_prefix = self.create_temp_file('') mergecontacts.run([ '--input_email=%s' % path_email, '--input_phone=%s' % path_phone, '--input_snailmail=%s' % path_snailmail, '--output_tsv=%s.tsv' % result_prefix, '--output_stats=%s.stats' % result_prefix], assert_results=(2, 1, 3)) with open('%s.tsv-00000-of-00001' % result_prefix) as f: contents = f.read() self.assertEqual(self.EXPECTED_TSV, self.normalize_tsv_results(contents)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/multiple_output_pardo.py000066400000000000000000000141051272340251000316410ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A workflow demonstrating a DoFn with multiple outputs. DoFns may produce a main output and additional side outputs. These side outputs are marked with a tag at output time and later the same tag will be used to get the corresponding result (a PCollection) for that side output. This is a slightly modified version of the basic wordcount example. In this example words are divided into 2 buckets as shorts words (3 characters in length or less) and words (all other words). There will be 3 output files: [OUTPUT]-chars : Character count for the input. [OUTPUT]-short-words : Word count for short words only. [OUTPUT]-words : Word count for all other words. To execute this pipeline locally, specify a local output file or output prefix on GCS: --output [YOUR_LOCAL_FILE | gs://YOUR_OUTPUT_PREFIX] To execute this pipeline using the Google Cloud Dataflow service, specify pipeline configuration: --project YOUR_PROJECT_ID --staging_location gs://YOUR_STAGING_DIRECTORY --temp_location gs://YOUR_TEMP_DIRECTORY --job_name YOUR_JOB_NAME --runner BlockingDataflowPipelineRunner and an output prefix on GCS: --output gs://YOUR_OUTPUT_PREFIX """ from __future__ import absolute_import import argparse import logging import re import google.cloud.dataflow as df from google.cloud.dataflow import pvalue class SplitLinesToWordsFn(df.DoFn): """A transform to split a line of text into individual words. This transform will have 3 outputs: - main output: all words that are longer than 3 characters. - short words side output: all other words. - character count side output: Number of characters in each processed line. """ # These tags will be used to tag the side outputs of this DoFn. SIDE_OUTPUT_TAG_SHORT_WORDS = 'tag_short_words' SIDE_OUTPUT_TAG_CHARACTER_COUNT = 'tag_character_count' def process(self, context): """Receives a single element (a line) and produces words and side outputs. Important things to note here: - For a single element you may produce multiple main outputs: words of a single line. - For that same input you may produce multiple side outputs, along with multiple main outputs. - Side outputs may have different types (count) or may share the same type (words) as with the main output. Args: context: processing context. Yields: words as main output, short words as side output, line character count as side output. """ # yield a count (integer) to the SIDE_OUTPUT_TAG_CHARACTER_COUNT tagged # collection. yield pvalue.SideOutputValue(self.SIDE_OUTPUT_TAG_CHARACTER_COUNT, len(context.element)) words = re.findall(r'[A-Za-z\']+', context.element) for word in words: if len(word) <= 3: # yield word as a side output to the SIDE_OUTPUT_TAG_SHORT_WORDS tagged # collection. yield pvalue.SideOutputValue(self.SIDE_OUTPUT_TAG_SHORT_WORDS, word) else: # yield word to add it to the main collection. yield word class CountWords(df.PTransform): """A transform to count the occurrences of each word. A PTransform that converts a PCollection containing words into a PCollection of "word: count" strings. """ def apply(self, pcoll): return (pcoll | df.Map('pair_with_one', lambda x: (x, 1)) | df.GroupByKey('group') | df.Map('count', lambda (word, ones): (word, sum(ones))) | df.Map('format', lambda (word, c): '%s: %s' % (word, c))) def run(argv=None): """Runs the workflow counting the long words and short words separately.""" parser = argparse.ArgumentParser() parser.add_argument('--input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument('--output', required=True, help='Output prefix for files to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) lines = p | df.Read('read', df.io.TextFileSource(known_args.input)) # with_outputs allows accessing the side outputs of a DoFn. split_lines_result = (lines | df.ParDo(SplitLinesToWordsFn()).with_outputs( SplitLinesToWordsFn.SIDE_OUTPUT_TAG_SHORT_WORDS, SplitLinesToWordsFn.SIDE_OUTPUT_TAG_CHARACTER_COUNT, main='words')) # split_lines_result is an object of type DoOutputsTuple. It supports # accessing result in alternative ways. words, _, _ = split_lines_result short_words = split_lines_result[ SplitLinesToWordsFn.SIDE_OUTPUT_TAG_SHORT_WORDS] character_count = split_lines_result.tag_character_count # pylint: disable=expression-not-assigned (character_count | df.Map('pair_with_key', lambda x: ('chars_temp_key', x)) | df.GroupByKey() | df.Map('count chars', lambda (_, counts): sum(counts)) | df.Write('write chars', df.io.TextFileSink(known_args.output + '-chars'))) # pylint: disable=expression-not-assigned (short_words | CountWords('count short words') | df.Write('write short words', df.io.TextFileSink(known_args.output + '-short-words'))) # pylint: disable=expression-not-assigned (words | CountWords('count words') | df.Write('write words', df.io.TextFileSink(known_args.output + '-words'))) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/cookbook/multiple_output_pardo_test.py000066400000000000000000000045211272340251000327010ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the multiple_output_pardo example.""" import logging import re import tempfile import unittest from google.cloud.dataflow.examples.cookbook import multiple_output_pardo class MultipleOutputParDo(unittest.TestCase): SAMPLE_TEXT = 'A whole new world\nA new fantastic point of view' EXPECTED_SHORT_WORDS = [('A', 2), ('new', 2), ('of', 1)] EXPECTED_WORDS = [ ('whole', 1), ('world', 1), ('fantastic', 1), ('point', 1), ('view', 1)] def create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def get_wordcount_results(self, temp_path): results = [] with open(temp_path) as result_file: for line in result_file: match = re.search(r'([A-Za-z]+): ([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) return results def test_multiple_output_pardo(self): temp_path = self.create_temp_file(self.SAMPLE_TEXT) result_prefix = temp_path + '.result' multiple_output_pardo.run([ '--input=%s*' % temp_path, '--output=%s' % result_prefix]) expected_char_count = len(''.join(self.SAMPLE_TEXT.split('\n'))) with open(result_prefix + '-chars-00000-of-00001') as f: contents = f.read() self.assertEqual(expected_char_count, int(contents)) short_words = self.get_wordcount_results( result_prefix + '-short-words-00000-of-00001') self.assertEqual(sorted(short_words), sorted(self.EXPECTED_SHORT_WORDS)) words = self.get_wordcount_results(result_prefix + '-words-00000-of-00001') self.assertEqual(sorted(words), sorted(self.EXPECTED_WORDS)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/snippets/000077500000000000000000000000001272340251000246655ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/snippets/snippets.py000066400000000000000000000717741272340251000271240ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Code snippets used in Cloud Dataflow webdocs. The examples here are written specifically to read well with the accompanying web docs from https://cloud.google.com/dataflow. Do not rewrite them until you make sure the webdocs still read well and the rewritten code supports the concept being described. For example, there are snippets that could be shorter but they are written like this to make a specific point in the docs. The code snippets are all organized as self contained functions. Parts of the function body delimited by [START tag] and [END tag] will be included automatically in the web docs. The naming convention for the tags is to have as prefix the PATH_TO_HTML where they are included followed by a descriptive string. For instance a code snippet that will be used as a code example at https://cloud.google.com/dataflow/model/pipelines will have the tag model_pipelines_DESCRIPTION. The tags can contain only letters, digits and _. """ import google.cloud.dataflow as df # Quiet some pylint warnings that happen because of the somewhat special # format for the code snippets. # pylint:disable=invalid-name # pylint:disable=expression-not-assigned # pylint:disable=redefined-outer-name # pylint:disable=unused-variable # pylint:disable=g-doc-args # pylint:disable=g-import-not-at-top class SnippetUtils(object): from google.cloud.dataflow.pipeline import PipelineVisitor class RenameFiles(PipelineVisitor): """RenameFiles will rewire source and sink for unit testing. RenameFiles will rewire the GCS files specified in the source and sink in the snippet pipeline to local files so the pipeline can be run as a unit test. This is as close as we can get to have code snippets that are executed and are also ready to presented in webdocs. """ def __init__(self, renames): self.renames = renames def visit_transform(self, transform_node): if hasattr(transform_node.transform, 'source'): source = transform_node.transform.source source.file_path = self.renames['read'] source.is_gcs_source = False elif hasattr(transform_node.transform, 'sink'): sink = transform_node.transform.sink sink.file_path = self.renames['write'] sink.is_gcs_sink = False def construct_pipeline(renames): """A reverse words snippet as an example for constructing a pipeline. URL: https://cloud.google.com/dataflow/pipelines/constructing-your-pipeline """ import re class ReverseWords(df.PTransform): """A PTransform that reverses individual elements in a PCollection.""" def apply(self, pcoll): return pcoll | df.Map(lambda e: e[::-1]) def filter_words(unused_x): """Pass through filter to select everything.""" return True # [START pipelines_constructing_creating] from google.cloud.dataflow.utils.options import PipelineOptions p = df.Pipeline(options=PipelineOptions()) # [END pipelines_constructing_creating] # [START pipelines_constructing_reading] lines = p | df.io.Read('ReadMyFile', df.io.TextFileSource('gs://some/inputData.txt')) # [END pipelines_constructing_reading] # [START pipelines_constructing_applying] words = lines | df.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) reversed_words = words | ReverseWords() # [END pipelines_constructing_applying] # [START pipelines_constructing_writing] filtered_words = reversed_words | df.Filter('FilterWords', filter_words) filtered_words | df.io.Write('WriteMyFile', df.io.TextFileSink('gs://some/outputData.txt')) # [END pipelines_constructing_writing] p.visit(SnippetUtils.RenameFiles(renames)) # [START pipelines_constructing_running] p.run() # [END pipelines_constructing_running] def model_pipelines(argv): """A wordcount snippet as a simple pipeline example. URL: https://cloud.google.com/dataflow/model/pipelines """ # [START model_pipelines] import re import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions class MyOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear' '.txt', help='Input file to process.') parser.add_argument('--output', dest='output', required=True, help='Output file to write results to.') pipeline_options = PipelineOptions(argv) my_options = pipeline_options.view_as(MyOptions) p = df.Pipeline(options=pipeline_options) (p | df.io.Read(df.io.TextFileSource(my_options.input)) | df.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) | df.Map(lambda x: (x, 1)) | df.combiners.Count.PerKey() | df.io.Write(df.io.TextFileSink(my_options.output))) p.run() # [END model_pipelines] def model_pcollection(argv): """Creating a PCollection from data in local memory. URL: https://cloud.google.com/dataflow/model/pcollection """ from google.cloud.dataflow.utils.options import PipelineOptions class MyOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--output', dest='output', required=True, help='Output file to write results to.') pipeline_options = PipelineOptions(argv) my_options = pipeline_options.view_as(MyOptions) # [START model_pcollection] p = df.Pipeline(options=pipeline_options) (p | df.Create([ 'To be, or not to be: that is the question: ', 'Whether \'tis nobler in the mind to suffer ', 'The slings and arrows of outrageous fortune, ', 'Or to take arms against a sea of troubles, ']) | df.io.Write(df.io.TextFileSink(my_options.output))) p.run() # [END model_pcollection] def pipeline_options_remote(argv): """"Creating a Pipeline using a PipelineOptions object for remote execution. URL: https://cloud.google.com/dataflow/pipelines/specifying-exec-params """ from google.cloud.dataflow import Pipeline from google.cloud.dataflow.utils.options import PipelineOptions # [START pipeline_options_create] options = PipelineOptions(flags=argv) # [END pipeline_options_create] # [START pipeline_options_define_custom] class MyOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--input') parser.add_argument('--output') # [END pipeline_options_define_custom] from google.cloud.dataflow.utils.options import GoogleCloudOptions from google.cloud.dataflow.utils.options import StandardOptions # [START pipeline_options_dataflow_service] # Create and set your PipelineOptions. options = PipelineOptions(flags=argv) # For Cloud execution, set the Cloud Platform project, job_name, # staging location, temp_location and specify DataflowPipelineRunner or # BlockingDataflowPipelineRunner. google_cloud_options = options.view_as(GoogleCloudOptions) google_cloud_options.project = 'my-project-id' google_cloud_options.job_name = 'myjob' google_cloud_options.staging_location = 'gs://my-bucket/binaries' google_cloud_options.temp_location = 'gs://my-bucket/temp' options.view_as(StandardOptions).runner = 'DataflowPipelineRunner' # Create the Pipeline with the specified options. p = Pipeline(options=options) # [END pipeline_options_dataflow_service] my_options = options.view_as(MyOptions) my_input = my_options.input my_output = my_options.output # Overriding the runner for tests. options.view_as(StandardOptions).runner = 'DirectPipelineRunner' p = Pipeline(options=options) lines = p | df.io.Read('ReadFromText', df.io.TextFileSource(my_input)) lines | df.io.Write('WriteToText', df.io.TextFileSink(my_output)) p.run() def pipeline_options_local(argv): """"Creating a Pipeline using a PipelineOptions object for local execution. URL: https://cloud.google.com/dataflow/pipelines/specifying-exec-params """ from google.cloud.dataflow import Pipeline from google.cloud.dataflow.utils.options import PipelineOptions options = PipelineOptions(flags=argv) # [START pipeline_options_define_custom_with_help_and_default] class MyOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--input', help='Input for the dataflow pipeline', default='gs://my-bucket/input') parser.add_argument('--output', help='Output for the dataflow pipeline', default='gs://my-bucket/output') # [END pipeline_options_define_custom_with_help_and_default] my_options = options.view_as(MyOptions) my_input = my_options.input my_output = my_options.output # [START pipeline_options_local] # Create and set your Pipeline Options. options = PipelineOptions() p = Pipeline(options=options) # [END pipeline_options_local] lines = p | df.io.Read('ReadFromText', df.io.TextFileSource(my_input)) lines | df.io.Write('WriteToText', df.io.TextFileSink(my_output)) p.run() def pipeline_options_command_line(argv): """Creating a Pipeline by passing a list of arguments. URL: https://cloud.google.com/dataflow/pipelines/specifying-exec-params """ # [START pipeline_options_command_line] # Use Python argparse module to parse custom arguments import argparse parser = argparse.ArgumentParser() parser.add_argument('--input') parser.add_argument('--output') known_args, pipeline_args = parser.parse_known_args(argv) # Create the Pipeline with remaining arguments. p = df.Pipeline(argv=pipeline_args) lines = p | df.io.Read('ReadFromText', df.io.TextFileSource(known_args.input)) lines | df.io.Write('WriteToText', df.io.TextFileSink(known_args.output)) # [END pipeline_options_command_line] p.run() def pipeline_logging(lines, output): """Logging Pipeline Messages. URL: https://cloud.google.com/dataflow/pipelines/logging """ import re import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions # [START pipeline_logging] # import Python logging module. import logging class ExtractWordsFn(df.DoFn): def process(self, context): words = re.findall(r'[A-Za-z\']+', context.element) for word in words: yield word if word.lower() == 'love': # Log using the root logger at info or higher levels logging.info('Found : %s', word.lower()) # Remaining WordCount example code ... # [END pipeline_logging] p = df.Pipeline(options=PipelineOptions()) (p | df.Create(lines) | df.ParDo('ExtractWords', ExtractWordsFn()) | df.io.Write('WriteToText', df.io.TextFileSink(output))) p.run() def pipeline_monitoring(renames): """Using monitoring interface snippets. URL: https://cloud.google.com/dataflow/pipelines/dataflow-monitoring-intf """ import re import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions class WordCountOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--input', help='Input for the dataflow pipeline', default='gs://my-bucket/input') parser.add_argument('--output', help='output for the dataflow pipeline', default='gs://my-bucket/output') class ExtractWordsFn(df.DoFn): def process(self, context): words = re.findall(r'[A-Za-z\']+', context.element) for word in words: yield word class FormatCountsFn(df.DoFn): def process(self, context): word, count = context.element yield '%s: %s' % (word, count) # [START pipeline_monitoring_composite] # The CountWords Composite Transform inside the WordCount pipeline. class CountWords(df.PTransform): def apply(self, pcoll): return (pcoll # Convert lines of text into individual words. | df.ParDo('ExtractWords', ExtractWordsFn()) # Count the number of times each word occurs. | df.combiners.Count.PerElement() # Format each word and count into a printable string. | df.ParDo('FormatCounts', FormatCountsFn())) # [END pipeline_monitoring_composite] pipeline_options = PipelineOptions() options = pipeline_options.view_as(WordCountOptions) p = df.Pipeline(options=pipeline_options) # [START pipeline_monitoring_execution] (p # Read the lines of the input text. | df.io.Read('ReadLines', df.io.TextFileSource(options.input)) # Count the words. | CountWords() # Write the formatted word counts to output. | df.io.Write('WriteCounts', df.io.TextFileSink(options.output))) # [END pipeline_monitoring_execution] p.visit(SnippetUtils.RenameFiles(renames)) p.run() def examples_wordcount_minimal(renames): """MinimalWordCount example snippets. URL: https://cloud.google.com/dataflow/examples/wordcount-example#MinimalWordCount """ import re import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import GoogleCloudOptions from google.cloud.dataflow.utils.options import StandardOptions from google.cloud.dataflow.utils.options import PipelineOptions # [START examples_wordcount_minimal_options] options = PipelineOptions() google_cloud_options = options.view_as(GoogleCloudOptions) google_cloud_options.project = 'my-project-id' google_cloud_options.job_name = 'myjob' google_cloud_options.staging_location = 'gs://your-bucket-name-here/staging' google_cloud_options.temp_location = 'gs://your-bucket-name-here/temp' options.view_as(StandardOptions).runner = 'BlockingDataflowPipelineRunner' # [END examples_wordcount_minimal_options] # Run it locally for testing. options = PipelineOptions() # [START examples_wordcount_minimal_create] p = df.Pipeline(options=options) # [END examples_wordcount_minimal_create] ( # [START examples_wordcount_minimal_read] p | df.io.Read(df.io.TextFileSource( 'gs://dataflow-samples/shakespeare/kinglear.txt')) # [END examples_wordcount_minimal_read] # [START examples_wordcount_minimal_pardo] | df.FlatMap('ExtractWords', lambda x: re.findall(r'[A-Za-z\']+', x)) # [END examples_wordcount_minimal_pardo] # [START examples_wordcount_minimal_count] | df.combiners.Count.PerElement() # [END examples_wordcount_minimal_count] # [START examples_wordcount_minimal_map] | df.Map(lambda (word, count): '%s: %s' % (word, count)) # [END examples_wordcount_minimal_map] # [START examples_wordcount_minimal_write] | df.io.Write(df.io.TextFileSink('gs://my-bucket/counts.txt')) # [END examples_wordcount_minimal_write] ) p.visit(SnippetUtils.RenameFiles(renames)) # [START examples_wordcount_minimal_run] p.run() # [END examples_wordcount_minimal_run] def examples_wordcount_wordcount(renames): """WordCount example snippets. URL: https://cloud.google.com/dataflow/examples/wordcount-example#WordCount """ import re import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions argv = [] # [START examples_wordcount_wordcount_options] class WordCountOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--input', help='Input for the dataflow pipeline', default='gs://my-bucket/input') options = PipelineOptions(argv) p = df.Pipeline(options=options) # [END examples_wordcount_wordcount_options] lines = p | df.io.Read(df.io.TextFileSource( 'gs://dataflow-samples/shakespeare/kinglear.txt')) # [START examples_wordcount_wordcount_composite] class CountWords(df.PTransform): def apply(self, pcoll): return (pcoll # Convert lines of text into individual words. | df.FlatMap( 'ExtractWords', lambda x: re.findall(r'[A-Za-z\']+', x)) # Count the number of times each word occurs. | df.combiners.Count.PerElement()) counts = lines | CountWords() # [END examples_wordcount_wordcount_composite] # [START examples_wordcount_wordcount_dofn] class FormatAsTextFn(df.DoFn): def process(self, context): word, count = context.element yield '%s: %s' % (word, count) formatted = counts | df.ParDo(FormatAsTextFn()) # [END examples_wordcount_wordcount_dofn] formatted | df.io.Write(df.io.TextFileSink('gs://my-bucket/counts.txt')) p.visit(SnippetUtils.RenameFiles(renames)) p.run() def examples_wordcount_debugging(renames): """DebuggingWordCount example snippets. URL: https://cloud.google.com/dataflow/examples/wordcount-example#DebuggingWordCount """ import re import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions # [START example_wordcount_debugging_logging] # [START example_wordcount_debugging_aggregators] import logging class FilterTextFn(df.DoFn): """A DoFn that filters for a specific key based on a regular expression.""" # A custom aggregator can track values in your pipeline as it runs. Create # custom aggregators matched_word and unmatched_words. matched_words = df.Aggregator('matched_words') umatched_words = df.Aggregator('umatched_words') def __init__(self, pattern): self.pattern = pattern def process(self, context): word, _ = context.element if re.match(self.pattern, word): # Log at INFO level each element we match. When executing this pipeline # using the Dataflow service, these log lines will appear in the Cloud # Logging UI. logging.info('Matched %s', word) # Add 1 to the custom aggregator matched_words context.aggregate_to(self.matched_words, 1) yield context.element else: # Log at the "DEBUG" level each element that is not matched. Different # log levels can be used to control the verbosity of logging providing # an effective mechanism to filter less important information. Note # currently only "INFO" and higher level logs are emitted to the Cloud # Logger. This log message will not be visible in the Cloud Logger. logging.debug('Did not match %s', word) # Add 1 to the custom aggregator umatched_words context.aggregate_to(self.umatched_words, 1) # [END example_wordcount_debugging_logging] # [END example_wordcount_debugging_aggregators] p = df.Pipeline(options=PipelineOptions()) filtered_words = ( p | df.io.Read(df.io.TextFileSource( 'gs://dataflow-samples/shakespeare/kinglear.txt')) | df.FlatMap('ExtractWords', lambda x: re.findall(r'[A-Za-z\']+', x)) | df.combiners.Count.PerElement() | df.ParDo('FilterText', FilterTextFn('Flourish|stomach'))) # [START example_wordcount_debugging_assert] df.assert_that(filtered_words, df.equal_to([('Flourish', 3), ('stomach', 1)])) # [END example_wordcount_debugging_assert] output = (filtered_words | df.Map('format', lambda (word, c): '%s: %s' % (word, c)) | df.io.Write( 'write', df.io.TextFileSink('gs://my-bucket/counts.txt'))) p.visit(SnippetUtils.RenameFiles(renames)) p.run() def model_textio(renames): """Using a Read and Write transform to read/write text files. URLs: https://cloud.google.com/dataflow/model/pipeline-io https://cloud.google.com/dataflow/model/text-io """ def filter_words(x): import re return re.findall(r'[A-Za-z\']+', x) import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions # [START model_textio_read] p = df.Pipeline(options=PipelineOptions()) # [START model_pipelineio_read] lines = p | df.io.Read( 'ReadFromText', df.io.TextFileSource('gs://my_bucket/path/to/input-*.csv')) # [END model_pipelineio_read] # [END model_textio_read] # [START model_textio_write] filtered_words = lines | df.FlatMap('FilterWords', filter_words) # [START model_pipelineio_write] filtered_words | df.io.Write( 'WriteToText', df.io.TextFileSink('gs://my_bucket/path/to/numbers', file_name_suffix='.csv')) # [END model_pipelineio_write] # [END model_textio_write] p.visit(SnippetUtils.RenameFiles(renames)) p.run() def model_bigqueryio(): """Using a Read and Write transform to read/write to BigQuery. URL: https://cloud.google.com/dataflow/model/bigquery-io """ import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions # [START model_bigqueryio_read] p = df.Pipeline(options=PipelineOptions()) weather_data = p | df.io.Read( 'ReadWeatherStations', df.io.BigQuerySource( 'clouddataflow-readonly:samples.weather_stations')) # [END model_bigqueryio_read] # [START model_bigqueryio_query] p = df.Pipeline(options=PipelineOptions()) weather_data = p | df.io.Read( 'ReadYearAndTemp', df.io.BigQuerySource( query='SELECT year, mean_temp FROM samples.weather_stations')) # [END model_bigqueryio_query] # [START model_bigqueryio_schema] schema = 'source:STRING, quote:STRING' # [END model_bigqueryio_schema] # [START model_bigqueryio_write] quotes = p | df.Create( [{'source': 'Mahatma Ghandi', 'quote': 'My life is my message.'}]) quotes | df.io.Write( 'Write', df.io.BigQuerySink( 'my-project:output.output_table', schema=schema, write_disposition=df.io.BigQueryDisposition.WRITE_TRUNCATE, create_disposition=df.io.BigQueryDisposition.CREATE_IF_NEEDED)) # [END model_bigqueryio_write] def model_composite_transform_example(contents, output_path): """Example of a composite transform. To declare a composite transform, define a subclass of PTransform. To override the apply method, define a method "apply" that takes a PCollection as its only parameter and returns a PCollection. URL: https://cloud.google.com/dataflow/model/composite-transforms """ import re import google.cloud.dataflow as df # [START composite_transform_example] # [START composite_ptransform_apply_method] # [START composite_ptransform_declare] class CountWords(df.PTransform): # [END composite_ptransform_declare] def apply(self, pcoll): return (pcoll | df.FlatMap(lambda x: re.findall(r'\w+', x)) | df.combiners.Count.PerElement() | df.Map(lambda (word, c): '%s: %s' % (word, c))) # [END composite_ptransform_apply_method] # [END composite_transform_example] from google.cloud.dataflow.utils.options import PipelineOptions p = df.Pipeline(options=PipelineOptions()) (p | df.Create(contents) | CountWords() | df.io.Write(df.io.TextFileSink(output_path))) p.run() def model_multiple_pcollections_flatten(contents, output_path): """Merging a PCollection with Flatten. URL: https://cloud.google.com/dataflow/model/multiple-pcollections """ some_hash_fn = lambda s: ord(s[0]) import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions p = df.Pipeline(options=PipelineOptions()) partition_fn = lambda element, partitions: some_hash_fn(element) % partitions # Partition into deciles partitioned = p | df.Create(contents) | df.Partition(partition_fn, 3) pcoll1 = partitioned[0] pcoll2 = partitioned[1] pcoll3 = partitioned[2] # Flatten them back into 1 # A collection of PCollection objects can be represented simply # as a tuple (or list) of PCollections. # (The SDK for Python has no separate type to store multiple # PCollection objects, whether containing the same or different # types.) # [START model_multiple_pcollections_flatten] merged = ( # [START model_multiple_pcollections_tuple] (pcoll1, pcoll2, pcoll3) # [END model_multiple_pcollections_tuple] # A list of tuples can be "piped" directly into a Flatten transform. | df.Flatten()) # [END model_multiple_pcollections_flatten] merged | df.io.Write(df.io.TextFileSink(output_path)) p.run() def model_multiple_pcollections_partition(contents, output_path): """Splitting a PCollection with Partition. URL: https://cloud.google.com/dataflow/model/multiple-pcollections """ some_hash_fn = lambda s: ord(s[0]) def get_percentile(i): """Assume i in [0,100).""" return i import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions p = df.Pipeline(options=PipelineOptions()) students = p | df.Create(contents) # [START model_multiple_pcollections_partition] def partition_fn(student, num_partitions): return int(get_percentile(student) * num_partitions / 100) by_decile = students | df.Partition(partition_fn, 10) # [END model_multiple_pcollections_partition] # [START model_multiple_pcollections_partition_40th] fortieth_percentile = by_decile[4] # [END model_multiple_pcollections_partition_40th] ([by_decile[d] for d in xrange(10) if d != 4] + [fortieth_percentile] | df.Flatten() | df.io.Write(df.io.TextFileSink(output_path))) p.run() def model_group_by_key(contents, output_path): """Applying a GroupByKey Transform. URL: https://cloud.google.com/dataflow/model/group-by-key """ import re import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions p = df.Pipeline(options=PipelineOptions()) words_and_counts = ( p | df.Create(contents) | df.FlatMap(lambda x: re.findall(r'\w+', x)) | df.Map('one word', lambda w: (w, 1))) # GroupByKey accepts a PCollection of (w, 1) and # outputs a PCollection of (w, (1, 1, ...)). # (A key/value pair is just a tuple in Python.) # This is a somewhat forced example, since one could # simply use df.combiners.Count.PerElement here. # [START model_group_by_key_transform] grouped_words = words_and_counts | df.GroupByKey() # [END model_group_by_key_transform] (grouped_words | df.Map('count words', lambda (word, counts): (word, len(counts))) | df.io.Write(df.io.TextFileSink(output_path))) p.run() def model_co_group_by_key_tuple(email_list, phone_list, output_path): """Applying a CoGroupByKey Transform to a tuple. URL: https://cloud.google.com/dataflow/model/group-by-key """ import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import PipelineOptions p = df.Pipeline(options=PipelineOptions()) # [START model_group_by_key_cogroupbykey_tuple] # Each data set is represented by key-value pairs in separate PCollections. # Both data sets share a common key type (in this example str). # The email_list contains values such as: ('joe', 'joe@example.com') with # multiple possible values for each key. # The phone_list contains values such as: ('mary': '111-222-3333') with # multiple possible values for each key. emails = p | df.Create('email', email_list) phones = p | df.Create('phone', phone_list) # The result PCollection contains one key-value element for each key in the # input PCollections. The key of the pair will be the key from the input and # the value will be a dictionary with two entries: 'emails' - an iterable of # all values for the current key in the emails PCollection and 'phones': an # iterable of all values for the current key in the phones PCollection. # For instance, if 'emails' contained ('joe', 'joe@example.com') and # ('joe', 'joe@gmail.com'), then 'result' will contain the element # ('joe', {'emails': ['joe@example.com', 'joe@gmail.com'], 'phones': ...}) result = {'emails': emails, 'phones': phones} | df.CoGroupByKey() def join_info((name, info)): return '; '.join(['%s' % name, '%s' % ','.join(info['emails']), '%s' % ','.join(info['phones'])]) contact_lines = result | df.Map(join_info) # [END model_group_by_key_cogroupbykey_tuple] contact_lines | df.io.Write(df.io.TextFileSink(output_path)) p.run() # [START model_library_transforms_keys] class Keys(df.PTransform): def apply(self, pcoll): return pcoll | df.Map('Keys', lambda (k, v): k) # [END model_library_transforms_keys] # pylint: enable=invalid-name # [START model_library_transforms_count] class Count(df.PTransform): def apply(self, pcoll): return ( pcoll | df.Map('Init', lambda v: (v, 1)) | df.CombinePerKey(sum)) # [END model_library_transforms_count] # pylint: enable=g-wrong-blank-lines DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/snippets/snippets_test.py000066400000000000000000000505071272340251000301520ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for all code snippets used in public docs.""" import logging import sys import tempfile import unittest import google.cloud.dataflow as df from google.cloud.dataflow import io from google.cloud.dataflow import pvalue from google.cloud.dataflow import typehints from google.cloud.dataflow.examples.snippets import snippets from google.cloud.dataflow.io import fileio from google.cloud.dataflow.utils.options import TypeOptions # Monky-patch to use native sink for file path re-writing. io.TextFileSink = fileio.NativeTextFileSink class ParDoTest(unittest.TestCase): """Tests for dataflow/model/par-do.""" def test_pardo(self): # Note: "words" and "ComputeWordLengthFn" are referenced by name in # the text of the doc. words = ['aa', 'bbb', 'c'] # [START model_pardo_pardo] class ComputeWordLengthFn(df.DoFn): def process(self, context): return [len(context.element)] # [END model_pardo_pardo] # [START model_pardo_apply] # Apply a ParDo to the PCollection "words" to compute lengths for each word. word_lengths = words | df.ParDo(ComputeWordLengthFn()) # [END model_pardo_apply] self.assertEqual({2, 3, 1}, set(word_lengths)) def test_pardo_yield(self): words = ['aa', 'bbb', 'c'] # [START model_pardo_yield] class ComputeWordLengthFn(df.DoFn): def process(self, context): yield len(context.element) # [END model_pardo_yield] word_lengths = words | df.ParDo(ComputeWordLengthFn()) self.assertEqual({2, 3, 1}, set(word_lengths)) def test_pardo_using_map(self): words = ['aa', 'bbb', 'c'] # [START model_pardo_using_map] word_lengths = words | df.Map(len) # [END model_pardo_using_map] self.assertEqual({2, 3, 1}, set(word_lengths)) def test_pardo_using_flatmap(self): words = ['aa', 'bbb', 'c'] # [START model_pardo_using_flatmap] word_lengths = words | df.FlatMap(lambda word: [len(word)]) # [END model_pardo_using_flatmap] self.assertEqual({2, 3, 1}, set(word_lengths)) def test_pardo_using_flatmap_yield(self): words = ['aA', 'bbb', 'C'] # [START model_pardo_using_flatmap_yield] def capitals(word): for letter in word: if 'A' <= letter <= 'Z': yield letter all_capitals = words | df.FlatMap(capitals) # [END model_pardo_using_flatmap_yield] self.assertEqual({'A', 'C'}, set(all_capitals)) def test_pardo_with_label(self): words = ['aa', 'bbc', 'defg'] # [START model_pardo_with_label] result = words | df.Map('CountUniqueLetters', lambda word: len(set(word))) # [END model_pardo_with_label] self.assertEqual({1, 2, 4}, set(result)) def test_pardo_side_input(self): p = df.Pipeline('DirectPipelineRunner') words = p | df.Create('start', ['a', 'bb', 'ccc', 'dddd']) # [START model_pardo_side_input] # Callable takes additional arguments. def filter_using_length(word, lower_bound, upper_bound=float('inf')): if lower_bound <= len(word) <= upper_bound: yield word # Construct a deferred side input. avg_word_len = words | df.Map(len) | df.CombineGlobally(df.combiners.MeanCombineFn()) # Call with explicit side inputs. small_words = words | df.FlatMap('small', filter_using_length, 0, 3) # A single deferred side input. larger_than_average = words | df.FlatMap('large', filter_using_length, lower_bound=pvalue.AsSingleton(avg_word_len)) # Mix and match. small_but_nontrivial = words | df.FlatMap(filter_using_length, lower_bound=2, upper_bound=pvalue.AsSingleton(avg_word_len)) # [END model_pardo_side_input] df.assert_that(small_words, df.equal_to(['a', 'bb', 'ccc'])) df.assert_that(larger_than_average, df.equal_to(['ccc', 'dddd']), label='larger_than_average') df.assert_that(small_but_nontrivial, df.equal_to(['bb']), label='small_but_not_trivial') p.run() def test_pardo_side_input_dofn(self): words = ['a', 'bb', 'ccc', 'dddd'] # [START model_pardo_side_input_dofn] class FilterUsingLength(df.DoFn): def process(self, context, lower_bound, upper_bound=float('inf')): if lower_bound <= len(context.element) <= upper_bound: yield context.element small_words = words | df.ParDo(FilterUsingLength(), 0, 3) # [END model_pardo_side_input_dofn] self.assertEqual({'a', 'bb', 'ccc'}, set(small_words)) def test_pardo_with_side_outputs(self): # [START model_pardo_emitting_values_on_side_outputs] class ProcessWords(df.DoFn): def process(self, context, cutoff_length, marker): if len(context.element) <= cutoff_length: # Emit this short word to the main output. yield context.element else: # Emit this word's long length to a side output. yield pvalue.SideOutputValue( 'above_cutoff_lengths', len(context.element)) if context.element.startswith(marker): # Emit this word to a different side output. yield pvalue.SideOutputValue('marked strings', context.element) # [END model_pardo_emitting_values_on_side_outputs] words = ['a', 'an', 'the', 'music', 'xyz'] # [START model_pardo_with_side_outputs] results = (words | df.ParDo(ProcessWords(), cutoff_length=2, marker='x') .with_outputs('above_cutoff_lengths', 'marked strings', main='below_cutoff_strings')) below = results.below_cutoff_strings above = results.above_cutoff_lengths marked = results['marked strings'] # indexing works as well # [END model_pardo_with_side_outputs] self.assertEqual({'a', 'an'}, set(below)) self.assertEqual({3, 5}, set(above)) self.assertEqual({'xyz'}, set(marked)) # [START model_pardo_with_side_outputs_iter] below, above, marked = (words | df.ParDo(ProcessWords(), cutoff_length=2, marker='x') .with_outputs('above_cutoff_lengths', 'marked strings', main='below_cutoff_strings')) # [END model_pardo_with_side_outputs_iter] self.assertEqual({'a', 'an'}, set(below)) self.assertEqual({3, 5}, set(above)) self.assertEqual({'xyz'}, set(marked)) def test_pardo_with_undeclared_side_outputs(self): numbers = [1, 2, 3, 4, 5, 10, 20] # [START model_pardo_with_side_outputs_undeclared] def even_odd(x): yield pvalue.SideOutputValue('odd' if x % 2 else 'even', x) if x % 10 == 0: yield x results = numbers | df.FlatMap(even_odd).with_outputs() evens = results.even odds = results.odd tens = results[None] # the undeclared main output # [END model_pardo_with_side_outputs_undeclared] self.assertEqual({2, 4, 10, 20}, set(evens)) self.assertEqual({1, 3, 5}, set(odds)) self.assertEqual({10, 20}, set(tens)) class TypeHintsTest(unittest.TestCase): def test_bad_types(self): p = df.Pipeline('DirectPipelineRunner', argv=sys.argv) # [START type_hints_missing_define_numbers] numbers = p | df.Create(['1', '2', '3']) # [END type_hints_missing_define_numbers] # Consider the following code. # [START type_hints_missing_apply] evens = numbers | df.Filter(lambda x: x % 2 == 0) # [END type_hints_missing_apply] # Now suppose numers was defined as [snippet above]. # When running this pipeline, you'd get a runtime error, # possibly on a remote machine, possibly very late. with self.assertRaises(TypeError): p.run() # To catch this early, we can assert what types we expect. with self.assertRaises(typehints.TypeCheckError): # [START type_hints_takes] p.options.view_as(TypeOptions).pipeline_type_check = True evens = numbers | df.Filter(lambda x: x % 2 == 0).with_input_types(int) # [END type_hints_takes] # Type hints can be declared on DoFns and callables as well, rather # than where they're used, to be more self contained. with self.assertRaises(typehints.TypeCheckError): # [START type_hints_do_fn] @df.typehints.with_input_types(int) class FilterEvensDoFn(df.DoFn): def process(self, context): if context.element % 2 == 0: yield context.element evens = numbers | df.ParDo(FilterEvensDoFn()) # [END type_hints_do_fn] words = p | df.Create('words', ['a', 'bb', 'c']) # One can assert outputs and apply them to transforms as well. # Helps document the contract and checks it at pipeline construction time. # [START type_hints_transform] T = df.typehints.TypeVariable('T') @df.typehints.with_input_types(T) @df.typehints.with_output_types(df.typehints.Tuple[int, T]) class MyTransform(df.PTransform): def apply(self, pcoll): return pcoll | df.Map(lambda x: (len(x), x)) words_with_lens = words | MyTransform() # [END type_hints_transform] with self.assertRaises(typehints.TypeCheckError): words_with_lens | df.Map(lambda x: x).with_input_types( df.typehints.Tuple[int, int]) def test_runtime_checks_off(self): p = df.Pipeline('DirectPipelineRunner', argv=sys.argv) # [START type_hints_runtime_off] p | df.Create(['a']) | df.Map(lambda x: 3).with_output_types(str) p.run() # [END type_hints_runtime_off] def test_runtime_checks_on(self): p = df.Pipeline('DirectPipelineRunner', argv=sys.argv) with self.assertRaises(typehints.TypeCheckError): # [START type_hints_runtime_on] p.options.view_as(TypeOptions).runtime_type_check = True p | df.Create(['a']) | df.Map(lambda x: 3).with_output_types(str) p.run() # [END type_hints_runtime_on] def test_deterministic_key(self): p = df.Pipeline('DirectPipelineRunner', argv=sys.argv) lines = ['banana,fruit,3', 'kiwi,fruit,2', 'kiwi,fruit,2', 'zucchini,veg,3'] # [START type_hints_deterministic_key] class Player(object): def __init__(self, team, name): self.team = team self.name = name class PlayerCoder(df.coders.Coder): def encode(self, player): return '%s:%s' % (player.team, player.name) def decode(self, s): return Player(*s.split(':')) def is_deterministic(self): return True df.coders.registry.register_coder(Player, PlayerCoder) def parse_player_and_score(csv): name, team, score = csv.split(',') return Player(team, name), int(score) totals = ( lines | df.Map(parse_player_and_score) | df.CombinePerKey(sum).with_input_types(df.typehints.Tuple[Player, int])) # [END type_hints_deterministic_key] self.assertEquals( {('banana', 3), ('kiwi', 4), ('zucchini', 3)}, set(totals | df.Map(lambda (k, v): (k.name, v)))) class SnippetsTest(unittest.TestCase): def create_temp_file(self, contents=''): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def get_output(self, path, sorted_output=True, suffix=''): with open(path + '-00000-of-00001' + suffix) as f: lines = f.readlines() if sorted_output: return sorted(s.rstrip('\n') for s in lines) else: return [s.rstrip('\n') for s in lines] def test_model_pipelines(self): temp_path = self.create_temp_file('aa bb cc\n bb cc\n cc') result_path = temp_path + '.result' snippets.model_pipelines([ '--input=%s*' % temp_path, '--output=%s' % result_path]) self.assertEqual( self.get_output(result_path), [str(s) for s in [(u'aa', 1), (u'bb', 2), (u'cc', 3)]]) def test_model_pcollection(self): temp_path = self.create_temp_file() snippets.model_pcollection(['--output=%s' % temp_path]) self.assertEqual(self.get_output(temp_path, sorted_output=False), [ 'To be, or not to be: that is the question: ', 'Whether \'tis nobler in the mind to suffer ', 'The slings and arrows of outrageous fortune, ', 'Or to take arms against a sea of troubles, ']) def test_construct_pipeline(self): temp_path = self.create_temp_file( 'abc def ghi\n jkl mno pqr\n stu vwx yz') result_path = self.create_temp_file() snippets.construct_pipeline({'read': temp_path, 'write': result_path}) self.assertEqual( self.get_output(result_path), ['cba', 'fed', 'ihg', 'lkj', 'onm', 'rqp', 'uts', 'xwv', 'zy']) def test_model_textio(self): temp_path = self.create_temp_file('aa bb cc\n bb cc\n cc') result_path = temp_path + '.result' snippets.model_textio({'read': temp_path, 'write': result_path}) self.assertEqual( ['aa', 'bb', 'bb', 'cc', 'cc', 'cc'], self.get_output(result_path, suffix='.csv')) def test_model_bigqueryio(self): # We cannot test BigQueryIO functionality in unit tests therefore we limit # ourselves to making sure the pipeline containing BigQuery sources and # sinks can be built. self.assertEqual(None, snippets.model_bigqueryio()) def _run_test_pipeline_for_options(self, fn): temp_path = self.create_temp_file('aa\nbb\ncc') result_path = temp_path + '.result' fn([ '--input=%s*' % temp_path, '--output=%s' % result_path]) self.assertEqual( ['aa', 'bb', 'cc'], self.get_output(result_path)) def test_pipeline_options_local(self): self._run_test_pipeline_for_options(snippets.pipeline_options_local) def test_pipeline_options_remote(self): self._run_test_pipeline_for_options(snippets.pipeline_options_remote) def test_pipeline_options_command_line(self): self._run_test_pipeline_for_options(snippets.pipeline_options_command_line) def test_pipeline_logging(self): result_path = self.create_temp_file() lines = ['we found love right where we are', 'we found love right from the start', 'we found love in a hopeless place'] snippets.pipeline_logging(lines, result_path) self.assertEqual( sorted(' '.join(lines).split(' ')), self.get_output(result_path)) def test_examples_wordcount(self): pipelines = [snippets.examples_wordcount_minimal, snippets.examples_wordcount_wordcount, snippets.pipeline_monitoring] for pipeline in pipelines: temp_path = self.create_temp_file( 'abc def ghi\n abc jkl') result_path = self.create_temp_file() pipeline({'read': temp_path, 'write': result_path}) self.assertEqual( self.get_output(result_path), ['abc: 2', 'def: 1', 'ghi: 1', 'jkl: 1']) def test_examples_wordcount_debugging(self): temp_path = self.create_temp_file( 'Flourish Flourish Flourish stomach abc def') result_path = self.create_temp_file() snippets.examples_wordcount_debugging( {'read': temp_path, 'write': result_path}) self.assertEqual( self.get_output(result_path), ['Flourish: 3', 'stomach: 1']) def test_model_composite_transform_example(self): contents = ['aa bb cc', 'bb cc', 'cc'] result_path = self.create_temp_file() snippets.model_composite_transform_example(contents, result_path) self.assertEqual(['aa: 1', 'bb: 2', 'cc: 3'], self.get_output(result_path)) def test_model_multiple_pcollections_flatten(self): contents = ['a', 'b', 'c', 'd', 'e', 'f'] result_path = self.create_temp_file() snippets.model_multiple_pcollections_flatten(contents, result_path) self.assertEqual(contents, self.get_output(result_path)) def test_model_multiple_pcollections_partition(self): contents = [17, 42, 64, 32, 0, 99, 53, 89] result_path = self.create_temp_file() snippets.model_multiple_pcollections_partition(contents, result_path) self.assertEqual(['0', '17', '32', '42', '53', '64', '89', '99'], self.get_output(result_path)) def test_model_group_by_key(self): contents = ['a bb ccc bb bb a'] result_path = self.create_temp_file() snippets.model_group_by_key(contents, result_path) expected = [('a', 2), ('bb', 3), ('ccc', 1)] self.assertEqual([str(s) for s in expected], self.get_output(result_path)) def test_model_co_group_by_key_tuple(self): email_list = [['a', 'a@example.com'], ['b', 'b@example.com']] phone_list = [['a', 'x4312'], ['b', 'x8452']] result_path = self.create_temp_file() snippets.model_co_group_by_key_tuple(email_list, phone_list, result_path) expect = ['a; a@example.com; x4312', 'b; b@example.com; x8452'] self.assertEqual(expect, self.get_output(result_path)) class CombineTest(unittest.TestCase): """Tests for dataflow/model/combine.""" def test_global_sum(self): pc = [1, 2, 3] # [START global_sum] result = pc | df.CombineGlobally(sum) # [END global_sum] self.assertEqual([6], result) def test_combine_values(self): occurences = [('cat', 1), ('cat', 5), ('cat', 9), ('dog', 5), ('dog', 2)] # [START combine_values] first_occurences = occurences | df.GroupByKey() | df.CombineValues(min) # [END combine_values] self.assertEqual({('cat', 1), ('dog', 2)}, set(first_occurences)) def test_combine_per_key(self): player_accuracies = [ ('cat', 1), ('cat', 5), ('cat', 9), ('cat', 1), ('dog', 5), ('dog', 2)] # [START combine_per_key] avg_accuracy_per_player = player_accuracies | df.CombinePerKey(df.combiners.MeanCombineFn()) # [END combine_per_key] self.assertEqual({('cat', 4.0), ('dog', 3.5)}, set(avg_accuracy_per_player)) def test_combine_concat(self): pc = ['a', 'b'] # [START combine_concat] def concat(values, separator=', '): return separator.join(values) with_commas = pc | df.CombineGlobally(concat) with_dashes = pc | df.CombineGlobally(concat, separator='-') # [END combine_concat] self.assertEqual(1, len(with_commas)) self.assertTrue(with_commas[0] in {'a, b', 'b, a'}) self.assertEqual(1, len(with_dashes)) self.assertTrue(with_dashes[0] in {'a-b', 'b-a'}) def test_bounded_sum(self): # [START combine_bounded_sum] pc = [1, 10, 100, 1000] def bounded_sum(values, bound=500): return min(sum(values), bound) small_sum = pc | df.CombineGlobally(bounded_sum) # [500] large_sum = pc | df.CombineGlobally(bounded_sum, bound=5000) # [1111] # [END combine_bounded_sum] self.assertEqual([500], small_sum) self.assertEqual([1111], large_sum) def test_combine_reduce(self): factors = [2, 3, 5, 7] # [START combine_reduce] import functools import operator product = factors | df.CombineGlobally(functools.partial(reduce, operator.mul), 1) # [END combine_reduce] self.assertEqual([210], product) def test_custom_average(self): pc = [2, 3, 5, 7] # [START combine_custom_average] class AverageFn(df.CombineFn): def create_accumulator(self): return (0.0, 0) def add_input(self, (sum, count), input): return sum + input, count + 1 def merge_accumulators(self, accumulators): sums, counts = zip(*accumulators) return sum(sums), sum(counts) def extract_output(self, (sum, count)): return sum / count if count else float('NaN') average = pc | df.CombineGlobally(AverageFn()) # [END combine_custom_average] self.assertEqual([4.25], average) def test_keys(self): occurrences = [('cat', 1), ('cat', 5), ('dog', 5), ('cat', 9), ('dog', 2)] unique_keys = occurrences | snippets.Keys() self.assertEqual({'cat', 'dog'}, set(unique_keys)) def test_count(self): occurrences = ['cat', 'dog', 'cat', 'cat', 'dog'] perkey_counts = occurrences | snippets.Count() self.assertEqual({('cat', 3), ('dog', 2)}, set(perkey_counts)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/streaming_wordcap.py000066400000000000000000000035511272340251000271060ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A streaming string-capitalization workflow. Important: streaming pipeline support in Python Dataflow is in development and is not yet available for use. """ from __future__ import absolute_import import argparse import logging import google.cloud.dataflow as df def run(argv=None): """Build and run the pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input_topic', dest='input_topic', required=True, help='Input PubSub topic of the form "/topics//".') parser.add_argument( '--output_topic', dest='output_topic', required=True, help='Output PubSub topic of the form "/topics//".') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Read the text file[pattern] into a PCollection. lines = p | df.io.Read( 'read', df.io.PubSubSource(known_args.input_topic)) # Capitalize the characters in each line. transformed = (lines | (df.Map('capitalize', lambda x: x.upper()))) # Write to PubSub. # pylint: disable=expression-not-assigned transformed | df.io.Write( 'pubsub_write', df.io.PubSubSink(known_args.output_topic)) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/streaming_wordcount.py000066400000000000000000000044151272340251000274730ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A streaming word-counting workflow. Important: streaming pipeline support in Python Dataflow is in development and is not yet available for use. """ from __future__ import absolute_import import argparse import logging import re import google.cloud.dataflow as df import google.cloud.dataflow.transforms.window as window def run(argv=None): """Build and run the pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input_topic', required=True, help='Input PubSub topic of the form "/topics//".') parser.add_argument( '--output_topic', required=True, help='Output PubSub topic of the form "/topics//".') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Read the text file[pattern] into a PCollection. lines = p | df.io.Read( 'read', df.io.PubSubSource(known_args.input_topic)) # Capitalize the characters in each line. transformed = (lines | (df.FlatMap('split', lambda x: re.findall(r'[A-Za-z\']+', x)) .with_output_types(unicode)) | df.Map('pair_with_one', lambda x: (x, 1)) | df.WindowInto(window.FixedWindows(15, 0)) | df.GroupByKey('group') | df.Map('count', lambda (word, ones): (word, sum(ones))) | df.Map('format', lambda tup: '%s: %d' % tup)) # Write to PubSub. # pylint: disable=expression-not-assigned transformed | df.io.Write( 'pubsub_write', df.io.PubSubSink(known_args.output_topic)) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/wordcount.py000066400000000000000000000066151272340251000254260ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A word-counting workflow.""" from __future__ import absolute_import import argparse import logging import re import google.cloud.dataflow as df empty_line_aggregator = df.Aggregator('emptyLines') average_word_size_aggregator = df.Aggregator('averageWordLength', df.combiners.MeanCombineFn(), float) class WordExtractingDoFn(df.DoFn): """Parse each line of input text into words.""" def process(self, context): """Returns an iterator over the words of this element. The element is a line of text. If the line is blank, note that, too. Args: context: the call-specific context: data and aggregator. Returns: The processed element. """ text_line = context.element.strip() if not text_line: context.aggregate_to(empty_line_aggregator, 1) words = re.findall(r'[A-Za-z\']+', text_line) for w in words: context.aggregate_to(average_word_size_aggregator, len(w)) return words def run(argv=None): """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() parser.add_argument('--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument('--output', dest='output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Read the text file[pattern] into a PCollection. lines = p | df.io.Read('read', df.io.TextFileSource(known_args.input)) # Count the occurrences of each word. counts = (lines | (df.ParDo('split', WordExtractingDoFn()) .with_output_types(unicode)) | df.Map('pair_with_one', lambda x: (x, 1)) | df.GroupByKey('group') | df.Map('count', lambda (word, ones): (word, sum(ones)))) # Format the counts into a PCollection of strings. output = counts | df.Map('format', lambda (word, c): '%s: %s' % (word, c)) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | df.io.Write('write', df.io.TextFileSink(known_args.output)) # Actually run the pipeline (all operations above are deferred). result = p.run() empty_line_values = result.aggregated_values(empty_line_aggregator) logging.info('number of empty lines: %d', sum(empty_line_values.values())) word_length_values = result.aggregated_values(average_word_size_aggregator) logging.info('average word lengths: %s', word_length_values.values()) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/wordcount_debugging.py000066400000000000000000000134351272340251000274370ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """An example that verifies the counts and includes Dataflow best practices. On top of the basic concepts in the wordcount example, this workflow introduces logging to Cloud Logging, and using assertions in a Dataflow pipeline. To execute this pipeline locally, specify a local output file or output prefix on GCS:: --output [YOUR_LOCAL_FILE | gs://YOUR_OUTPUT_PREFIX] To execute this pipeline using the Google Cloud Dataflow service, specify pipeline configuration:: --project YOUR_PROJECT_ID --staging_location gs://YOUR_STAGING_DIRECTORY --temp_location gs://YOUR_TEMP_DIRECTORY --job_name YOUR_JOB_NAME --runner BlockingDataflowPipelineRunner and an output prefix on GCS:: --output gs://YOUR_OUTPUT_PREFIX """ from __future__ import absolute_import import argparse import logging import re import google.cloud.dataflow as df class FilterTextFn(df.DoFn): """A DoFn that filters for a specific key based on a regular expression.""" # A custom aggregator can track values in your pipeline as it runs. Those # values will be displayed in the Dataflow Monitoring UI when this pipeline is # run using the Dataflow service. These aggregators below track the number of # matched and unmatched words. Learn more at # https://cloud.google.com/dataflow/pipelines/dataflow-monitoring-intf about # the Dataflow Monitoring UI. matched_words = df.Aggregator('matched_words') umatched_words = df.Aggregator('umatched_words') def __init__(self, pattern): super(FilterTextFn, self).__init__() self.pattern = pattern def process(self, context): word, _ = context.element if re.match(self.pattern, word): # Log at INFO level each element we match. When executing this pipeline # using the Dataflow service, these log lines will appear in the Cloud # Logging UI. logging.info('Matched %s', word) context.aggregate_to(self.matched_words, 1) yield context.element else: # Log at the "DEBUG" level each element that is not matched. Different log # levels can be used to control the verbosity of logging providing an # effective mechanism to filter less important information. # Note currently only "INFO" and higher level logs are emitted to the # Cloud Logger. This log message will not be visible in the Cloud Logger. logging.debug('Did not match %s', word) context.aggregate_to(self.umatched_words, 1) class CountWords(df.PTransform): """A transform to count the occurrences of each word. A PTransform that converts a PCollection containing lines of text into a PCollection of (word, count) tuples. """ def __init__(self): super(CountWords, self).__init__() def apply(self, pcoll): return (pcoll | (df.FlatMap('split', lambda x: re.findall(r'[A-Za-z\']+', x)) .with_output_types(unicode)) | df.Map('pair_with_one', lambda x: (x, 1)) | df.GroupByKey('group') | df.Map('count', lambda (word, ones): (word, sum(ones)))) def run(argv=None): """Runs the debugging wordcount pipeline.""" parser = argparse.ArgumentParser() parser.add_argument('--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument('--output', dest='output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) p = df.Pipeline(argv=pipeline_args) # Read the text file[pattern] into a PCollection, count the occurrences of # each word and filter by a list of words. filtered_words = ( p | df.io.Read('read', df.io.TextFileSource(known_args.input)) | CountWords() | df.ParDo('FilterText', FilterTextFn('Flourish|stomach'))) # assert_that is a convenient PTransform that checks a PCollection has an # expected value. Asserts are best used in unit tests with small data sets but # is demonstrated here as a teaching tool. # # Note assert_that does not provide any output and that successful completion # of the Pipeline implies that the expectations were met. Learn more at # https://cloud.google.com/dataflow/pipelines/testing-your-pipeline on how to # test your pipeline. df.assert_that(filtered_words, df.equal_to([('Flourish', 3), ('stomach', 1)])) # Format the counts into a PCollection of strings and write the output using a # "Write" transform that has side effects. # pylint: disable=unused-variable output = (filtered_words | df.Map('format', lambda (word, c): '%s: %s' % (word, c)) | df.io.Write('write', df.io.TextFileSink(known_args.output))) # Actually run the pipeline (all operations above are deferred). p.run() if __name__ == '__main__': # Cloud Logging would contain only logging.INFO and higher level logs logged # by the root logger. All log statements emitted by the root logger will be # visible in the Cloud Logging UI. Learn more at # https://cloud.google.com/logging about the Cloud Logging UI. # # You can set the default logging level to a different level when running # locally. logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/wordcount_debugging_test.py000066400000000000000000000034531272340251000304750ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the debugging wordcount example.""" import logging import re import tempfile import unittest from google.cloud.dataflow.examples import wordcount_debugging class WordCountTest(unittest.TestCase): SAMPLE_TEXT = 'xx yy Flourish\n zz Flourish Flourish stomach\n aa\n bb cc dd' def create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def get_results(self, temp_path): results = [] with open(temp_path + '.result-00000-of-00001') as result_file: for line in result_file: match = re.search(r'([A-Za-z]+): ([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) return results def test_basics(self): temp_path = self.create_temp_file(self.SAMPLE_TEXT) expected_words = [('Flourish', 3), ('stomach', 1)] wordcount_debugging.run([ '--input=%s*' % temp_path, '--output=%s.result' % temp_path]) # Parse result file and compare. results = self.get_results(temp_path) self.assertEqual(sorted(results), sorted(expected_words)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/wordcount_minimal.py000066400000000000000000000103711272340251000271260ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A minimalist word-counting workflow that counts words in Shakespeare. This is the first in a series of successively more detailed 'word count' examples. Next, see the wordcount pipeline, then the wordcount_debugging pipeline, for more detailed examples that introduce additional concepts. Concepts: 1. Reading data from text files 2. Specifying 'inline' transforms 3. Counting a PCollection 4. Writing data to Cloud Storage as text files To execute this pipeline locally, first edit the code to specify the output location. Output location could be a local file path or an output prefix on GCS. (Only update the output location marked with the first CHANGE comment.) To execute this pipeline remotely, first edit the code to set your project ID, runner type, the staging location, the temp location, and the output location. The specified GCS bucket(s) must already exist. (Update all the places marked with a CHANGE comment.) Then, run the pipeline as described in the README. It will be deployed and run using the Google Cloud Dataflow Service. No args are required to run the pipeline. You can see the results in your output bucket in the GCS browser. """ from __future__ import absolute_import import argparse import logging import re import google.cloud.dataflow as df def run(argv=None): """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() parser.add_argument('--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument('--output', dest='output', # CHANGE 1/5: The Google Cloud Storage path is required # for outputting the results. default='gs://YOUR_OUTPUT_BUCKET/AND_OUTPUT_PREFIX', help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) pipeline_args.extend([ # CHANGE 2/5: (OPTIONAL) Change this to BlockingDataflowPipelineRunner to # run your pipeline on the Google Cloud Dataflow Service. '--runner=DirectPipelineRunner', # CHANGE 3/5: Your project ID is required in order to run your pipeline on # the Google Cloud Dataflow Service. '--project=SET_YOUR_PROJECT_ID_HERE', # CHANGE 4/5: Your Google Cloud Storage path is required for staging local # files. '--staging_location=gs://YOUR_BUCKET_NAME/AND_STAGING_DIRECTORY', # CHANGE 5/5: Your Google Cloud Storage path is required for temporary # files. '--temp_location=gs://YOUR_BUCKET_NAME/AND_TEMP_DIRECTORY', '--job_name=your-wordcount-job', ]) p = df.Pipeline(argv=pipeline_args) # Read the text file[pattern] into a PCollection. lines = p | df.io.Read('read', df.io.TextFileSource(known_args.input)) # Count the occurrences of each word. counts = (lines | (df.FlatMap('split', lambda x: re.findall(r'[A-Za-z\']+', x)) .with_output_types(unicode)) | df.Map('pair_with_one', lambda x: (x, 1)) | df.GroupByKey('group') | df.Map('count', lambda (word, ones): (word, sum(ones)))) # Format the counts into a PCollection of strings. output = counts | df.Map('format', lambda (word, c): '%s: %s' % (word, c)) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | df.io.Write('write', df.io.TextFileSink(known_args.output)) # Actually run the pipeline (all operations above are deferred). p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/wordcount_minimal_test.py000066400000000000000000000035351272340251000301710ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the minimal wordcount example.""" import collections import logging import re import tempfile import unittest from google.cloud.dataflow.examples import wordcount_minimal class WordCountMinimalTest(unittest.TestCase): """Unit test for wordcount_minimal example with direct runner.""" SAMPLE_TEXT = 'a b c a b a\n aa bb cc aa bb aa' def create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def test_basics(self): temp_path = self.create_temp_file(self.SAMPLE_TEXT) expected_words = collections.defaultdict(int) for word in re.findall(r'\w+', self.SAMPLE_TEXT): expected_words[word] += 1 wordcount_minimal.run([ '--input=%s*' % temp_path, '--output=%s.result' % temp_path]) # Parse result file and compare. results = [] with open(temp_path + '.result-00000-of-00001') as result_file: for line in result_file: match = re.search(r'([a-z]+): ([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) self.assertEqual(sorted(results), sorted(expected_words.iteritems())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/examples/wordcount_test.py000066400000000000000000000033741272340251000264640ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Test for the wordcount example.""" import collections import logging import re import tempfile import unittest from google.cloud.dataflow.examples import wordcount class WordCountTest(unittest.TestCase): SAMPLE_TEXT = 'a b c a b a\n\n aa bb cc aa bb aa' def create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def test_basics(self): temp_path = self.create_temp_file(self.SAMPLE_TEXT) expected_words = collections.defaultdict(int) for word in re.findall(r'\w+', self.SAMPLE_TEXT): expected_words[word] += 1 wordcount.run([ '--input=%s*' % temp_path, '--output=%s.result' % temp_path]) # Parse result file and compare. results = [] with open(temp_path + '.result-00000-of-00001') as result_file: for line in result_file: match = re.search(r'([a-z]+): ([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) self.assertEqual(sorted(results), sorted(expected_words.iteritems())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/000077500000000000000000000000001272340251000230165ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/__init__.py000066400000000000000000000000001272340251000251150ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/apiclient.py000066400000000000000000001154451272340251000253520ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Dataflow client utility functions.""" import codecs import json import logging import os import re import time from google.cloud.dataflow import utils from google.cloud.dataflow import version from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.internal.auth import get_service_credentials from google.cloud.dataflow.internal.json_value import to_json_value from google.cloud.dataflow.io import iobase from google.cloud.dataflow.transforms import cy_combiners from google.cloud.dataflow.utils import dependency from google.cloud.dataflow.utils import names from google.cloud.dataflow.utils import retry from google.cloud.dataflow.utils.names import PropertyNames from google.cloud.dataflow.utils.options import GoogleCloudOptions from google.cloud.dataflow.utils.options import StandardOptions from google.cloud.dataflow.utils.options import WorkerOptions from apitools.base.py import encoding from apitools.base.py import exceptions from google.cloud.dataflow.internal.clients import storage import google.cloud.dataflow.internal.clients.dataflow as dataflow BIGQUERY_API_SERVICE = 'bigquery.googleapis.com' COMPUTE_API_SERVICE = 'compute.googleapis.com' STORAGE_API_SERVICE = 'storage.googleapis.com' def append_counter(status_object, counter, tentative): """Appends a counter to the status. Args: status_object: a work_item_status to which to add this counter counter: a counters.Counter object to append tentative: whether the value should be reported as tentative """ logging.debug('Appending counter%s %s', ' (tentative)' if tentative else '', counter) kind, setter = metric_translations[counter.combine_fn.__class__] append_metric( status_object, counter.name, kind, counter.accumulator, setter, tentative=tentative) def append_metric(status_object, metric_name, kind, value, setter=None, step=None, output_user_name=None, tentative=False, worker_id=None, cumulative=True): """Creates and adds a MetricUpdate field to the passed-in protobuf. Args: status_object: a work_item_status to which to add this metric metric_name: a string naming this metric kind: dataflow counter kind (e.g. 'sum') value: accumulator value to encode setter: if not None, a lambda to use to update metric_update with value step: the name of the associated step output_user_name: the user-visible name to use tentative: whether this should be labeled as a tentative metric worker_id: the id of this worker. Specifying a worker_id also causes this to be encoded as a metric, not a counter. cumulative: Whether this metric is cumulative, default True. Set to False for a delta value. """ # Does this look like a counter or like a metric? is_counter = not worker_id metric_update = dataflow.MetricUpdate() metric_update.name = dataflow.MetricStructuredName() metric_update.name.name = metric_name # Handle attributes stored in the name context if step or output_user_name or tentative or worker_id: metric_update.name.context = dataflow.MetricStructuredName.ContextValue() def append_to_context(key, value): metric_update.name.context.additionalProperties.append( dataflow.MetricStructuredName.ContextValue.AdditionalProperty( key=key, value=value)) if step: append_to_context('step', step) if output_user_name: append_to_context('output_user_name', output_user_name) if tentative: append_to_context('tentative', 'true') if worker_id: append_to_context('workerId', worker_id) if cumulative and is_counter: metric_update.cumulative = cumulative if is_counter: # Counters are distinguished by having a kind; metrics do not. metric_update.kind = kind if setter: setter(value, metric_update) else: metric_update.scalar = to_json_value(value, with_type=True) logging.debug('Appending metric_update: %s', metric_update) status_object.metricUpdates.append(metric_update) class Step(object): """Wrapper for a dataflow Step protobuf.""" def __init__(self, step_kind, step_name): self.step_kind = step_kind self.step_name = step_name self.proto = dataflow.Step(kind=step_kind, name=step_name) self.proto.properties = {} def add_property(self, name, value, with_type=False): self.proto.properties.additionalProperties.append( dataflow.Step.PropertiesValue.AdditionalProperty( key=name, value=to_json_value(value, with_type=with_type))) def _get_outputs(self): """Returns a list of all output labels for a step.""" outputs = [] for p in self.proto.properties.additionalProperties: if p.key == PropertyNames.OUTPUT_INFO: for entry in p.value.array_value.entries: for entry_prop in entry.object_value.properties: if entry_prop.key == PropertyNames.OUTPUT_NAME: outputs.append(entry_prop.value.string_value) return outputs def get_output(self, tag=None): """Returns name if it is one of the outputs or first output if name is None. Args: tag: tag of the output as a string or None if we want to get the name of the first output. Returns: The name of the output associated with the tag or the first output if tag was None. Raises: ValueError: if the tag does not exist within outputs. """ outputs = self._get_outputs() if tag is None: return outputs[0] else: name = '%s_%s' % (PropertyNames.OUT, tag) if name not in outputs: raise ValueError( 'Cannot find named output: %s in %s.' % (name, outputs)) return name class Environment(object): """Wrapper for a dataflow Environment protobuf.""" def __init__(self, packages, options, environment_version): self.standard_options = options.view_as(StandardOptions) self.google_cloud_options = options.view_as(GoogleCloudOptions) self.worker_options = options.view_as(WorkerOptions) self.proto = dataflow.Environment() self.proto.clusterManagerApiService = COMPUTE_API_SERVICE self.proto.dataset = '%s/cloud_dataflow' % BIGQUERY_API_SERVICE self.proto.tempStoragePrefix = ( self.google_cloud_options.temp_location.replace('gs:/', STORAGE_API_SERVICE)) # User agent information. self.proto.userAgent = dataflow.Environment.UserAgentValue() self.local = 'localhost' in self.google_cloud_options.dataflow_endpoint version_string = version.__version__ self.proto.userAgent.additionalProperties.extend([ dataflow.Environment.UserAgentValue.AdditionalProperty( key='name', value=to_json_value('Google Cloud Dataflow SDK for Python')), dataflow.Environment.UserAgentValue.AdditionalProperty( key='version', value=to_json_value(version_string))]) # Version information. self.proto.version = dataflow.Environment.VersionValue() if self.standard_options.streaming: job_type = 'PYTHON_STREAMING' else: job_type = 'PYTHON_BATCH' self.proto.version.additionalProperties.extend([ dataflow.Environment.VersionValue.AdditionalProperty( key='job_type', value=to_json_value(job_type)), dataflow.Environment.VersionValue.AdditionalProperty( key='major', value=to_json_value(environment_version))]) # Worker pool(s) information. package_descriptors = [] for package in packages: package_descriptors.append( dataflow.Package( location='%s/%s' % ( self.google_cloud_options.staging_location.replace( 'gs:/', STORAGE_API_SERVICE), package), name=package)) pool = dataflow.WorkerPool( kind='local' if self.local else 'harness', packages=package_descriptors, taskrunnerSettings=dataflow.TaskRunnerSettings( parallelWorkerSettings=dataflow.WorkerSettings( baseUrl='https://dataflow.googleapis.com', servicePath=self.google_cloud_options.dataflow_endpoint))) # Set worker pool options received through command line. if self.worker_options.num_workers: pool.numWorkers = self.worker_options.num_workers if self.worker_options.machine_type: pool.machineType = self.worker_options.machine_type if self.worker_options.disk_size_gb: pool.diskSizeGb = self.worker_options.disk_size_gb if self.worker_options.disk_type: pool.diskType = self.worker_options.disk_type if self.worker_options.disk_source_image: pool.diskSourceImage = self.worker_options.disk_source_image if self.worker_options.zone: pool.zone = self.worker_options.zone if self.worker_options.network: pool.network = self.worker_options.network if self.worker_options.worker_harness_container_image: pool.workerHarnessContainerImage = ( self.worker_options.worker_harness_container_image) else: # Default to using the worker harness container image for the current SDK # version. pool.workerHarnessContainerImage = ( 'dataflow.gcr.io/v1beta3/python:%s' % version.__version__) if self.worker_options.teardown_policy: if self.worker_options.teardown_policy == 'TEARDOWN_NEVER': pool.teardownPolicy = ( dataflow.WorkerPool.TeardownPolicyValueValuesEnum.TEARDOWN_NEVER) elif self.worker_options.teardown_policy == 'TEARDOWN_ALWAYS': pool.teardownPolicy = ( dataflow.WorkerPool.TeardownPolicyValueValuesEnum.TEARDOWN_ALWAYS) elif self.worker_options.teardown_policy == 'TEARDOWN_ON_SUCCESS': pool.teardownPolicy = ( dataflow.WorkerPool .TeardownPolicyValueValuesEnum.TEARDOWN_ON_SUCCESS) if self.standard_options.streaming: # Use separate data disk for streaming. disk = dataflow.Disk() if self.local: disk.diskType = 'local' # TODO(ccy): allow customization of disk. pool.dataDisks.append(disk) self.proto.workerPools.append(pool) sdk_pipeline_options = options.get_all_options() if sdk_pipeline_options: self.proto.sdkPipelineOptions = ( dataflow.Environment.SdkPipelineOptionsValue()) for k, v in sdk_pipeline_options.iteritems(): if v is not None: self.proto.sdkPipelineOptions.additionalProperties.append( dataflow.Environment.SdkPipelineOptionsValue.AdditionalProperty( key=k, value=to_json_value(v))) class Job(object): """Wrapper for a dataflow Job protobuf.""" def __str__(self): def encode_shortstrings(input_buffer, errors='strict'): """Encoder (from Unicode) that suppresses long base64 strings.""" original_len = len(input_buffer) if original_len > 150: if self.base64_str_re.match(input_buffer): input_buffer = '' % original_len input_buffer = input_buffer.encode('ascii', errors=errors) else: matched = self.coder_str_re.match(input_buffer) if matched: input_buffer = '%s' % ( matched.group(1), matched.end(2) - matched.start(2)) input_buffer = input_buffer.encode('ascii', errors=errors) return input_buffer, original_len def decode_shortstrings(input_buffer, errors='strict'): """Decoder (to Unicode) that suppresses long base64 strings.""" shortened, length = encode_shortstrings(input_buffer, errors) return unicode(shortened), length def shortstrings_registerer(encoding_name): if encoding_name == 'shortstrings': return codecs.CodecInfo(name='shortstrings', encode=encode_shortstrings, decode=decode_shortstrings) return None codecs.register(shortstrings_registerer) # Use json "dump string" method to get readable formatting; # further modify it to not output too-long strings, aimed at the # 10,000+ character hex-encoded "serialized_fn" values. return json.dumps( json.loads(encoding.MessageToJson(self.proto), encoding='shortstrings'), indent=2, sort_keys=True) def __init__(self, options): self.options = options self.google_cloud_options = options.view_as(GoogleCloudOptions) required_google_cloud_options = ['project', 'job_name', 'staging_location', 'temp_location'] missing = [ option for option in required_google_cloud_options if not getattr(self.google_cloud_options, option)] if missing: raise ValueError( 'Missing required configuration parameters: %s' % missing) # Make the staging and temp locations job name and time specific. This is # needed to avoid clashes between job submissions using the same staging # area or team members using same job names. This method is not entirely # foolproof since two job submissions with same name can happen at exactly # the same time. However the window is extremely small given that # time.time() has at least microseconds granularity. We add the suffix only # for GCS staging locations where the potential for such clashes is high. if self.google_cloud_options.staging_location.startswith('gs://'): path_suffix = '%s.%f' % (self.google_cloud_options.job_name, time.time()) self.google_cloud_options.staging_location = utils.path.join( self.google_cloud_options.staging_location, path_suffix) self.google_cloud_options.temp_location = utils.path.join( self.google_cloud_options.temp_location, path_suffix) self.proto = dataflow.Job(name=self.google_cloud_options.job_name) if self.options.view_as(StandardOptions).streaming: self.proto.type = dataflow.Job.TypeValueValuesEnum.JOB_TYPE_STREAMING else: self.proto.type = dataflow.Job.TypeValueValuesEnum.JOB_TYPE_BATCH self.base64_str_re = re.compile(r'^[A-Za-z0-9+/]*=*$') self.coder_str_re = re.compile(r'^([A-Za-z]+\$)([A-Za-z0-9+/]*=*)$') class DataflowApplicationClient(object): """A Dataflow API client used by application code to create and query jobs.""" def __init__(self, options, environment_version): """Initializes a Dataflow API client object.""" self.standard_options = options.view_as(StandardOptions) self.google_cloud_options = options.view_as(GoogleCloudOptions) self.environment_version = environment_version if self.google_cloud_options.no_auth: credentials = None else: credentials = get_service_credentials() self._client = dataflow.DataflowV1b3( url=self.google_cloud_options.dataflow_endpoint, credentials=credentials, get_credentials=(not self.google_cloud_options.no_auth)) self._storage_client = storage.StorageV1( url='https://www.googleapis.com/storage/v1', credentials=credentials, get_credentials=(not self.google_cloud_options.no_auth)) # TODO(silviuc): Refactor so that retry logic can be applied. @retry.no_retries # Using no_retries marks this as an integration point. def _gcs_file_copy(self, from_path, to_path): to_folder, to_name = os.path.split(to_path) with open(from_path, 'rb') as f: self.stage_file(to_folder, to_name, f) def stage_file(self, gcs_or_local_path, file_name, stream, mime_type='application/octet-stream'): """Stages a file at a GCS or local path with stream-supplied contents.""" if not gcs_or_local_path.startswith('gs://'): local_path = os.path.join(gcs_or_local_path, file_name) logging.info('Staging file locally to %s', local_path) with open(local_path, 'wb') as f: f.write(stream.read()) return gcs_location = gcs_or_local_path + '/' + file_name bucket, name = gcs_location[5:].split('/', 1) request = storage.StorageObjectsInsertRequest( bucket=bucket, name=name) logging.info('Starting GCS upload to %s...', gcs_location) upload = storage.Upload(stream, mime_type) try: response = self._storage_client.objects.Insert(request, upload=upload) except exceptions.HttpError as e: reportable_errors = { 403: 'access denied', 404: 'bucket not found', } if e.status_code in reportable_errors: raise IOError(('Could not upload to GCS path %s: %s. Please verify ' 'that credentials are valid and that you have write ' 'access to the specified path. Stale credentials can be ' 'refreshed by executing "gcloud auth login".') % (gcs_or_local_path, reportable_errors[e.status_code])) raise logging.info('Completed GCS upload to %s', gcs_location) return response # TODO(silviuc): Refactor so that retry logic can be applied. @retry.no_retries # Using no_retries marks this as an integration point. def create_job(self, job): """Submits for remote execution a job described by the workflow proto.""" # Stage job resources and add an environment proto with their paths. resources = dependency.stage_job_resources( job.options, file_copy=self._gcs_file_copy) job.proto.environment = Environment( packages=resources, options=job.options, environment_version=self.environment_version).proto # TODO(silviuc): Remove the debug logging eventually. logging.info('JOB: %s', job) request = dataflow.DataflowProjectsJobsCreateRequest() request.projectId = self.google_cloud_options.project request.job = job.proto try: response = self._client.projects_jobs.Create(request) except exceptions.BadStatusCodeError as e: logging.error('HTTP status %d trying to create job' ' at dataflow service endpoint %s', e.response.status, self.google_cloud_options.dataflow_endpoint) logging.fatal('details of server error: %s', e) raise logging.info('Create job: %s', response) # The response is a Job proto with the id for the new job. logging.info('Created job with id: [%s]', response.id) logging.info( 'To accesss the Dataflow monitoring console, please navigate to ' 'https://console.developers.google.com/project/%s/dataflow/job/%s', self.google_cloud_options.project, response.id) # Show the whitelisting warning. Projects should be whitelisted prior to # submitting jobs to Google Cloud Dataflow service. Please see documentation # for more information. # # TODO(altay): Remove once the whitelisting requirements are lifted. logging.warning( '\n\n***************************************************************\n' '* WARNING: PROJECT WHITELISTING REQUIRED. *' '\n***************************************************************\n' 'Please make sure your project is whitelisted for running\n' 'Python-based pipelines using the Google Cloud Dataflow service.\n\n' 'You may ignore this message if you have successfully ran\n' 'Python-based pipelines with this project on Google Cloud\n' 'Dataflow service before.\n\n' 'If your project is not whitelisted, your job will attempt to run\n' 'however it will fail to make any progress. Google Cloud Dataflow\n' 'service will automatically cancel your non-whitelisted job\n' 'after some time due to inactivity. You can also manually cancel\n' 'your job using the following command:\n\n' 'gcloud alpha dataflow jobs --project=%s cancel %s\n\n' 'Please refer to the documentation to learn more about whitelisting\n' 'your project at: %s' '\n***************************************************************\n\n', request.projectId, response.id, 'http://goo.gl/forms/o4w14whz9x' ) return response @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def modify_job_state(self, job_id, new_state): """Modify the run state of the job. Args: job_id: The id of the job. new_state: A string representing the new desired state. It could be set to either 'JOB_STATE_DONE', 'JOB_STATE_CANCELLED' or 'JOB_STATE_DRAINING'. Returns: True if the job was modified successfully. """ if new_state == 'JOB_STATE_DONE': new_state = dataflow.Job.RequestedStateValueValuesEnum.JOB_STATE_DONE elif new_state == 'JOB_STATE_CANCELLED': new_state = dataflow.Job.RequestedStateValueValuesEnum.JOB_STATE_CANCELLED elif new_state == 'JOB_STATE_DRAINING': new_state = dataflow.Job.RequestedStateValueValuesEnum.JOB_STATE_DRAINING else: # Other states could only be set by the service. return False request = dataflow.DataflowProjectsJobsUpdateRequest() request.jobId = job_id request.projectId = self.google_cloud_options.project request.job = dataflow.Job(requestedState=new_state) self._client.projects_jobs.Update(request) return True @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def get_job(self, job_id): """Gets the job status for a submitted job. Args: job_id: A string representing the job_id for the workflow as returned by the a create_job() request. Returns: A Job proto. See below for interesting fields. The Job proto returned from a get_job() request contains some interesting fields: currentState: An object representing the current state of the job. The string representation of the object (str() result) has the following possible values: JOB_STATE_UNKNONW, JOB_STATE_STOPPED, JOB_STATE_RUNNING, JOB_STATE_DONE, JOB_STATE_FAILED, JOB_STATE_CANCELLED. createTime: UTC time when the job was created (e.g. '2015-03-10T00:01:53.074Z') currentStateTime: UTC time for the current state of the job. """ request = dataflow.DataflowProjectsJobsGetRequest() request.jobId = job_id request.projectId = self.google_cloud_options.project response = self._client.projects_jobs.Get(request) return response @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def list_messages( self, job_id, start_time=None, end_time=None, page_token=None, minimum_importance=None): """List messages associated with the execution of a job. Args: job_id: A string representing the job_id for the workflow as returned by the a create_job() request. start_time: If specified, only messages generated after the start time will be returned, otherwise all messages since job started will be returned. The value is a string representing UTC time (e.g., '2015-08-18T21:03:50.644Z') end_time: If specified, only messages generated before the end time will be returned, otherwise all messages up to current time will be returned. The value is a string representing UTC time (e.g., '2015-08-18T21:03:50.644Z') page_token: A string to be used as next page token if the list call returned paginated results. minimum_importance: Filter for messages based on importance. The possible string values in increasing order of importance are: JOB_MESSAGE_DEBUG, JOB_MESSAGE_DETAILED, JOB_MESSAGE_BASIC, JOB_MESSAGE_WARNING, JOB_MESSAGE_ERROR. For example, a filter set on warning will allow only warnings and errors and exclude all others. Returns: A tuple consisting of a list of JobMessage instances and a next page token string. Raises: RuntimeError: if an unexpected value for the message_importance argument is used. The JobMessage objects returned by the call contain the following fields: id: A unique string identifier for the message. time: A string representing the UTC time of the message (e.g., '2015-08-18T21:03:50.644Z') messageImportance: An enumeration value for the message importance. The value if converted to string will have the following possible values: JOB_MESSAGE_DEBUG, JOB_MESSAGE_DETAILED, JOB_MESSAGE_BASIC, JOB_MESSAGE_WARNING, JOB_MESSAGE_ERROR. messageText: A message string. """ request = dataflow.DataflowProjectsJobsMessagesListRequest( jobId=job_id, projectId=self.google_cloud_options.project) if page_token is not None: request.pageToken = page_token if start_time is not None: request.startTime = start_time if end_time is not None: request.endTime = end_time if minimum_importance is not None: if minimum_importance == 'JOB_MESSAGE_DEBUG': request.minimumImportance = ( dataflow.DataflowProjectsJobsMessagesListRequest .MinimumImportanceValueValuesEnum .JOB_MESSAGE_DEBUG) elif minimum_importance == 'JOB_MESSAGE_DETAILED': request.minimumImportance = ( dataflow.DataflowProjectsJobsMessagesListRequest .MinimumImportanceValueValuesEnum .JOB_MESSAGE_DETAILED) elif minimum_importance == 'JOB_MESSAGE_BASIC': request.minimumImportance = ( dataflow.DataflowProjectsJobsMessagesListRequest .MinimumImportanceValueValuesEnum .JOB_MESSAGE_BASIC) elif minimum_importance == 'JOB_MESSAGE_WARNING': request.minimumImportance = ( dataflow.DataflowProjectsJobsMessagesListRequest .MinimumImportanceValueValuesEnum .JOB_MESSAGE_WARNING) elif minimum_importance == 'JOB_MESSAGE_ERROR': request.minimumImportance = ( dataflow.DataflowProjectsJobsMessagesListRequest .MinimumImportanceValueValuesEnum .JOB_MESSAGE_ERROR) else: raise RuntimeError( 'Unexpected value for minimum_importance argument: %r', minimum_importance) response = self._client.projects_jobs_messages.List(request) return response.jobMessages, response.nextPageToken class DataflowWorkerClient(object): """A Dataflow API client used by worker code to lease work items.""" def __init__(self, worker, skip_get_credentials=False): """Initializes a Dataflow API client object with worker functionality. Args: worker: A Worker instance. skip_get_credentials: If true disables credentials loading logic. """ self._client = ( dataflow.DataflowV1b3( url=worker.service_path, get_credentials=(not skip_get_credentials))) @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def lease_work(self, worker_info, desired_lease_duration): """Leases a work item from the service.""" work_request = dataflow.LeaseWorkItemRequest() work_request.workerId = worker_info.worker_id work_request.requestedLeaseDuration = desired_lease_duration work_request.currentWorkerTime = worker_info.formatted_current_time work_request.workerCapabilities.append(worker_info.worker_id) for value in worker_info.capabilities: work_request.workerCapabilities.append(value) for value in worker_info.work_types: work_request.workItemTypes.append(value) request = dataflow.DataflowProjectsJobsWorkItemsLeaseRequest() request.jobId = worker_info.job_id request.projectId = worker_info.project_id try: request.leaseWorkItemRequest = work_request except AttributeError: request.lease_work_item_request = work_request logging.debug('lease_work: %s', request) response = self._client.projects_jobs_workItems.Lease(request) logging.debug('lease_work: %s', response) return response def report_status(self, worker_info, desired_lease_duration, work_item, completed, progress, dynamic_split_result_to_report=None, source_operation_response=None, exception_details=None): """Reports status for a work item (success or failure). This is an integration point. The @retry decorator is used on callers of this method defined in google/cloud/dataflow/worker/worker.py because there are different retry strategies for a completed versus in progress work item. Args: worker_info: A batchworker.BatchWorkerInfo that contains information about the Worker instance executing the work item. desired_lease_duration: The duration for which the worker would like to extend the lease of the work item. Should be in seconds formatted as a string. work_item: The work item for which to report status. completed: True if there is no further work to be done on this work item either because it succeeded or because it failed. False if this is a progress report. progress: A SourceReaderProgress that gives the progress of worker handling the work item. dynamic_split_result_to_report: A successful dynamic split result that should be sent to the Dataflow service along with the status report. source_operation_response: Response to a source operation request from the service. This will be sent to the service along with the status report. exception_details: A string representation of the stack trace for an exception raised while executing the work item. The string is the output of the standard traceback.format_exc() function. Returns: A protobuf containing the response from the service for the status update (WorkItemServiceState). Raises: TypeError: if progress is of an unknown type RuntimeError: if dynamic split request is of an unknown type. """ work_item_status = dataflow.WorkItemStatus() work_item_status.completed = completed if not completed: work_item_status.requestedLeaseDuration = desired_lease_duration if progress is not None: work_item_progress = dataflow.ApproximateProgress() work_item_status.progress = work_item_progress if progress.position is not None: work_item_progress.position = ( reader_position_to_cloud_position(progress.position)) elif progress.percent_complete is not None: work_item_progress.percentComplete = progress.percent_complete elif progress.remaining_time is not None: work_item_progress.remainingTime = progress.remaining_time else: raise TypeError('Unknown type of progress') if dynamic_split_result_to_report is not None: assert isinstance(dynamic_split_result_to_report, iobase.DynamicSplitResult) if isinstance(dynamic_split_result_to_report, iobase.DynamicSplitResultWithPosition): work_item_status.stopPosition = ( dynamic_split_result_with_position_to_cloud_stop_position( dynamic_split_result_to_report)) else: raise RuntimeError('Unknown type of dynamic split result.') # The service keeps track of the report indexes in order to handle lost # and duplicate message. work_item_status.reportIndex = work_item.next_report_index work_item_status.workItemId = str(work_item.proto.id) # Add exception information if any. if exception_details is not None: status = dataflow.Status() # TODO(silviuc): Replace Code.UNKNOWN with a generated definition. status.code = 2 # TODO(silviuc): Attach the stack trace as exception details. status.message = exception_details work_item_status.errors.append(status) if source_operation_response is not None: work_item_status.sourceOperationResponse = source_operation_response # Look through the work item for metrics to send. if work_item.map_task: for counter in work_item.map_task.itercounters(): append_counter(work_item_status, counter, tentative=not completed) report_request = dataflow.ReportWorkItemStatusRequest() report_request.currentWorkerTime = worker_info.formatted_current_time report_request.workerId = worker_info.worker_id report_request.workItemStatuses.append(work_item_status) request = dataflow.DataflowProjectsJobsWorkItemsReportStatusRequest() request.jobId = worker_info.job_id request.projectId = worker_info.project_id try: request.reportWorkItemStatusRequest = report_request except AttributeError: request.report_work_item_status_request = report_request logging.debug('report_status: %s', request) response = self._client.projects_jobs_workItems.ReportStatus(request) logging.debug('report_status: %s', response) return response # Utility functions for translating cloud reader objects to corresponding SDK # reader objects and vice versa. def reader_progress_to_cloud_progress(reader_progress): """Converts a given 'ReaderProgress' to corresponding cloud format.""" cloud_progress = dataflow.ApproximateProgress() if reader_progress.position is not None: cloud_progress.position = reader_position_to_cloud_position( reader_progress.position) if reader_progress.percent_complete is not None: cloud_progress.percentComplete = reader_progress.percent_complete if reader_progress.remaining_time is not None: cloud_progress.remainingTime = reader_progress.remaining_time return cloud_progress def reader_position_to_cloud_position(reader_position): """Converts a given 'ReaderPosition' to corresponding cloud format.""" cloud_position = dataflow.Position() if reader_position.end is not None: cloud_position.end = reader_position.end if reader_position.key is not None: cloud_position.key = reader_position.key if reader_position.byte_offset is not None: cloud_position.byteOffset = reader_position.byte_offset if reader_position.record_index is not None: cloud_position.recordIndex = reader_position.record_index if reader_position.shuffle_position is not None: cloud_position.shufflePosition = reader_position.shuffle_position if reader_position.concat_position is not None: concat_position = dataflow.ConcatPosition() concat_position.index = reader_position.concat_position.index concat_position.position = reader_position_to_cloud_position( reader_position.concat_position.position) cloud_position.concatPosition = concat_position return cloud_position def dynamic_split_result_with_position_to_cloud_stop_position(split_result): """Converts a given 'DynamicSplitResultWithPosition' to cloud format.""" return reader_position_to_cloud_position(split_result.stop_position) def cloud_progress_to_reader_progress(cloud_progress): reader_position = None if cloud_progress.position is not None: reader_position = cloud_position_to_reader_position(cloud_progress.position) return iobase.ReaderProgress(reader_position, cloud_progress.percentComplete, cloud_progress.remainingTime) def cloud_position_to_reader_position(cloud_position): concat_position = None if cloud_position.concatPosition is not None: inner_position = cloud_position_to_reader_position( cloud_position.concatPosition.position) concat_position = iobase.ConcatPosition(cloud_position.index, inner_position) return iobase.ReaderPosition(cloud_position.end, cloud_position.key, cloud_position.byteOffset, cloud_position.recordIndex, cloud_position.shufflePosition, concat_position) def approximate_progress_to_dynamic_split_request(approximate_progress): return iobase.DynamicSplitRequest(cloud_progress_to_reader_progress( approximate_progress)) def set_scalar(accumulator, metric_update): metric_update.scalar = to_json_value(accumulator.value, with_type=True) def set_mean(accumulator, metric_update): if accumulator.count: metric_update.meanSum = to_json_value(accumulator.sum, with_type=True) metric_update.meanCount = to_json_value(accumulator.count, with_type=True) else: # A denominator of 0 will raise an error in the service. # What it means is we have nothing to report yet, so don't. metric_update.kind = None # To enable a counter on the service, add it to this dictionary. metric_translations = { cy_combiners.CountCombineFn: ('sum', set_scalar), cy_combiners.SumInt64Fn: ('sum', set_scalar), cy_combiners.MinInt64Fn: ('min', set_scalar), cy_combiners.MaxInt64Fn: ('max', set_scalar), cy_combiners.MeanInt64Fn: ('mean', set_mean), cy_combiners.SumFloatFn: ('sum', set_scalar), cy_combiners.MinFloatFn: ('min', set_scalar), cy_combiners.MaxFloatFn: ('max', set_scalar), cy_combiners.MeanFloatFn: ('mean', set_mean), cy_combiners.AllCombineFn: ('and', set_scalar), cy_combiners.AnyCombineFn: ('or', set_scalar), } def splits_to_split_response(bundles): """Generates a response to a custom source split request. Args: bundles: a set of bundles generated by a BoundedSource.split() invocation. Returns: a SourceOperationResponse object. """ derived_sources = [] for bundle in bundles: derived_source = dataflow.DerivedSource() derived_source.derivationMode = ( dataflow.DerivedSource.DerivationModeValueValuesEnum .SOURCE_DERIVATION_MODE_INDEPENDENT) derived_source.source = dataflow.Source() derived_source.source.doesNotNeedSplitting = True derived_source.source.spec = dataflow.Source.SpecValue() derived_source.source.spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty( key=names.SERIALIZED_SOURCE_KEY, value=to_json_value(pickler.dumps( (bundle.source, bundle.start_position, bundle.stop_position)), with_type=True))) derived_source.source.spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty(key='@type', value=to_json_value( names.SOURCE_TYPE))) derived_sources.append(derived_source) split_response = dataflow.SourceSplitResponse() split_response.bundles = derived_sources split_response.outcome = ( dataflow.SourceSplitResponse.OutcomeValueValuesEnum .SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED) response = dataflow.SourceOperationResponse() response.split = split_response return response DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/apiclient_test.py000066400000000000000000000107451272340251000264060ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the apiclient module.""" import unittest from google.cloud.dataflow.internal import apiclient from google.cloud.dataflow.io import iobase import google.cloud.dataflow.internal.clients.dataflow as dataflow class UtilTest(unittest.TestCase): def test_reader_progress_to_cloud_progress_position(self): reader_position = iobase.ReaderPosition(byte_offset=9999) reader_progress = iobase.ReaderProgress(position=reader_position) cloud_progress = apiclient.reader_progress_to_cloud_progress( reader_progress) self.assertIsNotNone(cloud_progress) self.assertIsInstance(cloud_progress, dataflow.ApproximateProgress) self.assertIsNotNone(cloud_progress.position) self.assertIsInstance(cloud_progress.position, dataflow.Position) self.assertEquals(9999, cloud_progress.position.byteOffset) def test_reader_progress_to_cloud_progress_percent_complete(self): reader_progress = iobase.ReaderProgress(percent_complete=0.123) cloud_progress = apiclient.reader_progress_to_cloud_progress( reader_progress) self.assertIsNotNone(cloud_progress) self.assertIsInstance(cloud_progress, dataflow.ApproximateProgress) self.assertIsNotNone(cloud_progress.percentComplete) self.assertEquals(0.123, cloud_progress.percentComplete) def test_reader_position_to_cloud_position(self): reader_position = iobase.ReaderPosition(byte_offset=9999) cloud_position = apiclient.reader_position_to_cloud_position( reader_position) self.assertIsNotNone(cloud_position) def test_dynamic_split_result_with_position_to_cloud_stop_position(self): position = iobase.ReaderPosition(byte_offset=9999) dynamic_split_result = iobase.DynamicSplitResultWithPosition(position) approximate_position = ( apiclient.dynamic_split_result_with_position_to_cloud_stop_position( dynamic_split_result)) self.assertIsNotNone(approximate_position) self.assertIsInstance(approximate_position, dataflow.Position) self.assertEqual(9999, approximate_position.byteOffset) def test_cloud_progress_to_reader_progress_index_position(self): cloud_progress = dataflow.ApproximateProgress() cloud_progress.position = dataflow.Position() cloud_progress.position.byteOffset = 9999 reader_progress = apiclient.cloud_progress_to_reader_progress( cloud_progress) self.assertIsNotNone(reader_progress.position) self.assertIsInstance(reader_progress.position, iobase.ReaderPosition) self.assertEqual(9999, reader_progress.position.byte_offset) def test_cloud_progress_to_reader_progress_percent_complete(self): cloud_progress = dataflow.ApproximateProgress() cloud_progress.percentComplete = 0.123 reader_progress = apiclient.cloud_progress_to_reader_progress( cloud_progress) self.assertIsNotNone(reader_progress.percent_complete) self.assertEqual(0.123, reader_progress.percent_complete) def test_cloud_position_to_reader_position_byte_offset(self): cloud_position = dataflow.Position() cloud_position.byteOffset = 9999 reader_position = apiclient.cloud_position_to_reader_position( cloud_position) self.assertIsNotNone(reader_position) self.assertIsInstance(reader_position, iobase.ReaderPosition) self.assertEqual(9999, reader_position.byte_offset) def test_approximate_progress_to_dynamic_split_request(self): approximate_progress = dataflow.ApproximateProgress() approximate_progress.percentComplete = 0.123 dynamic_split_request = ( apiclient.approximate_progress_to_dynamic_split_request( approximate_progress)) self.assertIsNotNone(dynamic_split_request) self.assertIsInstance(dynamic_split_request.progress, iobase.ReaderProgress) self.assertIsNotNone(dynamic_split_request.progress.percent_complete) self.assertEqual(dynamic_split_request.progress.percent_complete, 0.123) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/auth.py000066400000000000000000000136361272340251000243420ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Dataflow credentials and authentication.""" import datetime import json import logging import os import sys import urllib2 from oauth2client.client import OAuth2Credentials from google.cloud.dataflow.utils import processes from google.cloud.dataflow.utils import retry from google.cloud.dataflow.utils.options import GoogleCloudOptions from google.cloud.dataflow.utils.options import PipelineOptions # When we are running in GCE, we can authenticate with VM credentials. is_running_in_gce = False # When we are running in GCE, this value is set based on worker startup # information. executing_project = None def set_running_in_gce(worker_executing_project): """Informs the authentication library that we are running in GCE. When we are running in GCE, we have the option of using the VM metadata credentials for authentication to Google services. Args: worker_executing_project: The project running the workflow. This information comes from worker startup information. """ global is_running_in_gce global executing_project is_running_in_gce = True executing_project = worker_executing_project class AuthenticationException(retry.PermanentException): pass class GCEMetadataCredentials(OAuth2Credentials): """Credential object initialized using access token from GCE VM metadata.""" def __init__(self, user_agent=None): """Create an instance of GCEMetadataCredentials. These credentials are generated by contacting the metadata server on a GCE VM instance. Args: user_agent: string, The HTTP User-Agent to provide for this application. """ super(GCEMetadataCredentials, self).__init__( None, # access_token None, # client_id None, # client_secret None, # refresh_token datetime.datetime(2010, 1, 1), # token_expiry, set to time in past. None, # token_uri user_agent) def _refresh(self, http_request): refresh_time = datetime.datetime.now() req = urllib2.Request('http://metadata.google.internal/computeMetadata/v1/' 'instance/service-accounts/default/token', headers={'Metadata-Flavor': 'Google'}) token_data = json.loads(urllib2.urlopen(req).read()) self.access_token = token_data['access_token'] self.token_expiry = (refresh_time + datetime.timedelta(seconds=token_data['expires_in'])) class _GCloudWrapperCredentials(OAuth2Credentials): """Credentials class wrapping gcloud credentials via shell.""" def __init__(self, user_agent, **kwds): super(_GCloudWrapperCredentials, self).__init__( None, None, None, None, None, None, user_agent, **kwds) def _refresh(self, http_request): """Gets an access token using the gcloud client.""" try: gcloud_process = processes.Popen( ['gcloud', 'auth', 'print-access-token'], stdout=processes.PIPE) except OSError as exn: logging.error('The gcloud tool was not found.', exc_info=True) raise AuthenticationException('The gcloud tool was not found: %s' % exn) output, _ = gcloud_process.communicate() self.access_token = output.strip() def get_service_credentials(): """Get credentials to access Google services.""" user_agent = 'dataflow-python-sdk/1.0' if is_running_in_gce: # We are currently running as a GCE taskrunner worker. # # TODO(ccy): It's not entirely clear if these credentials are thread-safe. # If so, we can cache these credentials to save the overhead of creating # them again. return GCEMetadataCredentials(user_agent=user_agent) else: # We are currently being run from the command line. google_cloud_options = PipelineOptions( sys.argv).view_as(GoogleCloudOptions) if google_cloud_options.service_account_name: if not google_cloud_options.service_account_key_file: raise AuthenticationException( 'key file not provided for service account.') if not os.path.exists(google_cloud_options.service_account_key_file): raise AuthenticationException( 'Specified service account key file does not exist.') client_scopes = [ 'https://www.googleapis.com/auth/bigquery', 'https://www.googleapis.com/auth/cloud-platform', 'https://www.googleapis.com/auth/devstorage.full_control', 'https://www.googleapis.com/auth/userinfo.email', 'https://www.googleapis.com/auth/datastore' ] # The following code uses oauth2client >=2.0.0 functionality and if this # is not available due to import errors will use 1.5.2 functionality. try: from oauth2client.service_account import ServiceAccountCredentials return ServiceAccountCredentials.from_p12_keyfile( google_cloud_options.service_account_name, google_cloud_options.service_account_key_file, client_scopes, user_agent=user_agent) except ImportError: with file(google_cloud_options.service_account_key_file) as f: service_account_key = f.read() from oauth2client.client import SignedJwtAssertionCredentials return SignedJwtAssertionCredentials( google_cloud_options.service_account_name, service_account_key, client_scopes, user_agent=user_agent) else: return _GCloudWrapperCredentials(user_agent) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/000077500000000000000000000000001272340251000244575ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/__init__.py000066400000000000000000000000001272340251000265560ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/bigquery/000077500000000000000000000000001272340251000263065ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/bigquery/__init__.py000066400000000000000000000005451272340251000304230ustar00rootroot00000000000000"""Common imports for generated bigquery client library.""" # pylint:disable=wildcard-import import pkgutil from apitools.base.py import * from google.cloud.dataflow.internal.clients.bigquery.bigquery_v2_client import * from google.cloud.dataflow.internal.clients.bigquery.bigquery_v2_messages import * __path__ = pkgutil.extend_path(__path__, __name__) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/bigquery/bigquery_v2_client.py000066400000000000000000000656261272340251000324730ustar00rootroot00000000000000"""Generated client library for bigquery version v2.""" # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.py import base_api from google.cloud.dataflow.internal.clients.bigquery import bigquery_v2_messages as messages class BigqueryV2(base_api.BaseApiClient): """Generated client library for service bigquery version v2.""" MESSAGES_MODULE = messages _PACKAGE = u'bigquery' _SCOPES = [u'https://www.googleapis.com/auth/bigquery', u'https://www.googleapis.com/auth/bigquery.insertdata', u'https://www.googleapis.com/auth/cloud-platform', u'https://www.googleapis.com/auth/cloud-platform.read-only', u'https://www.googleapis.com/auth/devstorage.full_control', u'https://www.googleapis.com/auth/devstorage.read_only', u'https://www.googleapis.com/auth/devstorage.read_write'] _VERSION = u'v2' _CLIENT_ID = '1042881264118.apps.googleusercontent.com' _CLIENT_SECRET = 'x_Tw5K8nnjoRAqULM9PFAC2b' _USER_AGENT = 'x_Tw5K8nnjoRAqULM9PFAC2b' _CLIENT_CLASS_NAME = u'BigqueryV2' _URL_VERSION = u'v2' _API_KEY = None def __init__(self, url='', credentials=None, get_credentials=True, http=None, model=None, log_request=False, log_response=False, credentials_args=None, default_global_params=None, additional_http_headers=None): """Create a new bigquery handle.""" url = url or u'https://www.googleapis.com/bigquery/v2/' super(BigqueryV2, self).__init__( url, credentials=credentials, get_credentials=get_credentials, http=http, model=model, log_request=log_request, log_response=log_response, credentials_args=credentials_args, default_global_params=default_global_params, additional_http_headers=additional_http_headers) self.datasets = self.DatasetsService(self) self.jobs = self.JobsService(self) self.projects = self.ProjectsService(self) self.tabledata = self.TabledataService(self) self.tables = self.TablesService(self) class DatasetsService(base_api.BaseApiService): """Service class for the datasets resource.""" _NAME = u'datasets' def __init__(self, client): super(BigqueryV2.DatasetsService, self).__init__(client) self._method_configs = { 'Delete': base_api.ApiMethodInfo( http_method=u'DELETE', method_id=u'bigquery.datasets.delete', ordered_params=[u'projectId', u'datasetId'], path_params=[u'datasetId', u'projectId'], query_params=[u'deleteContents'], relative_path=u'projects/{projectId}/datasets/{datasetId}', request_field='', request_type_name=u'BigqueryDatasetsDeleteRequest', response_type_name=u'BigqueryDatasetsDeleteResponse', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.datasets.get', ordered_params=[u'projectId', u'datasetId'], path_params=[u'datasetId', u'projectId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}', request_field='', request_type_name=u'BigqueryDatasetsGetRequest', response_type_name=u'Dataset', supports_download=False, ), 'Insert': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'bigquery.datasets.insert', ordered_params=[u'projectId'], path_params=[u'projectId'], query_params=[], relative_path=u'projects/{projectId}/datasets', request_field=u'dataset', request_type_name=u'BigqueryDatasetsInsertRequest', response_type_name=u'Dataset', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.datasets.list', ordered_params=[u'projectId'], path_params=[u'projectId'], query_params=[u'all', u'maxResults', u'pageToken'], relative_path=u'projects/{projectId}/datasets', request_field='', request_type_name=u'BigqueryDatasetsListRequest', response_type_name=u'DatasetList', supports_download=False, ), 'Patch': base_api.ApiMethodInfo( http_method=u'PATCH', method_id=u'bigquery.datasets.patch', ordered_params=[u'projectId', u'datasetId'], path_params=[u'datasetId', u'projectId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}', request_field=u'dataset', request_type_name=u'BigqueryDatasetsPatchRequest', response_type_name=u'Dataset', supports_download=False, ), 'Update': base_api.ApiMethodInfo( http_method=u'PUT', method_id=u'bigquery.datasets.update', ordered_params=[u'projectId', u'datasetId'], path_params=[u'datasetId', u'projectId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}', request_field=u'dataset', request_type_name=u'BigqueryDatasetsUpdateRequest', response_type_name=u'Dataset', supports_download=False, ), } self._upload_configs = { } def Delete(self, request, global_params=None): """Deletes the dataset specified by the datasetId value. Before you can delete a dataset, you must delete all its tables, either manually or by specifying deleteContents. Immediately after deletion, you can create another dataset with the same name. Args: request: (BigqueryDatasetsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BigqueryDatasetsDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None): """Returns the dataset specified by datasetID. Args: request: (BigqueryDatasetsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Dataset) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) def Insert(self, request, global_params=None): """Creates a new empty dataset. Args: request: (BigqueryDatasetsInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Dataset) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod( config, request, global_params=global_params) def List(self, request, global_params=None): """Lists all datasets in the specified project to which you have been granted the READER dataset role. Args: request: (BigqueryDatasetsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (DatasetList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Patch(self, request, global_params=None): """Updates information in an existing dataset. The update method replaces the entire dataset resource, whereas the patch method only replaces fields that are provided in the submitted dataset resource. This method supports patch semantics. Args: request: (BigqueryDatasetsPatchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Dataset) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod( config, request, global_params=global_params) def Update(self, request, global_params=None): """Updates information in an existing dataset. The update method replaces the entire dataset resource, whereas the patch method only replaces fields that are provided in the submitted dataset resource. Args: request: (BigqueryDatasetsUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Dataset) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) class JobsService(base_api.BaseApiService): """Service class for the jobs resource.""" _NAME = u'jobs' def __init__(self, client): super(BigqueryV2.JobsService, self).__init__(client) self._method_configs = { 'Cancel': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'bigquery.jobs.cancel', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[], relative_path=u'project/{projectId}/jobs/{jobId}/cancel', request_field='', request_type_name=u'BigqueryJobsCancelRequest', response_type_name=u'JobCancelResponse', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.jobs.get', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[], relative_path=u'projects/{projectId}/jobs/{jobId}', request_field='', request_type_name=u'BigqueryJobsGetRequest', response_type_name=u'Job', supports_download=False, ), 'GetQueryResults': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.jobs.getQueryResults', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[u'maxResults', u'pageToken', u'startIndex', u'timeoutMs'], relative_path=u'projects/{projectId}/queries/{jobId}', request_field='', request_type_name=u'BigqueryJobsGetQueryResultsRequest', response_type_name=u'GetQueryResultsResponse', supports_download=False, ), 'Insert': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'bigquery.jobs.insert', ordered_params=[u'projectId'], path_params=[u'projectId'], query_params=[], relative_path=u'projects/{projectId}/jobs', request_field=u'job', request_type_name=u'BigqueryJobsInsertRequest', response_type_name=u'Job', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.jobs.list', ordered_params=[u'projectId'], path_params=[u'projectId'], query_params=[u'allUsers', u'maxResults', u'pageToken', u'projection', u'stateFilter'], relative_path=u'projects/{projectId}/jobs', request_field='', request_type_name=u'BigqueryJobsListRequest', response_type_name=u'JobList', supports_download=False, ), 'Query': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'bigquery.jobs.query', ordered_params=[u'projectId'], path_params=[u'projectId'], query_params=[], relative_path=u'projects/{projectId}/queries', request_field=u'queryRequest', request_type_name=u'BigqueryJobsQueryRequest', response_type_name=u'QueryResponse', supports_download=False, ), } self._upload_configs = { 'Insert': base_api.ApiUploadInfo( accept=['*/*'], max_size=None, resumable_multipart=True, resumable_path=u'/resumable/upload/bigquery/v2/projects/{projectId}/jobs', simple_multipart=True, simple_path=u'/upload/bigquery/v2/projects/{projectId}/jobs', ), } def Cancel(self, request, global_params=None): """Requests that a job be cancelled. This call will return immediately, and the client will need to poll for the job status to see if the cancel completed successfully. Cancelled jobs may still incur costs. Args: request: (BigqueryJobsCancelRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (JobCancelResponse) The response message. """ config = self.GetMethodConfig('Cancel') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None): """Returns information about a specific job. Job information is available for a six month period after creation. Requires that you're the person who ran the job, or have the Is Owner project role. Args: request: (BigqueryJobsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) def GetQueryResults(self, request, global_params=None): """Retrieves the results of a query job. Args: request: (BigqueryJobsGetQueryResultsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (GetQueryResultsResponse) The response message. """ config = self.GetMethodConfig('GetQueryResults') return self._RunMethod( config, request, global_params=global_params) def Insert(self, request, global_params=None, upload=None): """Starts a new asynchronous job. Requires the Can View project role. Args: request: (BigqueryJobsInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments upload: (Upload, default: None) If present, upload this stream with the request. Returns: (Job) The response message. """ config = self.GetMethodConfig('Insert') upload_config = self.GetUploadConfig('Insert') return self._RunMethod( config, request, global_params=global_params, upload=upload, upload_config=upload_config) def List(self, request, global_params=None): """Lists all jobs that you started in the specified project. Job information is available for a six month period after creation. The job list is sorted in reverse chronological order, by job creation time. Requires the Can View project role, or the Is Owner project role if you set the allUsers property. Args: request: (BigqueryJobsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (JobList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Query(self, request, global_params=None): """Runs a BigQuery SQL query synchronously and returns query results if the query completes within a specified timeout. Args: request: (BigqueryJobsQueryRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (QueryResponse) The response message. """ config = self.GetMethodConfig('Query') return self._RunMethod( config, request, global_params=global_params) class ProjectsService(base_api.BaseApiService): """Service class for the projects resource.""" _NAME = u'projects' def __init__(self, client): super(BigqueryV2.ProjectsService, self).__init__(client) self._method_configs = { 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.projects.list', ordered_params=[], path_params=[], query_params=[u'maxResults', u'pageToken'], relative_path=u'projects', request_field='', request_type_name=u'BigqueryProjectsListRequest', response_type_name=u'ProjectList', supports_download=False, ), } self._upload_configs = { } def List(self, request, global_params=None): """Lists all projects to which you have been granted any project role. Args: request: (BigqueryProjectsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ProjectList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) class TabledataService(base_api.BaseApiService): """Service class for the tabledata resource.""" _NAME = u'tabledata' def __init__(self, client): super(BigqueryV2.TabledataService, self).__init__(client) self._method_configs = { 'InsertAll': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'bigquery.tabledata.insertAll', ordered_params=[u'projectId', u'datasetId', u'tableId'], path_params=[u'datasetId', u'projectId', u'tableId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}/tables/{tableId}/insertAll', request_field=u'tableDataInsertAllRequest', request_type_name=u'BigqueryTabledataInsertAllRequest', response_type_name=u'TableDataInsertAllResponse', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.tabledata.list', ordered_params=[u'projectId', u'datasetId', u'tableId'], path_params=[u'datasetId', u'projectId', u'tableId'], query_params=[u'maxResults', u'pageToken', u'startIndex'], relative_path=u'projects/{projectId}/datasets/{datasetId}/tables/{tableId}/data', request_field='', request_type_name=u'BigqueryTabledataListRequest', response_type_name=u'TableDataList', supports_download=False, ), } self._upload_configs = { } def InsertAll(self, request, global_params=None): """Streams data into BigQuery one record at a time without needing to run a load job. Requires the WRITER dataset role. Args: request: (BigqueryTabledataInsertAllRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (TableDataInsertAllResponse) The response message. """ config = self.GetMethodConfig('InsertAll') return self._RunMethod( config, request, global_params=global_params) def List(self, request, global_params=None): """Retrieves table data from a specified set of rows. Requires the READER dataset role. Args: request: (BigqueryTabledataListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (TableDataList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) class TablesService(base_api.BaseApiService): """Service class for the tables resource.""" _NAME = u'tables' def __init__(self, client): super(BigqueryV2.TablesService, self).__init__(client) self._method_configs = { 'Delete': base_api.ApiMethodInfo( http_method=u'DELETE', method_id=u'bigquery.tables.delete', ordered_params=[u'projectId', u'datasetId', u'tableId'], path_params=[u'datasetId', u'projectId', u'tableId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}/tables/{tableId}', request_field='', request_type_name=u'BigqueryTablesDeleteRequest', response_type_name=u'BigqueryTablesDeleteResponse', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.tables.get', ordered_params=[u'projectId', u'datasetId', u'tableId'], path_params=[u'datasetId', u'projectId', u'tableId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}/tables/{tableId}', request_field='', request_type_name=u'BigqueryTablesGetRequest', response_type_name=u'Table', supports_download=False, ), 'Insert': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'bigquery.tables.insert', ordered_params=[u'projectId', u'datasetId'], path_params=[u'datasetId', u'projectId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}/tables', request_field=u'table', request_type_name=u'BigqueryTablesInsertRequest', response_type_name=u'Table', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'bigquery.tables.list', ordered_params=[u'projectId', u'datasetId'], path_params=[u'datasetId', u'projectId'], query_params=[u'maxResults', u'pageToken'], relative_path=u'projects/{projectId}/datasets/{datasetId}/tables', request_field='', request_type_name=u'BigqueryTablesListRequest', response_type_name=u'TableList', supports_download=False, ), 'Patch': base_api.ApiMethodInfo( http_method=u'PATCH', method_id=u'bigquery.tables.patch', ordered_params=[u'projectId', u'datasetId', u'tableId'], path_params=[u'datasetId', u'projectId', u'tableId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}/tables/{tableId}', request_field=u'table', request_type_name=u'BigqueryTablesPatchRequest', response_type_name=u'Table', supports_download=False, ), 'Update': base_api.ApiMethodInfo( http_method=u'PUT', method_id=u'bigquery.tables.update', ordered_params=[u'projectId', u'datasetId', u'tableId'], path_params=[u'datasetId', u'projectId', u'tableId'], query_params=[], relative_path=u'projects/{projectId}/datasets/{datasetId}/tables/{tableId}', request_field=u'table', request_type_name=u'BigqueryTablesUpdateRequest', response_type_name=u'Table', supports_download=False, ), } self._upload_configs = { } def Delete(self, request, global_params=None): """Deletes the table specified by tableId from the dataset. If the table contains data, all the data will be deleted. Args: request: (BigqueryTablesDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BigqueryTablesDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None): """Gets the specified table resource by table ID. This method does not return the data in the table, it only returns the table resource, which describes the structure of this table. Args: request: (BigqueryTablesGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Table) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) def Insert(self, request, global_params=None): """Creates a new, empty table in the dataset. Args: request: (BigqueryTablesInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Table) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod( config, request, global_params=global_params) def List(self, request, global_params=None): """Lists all tables in the specified dataset. Requires the READER dataset role. Args: request: (BigqueryTablesListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (TableList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Patch(self, request, global_params=None): """Updates information in an existing table. The update method replaces the entire table resource, whereas the patch method only replaces fields that are provided in the submitted table resource. This method supports patch semantics. Args: request: (BigqueryTablesPatchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Table) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod( config, request, global_params=global_params) def Update(self, request, global_params=None): """Updates information in an existing table. The update method replaces the entire table resource, whereas the patch method only replaces fields that are provided in the submitted table resource. Args: request: (BigqueryTablesUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Table) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/bigquery/bigquery_v2_messages.py000066400000000000000000002363101272340251000330120ustar00rootroot00000000000000"""Generated message classes for bigquery version v2. A data platform for customers to create, manage, share and query data. """ # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.protorpclite import messages as _messages from apitools.base.py import encoding from apitools.base.py import extra_types package = 'bigquery' class BigqueryDatasetsDeleteRequest(_messages.Message): """A BigqueryDatasetsDeleteRequest object. Fields: datasetId: Dataset ID of dataset being deleted deleteContents: If True, delete all the tables in the dataset. If False and the dataset contains tables, the request will fail. Default is False projectId: Project ID of the dataset being deleted """ datasetId = _messages.StringField(1, required=True) deleteContents = _messages.BooleanField(2) projectId = _messages.StringField(3, required=True) class BigqueryDatasetsDeleteResponse(_messages.Message): """An empty BigqueryDatasetsDelete response.""" class BigqueryDatasetsGetRequest(_messages.Message): """A BigqueryDatasetsGetRequest object. Fields: datasetId: Dataset ID of the requested dataset projectId: Project ID of the requested dataset """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) class BigqueryDatasetsInsertRequest(_messages.Message): """A BigqueryDatasetsInsertRequest object. Fields: dataset: A Dataset resource to be passed as the request body. projectId: Project ID of the new dataset """ dataset = _messages.MessageField('Dataset', 1) projectId = _messages.StringField(2, required=True) class BigqueryDatasetsListRequest(_messages.Message): """A BigqueryDatasetsListRequest object. Fields: all: Whether to list all datasets, including hidden ones maxResults: The maximum number of results to return pageToken: Page token, returned by a previous call, to request the next page of results projectId: Project ID of the datasets to be listed """ all = _messages.BooleanField(1) maxResults = _messages.IntegerField(2, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) class BigqueryDatasetsPatchRequest(_messages.Message): """A BigqueryDatasetsPatchRequest object. Fields: dataset: A Dataset resource to be passed as the request body. datasetId: Dataset ID of the dataset being updated projectId: Project ID of the dataset being updated """ dataset = _messages.MessageField('Dataset', 1) datasetId = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class BigqueryDatasetsUpdateRequest(_messages.Message): """A BigqueryDatasetsUpdateRequest object. Fields: dataset: A Dataset resource to be passed as the request body. datasetId: Dataset ID of the dataset being updated projectId: Project ID of the dataset being updated """ dataset = _messages.MessageField('Dataset', 1) datasetId = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class BigqueryJobsCancelRequest(_messages.Message): """A BigqueryJobsCancelRequest object. Fields: jobId: [Required] Job ID of the job to cancel projectId: [Required] Project ID of the job to cancel """ jobId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) class BigqueryJobsGetQueryResultsRequest(_messages.Message): """A BigqueryJobsGetQueryResultsRequest object. Fields: jobId: [Required] Job ID of the query job maxResults: Maximum number of results to read pageToken: Page token, returned by a previous call, to request the next page of results projectId: [Required] Project ID of the query job startIndex: Zero-based index of the starting row timeoutMs: How long to wait for the query to complete, in milliseconds, before returning. Default is 10 seconds. If the timeout passes before the job completes, the 'jobComplete' field in the response will be false """ jobId = _messages.StringField(1, required=True) maxResults = _messages.IntegerField(2, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) startIndex = _messages.IntegerField(5, variant=_messages.Variant.UINT64) timeoutMs = _messages.IntegerField(6, variant=_messages.Variant.UINT32) class BigqueryJobsGetRequest(_messages.Message): """A BigqueryJobsGetRequest object. Fields: jobId: [Required] Job ID of the requested job projectId: [Required] Project ID of the requested job """ jobId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) class BigqueryJobsInsertRequest(_messages.Message): """A BigqueryJobsInsertRequest object. Fields: job: A Job resource to be passed as the request body. projectId: Project ID of the project that will be billed for the job """ job = _messages.MessageField('Job', 1) projectId = _messages.StringField(2, required=True) class BigqueryJobsListRequest(_messages.Message): """A BigqueryJobsListRequest object. Enums: ProjectionValueValuesEnum: Restrict information returned to a set of selected fields StateFilterValueValuesEnum: Filter for job state Fields: allUsers: Whether to display jobs owned by all users in the project. Default false maxResults: Maximum number of results to return pageToken: Page token, returned by a previous call, to request the next page of results projectId: Project ID of the jobs to list projection: Restrict information returned to a set of selected fields stateFilter: Filter for job state """ class ProjectionValueValuesEnum(_messages.Enum): """Restrict information returned to a set of selected fields Values: full: Includes all job data minimal: Does not include the job configuration """ full = 0 minimal = 1 class StateFilterValueValuesEnum(_messages.Enum): """Filter for job state Values: done: Finished jobs pending: Pending jobs running: Running jobs """ done = 0 pending = 1 running = 2 allUsers = _messages.BooleanField(1) maxResults = _messages.IntegerField(2, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) projection = _messages.EnumField('ProjectionValueValuesEnum', 5) stateFilter = _messages.EnumField('StateFilterValueValuesEnum', 6, repeated=True) class BigqueryJobsQueryRequest(_messages.Message): """A BigqueryJobsQueryRequest object. Fields: projectId: Project ID of the project billed for the query queryRequest: A QueryRequest resource to be passed as the request body. """ projectId = _messages.StringField(1, required=True) queryRequest = _messages.MessageField('QueryRequest', 2) class BigqueryProjectsListRequest(_messages.Message): """A BigqueryProjectsListRequest object. Fields: maxResults: Maximum number of results to return pageToken: Page token, returned by a previous call, to request the next page of results """ maxResults = _messages.IntegerField(1, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(2) class BigqueryTabledataInsertAllRequest(_messages.Message): """A BigqueryTabledataInsertAllRequest object. Fields: datasetId: Dataset ID of the destination table. projectId: Project ID of the destination table. tableDataInsertAllRequest: A TableDataInsertAllRequest resource to be passed as the request body. tableId: Table ID of the destination table. """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) tableDataInsertAllRequest = _messages.MessageField('TableDataInsertAllRequest', 3) tableId = _messages.StringField(4, required=True) class BigqueryTabledataListRequest(_messages.Message): """A BigqueryTabledataListRequest object. Fields: datasetId: Dataset ID of the table to read maxResults: Maximum number of results to return pageToken: Page token, returned by a previous call, identifying the result set projectId: Project ID of the table to read startIndex: Zero-based index of the starting row to read tableId: Table ID of the table to read """ datasetId = _messages.StringField(1, required=True) maxResults = _messages.IntegerField(2, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) startIndex = _messages.IntegerField(5, variant=_messages.Variant.UINT64) tableId = _messages.StringField(6, required=True) class BigqueryTablesDeleteRequest(_messages.Message): """A BigqueryTablesDeleteRequest object. Fields: datasetId: Dataset ID of the table to delete projectId: Project ID of the table to delete tableId: Table ID of the table to delete """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) tableId = _messages.StringField(3, required=True) class BigqueryTablesDeleteResponse(_messages.Message): """An empty BigqueryTablesDelete response.""" class BigqueryTablesGetRequest(_messages.Message): """A BigqueryTablesGetRequest object. Fields: datasetId: Dataset ID of the requested table projectId: Project ID of the requested table tableId: Table ID of the requested table """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) tableId = _messages.StringField(3, required=True) class BigqueryTablesInsertRequest(_messages.Message): """A BigqueryTablesInsertRequest object. Fields: datasetId: Dataset ID of the new table projectId: Project ID of the new table table: A Table resource to be passed as the request body. """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) table = _messages.MessageField('Table', 3) class BigqueryTablesListRequest(_messages.Message): """A BigqueryTablesListRequest object. Fields: datasetId: Dataset ID of the tables to list maxResults: Maximum number of results to return pageToken: Page token, returned by a previous call, to request the next page of results projectId: Project ID of the tables to list """ datasetId = _messages.StringField(1, required=True) maxResults = _messages.IntegerField(2, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) class BigqueryTablesPatchRequest(_messages.Message): """A BigqueryTablesPatchRequest object. Fields: datasetId: Dataset ID of the table to update projectId: Project ID of the table to update table: A Table resource to be passed as the request body. tableId: Table ID of the table to update """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) table = _messages.MessageField('Table', 3) tableId = _messages.StringField(4, required=True) class BigqueryTablesUpdateRequest(_messages.Message): """A BigqueryTablesUpdateRequest object. Fields: datasetId: Dataset ID of the table to update projectId: Project ID of the table to update table: A Table resource to be passed as the request body. tableId: Table ID of the table to update """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) table = _messages.MessageField('Table', 3) tableId = _messages.StringField(4, required=True) class BigtableColumn(_messages.Message): """A BigtableColumn object. Fields: encoding: [Optional] The encoding of the values when the type is not STRING. Acceptable encoding values are: TEXT - indicates values are alphanumeric text strings. BINARY - indicates values are encoded using HBase Bytes.toBytes family of functions. 'encoding' can also be set at the column family level. However, the setting at this level takes precedence if 'encoding' is set at both levels. fieldName: [Optional] If the qualifier is not a valid BigQuery field identifier i.e. does not match [a-zA-Z][a-zA-Z0-9_]*, a valid identifier must be provided as the column field name and is used as field name in queries. onlyReadLatest: [Optional] If this is set, only the latest version of value in this column are exposed. 'onlyReadLatest' can also be set at the column family level. However, the setting at this level takes precedence if 'onlyReadLatest' is set at both levels. qualifierEncoded: [Required] Qualifier of the column. Columns in the parent column family that has this exact qualifier are exposed as . field. If the qualifier is valid UTF-8 string, it can be specified in the qualifier_string field. Otherwise, a base-64 encoded value must be set to qualifier_encoded. The column field name is the same as the column qualifier. However, if the qualifier is not a valid BigQuery field identifier i.e. does not match [a-zA-Z][a-zA-Z0-9_]*, a valid identifier must be provided as field_name. qualifierString: A string attribute. type: [Optional] The type to convert the value in cells of this column. The values are expected to be encoded using HBase Bytes.toBytes function when using the BINARY encoding value. Following BigQuery types are allowed (case-sensitive) - BYTES STRING INTEGER FLOAT BOOLEAN Defaut type is BYTES. 'type' can also be set at the column family level. However, the setting at this level takes precedence if 'type' is set at both levels. """ encoding = _messages.StringField(1) fieldName = _messages.StringField(2) onlyReadLatest = _messages.BooleanField(3) qualifierEncoded = _messages.BytesField(4) qualifierString = _messages.StringField(5) type = _messages.StringField(6) class BigtableColumnFamily(_messages.Message): """A BigtableColumnFamily object. Fields: columns: [Optional] Lists of columns that should be exposed as individual fields as opposed to a list of (column name, value) pairs. All columns whose qualifier matches a qualifier in this list can be accessed as .. Other columns can be accessed as a list through .Column field. encoding: [Optional] The encoding of the values when the type is not STRING. Acceptable encoding values are: TEXT - indicates values are alphanumeric text strings. BINARY - indicates values are encoded using HBase Bytes.toBytes family of functions. This can be overridden for a specific column by listing that column in 'columns' and specifying an encoding for it. familyId: Identifier of the column family. onlyReadLatest: [Optional] If this is set only the latest version of value are exposed for all columns in this column family. This can be overridden for a specific column by listing that column in 'columns' and specifying a different setting for that column. type: [Optional] The type to convert the value in cells of this column family. The values are expected to be encoded using HBase Bytes.toBytes function when using the BINARY encoding value. Following BigQuery types are allowed (case-sensitive) - BYTES STRING INTEGER FLOAT BOOLEAN Defaut type is BYTES. This can be overridden for a specific column by listing that column in 'columns' and specifying a type for it. """ columns = _messages.MessageField('BigtableColumn', 1, repeated=True) encoding = _messages.StringField(2) familyId = _messages.StringField(3) onlyReadLatest = _messages.BooleanField(4) type = _messages.StringField(5) class BigtableOptions(_messages.Message): """A BigtableOptions object. Fields: columnFamilies: [Optional] List of column families to expose in the table schema along with their types. This list restricts the column families that can be referenced in queries and specifies their value types. You can use this list to do type conversions - see the 'type' field for more details. If you leave this list empty, all column families are present in the table schema and their values are read as BYTES. During a query only the column families referenced in that query are read from Bigtable. ignoreUnspecifiedColumnFamilies: [Optional] If field is true, then the column families that are not specified in columnFamilies list are not exposed in the table schema. Otherwise, they are read with BYTES type values. The default value is false. """ columnFamilies = _messages.MessageField('BigtableColumnFamily', 1, repeated=True) ignoreUnspecifiedColumnFamilies = _messages.BooleanField(2) class CsvOptions(_messages.Message): """A CsvOptions object. Fields: allowJaggedRows: [Optional] Indicates if BigQuery should accept rows that are missing trailing optional columns. If true, BigQuery treats missing trailing columns as null values. If false, records with missing trailing columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. The default value is false. allowQuotedNewlines: [Optional] Indicates if BigQuery should allow quoted data sections that contain newline characters in a CSV file. The default value is false. encoding: [Optional] The character encoding of the data. The supported values are UTF-8 or ISO-8859-1. The default value is UTF-8. BigQuery decodes the data after the raw, binary data has been split using the values of the quote and fieldDelimiter properties. fieldDelimiter: [Optional] The separator for fields in a CSV file. BigQuery converts the string to ISO-8859-1 encoding, and then uses the first byte of the encoded string to split the data in its raw, binary state. BigQuery also supports the escape sequence "\t" to specify a tab separator. The default value is a comma (','). quote: [Optional] The value that is used to quote data sections in a CSV file. BigQuery converts the string to ISO-8859-1 encoding, and then uses the first byte of the encoded string to split the data in its raw, binary state. The default value is a double-quote ('"'). If your data does not contain quoted sections, set the property value to an empty string. If your data contains quoted newline characters, you must also set the allowQuotedNewlines property to true. skipLeadingRows: [Optional] The number of rows at the top of a CSV file that BigQuery will skip when reading the data. The default value is 0. This property is useful if you have header rows in the file that should be skipped. """ allowJaggedRows = _messages.BooleanField(1) allowQuotedNewlines = _messages.BooleanField(2) encoding = _messages.StringField(3) fieldDelimiter = _messages.StringField(4) quote = _messages.StringField(5, default=u'"') skipLeadingRows = _messages.IntegerField(6, variant=_messages.Variant.INT32) class Dataset(_messages.Message): """A Dataset object. Messages: AccessValueListEntry: A AccessValueListEntry object. Fields: access: [Optional] An array of objects that define dataset access for one or more entities. You can set this property when inserting or updating a dataset in order to control who is allowed to access the data. If unspecified at dataset creation time, BigQuery adds default dataset access for the following entities: access.specialGroup: projectReaders; access.role: READER; access.specialGroup: projectWriters; access.role: WRITER; access.specialGroup: projectOwners; access.role: OWNER; access.userByEmail: [dataset creator email]; access.role: OWNER; creationTime: [Output-only] The time when this dataset was created, in milliseconds since the epoch. datasetReference: [Required] A reference that identifies the dataset. defaultTableExpirationMs: [Optional] The default lifetime of all tables in the dataset, in milliseconds. The minimum value is 3600000 milliseconds (one hour). Once this property is set, all newly-created tables in the dataset will have an expirationTime property set to the creation time plus the value in this property, and changing the value will only affect new tables, not existing ones. When the expirationTime for a given table is reached, that table will be deleted automatically. If a table's expirationTime is modified or removed before the table expires, or if you provide an explicit expirationTime when creating a table, that value takes precedence over the default expiration time indicated by this property. description: [Optional] A user-friendly description of the dataset. etag: [Output-only] A hash of the resource. friendlyName: [Optional] A descriptive name for the dataset. id: [Output-only] The fully-qualified unique name of the dataset in the format projectId:datasetId. The dataset name without the project name is given in the datasetId field. When creating a new dataset, leave this field blank, and instead specify the datasetId field. kind: [Output-only] The resource type. lastModifiedTime: [Output-only] The date when this dataset or any of its tables was last modified, in milliseconds since the epoch. location: [Experimental] The geographic location where the dataset should reside. Possible values include EU and US. The default value is US. selfLink: [Output-only] A URL that can be used to access the resource again. You can use this URL in Get or Update requests to the resource. """ class AccessValueListEntry(_messages.Message): """A AccessValueListEntry object. Fields: domain: [Pick one] A domain to grant access to. Any users signed in with the domain specified will be granted the specified access. Example: "example.com". groupByEmail: [Pick one] An email address of a Google Group to grant access to. role: [Required] Describes the rights granted to the user specified by the other member of the access object. The following string values are supported: READER, WRITER, OWNER. specialGroup: [Pick one] A special group to grant access to. Possible values include: projectOwners: Owners of the enclosing project. projectReaders: Readers of the enclosing project. projectWriters: Writers of the enclosing project. allAuthenticatedUsers: All authenticated BigQuery users. userByEmail: [Pick one] An email address of a user to grant access to. For example: fred@example.com. view: [Pick one] A view from a different dataset to grant access to. Queries executed against that view will have read access to tables in this dataset. The role field is not required when this field is set. If that view is updated by any user, access to the view needs to be granted again via an update operation. """ domain = _messages.StringField(1) groupByEmail = _messages.StringField(2) role = _messages.StringField(3) specialGroup = _messages.StringField(4) userByEmail = _messages.StringField(5) view = _messages.MessageField('TableReference', 6) access = _messages.MessageField('AccessValueListEntry', 1, repeated=True) creationTime = _messages.IntegerField(2) datasetReference = _messages.MessageField('DatasetReference', 3) defaultTableExpirationMs = _messages.IntegerField(4) description = _messages.StringField(5) etag = _messages.StringField(6) friendlyName = _messages.StringField(7) id = _messages.StringField(8) kind = _messages.StringField(9, default=u'bigquery#dataset') lastModifiedTime = _messages.IntegerField(10) location = _messages.StringField(11) selfLink = _messages.StringField(12) class DatasetList(_messages.Message): """A DatasetList object. Messages: DatasetsValueListEntry: A DatasetsValueListEntry object. Fields: datasets: An array of the dataset resources in the project. Each resource contains basic information. For full information about a particular dataset resource, use the Datasets: get method. This property is omitted when there are no datasets in the project. etag: A hash value of the results page. You can use this property to determine if the page has changed since the last request. kind: The list type. This property always returns the value "bigquery#datasetList". nextPageToken: A token that can be used to request the next results page. This property is omitted on the final results page. """ class DatasetsValueListEntry(_messages.Message): """A DatasetsValueListEntry object. Fields: datasetReference: The dataset reference. Use this property to access specific parts of the dataset's ID, such as project ID or dataset ID. friendlyName: A descriptive name for the dataset, if one exists. id: The fully-qualified, unique, opaque ID of the dataset. kind: The resource type. This property always returns the value "bigquery#dataset". """ datasetReference = _messages.MessageField('DatasetReference', 1) friendlyName = _messages.StringField(2) id = _messages.StringField(3) kind = _messages.StringField(4, default=u'bigquery#dataset') datasets = _messages.MessageField('DatasetsValueListEntry', 1, repeated=True) etag = _messages.StringField(2) kind = _messages.StringField(3, default=u'bigquery#datasetList') nextPageToken = _messages.StringField(4) class DatasetReference(_messages.Message): """A DatasetReference object. Fields: datasetId: [Required] A unique ID for this dataset, without the project name. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). The maximum length is 1,024 characters. projectId: [Optional] The ID of the project containing this dataset. """ datasetId = _messages.StringField(1) projectId = _messages.StringField(2) class ErrorProto(_messages.Message): """A ErrorProto object. Fields: debugInfo: Debugging information. This property is internal to Google and should not be used. location: Specifies where the error occurred, if present. message: A human-readable description of the error. reason: A short error code that summarizes the error. """ debugInfo = _messages.StringField(1) location = _messages.StringField(2) message = _messages.StringField(3) reason = _messages.StringField(4) class ExplainQueryStage(_messages.Message): """A ExplainQueryStage object. Fields: computeRatioAvg: Relative amount of time the average shard spent on CPU- bound tasks. computeRatioMax: Relative amount of time the slowest shard spent on CPU- bound tasks. id: Unique ID for stage within plan. name: Human-readable name for stage. readRatioAvg: Relative amount of time the average shard spent reading input. readRatioMax: Relative amount of time the slowest shard spent reading input. recordsRead: Number of records read into the stage. recordsWritten: Number of records written by the stage. steps: List of operations within the stage in dependency order (approximately chronological). waitRatioAvg: Relative amount of time the average shard spent waiting to be scheduled. waitRatioMax: Relative amount of time the slowest shard spent waiting to be scheduled. writeRatioAvg: Relative amount of time the average shard spent on writing output. writeRatioMax: Relative amount of time the slowest shard spent on writing output. """ computeRatioAvg = _messages.FloatField(1) computeRatioMax = _messages.FloatField(2) id = _messages.IntegerField(3) name = _messages.StringField(4) readRatioAvg = _messages.FloatField(5) readRatioMax = _messages.FloatField(6) recordsRead = _messages.IntegerField(7) recordsWritten = _messages.IntegerField(8) steps = _messages.MessageField('ExplainQueryStep', 9, repeated=True) waitRatioAvg = _messages.FloatField(10) waitRatioMax = _messages.FloatField(11) writeRatioAvg = _messages.FloatField(12) writeRatioMax = _messages.FloatField(13) class ExplainQueryStep(_messages.Message): """A ExplainQueryStep object. Fields: kind: Machine-readable operation type. substeps: Human-readable stage descriptions. """ kind = _messages.StringField(1) substeps = _messages.StringField(2, repeated=True) class ExternalDataConfiguration(_messages.Message): """A ExternalDataConfiguration object. Fields: autodetect: [Experimental] Try to detect schema and format options automatically. Any option specified explicitly will be honored. bigtableOptions: [Optional] Additional options if sourceFormat is set to BIGTABLE. compression: [Optional] The compression type of the data source. Possible values include GZIP and NONE. The default value is NONE. This setting is ignored for Google Cloud Bigtable, Google Cloud Datastore backups and Avro formats. csvOptions: Additional properties to set if sourceFormat is set to CSV. ignoreUnknownValues: [Optional] Indicates if BigQuery should allow extra values that are not represented in the table schema. If true, the extra values are ignored. If false, records with extra columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. The default value is false. The sourceFormat property determines what BigQuery treats as an extra value: CSV: Trailing columns JSON: Named values that don't match any column names Google Cloud Bigtable: This setting is ignored. Google Cloud Datastore backups: This setting is ignored. Avro: This setting is ignored. maxBadRecords: [Optional] The maximum number of bad records that BigQuery can ignore when reading data. If the number of bad records exceeds this value, an invalid error is returned in the job result. The default value is 0, which requires that all records are valid. This setting is ignored for Google Cloud Bigtable, Google Cloud Datastore backups and Avro formats. schema: [Optional] The schema for the data. Schema is required for CSV and JSON formats. Schema is disallowed for Google Cloud Bigtable, Cloud Datastore backups, and Avro formats. sourceFormat: [Required] The data format. For CSV files, specify "CSV". For newline-delimited JSON, specify "NEWLINE_DELIMITED_JSON". For Avro files, specify "AVRO". For Google Cloud Datastore backups, specify "DATASTORE_BACKUP". [Experimental] For Google Cloud Bigtable, specify "BIGTABLE". Please note that reading from Google Cloud Bigtable is experimental and has to be enabled for your project. Please contact Google Cloud Support to enable this for your project. sourceUris: [Required] The fully-qualified URIs that point to your data in Google Cloud. For Google Cloud Storage URIs: Each URI can contain one '*' wildcard character and it must come after the 'bucket' name. Size limits related to load jobs apply to external data sources, plus an additional limit of 10 GB maximum size across all URIs. For Google Cloud Bigtable URIs: Exactly one URI can be specified and it has be a fully specified and valid HTTPS URL for a Google Cloud Bigtable table. For Google Cloud Datastore backups, exactly one URI can be specified, and it must end with '.backup_info'. Also, the '*' wildcard character is not allowed. """ autodetect = _messages.BooleanField(1) bigtableOptions = _messages.MessageField('BigtableOptions', 2) compression = _messages.StringField(3) csvOptions = _messages.MessageField('CsvOptions', 4) ignoreUnknownValues = _messages.BooleanField(5) maxBadRecords = _messages.IntegerField(6, variant=_messages.Variant.INT32) schema = _messages.MessageField('TableSchema', 7) sourceFormat = _messages.StringField(8) sourceUris = _messages.StringField(9, repeated=True) class GetQueryResultsResponse(_messages.Message): """A GetQueryResultsResponse object. Fields: cacheHit: Whether the query result was fetched from the query cache. errors: [Output-only] All errors and warnings encountered during the running of the job. Errors here do not necessarily mean that the job has completed or was unsuccessful. etag: A hash of this response. jobComplete: Whether the query has completed or not. If rows or totalRows are present, this will always be true. If this is false, totalRows will not be available. jobReference: Reference to the BigQuery Job that was created to run the query. This field will be present even if the original request timed out, in which case GetQueryResults can be used to read the results once the query has completed. Since this API only returns the first page of results, subsequent pages can be fetched via the same mechanism (GetQueryResults). kind: The resource type of the response. pageToken: A token used for paging results. rows: An object with as many results as can be contained within the maximum permitted reply size. To get any additional rows, you can call GetQueryResults and specify the jobReference returned above. Present only when the query completes successfully. schema: The schema of the results. Present only when the query completes successfully. totalBytesProcessed: The total number of bytes processed for this query. totalRows: The total number of rows in the complete query result set, which can be more than the number of rows in this single page of results. Present only when the query completes successfully. """ cacheHit = _messages.BooleanField(1) errors = _messages.MessageField('ErrorProto', 2, repeated=True) etag = _messages.StringField(3) jobComplete = _messages.BooleanField(4) jobReference = _messages.MessageField('JobReference', 5) kind = _messages.StringField(6, default=u'bigquery#getQueryResultsResponse') pageToken = _messages.StringField(7) rows = _messages.MessageField('TableRow', 8, repeated=True) schema = _messages.MessageField('TableSchema', 9) totalBytesProcessed = _messages.IntegerField(10) totalRows = _messages.IntegerField(11, variant=_messages.Variant.UINT64) class IntervalPartitionConfiguration(_messages.Message): """A IntervalPartitionConfiguration object. Fields: expirationMs: A string attribute. type: A string attribute. """ expirationMs = _messages.IntegerField(1) type = _messages.StringField(2) class Job(_messages.Message): """A Job object. Fields: configuration: [Required] Describes the job configuration. etag: [Output-only] A hash of this resource. id: [Output-only] Opaque ID field of the job jobReference: [Optional] Reference describing the unique-per-user name of the job. kind: [Output-only] The type of the resource. selfLink: [Output-only] A URL that can be used to access this resource again. statistics: [Output-only] Information about the job, including starting time and ending time of the job. status: [Output-only] The status of this job. Examine this value when polling an asynchronous job to see if the job is complete. user_email: [Output-only] Email address of the user who ran the job. """ configuration = _messages.MessageField('JobConfiguration', 1) etag = _messages.StringField(2) id = _messages.StringField(3) jobReference = _messages.MessageField('JobReference', 4) kind = _messages.StringField(5, default=u'bigquery#job') selfLink = _messages.StringField(6) statistics = _messages.MessageField('JobStatistics', 7) status = _messages.MessageField('JobStatus', 8) user_email = _messages.StringField(9) class JobCancelResponse(_messages.Message): """A JobCancelResponse object. Fields: job: The final state of the job. kind: The resource type of the response. """ job = _messages.MessageField('Job', 1) kind = _messages.StringField(2, default=u'bigquery#jobCancelResponse') class JobConfiguration(_messages.Message): """A JobConfiguration object. Fields: copy: [Pick one] Copies a table. dryRun: [Optional] If set, don't actually run this job. A valid query will return a mostly empty response with some processing statistics, while an invalid query will return the same error it would if it wasn't a dry run. Behavior of non-query jobs is undefined. extract: [Pick one] Configures an extract job. load: [Pick one] Configures a load job. query: [Pick one] Configures a query job. """ copy = _messages.MessageField('JobConfigurationTableCopy', 1) dryRun = _messages.BooleanField(2) extract = _messages.MessageField('JobConfigurationExtract', 3) load = _messages.MessageField('JobConfigurationLoad', 4) query = _messages.MessageField('JobConfigurationQuery', 5) class JobConfigurationExtract(_messages.Message): """A JobConfigurationExtract object. Fields: compression: [Optional] The compression type to use for exported files. Possible values include GZIP and NONE. The default value is NONE. destinationFormat: [Optional] The exported file format. Possible values include CSV, NEWLINE_DELIMITED_JSON and AVRO. The default value is CSV. Tables with nested or repeated fields cannot be exported as CSV. destinationUri: [Pick one] DEPRECATED: Use destinationUris instead, passing only one URI as necessary. The fully-qualified Google Cloud Storage URI where the extracted table should be written. destinationUris: [Pick one] A list of fully-qualified Google Cloud Storage URIs where the extracted table should be written. fieldDelimiter: [Optional] Delimiter to use between fields in the exported data. Default is ',' printHeader: [Optional] Whether to print out a header row in the results. Default is true. sourceTable: [Required] A reference to the table being exported. """ compression = _messages.StringField(1) destinationFormat = _messages.StringField(2) destinationUri = _messages.StringField(3) destinationUris = _messages.StringField(4, repeated=True) fieldDelimiter = _messages.StringField(5) printHeader = _messages.BooleanField(6, default=True) sourceTable = _messages.MessageField('TableReference', 7) class JobConfigurationLoad(_messages.Message): """A JobConfigurationLoad object. Fields: allowJaggedRows: [Optional] Accept rows that are missing trailing optional columns. The missing values are treated as nulls. If false, records with missing trailing columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. The default value is false. Only applicable to CSV, ignored for other formats. allowQuotedNewlines: Indicates if BigQuery should allow quoted data sections that contain newline characters in a CSV file. The default value is false. createDisposition: [Optional] Specifies whether the job is allowed to create new tables. The following values are supported: CREATE_IF_NEEDED: If the table does not exist, BigQuery creates the table. CREATE_NEVER: The table must already exist. If it does not, a 'notFound' error is returned in the job result. The default value is CREATE_IF_NEEDED. Creation, truncation and append actions occur as one atomic update upon job completion. destinationTable: [Required] The destination table to load the data into. encoding: [Optional] The character encoding of the data. The supported values are UTF-8 or ISO-8859-1. The default value is UTF-8. BigQuery decodes the data after the raw, binary data has been split using the values of the quote and fieldDelimiter properties. fieldDelimiter: [Optional] The separator for fields in a CSV file. The separator can be any ISO-8859-1 single-byte character. To use a character in the range 128-255, you must encode the character as UTF8. BigQuery converts the string to ISO-8859-1 encoding, and then uses the first byte of the encoded string to split the data in its raw, binary state. BigQuery also supports the escape sequence "\t" to specify a tab separator. The default value is a comma (','). ignoreUnknownValues: [Optional] Indicates if BigQuery should allow extra values that are not represented in the table schema. If true, the extra values are ignored. If false, records with extra columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. The default value is false. The sourceFormat property determines what BigQuery treats as an extra value: CSV: Trailing columns JSON: Named values that don't match any column names maxBadRecords: [Optional] The maximum number of bad records that BigQuery can ignore when running the job. If the number of bad records exceeds this value, an invalid error is returned in the job result. The default value is 0, which requires that all records are valid. projectionFields: [Experimental] If sourceFormat is set to "DATASTORE_BACKUP", indicates which entity properties to load into BigQuery from a Cloud Datastore backup. Property names are case sensitive and must be top-level properties. If no properties are specified, BigQuery loads all properties. If any named property isn't found in the Cloud Datastore backup, an invalid error is returned in the job result. quote: [Optional] The value that is used to quote data sections in a CSV file. BigQuery converts the string to ISO-8859-1 encoding, and then uses the first byte of the encoded string to split the data in its raw, binary state. The default value is a double-quote ('"'). If your data does not contain quoted sections, set the property value to an empty string. If your data contains quoted newline characters, you must also set the allowQuotedNewlines property to true. schema: [Optional] The schema for the destination table. The schema can be omitted if the destination table already exists, or if you're loading data from Google Cloud Datastore. schemaInline: [Deprecated] The inline schema. For CSV schemas, specify as "Field1:Type1[,Field2:Type2]*". For example, "foo:STRING, bar:INTEGER, baz:FLOAT". schemaInlineFormat: [Deprecated] The format of the schemaInline property. skipLeadingRows: [Optional] The number of rows at the top of a CSV file that BigQuery will skip when loading the data. The default value is 0. This property is useful if you have header rows in the file that should be skipped. sourceFormat: [Optional] The format of the data files. For CSV files, specify "CSV". For datastore backups, specify "DATASTORE_BACKUP". For newline-delimited JSON, specify "NEWLINE_DELIMITED_JSON". For Avro, specify "AVRO". The default value is CSV. sourceUris: [Required] The fully-qualified URIs that point to your data in Google Cloud Storage. Each URI can contain one '*' wildcard character and it must come after the 'bucket' name. writeDisposition: [Optional] Specifies the action that occurs if the destination table already exists. The following values are supported: WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the table data. WRITE_APPEND: If the table already exists, BigQuery appends the data to the table. WRITE_EMPTY: If the table already exists and contains data, a 'duplicate' error is returned in the job result. The default value is WRITE_APPEND. Each action is atomic and only occurs if BigQuery is able to complete the job successfully. Creation, truncation and append actions occur as one atomic update upon job completion. """ allowJaggedRows = _messages.BooleanField(1) allowQuotedNewlines = _messages.BooleanField(2) createDisposition = _messages.StringField(3) destinationTable = _messages.MessageField('TableReference', 4) encoding = _messages.StringField(5) fieldDelimiter = _messages.StringField(6) ignoreUnknownValues = _messages.BooleanField(7) maxBadRecords = _messages.IntegerField(8, variant=_messages.Variant.INT32) projectionFields = _messages.StringField(9, repeated=True) quote = _messages.StringField(10, default=u'"') schema = _messages.MessageField('TableSchema', 11) schemaInline = _messages.StringField(12) schemaInlineFormat = _messages.StringField(13) skipLeadingRows = _messages.IntegerField(14, variant=_messages.Variant.INT32) sourceFormat = _messages.StringField(15) sourceUris = _messages.StringField(16, repeated=True) writeDisposition = _messages.StringField(17) class JobConfigurationQuery(_messages.Message): """A JobConfigurationQuery object. Messages: TableDefinitionsValue: [Optional] If querying an external data source outside of BigQuery, describes the data format, location and other properties of the data source. By defining these properties, the data source can then be queried as if it were a standard BigQuery table. Fields: allowLargeResults: If true, allows the query to produce arbitrarily large result tables at a slight cost in performance. Requires destinationTable to be set. createDisposition: [Optional] Specifies whether the job is allowed to create new tables. The following values are supported: CREATE_IF_NEEDED: If the table does not exist, BigQuery creates the table. CREATE_NEVER: The table must already exist. If it does not, a 'notFound' error is returned in the job result. The default value is CREATE_IF_NEEDED. Creation, truncation and append actions occur as one atomic update upon job completion. defaultDataset: [Optional] Specifies the default dataset to use for unqualified table names in the query. destinationTable: [Optional] Describes the table where the query results should be stored. If not present, a new table will be created to store the results. flattenResults: [Optional] Flattens all nested and repeated fields in the query results. The default value is true. allowLargeResults must be true if this is set to false. maximumBillingTier: [Optional] Limits the billing tier for this job. Queries that have resource usage beyond this tier will fail (without incurring a charge). If unspecified, this will be set to your project default. preserveNulls: [Deprecated] This property is deprecated. priority: [Optional] Specifies a priority for the query. Possible values include INTERACTIVE and BATCH. The default value is INTERACTIVE. query: [Required] BigQuery SQL query to execute. tableDefinitions: [Optional] If querying an external data source outside of BigQuery, describes the data format, location and other properties of the data source. By defining these properties, the data source can then be queried as if it were a standard BigQuery table. useLegacySql: [Experimental] Specifies whether to use BigQuery's legacy SQL dialect for this query. The default value is true. If set to false, the query will use BigQuery's updated SQL dialect with improved standards compliance. When using BigQuery's updated SQL, the values of allowLargeResults and flattenResults are ignored. Queries with useLegacySql set to false will be run as if allowLargeResults is true and flattenResults is false. useQueryCache: [Optional] Whether to look for the result in the query cache. The query cache is a best-effort cache that will be flushed whenever tables in the query are modified. Moreover, the query cache is only available when a query does not have a destination table specified. The default value is true. userDefinedFunctionResources: [Experimental] Describes user-defined function resources used in the query. writeDisposition: [Optional] Specifies the action that occurs if the destination table already exists. The following values are supported: WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the table data. WRITE_APPEND: If the table already exists, BigQuery appends the data to the table. WRITE_EMPTY: If the table already exists and contains data, a 'duplicate' error is returned in the job result. The default value is WRITE_EMPTY. Each action is atomic and only occurs if BigQuery is able to complete the job successfully. Creation, truncation and append actions occur as one atomic update upon job completion. """ @encoding.MapUnrecognizedFields('additionalProperties') class TableDefinitionsValue(_messages.Message): """[Optional] If querying an external data source outside of BigQuery, describes the data format, location and other properties of the data source. By defining these properties, the data source can then be queried as if it were a standard BigQuery table. Messages: AdditionalProperty: An additional property for a TableDefinitionsValue object. Fields: additionalProperties: Additional properties of type TableDefinitionsValue """ class AdditionalProperty(_messages.Message): """An additional property for a TableDefinitionsValue object. Fields: key: Name of the additional property. value: A ExternalDataConfiguration attribute. """ key = _messages.StringField(1) value = _messages.MessageField('ExternalDataConfiguration', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) allowLargeResults = _messages.BooleanField(1) createDisposition = _messages.StringField(2) defaultDataset = _messages.MessageField('DatasetReference', 3) destinationTable = _messages.MessageField('TableReference', 4) flattenResults = _messages.BooleanField(5, default=True) maximumBillingTier = _messages.IntegerField(6, variant=_messages.Variant.INT32, default=1) preserveNulls = _messages.BooleanField(7) priority = _messages.StringField(8) query = _messages.StringField(9) tableDefinitions = _messages.MessageField('TableDefinitionsValue', 10) useLegacySql = _messages.BooleanField(11) useQueryCache = _messages.BooleanField(12, default=True) userDefinedFunctionResources = _messages.MessageField('UserDefinedFunctionResource', 13, repeated=True) writeDisposition = _messages.StringField(14) class JobConfigurationTableCopy(_messages.Message): """A JobConfigurationTableCopy object. Fields: createDisposition: [Optional] Specifies whether the job is allowed to create new tables. The following values are supported: CREATE_IF_NEEDED: If the table does not exist, BigQuery creates the table. CREATE_NEVER: The table must already exist. If it does not, a 'notFound' error is returned in the job result. The default value is CREATE_IF_NEEDED. Creation, truncation and append actions occur as one atomic update upon job completion. destinationTable: [Required] The destination table sourceTable: [Pick one] Source table to copy. sourceTables: [Pick one] Source tables to copy. writeDisposition: [Optional] Specifies the action that occurs if the destination table already exists. The following values are supported: WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the table data. WRITE_APPEND: If the table already exists, BigQuery appends the data to the table. WRITE_EMPTY: If the table already exists and contains data, a 'duplicate' error is returned in the job result. The default value is WRITE_EMPTY. Each action is atomic and only occurs if BigQuery is able to complete the job successfully. Creation, truncation and append actions occur as one atomic update upon job completion. """ createDisposition = _messages.StringField(1) destinationTable = _messages.MessageField('TableReference', 2) sourceTable = _messages.MessageField('TableReference', 3) sourceTables = _messages.MessageField('TableReference', 4, repeated=True) writeDisposition = _messages.StringField(5) class JobList(_messages.Message): """A JobList object. Messages: JobsValueListEntry: A JobsValueListEntry object. Fields: etag: A hash of this page of results. jobs: List of jobs that were requested. kind: The resource type of the response. nextPageToken: A token to request the next page of results. """ class JobsValueListEntry(_messages.Message): """A JobsValueListEntry object. Fields: configuration: [Full-projection-only] Specifies the job configuration. errorResult: A result object that will be present only if the job has failed. id: Unique opaque ID of the job. jobReference: Job reference uniquely identifying the job. kind: The resource type. state: Running state of the job. When the state is DONE, errorResult can be checked to determine whether the job succeeded or failed. statistics: [Output-only] Information about the job, including starting time and ending time of the job. status: [Full-projection-only] Describes the state of the job. user_email: [Full-projection-only] Email address of the user who ran the job. """ configuration = _messages.MessageField('JobConfiguration', 1) errorResult = _messages.MessageField('ErrorProto', 2) id = _messages.StringField(3) jobReference = _messages.MessageField('JobReference', 4) kind = _messages.StringField(5, default=u'bigquery#job') state = _messages.StringField(6) statistics = _messages.MessageField('JobStatistics', 7) status = _messages.MessageField('JobStatus', 8) user_email = _messages.StringField(9) etag = _messages.StringField(1) jobs = _messages.MessageField('JobsValueListEntry', 2, repeated=True) kind = _messages.StringField(3, default=u'bigquery#jobList') nextPageToken = _messages.StringField(4) class JobReference(_messages.Message): """A JobReference object. Fields: jobId: [Required] The ID of the job. The ID must contain only letters (a-z, A-Z), numbers (0-9), underscores (_), or dashes (-). The maximum length is 1,024 characters. projectId: [Required] The ID of the project containing this job. """ jobId = _messages.StringField(1) projectId = _messages.StringField(2) class JobStatistics(_messages.Message): """A JobStatistics object. Fields: creationTime: [Output-only] Creation time of this job, in milliseconds since the epoch. This field will be present on all jobs. endTime: [Output-only] End time of this job, in milliseconds since the epoch. This field will be present whenever a job is in the DONE state. extract: [Output-only] Statistics for an extract job. load: [Output-only] Statistics for a load job. query: [Output-only] Statistics for a query job. startTime: [Output-only] Start time of this job, in milliseconds since the epoch. This field will be present when the job transitions from the PENDING state to either RUNNING or DONE. totalBytesProcessed: [Output-only] [Deprecated] Use the bytes processed in the query statistics instead. """ creationTime = _messages.IntegerField(1) endTime = _messages.IntegerField(2) extract = _messages.MessageField('JobStatistics4', 3) load = _messages.MessageField('JobStatistics3', 4) query = _messages.MessageField('JobStatistics2', 5) startTime = _messages.IntegerField(6) totalBytesProcessed = _messages.IntegerField(7) class JobStatistics2(_messages.Message): """A JobStatistics2 object. Fields: billingTier: [Output-only] Billing tier for the job. cacheHit: [Output-only] Whether the query result was fetched from the query cache. queryPlan: [Output-only, Experimental] Describes execution plan for the query as a list of stages. referencedTables: [Output-only, Experimental] Referenced tables for the job. Queries that reference more than 50 tables will not have a complete list. totalBytesBilled: [Output-only] Total bytes billed for the job. totalBytesProcessed: [Output-only] Total bytes processed for the job. """ billingTier = _messages.IntegerField(1, variant=_messages.Variant.INT32) cacheHit = _messages.BooleanField(2) queryPlan = _messages.MessageField('ExplainQueryStage', 3, repeated=True) referencedTables = _messages.MessageField('TableReference', 4, repeated=True) totalBytesBilled = _messages.IntegerField(5) totalBytesProcessed = _messages.IntegerField(6) class JobStatistics3(_messages.Message): """A JobStatistics3 object. Fields: inputFileBytes: [Output-only] Number of bytes of source data in a load job. inputFiles: [Output-only] Number of source files in a load job. outputBytes: [Output-only] Size of the loaded data in bytes. Note that while a load job is in the running state, this value may change. outputRows: [Output-only] Number of rows imported in a load job. Note that while an import job is in the running state, this value may change. """ inputFileBytes = _messages.IntegerField(1) inputFiles = _messages.IntegerField(2) outputBytes = _messages.IntegerField(3) outputRows = _messages.IntegerField(4) class JobStatistics4(_messages.Message): """A JobStatistics4 object. Fields: destinationUriFileCounts: [Output-only] Number of files per destination URI or URI pattern specified in the extract configuration. These values will be in the same order as the URIs specified in the 'destinationUris' field. """ destinationUriFileCounts = _messages.IntegerField(1, repeated=True) class JobStatus(_messages.Message): """A JobStatus object. Fields: errorResult: [Output-only] Final error result of the job. If present, indicates that the job has completed and was unsuccessful. errors: [Output-only] All errors encountered during the running of the job. Errors here do not necessarily mean that the job has completed or was unsuccessful. state: [Output-only] Running state of the job. """ errorResult = _messages.MessageField('ErrorProto', 1) errors = _messages.MessageField('ErrorProto', 2, repeated=True) state = _messages.StringField(3) @encoding.MapUnrecognizedFields('additionalProperties') class JsonObject(_messages.Message): """Represents a single JSON object. Messages: AdditionalProperty: An additional property for a JsonObject object. Fields: additionalProperties: Additional properties of type JsonObject """ class AdditionalProperty(_messages.Message): """An additional property for a JsonObject object. Fields: key: Name of the additional property. value: A JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) JsonValue = extra_types.JsonValue class ProjectList(_messages.Message): """A ProjectList object. Messages: ProjectsValueListEntry: A ProjectsValueListEntry object. Fields: etag: A hash of the page of results kind: The type of list. nextPageToken: A token to request the next page of results. projects: Projects to which you have at least READ access. totalItems: The total number of projects in the list. """ class ProjectsValueListEntry(_messages.Message): """A ProjectsValueListEntry object. Fields: friendlyName: A descriptive name for this project. id: An opaque ID of this project. kind: The resource type. numericId: The numeric ID of this project. projectReference: A unique reference to this project. """ friendlyName = _messages.StringField(1) id = _messages.StringField(2) kind = _messages.StringField(3, default=u'bigquery#project') numericId = _messages.IntegerField(4, variant=_messages.Variant.UINT64) projectReference = _messages.MessageField('ProjectReference', 5) etag = _messages.StringField(1) kind = _messages.StringField(2, default=u'bigquery#projectList') nextPageToken = _messages.StringField(3) projects = _messages.MessageField('ProjectsValueListEntry', 4, repeated=True) totalItems = _messages.IntegerField(5, variant=_messages.Variant.INT32) class ProjectReference(_messages.Message): """A ProjectReference object. Fields: projectId: [Required] ID of the project. Can be either the numeric ID or the assigned ID of the project. """ projectId = _messages.StringField(1) class QueryRequest(_messages.Message): """A QueryRequest object. Fields: defaultDataset: [Optional] Specifies the default datasetId and projectId to assume for any unqualified table names in the query. If not set, all table names in the query string must be qualified in the format 'datasetId.tableId'. dryRun: [Optional] If set to true, BigQuery doesn't run the job. Instead, if the query is valid, BigQuery returns statistics about the job such as how many bytes would be processed. If the query is invalid, an error returns. The default value is false. kind: The resource type of the request. maxResults: [Optional] The maximum number of rows of data to return per page of results. Setting this flag to a small value such as 1000 and then paging through results might improve reliability when the query result set is large. In addition to this limit, responses are also limited to 10 MB. By default, there is no maximum row count, and only the byte limit applies. preserveNulls: [Deprecated] This property is deprecated. query: [Required] A query string, following the BigQuery query syntax, of the query to execute. Example: "SELECT count(f1) FROM [myProjectId:myDatasetId.myTableId]". timeoutMs: [Optional] How long to wait for the query to complete, in milliseconds, before the request times out and returns. Note that this is only a timeout for the request, not the query. If the query takes longer to run than the timeout value, the call returns without any results and with the 'jobComplete' flag set to false. You can call GetQueryResults() to wait for the query to complete and read the results. The default value is 10000 milliseconds (10 seconds). useLegacySql: [Experimental] Specifies whether to use BigQuery's legacy SQL dialect for this query. The default value is true. If set to false, the query will use BigQuery's updated SQL dialect with improved standards compliance. When using BigQuery's updated SQL, the values of allowLargeResults and flattenResults are ignored. Queries with useLegacySql set to false will be run as if allowLargeResults is true and flattenResults is false. useQueryCache: [Optional] Whether to look for the result in the query cache. The query cache is a best-effort cache that will be flushed whenever tables in the query are modified. The default value is true. """ defaultDataset = _messages.MessageField('DatasetReference', 1) dryRun = _messages.BooleanField(2) kind = _messages.StringField(3, default=u'bigquery#queryRequest') maxResults = _messages.IntegerField(4, variant=_messages.Variant.UINT32) preserveNulls = _messages.BooleanField(5) query = _messages.StringField(6) timeoutMs = _messages.IntegerField(7, variant=_messages.Variant.UINT32) useLegacySql = _messages.BooleanField(8) useQueryCache = _messages.BooleanField(9, default=True) class QueryResponse(_messages.Message): """A QueryResponse object. Fields: cacheHit: Whether the query result was fetched from the query cache. errors: [Output-only] All errors and warnings encountered during the running of the job. Errors here do not necessarily mean that the job has completed or was unsuccessful. jobComplete: Whether the query has completed or not. If rows or totalRows are present, this will always be true. If this is false, totalRows will not be available. jobReference: Reference to the Job that was created to run the query. This field will be present even if the original request timed out, in which case GetQueryResults can be used to read the results once the query has completed. Since this API only returns the first page of results, subsequent pages can be fetched via the same mechanism (GetQueryResults). kind: The resource type. pageToken: A token used for paging results. rows: An object with as many results as can be contained within the maximum permitted reply size. To get any additional rows, you can call GetQueryResults and specify the jobReference returned above. schema: The schema of the results. Present only when the query completes successfully. totalBytesProcessed: The total number of bytes processed for this query. If this query was a dry run, this is the number of bytes that would be processed if the query were run. totalRows: The total number of rows in the complete query result set, which can be more than the number of rows in this single page of results. """ cacheHit = _messages.BooleanField(1) errors = _messages.MessageField('ErrorProto', 2, repeated=True) jobComplete = _messages.BooleanField(3) jobReference = _messages.MessageField('JobReference', 4) kind = _messages.StringField(5, default=u'bigquery#queryResponse') pageToken = _messages.StringField(6) rows = _messages.MessageField('TableRow', 7, repeated=True) schema = _messages.MessageField('TableSchema', 8) totalBytesProcessed = _messages.IntegerField(9) totalRows = _messages.IntegerField(10, variant=_messages.Variant.UINT64) class StandardQueryParameters(_messages.Message): """Query parameters accepted by all methods. Enums: AltValueValuesEnum: Data format for the response. Fields: alt: Data format for the response. fields: Selector specifying which fields to include in a partial response. key: API key. Your API key identifies your project and provides you with API access, quota, and reports. Required unless you provide an OAuth 2.0 token. oauth_token: OAuth 2.0 token for the current user. prettyPrint: Returns response with indentations and line breaks. quotaUser: Available to use for quota purposes for server-side applications. Can be any arbitrary string assigned to a user, but should not exceed 40 characters. Overrides userIp if both are provided. trace: A tracing token of the form "token:" to include in api requests. userIp: IP address of the site where the request originates. Use this if you want to enforce per-user limits. """ class AltValueValuesEnum(_messages.Enum): """Data format for the response. Values: csv: Responses with Content-Type of text/csv json: Responses with Content-Type of application/json """ csv = 0 json = 1 alt = _messages.EnumField('AltValueValuesEnum', 1, default=u'json') fields = _messages.StringField(2) key = _messages.StringField(3) oauth_token = _messages.StringField(4) prettyPrint = _messages.BooleanField(5, default=True) quotaUser = _messages.StringField(6) trace = _messages.StringField(7) userIp = _messages.StringField(8) class Streamingbuffer(_messages.Message): """A Streamingbuffer object. Fields: estimatedBytes: [Output-only] A lower-bound estimate of the number of bytes currently in the streaming buffer. estimatedRows: [Output-only] A lower-bound estimate of the number of rows currently in the streaming buffer. oldestEntryTime: [Output-only] Contains the timestamp of the oldest entry in the streaming buffer, in milliseconds since the epoch, if the streaming buffer is available. """ estimatedBytes = _messages.IntegerField(1, variant=_messages.Variant.UINT64) estimatedRows = _messages.IntegerField(2, variant=_messages.Variant.UINT64) oldestEntryTime = _messages.IntegerField(3, variant=_messages.Variant.UINT64) class Table(_messages.Message): """A Table object. Fields: creationTime: [Output-only] The time when this table was created, in milliseconds since the epoch. description: [Optional] A user-friendly description of this table. etag: [Output-only] A hash of this resource. expirationTime: [Optional] The time when this table expires, in milliseconds since the epoch. If not present, the table will persist indefinitely. Expired tables will be deleted and their storage reclaimed. externalDataConfiguration: [Optional] Describes the data format, location, and other properties of a table stored outside of BigQuery. By defining these properties, the data source can then be queried as if it were a standard BigQuery table. friendlyName: [Optional] A descriptive name for this table. id: [Output-only] An opaque ID uniquely identifying the table. kind: [Output-only] The type of the resource. lastModifiedTime: [Output-only] The time when this table was last modified, in milliseconds since the epoch. location: [Output-only] The geographic location where the table resides. This value is inherited from the dataset. numBytes: [Output-only] The size of this table in bytes, excluding any data in the streaming buffer. numRows: [Output-only] The number of rows of data in this table, excluding any data in the streaming buffer. partitionConfigurations: [Experimental] List of partition configurations for this table. Currently only one configuration can be specified and it can only be an interval partition with type daily. schema: [Optional] Describes the schema of this table. selfLink: [Output-only] A URL that can be used to access this resource again. streamingBuffer: [Output-only] Contains information regarding this table's streaming buffer, if one is present. This field will be absent if the table is not being streamed to or if there is no data in the streaming buffer. tableReference: [Required] Reference describing the ID of this table. type: [Output-only] Describes the table type. The following values are supported: TABLE: A normal BigQuery table. VIEW: A virtual table defined by a SQL query. EXTERNAL: A table that references data stored in an external storage system, such as Google Cloud Storage. The default value is TABLE. view: [Optional] The view definition. """ creationTime = _messages.IntegerField(1) description = _messages.StringField(2) etag = _messages.StringField(3) expirationTime = _messages.IntegerField(4) externalDataConfiguration = _messages.MessageField('ExternalDataConfiguration', 5) friendlyName = _messages.StringField(6) id = _messages.StringField(7) kind = _messages.StringField(8, default=u'bigquery#table') lastModifiedTime = _messages.IntegerField(9, variant=_messages.Variant.UINT64) location = _messages.StringField(10) numBytes = _messages.IntegerField(11) numRows = _messages.IntegerField(12, variant=_messages.Variant.UINT64) partitionConfigurations = _messages.MessageField('TablePartitionConfiguration', 13, repeated=True) schema = _messages.MessageField('TableSchema', 14) selfLink = _messages.StringField(15) streamingBuffer = _messages.MessageField('Streamingbuffer', 16) tableReference = _messages.MessageField('TableReference', 17) type = _messages.StringField(18) view = _messages.MessageField('ViewDefinition', 19) class TableCell(_messages.Message): """A TableCell object. Fields: v: A extra_types.JsonValue attribute. """ v = _messages.MessageField('extra_types.JsonValue', 1) class TableDataInsertAllRequest(_messages.Message): """A TableDataInsertAllRequest object. Messages: RowsValueListEntry: A RowsValueListEntry object. Fields: ignoreUnknownValues: [Optional] Accept rows that contain values that do not match the schema. The unknown values are ignored. Default is false, which treats unknown values as errors. kind: The resource type of the response. rows: The rows to insert. skipInvalidRows: [Optional] Insert all valid rows of a request, even if invalid rows exist. The default value is false, which causes the entire request to fail if any invalid rows exist. templateSuffix: [Experimental] If specified, treats the destination table as a base template, and inserts the rows into an instance table named "{destination}{templateSuffix}". BigQuery will manage creation of the instance table, using the schema of the base template table. See https://cloud.google.com/bigquery/streaming-data-into-bigquery#template- tables for considerations when working with templates tables. """ class RowsValueListEntry(_messages.Message): """A RowsValueListEntry object. Fields: insertId: [Optional] A unique ID for each row. BigQuery uses this property to detect duplicate insertion requests on a best-effort basis. json: [Required] A JSON object that contains a row of data. The object's properties and values must match the destination table's schema. """ insertId = _messages.StringField(1) json = _messages.MessageField('JsonObject', 2) ignoreUnknownValues = _messages.BooleanField(1) kind = _messages.StringField(2, default=u'bigquery#tableDataInsertAllRequest') rows = _messages.MessageField('RowsValueListEntry', 3, repeated=True) skipInvalidRows = _messages.BooleanField(4) templateSuffix = _messages.StringField(5) class TableDataInsertAllResponse(_messages.Message): """A TableDataInsertAllResponse object. Messages: InsertErrorsValueListEntry: A InsertErrorsValueListEntry object. Fields: insertErrors: An array of errors for rows that were not inserted. kind: The resource type of the response. """ class InsertErrorsValueListEntry(_messages.Message): """A InsertErrorsValueListEntry object. Fields: errors: Error information for the row indicated by the index property. index: The index of the row that error applies to. """ errors = _messages.MessageField('ErrorProto', 1, repeated=True) index = _messages.IntegerField(2, variant=_messages.Variant.UINT32) insertErrors = _messages.MessageField('InsertErrorsValueListEntry', 1, repeated=True) kind = _messages.StringField(2, default=u'bigquery#tableDataInsertAllResponse') class TableDataList(_messages.Message): """A TableDataList object. Fields: etag: A hash of this page of results. kind: The resource type of the response. pageToken: A token used for paging results. Providing this token instead of the startIndex parameter can help you retrieve stable results when an underlying table is changing. rows: Rows of results. totalRows: The total number of rows in the complete table. """ etag = _messages.StringField(1) kind = _messages.StringField(2, default=u'bigquery#tableDataList') pageToken = _messages.StringField(3) rows = _messages.MessageField('TableRow', 4, repeated=True) totalRows = _messages.IntegerField(5) class TableFieldSchema(_messages.Message): """A TableFieldSchema object. Fields: description: [Optional] The field description. The maximum length is 16K characters. fields: [Optional] Describes the nested schema fields if the type property is set to RECORD. mode: [Optional] The field mode. Possible values include NULLABLE, REQUIRED and REPEATED. The default value is NULLABLE. name: [Required] The field name. The name must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_), and must start with a letter or underscore. The maximum length is 128 characters. type: [Required] The field data type. Possible values include STRING, BYTES, INTEGER, FLOAT, BOOLEAN, TIMESTAMP or RECORD (where RECORD indicates that the field contains a nested schema). """ description = _messages.StringField(1) fields = _messages.MessageField('TableFieldSchema', 2, repeated=True) mode = _messages.StringField(3) name = _messages.StringField(4) type = _messages.StringField(5) class TableList(_messages.Message): """A TableList object. Messages: TablesValueListEntry: A TablesValueListEntry object. Fields: etag: A hash of this page of results. kind: The type of list. nextPageToken: A token to request the next page of results. tables: Tables in the requested dataset. totalItems: The total number of tables in the dataset. """ class TablesValueListEntry(_messages.Message): """A TablesValueListEntry object. Fields: friendlyName: The user-friendly name for this table. id: An opaque ID of the table kind: The resource type. tableReference: A reference uniquely identifying the table. type: The type of table. Possible values are: TABLE, VIEW. """ friendlyName = _messages.StringField(1) id = _messages.StringField(2) kind = _messages.StringField(3, default=u'bigquery#table') tableReference = _messages.MessageField('TableReference', 4) type = _messages.StringField(5) etag = _messages.StringField(1) kind = _messages.StringField(2, default=u'bigquery#tableList') nextPageToken = _messages.StringField(3) tables = _messages.MessageField('TablesValueListEntry', 4, repeated=True) totalItems = _messages.IntegerField(5, variant=_messages.Variant.INT32) class TablePartitionConfiguration(_messages.Message): """[Required] A partition configuration. Only one type of partition should be configured. Fields: interval: [Pick one] Configures an interval partition. """ interval = _messages.MessageField('IntervalPartitionConfiguration', 1) class TableReference(_messages.Message): """A TableReference object. Fields: datasetId: [Required] The ID of the dataset containing this table. projectId: [Required] The ID of the project containing this table. tableId: [Required] The ID of the table. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). The maximum length is 1,024 characters. """ datasetId = _messages.StringField(1) projectId = _messages.StringField(2) tableId = _messages.StringField(3) class TableRow(_messages.Message): """A TableRow object. Fields: f: Represents a single row in the result set, consisting of one or more fields. """ f = _messages.MessageField('TableCell', 1, repeated=True) class TableSchema(_messages.Message): """A TableSchema object. Fields: fields: Describes the fields in a table. """ fields = _messages.MessageField('TableFieldSchema', 1, repeated=True) class UserDefinedFunctionResource(_messages.Message): """A UserDefinedFunctionResource object. Fields: inlineCode: [Pick one] An inline resource that contains code for a user- defined function (UDF). Providing a inline code resource is equivalent to providing a URI for a file containing the same code. resourceUri: [Pick one] A code resource to load from a Google Cloud Storage URI (gs://bucket/path). """ inlineCode = _messages.StringField(1) resourceUri = _messages.StringField(2) class ViewDefinition(_messages.Message): """A ViewDefinition object. Fields: query: [Required] A query that BigQuery executes when the view is referenced. userDefinedFunctionResources: [Experimental] Describes user-defined function resources used in the query. """ query = _messages.StringField(1) userDefinedFunctionResources = _messages.MessageField('UserDefinedFunctionResource', 2, repeated=True) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/dataflow/000077500000000000000000000000001272340251000262605ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/dataflow/__init__.py000066400000000000000000000005511272340251000303720ustar00rootroot00000000000000"""Common imports for generated dataflow client library.""" # pylint:disable=wildcard-import import pkgutil from apitools.base.py import * from google.cloud.dataflow.internal.clients.dataflow.dataflow_v1b3_client import * from google.cloud.dataflow.internal.clients.dataflow.dataflow_v1b3_messages import * __path__ = pkgutil.extend_path(__path__, __name__) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/dataflow/dataflow_v1b3_client.py000066400000000000000000000305331272340251000326300ustar00rootroot00000000000000"""Generated client library for dataflow version v1b3.""" # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.py import base_api from google.cloud.dataflow.internal.clients.dataflow import dataflow_v1b3_messages as messages class DataflowV1b3(base_api.BaseApiClient): """Generated client library for service dataflow version v1b3.""" MESSAGES_MODULE = messages _PACKAGE = u'dataflow' _SCOPES = [u'https://www.googleapis.com/auth/cloud-platform', u'https://www.googleapis.com/auth/userinfo.email'] _VERSION = u'v1b3' _CLIENT_ID = '1042881264118.apps.googleusercontent.com' _CLIENT_SECRET = 'x_Tw5K8nnjoRAqULM9PFAC2b' _USER_AGENT = 'x_Tw5K8nnjoRAqULM9PFAC2b' _CLIENT_CLASS_NAME = u'DataflowV1b3' _URL_VERSION = u'v1b3' _API_KEY = None def __init__(self, url='', credentials=None, get_credentials=True, http=None, model=None, log_request=False, log_response=False, credentials_args=None, default_global_params=None, additional_http_headers=None): """Create a new dataflow handle.""" url = url or u'https://dataflow.googleapis.com/' super(DataflowV1b3, self).__init__( url, credentials=credentials, get_credentials=get_credentials, http=http, model=model, log_request=log_request, log_response=log_response, credentials_args=credentials_args, default_global_params=default_global_params, additional_http_headers=additional_http_headers) self.projects_jobs_messages = self.ProjectsJobsMessagesService(self) self.projects_jobs_workItems = self.ProjectsJobsWorkItemsService(self) self.projects_jobs = self.ProjectsJobsService(self) self.projects = self.ProjectsService(self) class ProjectsJobsMessagesService(base_api.BaseApiService): """Service class for the projects_jobs_messages resource.""" _NAME = u'projects_jobs_messages' def __init__(self, client): super(DataflowV1b3.ProjectsJobsMessagesService, self).__init__(client) self._method_configs = { 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'dataflow.projects.jobs.messages.list', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[u'endTime', u'minimumImportance', u'pageSize', u'pageToken', u'startTime'], relative_path=u'v1b3/projects/{projectId}/jobs/{jobId}/messages', request_field='', request_type_name=u'DataflowProjectsJobsMessagesListRequest', response_type_name=u'ListJobMessagesResponse', supports_download=False, ), } self._upload_configs = { } def List(self, request, global_params=None): """Request the job status. Args: request: (DataflowProjectsJobsMessagesListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListJobMessagesResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) class ProjectsJobsWorkItemsService(base_api.BaseApiService): """Service class for the projects_jobs_workItems resource.""" _NAME = u'projects_jobs_workItems' def __init__(self, client): super(DataflowV1b3.ProjectsJobsWorkItemsService, self).__init__(client) self._method_configs = { 'Lease': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'dataflow.projects.jobs.workItems.lease', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[], relative_path=u'v1b3/projects/{projectId}/jobs/{jobId}/workItems:lease', request_field=u'leaseWorkItemRequest', request_type_name=u'DataflowProjectsJobsWorkItemsLeaseRequest', response_type_name=u'LeaseWorkItemResponse', supports_download=False, ), 'ReportStatus': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'dataflow.projects.jobs.workItems.reportStatus', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[], relative_path=u'v1b3/projects/{projectId}/jobs/{jobId}/workItems:reportStatus', request_field=u'reportWorkItemStatusRequest', request_type_name=u'DataflowProjectsJobsWorkItemsReportStatusRequest', response_type_name=u'ReportWorkItemStatusResponse', supports_download=False, ), } self._upload_configs = { } def Lease(self, request, global_params=None): """Leases a dataflow WorkItem to run. Args: request: (DataflowProjectsJobsWorkItemsLeaseRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (LeaseWorkItemResponse) The response message. """ config = self.GetMethodConfig('Lease') return self._RunMethod( config, request, global_params=global_params) def ReportStatus(self, request, global_params=None): """Reports the status of dataflow WorkItems leased by a worker. Args: request: (DataflowProjectsJobsWorkItemsReportStatusRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ReportWorkItemStatusResponse) The response message. """ config = self.GetMethodConfig('ReportStatus') return self._RunMethod( config, request, global_params=global_params) class ProjectsJobsService(base_api.BaseApiService): """Service class for the projects_jobs resource.""" _NAME = u'projects_jobs' def __init__(self, client): super(DataflowV1b3.ProjectsJobsService, self).__init__(client) self._method_configs = { 'Create': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'dataflow.projects.jobs.create', ordered_params=[u'projectId'], path_params=[u'projectId'], query_params=[u'replaceJobId', u'view'], relative_path=u'v1b3/projects/{projectId}/jobs', request_field=u'job', request_type_name=u'DataflowProjectsJobsCreateRequest', response_type_name=u'Job', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'dataflow.projects.jobs.get', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[u'view'], relative_path=u'v1b3/projects/{projectId}/jobs/{jobId}', request_field='', request_type_name=u'DataflowProjectsJobsGetRequest', response_type_name=u'Job', supports_download=False, ), 'GetMetrics': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'dataflow.projects.jobs.getMetrics', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[u'startTime'], relative_path=u'v1b3/projects/{projectId}/jobs/{jobId}/metrics', request_field='', request_type_name=u'DataflowProjectsJobsGetMetricsRequest', response_type_name=u'JobMetrics', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'dataflow.projects.jobs.list', ordered_params=[u'projectId'], path_params=[u'projectId'], query_params=[u'filter', u'pageSize', u'pageToken', u'view'], relative_path=u'v1b3/projects/{projectId}/jobs', request_field='', request_type_name=u'DataflowProjectsJobsListRequest', response_type_name=u'ListJobsResponse', supports_download=False, ), 'Update': base_api.ApiMethodInfo( http_method=u'PUT', method_id=u'dataflow.projects.jobs.update', ordered_params=[u'projectId', u'jobId'], path_params=[u'jobId', u'projectId'], query_params=[], relative_path=u'v1b3/projects/{projectId}/jobs/{jobId}', request_field=u'job', request_type_name=u'DataflowProjectsJobsUpdateRequest', response_type_name=u'Job', supports_download=False, ), } self._upload_configs = { } def Create(self, request, global_params=None): """Creates a dataflow job. Args: request: (DataflowProjectsJobsCreateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Create') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None): """Gets the state of the specified dataflow job. Args: request: (DataflowProjectsJobsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) def GetMetrics(self, request, global_params=None): """Request the job status. Args: request: (DataflowProjectsJobsGetMetricsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (JobMetrics) The response message. """ config = self.GetMethodConfig('GetMetrics') return self._RunMethod( config, request, global_params=global_params) def List(self, request, global_params=None): """List the jobs of a project. Args: request: (DataflowProjectsJobsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListJobsResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Update(self, request, global_params=None): """Updates the state of an existing dataflow job. Args: request: (DataflowProjectsJobsUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) class ProjectsService(base_api.BaseApiService): """Service class for the projects resource.""" _NAME = u'projects' def __init__(self, client): super(DataflowV1b3.ProjectsService, self).__init__(client) self._method_configs = { 'WorkerMessages': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'dataflow.projects.workerMessages', ordered_params=[u'projectId'], path_params=[u'projectId'], query_params=[], relative_path=u'v1b3/projects/{projectId}/WorkerMessages', request_field=u'sendWorkerMessagesRequest', request_type_name=u'DataflowProjectsWorkerMessagesRequest', response_type_name=u'SendWorkerMessagesResponse', supports_download=False, ), } self._upload_configs = { } def WorkerMessages(self, request, global_params=None): """Send a worker_message to the service. Args: request: (DataflowProjectsWorkerMessagesRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (SendWorkerMessagesResponse) The response message. """ config = self.GetMethodConfig('WorkerMessages') return self._RunMethod( config, request, global_params=global_params) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/dataflow/dataflow_v1b3_messages.py000066400000000000000000003554531272340251000331740ustar00rootroot00000000000000"""Generated message classes for dataflow version v1b3. Develops and executes data processing patterns like ETL, batch computation, and continuous computation. """ # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.protorpclite import messages as _messages from apitools.base.py import encoding package = 'dataflow' class ApproximateProgress(_messages.Message): """Obsolete in favor of ApproximateReportedProgress and ApproximateSplitRequest. Fields: percentComplete: Obsolete. position: Obsolete. remainingTime: Obsolete. """ percentComplete = _messages.FloatField(1, variant=_messages.Variant.FLOAT) position = _messages.MessageField('Position', 2) remainingTime = _messages.StringField(3) class ApproximateReportedProgress(_messages.Message): """A progress measurement of a WorkItem by a worker. Fields: consumedParallelism: Total amount of parallelism in the portion of input of this work item that has already been consumed. In the first two examples above (see remaining_parallelism), the value should be 30 or 3 respectively. The sum of remaining_parallelism and consumed_parallelism should equal the total amount of parallelism in this work item. If specified, must be finite. fractionConsumed: Completion as fraction of the input consumed, from 0.0 (beginning, nothing consumed), to 1.0 (end of the input, entire input consumed). position: A Position within the work to represent a progress. remainingParallelism: Total amount of parallelism in the input of this WorkItem that has not been consumed yet (i.e. can be delegated to a new WorkItem via dynamic splitting). "Amount of parallelism" refers to how many non-empty parts of the input can be read in parallel. This does not necessarily equal number of records. An input that can be read in parallel down to the individual records is called "perfectly splittable". An example of non-perfectly parallelizable input is a block-compressed file format where a block of records has to be read as a whole, but different blocks can be read in parallel. Examples: * If we have read 30 records out of 50 in a perfectly splittable 50-record input, this value should be 20. * If we are reading through block 3 in a block-compressed file consisting of 5 blocks, this value should be 2 (since blocks 4 and 5 can be processed in parallel by new work items via dynamic splitting). * If we are reading through the last block in a block-compressed file, or reading or processing the last record in a perfectly splittable input, this value should be 0, because the remainder of the work item cannot be further split. """ consumedParallelism = _messages.MessageField('ReportedParallelism', 1) fractionConsumed = _messages.FloatField(2) position = _messages.MessageField('Position', 3) remainingParallelism = _messages.MessageField('ReportedParallelism', 4) class ApproximateSplitRequest(_messages.Message): """A suggestion by the service to the worker to dynamically split the WorkItem. Fields: fractionConsumed: A fraction at which to split the work item, from 0.0 (beginning of the input) to 1.0 (end of the input). position: A Position at which to split the work item. """ fractionConsumed = _messages.FloatField(1) position = _messages.MessageField('Position', 2) class AutoscalingSettings(_messages.Message): """Settings for WorkerPool autoscaling. Enums: AlgorithmValueValuesEnum: The algorithm to use for autoscaling. Fields: algorithm: The algorithm to use for autoscaling. maxNumWorkers: The maximum number of workers to cap scaling at. """ class AlgorithmValueValuesEnum(_messages.Enum): """The algorithm to use for autoscaling. Values: AUTOSCALING_ALGORITHM_UNKNOWN: AUTOSCALING_ALGORITHM_NONE: AUTOSCALING_ALGORITHM_BASIC: """ AUTOSCALING_ALGORITHM_UNKNOWN = 0 AUTOSCALING_ALGORITHM_NONE = 1 AUTOSCALING_ALGORITHM_BASIC = 2 algorithm = _messages.EnumField('AlgorithmValueValuesEnum', 1) maxNumWorkers = _messages.IntegerField(2, variant=_messages.Variant.INT32) class ComputationTopology(_messages.Message): """All configuration data for a particular Computation. Fields: computationId: The ID of the computation. inputs: The inputs to the computation. keyRanges: The key ranges processed by the computation. outputs: The outputs from the computation. stateFamilies: The state family values. systemStageName: The system stage name. userStageName: The user stage name. """ computationId = _messages.StringField(1) inputs = _messages.MessageField('StreamLocation', 2, repeated=True) keyRanges = _messages.MessageField('KeyRangeLocation', 3, repeated=True) outputs = _messages.MessageField('StreamLocation', 4, repeated=True) stateFamilies = _messages.MessageField('StateFamilyConfig', 5, repeated=True) systemStageName = _messages.StringField(6) userStageName = _messages.StringField(7) class ConcatPosition(_messages.Message): """A position that encapsulates an inner position and an index for the inner position. A ConcatPosition can be used by a reader of a source that encapsulates a set of other sources. Fields: index: Index of the inner source. position: Position within the inner source. """ index = _messages.IntegerField(1, variant=_messages.Variant.INT32) position = _messages.MessageField('Position', 2) class CustomSourceLocation(_messages.Message): """Identifies the location of a custom souce. Fields: stateful: Whether this source is stateful. """ stateful = _messages.BooleanField(1) class DataDiskAssignment(_messages.Message): """Data disk assignment for a given VM instance. Fields: dataDisks: Mounted data disks. The order is important a data disk's 0-based index in this list defines which persistent directory the disk is mounted to, for example the list of { "myproject-1014-104817-4c2-harness-0-disk-0" }, { "myproject-1014-104817-4c2-harness-0-disk-1" }. vmInstance: VM instance name the data disks mounted to, for example "myproject-1014-104817-4c2-harness-0". """ dataDisks = _messages.StringField(1, repeated=True) vmInstance = _messages.StringField(2) class DataflowProjectsJobsCreateRequest(_messages.Message): """A DataflowProjectsJobsCreateRequest object. Enums: ViewValueValuesEnum: Level of information requested in response. Fields: job: A Job resource to be passed as the request body. projectId: The project which owns the job. replaceJobId: DEPRECATED. This field is now on the Job message. view: Level of information requested in response. """ class ViewValueValuesEnum(_messages.Enum): """Level of information requested in response. Values: JOB_VIEW_UNKNOWN: JOB_VIEW_SUMMARY: JOB_VIEW_ALL: """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 job = _messages.MessageField('Job', 1) projectId = _messages.StringField(2, required=True) replaceJobId = _messages.StringField(3) view = _messages.EnumField('ViewValueValuesEnum', 4) class DataflowProjectsJobsGetMetricsRequest(_messages.Message): """A DataflowProjectsJobsGetMetricsRequest object. Fields: jobId: The job to get messages for. projectId: A project id. startTime: Return only metric data that has changed since this time. Default is to return all information about all metrics for the job. """ jobId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) startTime = _messages.StringField(3) class DataflowProjectsJobsGetRequest(_messages.Message): """A DataflowProjectsJobsGetRequest object. Enums: ViewValueValuesEnum: Level of information requested in response. Fields: jobId: Identifies a single job. projectId: The project which owns the job. view: Level of information requested in response. """ class ViewValueValuesEnum(_messages.Enum): """Level of information requested in response. Values: JOB_VIEW_UNKNOWN: JOB_VIEW_SUMMARY: JOB_VIEW_ALL: """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 jobId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) view = _messages.EnumField('ViewValueValuesEnum', 3) class DataflowProjectsJobsListRequest(_messages.Message): """A DataflowProjectsJobsListRequest object. Enums: FilterValueValuesEnum: The kind of filter to use. ViewValueValuesEnum: Level of information requested in response. Default is SUMMARY. Fields: filter: The kind of filter to use. pageSize: If there are many jobs, limit response to at most this many. The actual number of jobs returned will be the lesser of max_responses and an unspecified server-defined limit. pageToken: Set this to the 'next_page_token' field of a previous response to request additional results in a long list. projectId: The project which owns the jobs. view: Level of information requested in response. Default is SUMMARY. """ class FilterValueValuesEnum(_messages.Enum): """The kind of filter to use. Values: ALL: TERMINATED: ACTIVE: """ ALL = 0 TERMINATED = 1 ACTIVE = 2 class ViewValueValuesEnum(_messages.Enum): """Level of information requested in response. Default is SUMMARY. Values: JOB_VIEW_UNKNOWN: JOB_VIEW_SUMMARY: JOB_VIEW_ALL: """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 filter = _messages.EnumField('FilterValueValuesEnum', 1) pageSize = _messages.IntegerField(2, variant=_messages.Variant.INT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) view = _messages.EnumField('ViewValueValuesEnum', 5) class DataflowProjectsJobsMessagesListRequest(_messages.Message): """A DataflowProjectsJobsMessagesListRequest object. Enums: MinimumImportanceValueValuesEnum: Filter to only get messages with importance >= level Fields: endTime: Return only messages with timestamps < end_time. The default is now (i.e. return up to the latest messages available). jobId: The job to get messages about. minimumImportance: Filter to only get messages with importance >= level pageSize: If specified, determines the maximum number of messages to return. If unspecified, the service may choose an appropriate default, or may return an arbitrarily large number of results. pageToken: If supplied, this should be the value of next_page_token returned by an earlier call. This will cause the next page of results to be returned. projectId: A project id. startTime: If specified, return only messages with timestamps >= start_time. The default is the job creation time (i.e. beginning of messages). """ class MinimumImportanceValueValuesEnum(_messages.Enum): """Filter to only get messages with importance >= level Values: JOB_MESSAGE_IMPORTANCE_UNKNOWN: JOB_MESSAGE_DEBUG: JOB_MESSAGE_DETAILED: JOB_MESSAGE_BASIC: JOB_MESSAGE_WARNING: JOB_MESSAGE_ERROR: """ JOB_MESSAGE_IMPORTANCE_UNKNOWN = 0 JOB_MESSAGE_DEBUG = 1 JOB_MESSAGE_DETAILED = 2 JOB_MESSAGE_BASIC = 3 JOB_MESSAGE_WARNING = 4 JOB_MESSAGE_ERROR = 5 endTime = _messages.StringField(1) jobId = _messages.StringField(2, required=True) minimumImportance = _messages.EnumField('MinimumImportanceValueValuesEnum', 3) pageSize = _messages.IntegerField(4, variant=_messages.Variant.INT32) pageToken = _messages.StringField(5) projectId = _messages.StringField(6, required=True) startTime = _messages.StringField(7) class DataflowProjectsJobsUpdateRequest(_messages.Message): """A DataflowProjectsJobsUpdateRequest object. Fields: job: A Job resource to be passed as the request body. jobId: Identifies a single job. projectId: The project which owns the job. """ job = _messages.MessageField('Job', 1) jobId = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class DataflowProjectsJobsWorkItemsLeaseRequest(_messages.Message): """A DataflowProjectsJobsWorkItemsLeaseRequest object. Fields: jobId: Identifies the workflow job this worker belongs to. leaseWorkItemRequest: A LeaseWorkItemRequest resource to be passed as the request body. projectId: Identifies the project this worker belongs to. """ jobId = _messages.StringField(1, required=True) leaseWorkItemRequest = _messages.MessageField('LeaseWorkItemRequest', 2) projectId = _messages.StringField(3, required=True) class DataflowProjectsJobsWorkItemsReportStatusRequest(_messages.Message): """A DataflowProjectsJobsWorkItemsReportStatusRequest object. Fields: jobId: The job which the WorkItem is part of. projectId: The project which owns the WorkItem's job. reportWorkItemStatusRequest: A ReportWorkItemStatusRequest resource to be passed as the request body. """ jobId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) reportWorkItemStatusRequest = _messages.MessageField('ReportWorkItemStatusRequest', 3) class DataflowProjectsWorkerMessagesRequest(_messages.Message): """A DataflowProjectsWorkerMessagesRequest object. Fields: projectId: The project to send the WorkerMessages to. sendWorkerMessagesRequest: A SendWorkerMessagesRequest resource to be passed as the request body. """ projectId = _messages.StringField(1, required=True) sendWorkerMessagesRequest = _messages.MessageField('SendWorkerMessagesRequest', 2) class DerivedSource(_messages.Message): """Specification of one of the bundles produced as a result of splitting a Source (e.g. when executing a SourceSplitRequest, or when splitting an active task using WorkItemStatus.dynamic_source_split), relative to the source being split. Enums: DerivationModeValueValuesEnum: What source to base the produced source on (if any). Fields: derivationMode: What source to base the produced source on (if any). source: Specification of the source. """ class DerivationModeValueValuesEnum(_messages.Enum): """What source to base the produced source on (if any). Values: SOURCE_DERIVATION_MODE_UNKNOWN: SOURCE_DERIVATION_MODE_INDEPENDENT: SOURCE_DERIVATION_MODE_CHILD_OF_CURRENT: SOURCE_DERIVATION_MODE_SIBLING_OF_CURRENT: """ SOURCE_DERIVATION_MODE_UNKNOWN = 0 SOURCE_DERIVATION_MODE_INDEPENDENT = 1 SOURCE_DERIVATION_MODE_CHILD_OF_CURRENT = 2 SOURCE_DERIVATION_MODE_SIBLING_OF_CURRENT = 3 derivationMode = _messages.EnumField('DerivationModeValueValuesEnum', 1) source = _messages.MessageField('Source', 2) class Disk(_messages.Message): """Describes the data disk used by a workflow job. Fields: diskType: Disk storage type, as defined by Google Compute Engine. This must be a disk type appropriate to the project and zone in which the workers will run. If unknown or unspecified, the service will attempt to choose a reasonable default. For example, the standard persistent disk type is a resource name typically ending in "pd-standard". If SSD persistent disks are available, the resource name typically ends with "pd-ssd". The actual valid values are defined the Google Compute Engine API, not by the Dataflow API; consult the Google Compute Engine documentation for more information about determining the set of available disk types for a particular project and zone. Google Compute Engine Disk types are local to a particular project in a particular zone, and so the resource name will typically look something like this: compute.googleapis.com/projects/ /zones//diskTypes/pd-standard mountPoint: Directory in a VM where disk is mounted. sizeGb: Size of disk in GB. If zero or unspecified, the service will attempt to choose a reasonable default. """ diskType = _messages.StringField(1) mountPoint = _messages.StringField(2) sizeGb = _messages.IntegerField(3, variant=_messages.Variant.INT32) class DynamicSourceSplit(_messages.Message): """When a task splits using WorkItemStatus.dynamic_source_split, this message describes the two parts of the split relative to the description of the current task's input. Fields: primary: Primary part (continued to be processed by worker). Specified relative to the previously-current source. Becomes current. residual: Residual part (returned to the pool of work). Specified relative to the previously-current source. """ primary = _messages.MessageField('DerivedSource', 1) residual = _messages.MessageField('DerivedSource', 2) class Environment(_messages.Message): """Describes the environment in which a Dataflow Job runs. Messages: InternalExperimentsValue: Experimental settings. SdkPipelineOptionsValue: The Dataflow SDK pipeline options specified by the user. These options are passed through the service and are used to recreate the SDK pipeline options on the worker in a language agnostic and platform independent way. UserAgentValue: A description of the process that generated the request. VersionValue: A structure describing which components and their versions of the service are required in order to run the job. Fields: clusterManagerApiService: The type of cluster manager API to use. If unknown or unspecified, the service will attempt to choose a reasonable default. This should be in the form of the API service name, e.g. "compute.googleapis.com". dataset: The dataset for the current project where various workflow related tables are stored. The supported resource type is: Google BigQuery: bigquery.googleapis.com/{dataset} experiments: The list of experiments to enable. internalExperiments: Experimental settings. sdkPipelineOptions: The Dataflow SDK pipeline options specified by the user. These options are passed through the service and are used to recreate the SDK pipeline options on the worker in a language agnostic and platform independent way. tempStoragePrefix: The prefix of the resources the system should use for temporary storage. The system will append the suffix "/temp-{JOBNAME} to this resource prefix, where {JOBNAME} is the value of the job_name field. The resulting bucket and object prefix is used as the prefix of the resources used to store temporary data needed during the job execution. NOTE: This will override the value in taskrunner_settings. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} userAgent: A description of the process that generated the request. version: A structure describing which components and their versions of the service are required in order to run the job. workerPools: Worker pools. At least one "harness" worker pool must be specified in order for the job to have workers. """ @encoding.MapUnrecognizedFields('additionalProperties') class InternalExperimentsValue(_messages.Message): """Experimental settings. Messages: AdditionalProperty: An additional property for a InternalExperimentsValue object. Fields: additionalProperties: Properties of the object. Contains field @ype with type URL. """ class AdditionalProperty(_messages.Message): """An additional property for a InternalExperimentsValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class SdkPipelineOptionsValue(_messages.Message): """The Dataflow SDK pipeline options specified by the user. These options are passed through the service and are used to recreate the SDK pipeline options on the worker in a language agnostic and platform independent way. Messages: AdditionalProperty: An additional property for a SdkPipelineOptionsValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a SdkPipelineOptionsValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class UserAgentValue(_messages.Message): """A description of the process that generated the request. Messages: AdditionalProperty: An additional property for a UserAgentValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a UserAgentValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class VersionValue(_messages.Message): """A structure describing which components and their versions of the service are required in order to run the job. Messages: AdditionalProperty: An additional property for a VersionValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a VersionValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) clusterManagerApiService = _messages.StringField(1) dataset = _messages.StringField(2) experiments = _messages.StringField(3, repeated=True) internalExperiments = _messages.MessageField('InternalExperimentsValue', 4) sdkPipelineOptions = _messages.MessageField('SdkPipelineOptionsValue', 5) tempStoragePrefix = _messages.StringField(6) userAgent = _messages.MessageField('UserAgentValue', 7) version = _messages.MessageField('VersionValue', 8) workerPools = _messages.MessageField('WorkerPool', 9, repeated=True) class FlattenInstruction(_messages.Message): """An instruction that copies its inputs (zero or more) to its (single) output. Fields: inputs: Describes the inputs to the flatten instruction. """ inputs = _messages.MessageField('InstructionInput', 1, repeated=True) class InstructionInput(_messages.Message): """An input of an instruction, as a reference to an output of a producer instruction. Fields: outputNum: The output index (origin zero) within the producer. producerInstructionIndex: The index (origin zero) of the parallel instruction that produces the output to be consumed by this input. This index is relative to the list of instructions in this input's instruction's containing MapTask. """ outputNum = _messages.IntegerField(1, variant=_messages.Variant.INT32) producerInstructionIndex = _messages.IntegerField(2, variant=_messages.Variant.INT32) class InstructionOutput(_messages.Message): """An output of an instruction. Messages: CodecValue: The codec to use to encode data being written via this output. Fields: codec: The codec to use to encode data being written via this output. name: The user-provided name of this output. systemName: System-defined name of this output. Unique across the workflow. """ @encoding.MapUnrecognizedFields('additionalProperties') class CodecValue(_messages.Message): """The codec to use to encode data being written via this output. Messages: AdditionalProperty: An additional property for a CodecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a CodecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) codec = _messages.MessageField('CodecValue', 1) name = _messages.StringField(2) systemName = _messages.StringField(3) class Job(_messages.Message): """Defines a job to be run by the Dataflow service. Enums: CurrentStateValueValuesEnum: The current state of the job. Jobs are created in the JOB_STATE_STOPPED state unless otherwise specified. A job in the JOB_STATE_RUNNING state may asynchronously enter a terminal state. Once a job has reached a terminal state, no further state updates may be made. This field may be mutated by the Dataflow service; callers cannot mutate it. RequestedStateValueValuesEnum: The job's requested state. UpdateJob may be used to switch between the JOB_STATE_STOPPED and JOB_STATE_RUNNING states, by setting requested_state. UpdateJob may also be used to directly set a job's requested state to JOB_STATE_CANCELLED or JOB_STATE_DONE, irrevocably terminating the job if it has not already reached a terminal state. TypeValueValuesEnum: The type of dataflow job. Messages: TransformNameMappingValue: Map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. Fields: clientRequestId: Client's unique identifier of the job, re-used by SDK across retried attempts. If this field is set, the service will ensure its uniqueness. That is, the request to create a job will fail if the service has knowledge of a previously submitted job with the same client's id and job name. The caller may, for example, use this field to ensure idempotence of job creation across retried attempts to create a job. By default, the field is empty and, in that case, the service ignores it. createTime: Timestamp when job was initially created. Immutable, set by the Dataflow service. currentState: The current state of the job. Jobs are created in the JOB_STATE_STOPPED state unless otherwise specified. A job in the JOB_STATE_RUNNING state may asynchronously enter a terminal state. Once a job has reached a terminal state, no further state updates may be made. This field may be mutated by the Dataflow service; callers cannot mutate it. currentStateTime: The timestamp associated with the current state. environment: Environment for the job. executionInfo: Information about how the Dataflow service will actually run the job. id: The unique ID of this job. This field is set by the Dataflow service when the Job is created, and is immutable for the life of the Job. name: The user-specified Dataflow job name. Only one Job with a given name may exist in a project at any given time. If a caller attempts to create a Job with the same name as an already-existing Job, the attempt will return the existing Job. The name must match the regular expression [a-z]([-a-z0-9]{0,38}[a-z0-9])? projectId: The project which owns the job. replaceJobId: If this job is an update of an existing job, this field will be the ID of the job it replaced. When sending a CreateJobRequest, you can update a job by specifying it here. The job named here will be stopped, and its intermediate state transferred to this job. replacedByJobId: If another job is an update of this job (and thus, this job is in JOB_STATE_UPDATED), this field will contain the ID of that job. requestedState: The job's requested state. UpdateJob may be used to switch between the JOB_STATE_STOPPED and JOB_STATE_RUNNING states, by setting requested_state. UpdateJob may also be used to directly set a job's requested state to JOB_STATE_CANCELLED or JOB_STATE_DONE, irrevocably terminating the job if it has not already reached a terminal state. steps: The top-level steps that constitute the entire job. tempFiles: A set of files the system should be aware of that are used for temporary storage. These temporary files will be removed on job completion. No duplicates are allowed. No file patterns are supported. The supported files are: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} transformNameMapping: Map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. type: The type of dataflow job. """ class CurrentStateValueValuesEnum(_messages.Enum): """The current state of the job. Jobs are created in the JOB_STATE_STOPPED state unless otherwise specified. A job in the JOB_STATE_RUNNING state may asynchronously enter a terminal state. Once a job has reached a terminal state, no further state updates may be made. This field may be mutated by the Dataflow service; callers cannot mutate it. Values: JOB_STATE_UNKNOWN: JOB_STATE_STOPPED: JOB_STATE_RUNNING: JOB_STATE_DONE: JOB_STATE_FAILED: JOB_STATE_CANCELLED: JOB_STATE_UPDATED: JOB_STATE_DRAINING: JOB_STATE_DRAINED: """ JOB_STATE_UNKNOWN = 0 JOB_STATE_STOPPED = 1 JOB_STATE_RUNNING = 2 JOB_STATE_DONE = 3 JOB_STATE_FAILED = 4 JOB_STATE_CANCELLED = 5 JOB_STATE_UPDATED = 6 JOB_STATE_DRAINING = 7 JOB_STATE_DRAINED = 8 class RequestedStateValueValuesEnum(_messages.Enum): """The job's requested state. UpdateJob may be used to switch between the JOB_STATE_STOPPED and JOB_STATE_RUNNING states, by setting requested_state. UpdateJob may also be used to directly set a job's requested state to JOB_STATE_CANCELLED or JOB_STATE_DONE, irrevocably terminating the job if it has not already reached a terminal state. Values: JOB_STATE_UNKNOWN: JOB_STATE_STOPPED: JOB_STATE_RUNNING: JOB_STATE_DONE: JOB_STATE_FAILED: JOB_STATE_CANCELLED: JOB_STATE_UPDATED: JOB_STATE_DRAINING: JOB_STATE_DRAINED: """ JOB_STATE_UNKNOWN = 0 JOB_STATE_STOPPED = 1 JOB_STATE_RUNNING = 2 JOB_STATE_DONE = 3 JOB_STATE_FAILED = 4 JOB_STATE_CANCELLED = 5 JOB_STATE_UPDATED = 6 JOB_STATE_DRAINING = 7 JOB_STATE_DRAINED = 8 class TypeValueValuesEnum(_messages.Enum): """The type of dataflow job. Values: JOB_TYPE_UNKNOWN: JOB_TYPE_BATCH: JOB_TYPE_STREAMING: """ JOB_TYPE_UNKNOWN = 0 JOB_TYPE_BATCH = 1 JOB_TYPE_STREAMING = 2 @encoding.MapUnrecognizedFields('additionalProperties') class TransformNameMappingValue(_messages.Message): """Map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. Messages: AdditionalProperty: An additional property for a TransformNameMappingValue object. Fields: additionalProperties: Additional properties of type TransformNameMappingValue """ class AdditionalProperty(_messages.Message): """An additional property for a TransformNameMappingValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) clientRequestId = _messages.StringField(1) createTime = _messages.StringField(2) currentState = _messages.EnumField('CurrentStateValueValuesEnum', 3) currentStateTime = _messages.StringField(4) environment = _messages.MessageField('Environment', 5) executionInfo = _messages.MessageField('JobExecutionInfo', 6) id = _messages.StringField(7) name = _messages.StringField(8) projectId = _messages.StringField(9) replaceJobId = _messages.StringField(10) replacedByJobId = _messages.StringField(11) requestedState = _messages.EnumField('RequestedStateValueValuesEnum', 12) steps = _messages.MessageField('Step', 13, repeated=True) tempFiles = _messages.StringField(14, repeated=True) transformNameMapping = _messages.MessageField('TransformNameMappingValue', 15) type = _messages.EnumField('TypeValueValuesEnum', 16) class JobExecutionInfo(_messages.Message): """Additional information about how a Dataflow job will be executed which isn\u2019t contained in the submitted job. Messages: StagesValue: A mapping from each stage to the information about that stage. Fields: stages: A mapping from each stage to the information about that stage. """ @encoding.MapUnrecognizedFields('additionalProperties') class StagesValue(_messages.Message): """A mapping from each stage to the information about that stage. Messages: AdditionalProperty: An additional property for a StagesValue object. Fields: additionalProperties: Additional properties of type StagesValue """ class AdditionalProperty(_messages.Message): """An additional property for a StagesValue object. Fields: key: Name of the additional property. value: A JobExecutionStageInfo attribute. """ key = _messages.StringField(1) value = _messages.MessageField('JobExecutionStageInfo', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) stages = _messages.MessageField('StagesValue', 1) class JobExecutionStageInfo(_messages.Message): """Contains information about how a particular google.dataflow.v1beta3.Step will be executed. Fields: stepName: The steps associated with the execution stage. Note that stages may have several steps, and that a given step might be run by more than one stage. """ stepName = _messages.StringField(1, repeated=True) class JobMessage(_messages.Message): """A particular message pertaining to a Dataflow job. Enums: MessageImportanceValueValuesEnum: Importance level of the message. Fields: id: Identifies the message. This is automatically generated by the service; the caller should treat it as an opaque string. messageImportance: Importance level of the message. messageText: The text of the message. time: The timestamp of the message. """ class MessageImportanceValueValuesEnum(_messages.Enum): """Importance level of the message. Values: JOB_MESSAGE_IMPORTANCE_UNKNOWN: JOB_MESSAGE_DEBUG: JOB_MESSAGE_DETAILED: JOB_MESSAGE_BASIC: JOB_MESSAGE_WARNING: JOB_MESSAGE_ERROR: """ JOB_MESSAGE_IMPORTANCE_UNKNOWN = 0 JOB_MESSAGE_DEBUG = 1 JOB_MESSAGE_DETAILED = 2 JOB_MESSAGE_BASIC = 3 JOB_MESSAGE_WARNING = 4 JOB_MESSAGE_ERROR = 5 id = _messages.StringField(1) messageImportance = _messages.EnumField('MessageImportanceValueValuesEnum', 2) messageText = _messages.StringField(3) time = _messages.StringField(4) class JobMetrics(_messages.Message): """JobMetrics contains a collection of metrics descibing the detailed progress of a Dataflow job. Metrics correspond to user-defined and system- defined metrics in the job. This resource captures only the most recent values of each metric; time-series data can be queried for them (under the same metric names) from Cloud Monitoring. Fields: metricTime: Timestamp as of which metric values are current. metrics: All metrics for this job. """ metricTime = _messages.StringField(1) metrics = _messages.MessageField('MetricUpdate', 2, repeated=True) class KeyRangeDataDiskAssignment(_messages.Message): """Data disk assignment information for a specific key-range of a sharded computation. Currently we only support UTF-8 character splits to simplify encoding into JSON. Fields: dataDisk: The name of the data disk where data for this range is stored. This name is local to the Google Cloud Platform project and uniquely identifies the disk within that project, for example "myproject-1014-104817-4c2-harness-0-disk-1". end: The end (exclusive) of the key range. start: The start (inclusive) of the key range. """ dataDisk = _messages.StringField(1) end = _messages.StringField(2) start = _messages.StringField(3) class KeyRangeLocation(_messages.Message): """Location information for a specific key-range of a sharded computation. Currently we only support UTF-8 character splits to simplify encoding into JSON. Fields: dataDisk: The name of the data disk where data for this range is stored. This name is local to the Google Cloud Platform project and uniquely identifies the disk within that project, for example "myproject-1014-104817-4c2-harness-0-disk-1". deliveryEndpoint: The physical location of this range assignment to be used for streaming computation cross-worker message delivery. end: The end (exclusive) of the key range. persistentDirectory: The location of the persistent state for this range, as a persistent directory in the worker local filesystem. start: The start (inclusive) of the key range. """ dataDisk = _messages.StringField(1) deliveryEndpoint = _messages.StringField(2) end = _messages.StringField(3) persistentDirectory = _messages.StringField(4) start = _messages.StringField(5) class LeaseWorkItemRequest(_messages.Message): """Request to lease WorkItems. Fields: currentWorkerTime: The current timestamp at the worker. requestedLeaseDuration: The initial lease period. workItemTypes: Filter for WorkItem type. workerCapabilities: Worker capabilities. WorkItems might be limited to workers with specific capabilities. workerId: Identifies the worker leasing work -- typically the ID of the virtual machine running the worker. """ currentWorkerTime = _messages.StringField(1) requestedLeaseDuration = _messages.StringField(2) workItemTypes = _messages.StringField(3, repeated=True) workerCapabilities = _messages.StringField(4, repeated=True) workerId = _messages.StringField(5) class LeaseWorkItemResponse(_messages.Message): """Response to a request to lease WorkItems. Fields: workItems: A list of the leased WorkItems. """ workItems = _messages.MessageField('WorkItem', 1, repeated=True) class ListJobMessagesResponse(_messages.Message): """Response to a request to list job messages. Fields: jobMessages: Messages in ascending timestamp order. nextPageToken: The token to obtain the next page of results if there are more. """ jobMessages = _messages.MessageField('JobMessage', 1, repeated=True) nextPageToken = _messages.StringField(2) class ListJobsResponse(_messages.Message): """Response to a request to list Dataflow jobs. This may be a partial response, depending on the page size in the ListJobsRequest. Fields: jobs: A subset of the requested job information. nextPageToken: Set if there may be more results than fit in this response. """ jobs = _messages.MessageField('Job', 1, repeated=True) nextPageToken = _messages.StringField(2) class MapTask(_messages.Message): """MapTask consists of an ordered set of instructions, each of which describes one particular low-level operation for the worker to perform in order to accomplish the MapTask's WorkItem. Each instruction must appear in the list before any instructions which depends on its output. Fields: instructions: The instructions in the MapTask. stageName: System-defined name of the stage containing this MapTask. Unique across the workflow. systemName: System-defined name of this MapTask. Unique across the workflow. """ instructions = _messages.MessageField('ParallelInstruction', 1, repeated=True) stageName = _messages.StringField(2) systemName = _messages.StringField(3) class MetricStructuredName(_messages.Message): """Identifies a metric, by describing the source which generated the metric. Messages: ContextValue: Zero or more labeled fields which identify the part of the job this metric is associated with, such as the name of a step or collection. For example, built-in counters associated with steps will have context['step'] = . Counters associated with PCollections in the SDK will have context['pcollection'] = . Fields: context: Zero or more labeled fields which identify the part of the job this metric is associated with, such as the name of a step or collection. For example, built-in counters associated with steps will have context['step'] = . Counters associated with PCollections in the SDK will have context['pcollection'] = . name: Worker-defined metric name. origin: Origin (namespace) of metric name. May be blank for user-define metrics; will be "dataflow" for metrics defined by the Dataflow service or SDK. """ @encoding.MapUnrecognizedFields('additionalProperties') class ContextValue(_messages.Message): """Zero or more labeled fields which identify the part of the job this metric is associated with, such as the name of a step or collection. For example, built-in counters associated with steps will have context['step'] = . Counters associated with PCollections in the SDK will have context['pcollection'] = . Messages: AdditionalProperty: An additional property for a ContextValue object. Fields: additionalProperties: Additional properties of type ContextValue """ class AdditionalProperty(_messages.Message): """An additional property for a ContextValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) context = _messages.MessageField('ContextValue', 1) name = _messages.StringField(2) origin = _messages.StringField(3) class MetricUpdate(_messages.Message): """Describes the state of a metric. Fields: cumulative: True if this metric is reported as the total cumulative aggregate value accumulated since the worker started working on this WorkItem. By default this is false, indicating that this metric is reported as a delta that is not associated with any WorkItem. internal: Worker-computed aggregate value for internal use by the Dataflow service. kind: Metric aggregation kind. The possible metric aggregation kinds are "Sum", "Max", "Min", "Mean", "Set", "And", and "Or". The specified aggregation kind is case-insensitive. If omitted, this is not an aggregated value but instead a single metric sample value. meanCount: Worker-computed aggregate value for the "Mean" aggregation kind. This holds the count of the aggregated values and is used in combination with mean_sum above to obtain the actual mean aggregate value. The only possible value type is Long. meanSum: Worker-computed aggregate value for the "Mean" aggregation kind. This holds the sum of the aggregated values and is used in combination with mean_count below to obtain the actual mean aggregate value. The only possible value types are Long and Double. name: Name of the metric. scalar: Worker-computed aggregate value for aggregation kinds "Sum", "Max", "Min", "And", and "Or". The possible value types are Long, Double, and Boolean. set: Worker-computed aggregate value for the "Set" aggregation kind. The only possible value type is a list of Values whose type can be Long, Double, or String, according to the metric's type. All Values in the list must be of the same type. updateTime: Timestamp associated with the metric value. Optional when workers are reporting work progress; it will be filled in responses from the metrics API. """ cumulative = _messages.BooleanField(1) internal = _messages.MessageField('extra_types.JsonValue', 2) kind = _messages.StringField(3) meanCount = _messages.MessageField('extra_types.JsonValue', 4) meanSum = _messages.MessageField('extra_types.JsonValue', 5) name = _messages.MessageField('MetricStructuredName', 6) scalar = _messages.MessageField('extra_types.JsonValue', 7) set = _messages.MessageField('extra_types.JsonValue', 8) updateTime = _messages.StringField(9) class MountedDataDisk(_messages.Message): """Describes mounted data disk. Fields: dataDisk: The name of the data disk. This name is local to the Google Cloud Platform project and uniquely identifies the disk within that project, for example "myproject-1014-104817-4c2-harness-0-disk-1". """ dataDisk = _messages.StringField(1) class MultiOutputInfo(_messages.Message): """Information about an output of a multi-output DoFn. Fields: tag: The id of the tag the user code will emit to this output by; this should correspond to the tag of some SideInputInfo. """ tag = _messages.StringField(1) class Package(_messages.Message): """Packages that need to be installed in order for a worker to run the steps of the Dataflow job which will be assigned to its worker pool. This is the mechanism by which the SDK causes code to be loaded onto the workers. For example, the Dataflow Java SDK might use this to install jars containing the user's code and all of the various dependencies (libraries, data files, etc) required in order for that code to run. Fields: location: The resource to read the package from. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket} bucket.storage.googleapis.com/ name: The name of the package. """ location = _messages.StringField(1) name = _messages.StringField(2) class ParDoInstruction(_messages.Message): """An instruction that does a ParDo operation. Takes one main input and zero or more side inputs, and produces zero or more outputs. Runs user code. Messages: UserFnValue: The user function to invoke. Fields: input: The input. multiOutputInfos: Information about each of the outputs, if user_fn is a MultiDoFn. numOutputs: The number of outputs. sideInputs: Zero or more side inputs. userFn: The user function to invoke. """ @encoding.MapUnrecognizedFields('additionalProperties') class UserFnValue(_messages.Message): """The user function to invoke. Messages: AdditionalProperty: An additional property for a UserFnValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a UserFnValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) input = _messages.MessageField('InstructionInput', 1) multiOutputInfos = _messages.MessageField('MultiOutputInfo', 2, repeated=True) numOutputs = _messages.IntegerField(3, variant=_messages.Variant.INT32) sideInputs = _messages.MessageField('SideInputInfo', 4, repeated=True) userFn = _messages.MessageField('UserFnValue', 5) class ParallelInstruction(_messages.Message): """Describes a particular operation comprising a MapTask. Fields: flatten: Additional information for Flatten instructions. name: User-provided name of this operation. outputs: Describes the outputs of the instruction. parDo: Additional information for ParDo instructions. partialGroupByKey: Additional information for PartialGroupByKey instructions. read: Additional information for Read instructions. systemName: System-defined name of this operation. Unique across the workflow. write: Additional information for Write instructions. """ flatten = _messages.MessageField('FlattenInstruction', 1) name = _messages.StringField(2) outputs = _messages.MessageField('InstructionOutput', 3, repeated=True) parDo = _messages.MessageField('ParDoInstruction', 4) partialGroupByKey = _messages.MessageField('PartialGroupByKeyInstruction', 5) read = _messages.MessageField('ReadInstruction', 6) systemName = _messages.StringField(7) write = _messages.MessageField('WriteInstruction', 8) class PartialGroupByKeyInstruction(_messages.Message): """An instruction that does a partial group-by-key. One input and one output. Messages: InputElementCodecValue: The codec to use for interpreting an element in the input PTable. ValueCombiningFnValue: The value combining function to invoke. Fields: input: Describes the input to the partial group-by-key instruction. inputElementCodec: The codec to use for interpreting an element in the input PTable. sideInputs: Zero or more side inputs. valueCombiningFn: The value combining function to invoke. """ @encoding.MapUnrecognizedFields('additionalProperties') class InputElementCodecValue(_messages.Message): """The codec to use for interpreting an element in the input PTable. Messages: AdditionalProperty: An additional property for a InputElementCodecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a InputElementCodecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class ValueCombiningFnValue(_messages.Message): """The value combining function to invoke. Messages: AdditionalProperty: An additional property for a ValueCombiningFnValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a ValueCombiningFnValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) input = _messages.MessageField('InstructionInput', 1) inputElementCodec = _messages.MessageField('InputElementCodecValue', 2) sideInputs = _messages.MessageField('SideInputInfo', 3, repeated=True) valueCombiningFn = _messages.MessageField('ValueCombiningFnValue', 4) class Position(_messages.Message): """Position defines a position within a collection of data. The value can be either the end position, a key (used with ordered collections), a byte offset, or a record index. Fields: byteOffset: Position is a byte offset. concatPosition: CloudPosition is a concat position. end: Position is past all other positions. Also useful for the end position of an unbounded range. key: Position is a string key, ordered lexicographically. recordIndex: Position is a record index. shufflePosition: CloudPosition is a base64 encoded BatchShufflePosition (with FIXED sharding). """ byteOffset = _messages.IntegerField(1) concatPosition = _messages.MessageField('ConcatPosition', 2) end = _messages.BooleanField(3) key = _messages.StringField(4) recordIndex = _messages.IntegerField(5) shufflePosition = _messages.StringField(6) class PubsubLocation(_messages.Message): """Identifies a pubsub location to use for transferring data into or out of a streaming Dataflow job. Fields: dropLateData: Indicates whether the pipeline allows late-arriving data. idLabel: If set, contains a pubsub label from which to extract record ids. If left empty, record deduplication will be strictly best effort. subscription: A pubsub subscription, in the form of "pubsub.googleapis.com/subscriptions/ /" timestampLabel: If set, contains a pubsub label from which to extract record timestamps. If left empty, record timestamps will be generated upon arrival. topic: A pubsub topic, in the form of "pubsub.googleapis.com/topics/ /" trackingSubscription: If set, specifies the pubsub subscription that will be used for tracking custom time timestamps for watermark estimation. """ dropLateData = _messages.BooleanField(1) idLabel = _messages.StringField(2) subscription = _messages.StringField(3) timestampLabel = _messages.StringField(4) topic = _messages.StringField(5) trackingSubscription = _messages.StringField(6) class ReadInstruction(_messages.Message): """An instruction that reads records. Takes no inputs, produces one output. Fields: source: The source to read from. """ source = _messages.MessageField('Source', 1) class ReportWorkItemStatusRequest(_messages.Message): """Request to report the status of WorkItems. Fields: currentWorkerTime: The current timestamp at the worker. workItemStatuses: The order is unimportant, except that the order of the WorkItemServiceState messages in the ReportWorkItemStatusResponse corresponds to the order of WorkItemStatus messages here. workerId: The ID of the worker reporting the WorkItem status. If this does not match the ID of the worker which the Dataflow service believes currently has the lease on the WorkItem, the report will be dropped (with an error response). """ currentWorkerTime = _messages.StringField(1) workItemStatuses = _messages.MessageField('WorkItemStatus', 2, repeated=True) workerId = _messages.StringField(3) class ReportWorkItemStatusResponse(_messages.Message): """Response from a request to report the status of WorkItems. Fields: workItemServiceStates: A set of messages indicating the service-side state for each WorkItem whose status was reported, in the same order as the WorkItemStatus messages in the ReportWorkItemStatusRequest which resulting in this response. """ workItemServiceStates = _messages.MessageField('WorkItemServiceState', 1, repeated=True) class ReportedParallelism(_messages.Message): """Represents the level of parallelism in a WorkItem's input, reported by the worker. Fields: isInfinite: Specifies whether the parallelism is infinite. If true, "value" is ignored. Infinite parallelism means the service will assume that the work item can always be split into more non-empty work items by dynamic splitting. This is a work-around for lack of support for infinity by the current JSON-based Java RPC stack. value: Specifies the level of parallelism in case it is finite. """ isInfinite = _messages.BooleanField(1) value = _messages.FloatField(2) class SendWorkerMessagesRequest(_messages.Message): """A request for sending worker messages to the service. Fields: workerMessages: The WorkerMessages to send. """ workerMessages = _messages.MessageField('WorkerMessage', 1, repeated=True) class SendWorkerMessagesResponse(_messages.Message): """The response to the worker messages. Fields: workerMessageResponses: The servers response to the worker messages. """ workerMessageResponses = _messages.MessageField('WorkerMessageResponse', 1, repeated=True) class SeqMapTask(_messages.Message): """Describes a particular function to invoke. Messages: UserFnValue: The user function to invoke. Fields: inputs: Information about each of the inputs. name: The user-provided name of the SeqDo operation. outputInfos: Information about each of the outputs. stageName: System-defined name of the stage containing the SeqDo operation. Unique across the workflow. systemName: System-defined name of the SeqDo operation. Unique across the workflow. userFn: The user function to invoke. """ @encoding.MapUnrecognizedFields('additionalProperties') class UserFnValue(_messages.Message): """The user function to invoke. Messages: AdditionalProperty: An additional property for a UserFnValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a UserFnValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) inputs = _messages.MessageField('SideInputInfo', 1, repeated=True) name = _messages.StringField(2) outputInfos = _messages.MessageField('SeqMapTaskOutputInfo', 3, repeated=True) stageName = _messages.StringField(4) systemName = _messages.StringField(5) userFn = _messages.MessageField('UserFnValue', 6) class SeqMapTaskOutputInfo(_messages.Message): """Information about an output of a SeqMapTask. Fields: sink: The sink to write the output value to. tag: The id of the TupleTag the user code will tag the output value by. """ sink = _messages.MessageField('Sink', 1) tag = _messages.StringField(2) class ShellTask(_messages.Message): """A task which consists of a shell command for the worker to execute. Fields: command: The shell command to run. exitCode: Exit code for the task. """ command = _messages.StringField(1) exitCode = _messages.IntegerField(2, variant=_messages.Variant.INT32) class SideInputInfo(_messages.Message): """Information about a side input of a DoFn or an input of a SeqDoFn. Messages: KindValue: How to interpret the source element(s) as a side input value. Fields: kind: How to interpret the source element(s) as a side input value. sources: The source(s) to read element(s) from to get the value of this side input. If more than one source, then the elements are taken from the sources, in the specified order if order matters. At least one source is required. tag: The id of the tag the user code will access this side input by; this should correspond to the tag of some MultiOutputInfo. """ @encoding.MapUnrecognizedFields('additionalProperties') class KindValue(_messages.Message): """How to interpret the source element(s) as a side input value. Messages: AdditionalProperty: An additional property for a KindValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a KindValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) kind = _messages.MessageField('KindValue', 1) sources = _messages.MessageField('Source', 2, repeated=True) tag = _messages.StringField(3) class Sink(_messages.Message): """A sink that records can be encoded and written to. Messages: CodecValue: The codec to use to encode data written to the sink. SpecValue: The sink to write to, plus its parameters. Fields: codec: The codec to use to encode data written to the sink. spec: The sink to write to, plus its parameters. """ @encoding.MapUnrecognizedFields('additionalProperties') class CodecValue(_messages.Message): """The codec to use to encode data written to the sink. Messages: AdditionalProperty: An additional property for a CodecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a CodecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class SpecValue(_messages.Message): """The sink to write to, plus its parameters. Messages: AdditionalProperty: An additional property for a SpecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a SpecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) codec = _messages.MessageField('CodecValue', 1) spec = _messages.MessageField('SpecValue', 2) class Source(_messages.Message): """A source that records can be read and decoded from. Messages: BaseSpecsValueListEntry: A BaseSpecsValueListEntry object. CodecValue: The codec to use to decode data read from the source. SpecValue: The source to read from, plus its parameters. Fields: baseSpecs: While splitting, sources may specify the produced bundles as differences against another source, in order to save backend-side memory and allow bigger jobs. For details, see SourceSplitRequest. To support this use case, the full set of parameters of the source is logically obtained by taking the latest explicitly specified value of each parameter in the order: base_specs (later items win), spec (overrides anything in base_specs). codec: The codec to use to decode data read from the source. doesNotNeedSplitting: Setting this value to true hints to the framework that the source doesn't need splitting, and using SourceSplitRequest on it would yield SOURCE_SPLIT_OUTCOME_USE_CURRENT. E.g. a file splitter may set this to true when splitting a single file into a set of byte ranges of appropriate size, and set this to false when splitting a filepattern into individual files. However, for efficiency, a file splitter may decide to produce file subranges directly from the filepattern to avoid a splitting round-trip. See SourceSplitRequest for an overview of the splitting process. This field is meaningful only in the Source objects populated by the user (e.g. when filling in a DerivedSource). Source objects supplied by the framework to the user don't have this field populated. metadata: Optionally, metadata for this source can be supplied right away, avoiding a SourceGetMetadataOperation roundtrip (see SourceOperationRequest). This field is meaningful only in the Source objects populated by the user (e.g. when filling in a DerivedSource). Source objects supplied by the framework to the user don't have this field populated. spec: The source to read from, plus its parameters. """ @encoding.MapUnrecognizedFields('additionalProperties') class BaseSpecsValueListEntry(_messages.Message): """A BaseSpecsValueListEntry object. Messages: AdditionalProperty: An additional property for a BaseSpecsValueListEntry object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a BaseSpecsValueListEntry object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class CodecValue(_messages.Message): """The codec to use to decode data read from the source. Messages: AdditionalProperty: An additional property for a CodecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a CodecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class SpecValue(_messages.Message): """The source to read from, plus its parameters. Messages: AdditionalProperty: An additional property for a SpecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a SpecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) baseSpecs = _messages.MessageField('BaseSpecsValueListEntry', 1, repeated=True) codec = _messages.MessageField('CodecValue', 2) doesNotNeedSplitting = _messages.BooleanField(3) metadata = _messages.MessageField('SourceMetadata', 4) spec = _messages.MessageField('SpecValue', 5) class SourceFork(_messages.Message): """DEPRECATED in favor of DynamicSourceSplit. Fields: primary: DEPRECATED primarySource: DEPRECATED residual: DEPRECATED residualSource: DEPRECATED """ primary = _messages.MessageField('SourceSplitShard', 1) primarySource = _messages.MessageField('DerivedSource', 2) residual = _messages.MessageField('SourceSplitShard', 3) residualSource = _messages.MessageField('DerivedSource', 4) class SourceGetMetadataRequest(_messages.Message): """A request to compute the SourceMetadata of a Source. Fields: source: Specification of the source whose metadata should be computed. """ source = _messages.MessageField('Source', 1) class SourceGetMetadataResponse(_messages.Message): """The result of a SourceGetMetadataOperation. Fields: metadata: The computed metadata. """ metadata = _messages.MessageField('SourceMetadata', 1) class SourceMetadata(_messages.Message): """Metadata about a Source useful for automatically optimizing and tuning the pipeline, etc. Fields: estimatedSizeBytes: An estimate of the total size (in bytes) of the data that would be read from this source. This estimate is in terms of external storage size, before any decompression or other processing done by the reader. infinite: Specifies that the size of this source is known to be infinite (this is a streaming source). producesSortedKeys: Whether this source is known to produce key/value pairs with the (encoded) keys in lexicographically sorted order. """ estimatedSizeBytes = _messages.IntegerField(1) infinite = _messages.BooleanField(2) producesSortedKeys = _messages.BooleanField(3) class SourceOperationRequest(_messages.Message): """A work item that represents the different operations that can be performed on a user-defined Source specification. Fields: getMetadata: Information about a request to get metadata about a source. split: Information about a request to split a source. """ getMetadata = _messages.MessageField('SourceGetMetadataRequest', 1) split = _messages.MessageField('SourceSplitRequest', 2) class SourceOperationResponse(_messages.Message): """The result of a SourceOperationRequest, specified in ReportWorkItemStatusRequest.source_operation when the work item is completed. Fields: getMetadata: A response to a request to get metadata about a source. split: A response to a request to split a source. """ getMetadata = _messages.MessageField('SourceGetMetadataResponse', 1) split = _messages.MessageField('SourceSplitResponse', 2) class SourceSplitOptions(_messages.Message): """Hints for splitting a Source into bundles (parts for parallel processing) using SourceSplitRequest. Fields: desiredBundleSizeBytes: The source should be split into a set of bundles where the estimated size of each is approximately this many bytes. desiredShardSizeBytes: DEPRECATED in favor of desired_bundle_size_bytes. """ desiredBundleSizeBytes = _messages.IntegerField(1) desiredShardSizeBytes = _messages.IntegerField(2) class SourceSplitRequest(_messages.Message): """Represents the operation to split a high-level Source specification into bundles (parts for parallel processing). At a high level, splitting of a source into bundles happens as follows: SourceSplitRequest is applied to the source. If it returns SOURCE_SPLIT_OUTCOME_USE_CURRENT, no further splitting happens and the source is used "as is". Otherwise, splitting is applied recursively to each produced DerivedSource. As an optimization, for any Source, if its does_not_need_splitting is true, the framework assumes that splitting this source would return SOURCE_SPLIT_OUTCOME_USE_CURRENT, and doesn't initiate a SourceSplitRequest. This applies both to the initial source being split and to bundles produced from it. Fields: options: Hints for tuning the splitting process. source: Specification of the source to be split. """ options = _messages.MessageField('SourceSplitOptions', 1) source = _messages.MessageField('Source', 2) class SourceSplitResponse(_messages.Message): """The response to a SourceSplitRequest. Enums: OutcomeValueValuesEnum: Indicates whether splitting happened and produced a list of bundles. If this is USE_CURRENT_SOURCE_AS_IS, the current source should be processed "as is" without splitting. "bundles" is ignored in this case. If this is SPLITTING_HAPPENED, then "bundles" contains a list of bundles into which the source was split. Fields: bundles: If outcome is SPLITTING_HAPPENED, then this is a list of bundles into which the source was split. Otherwise this field is ignored. This list can be empty, which means the source represents an empty input. outcome: Indicates whether splitting happened and produced a list of bundles. If this is USE_CURRENT_SOURCE_AS_IS, the current source should be processed "as is" without splitting. "bundles" is ignored in this case. If this is SPLITTING_HAPPENED, then "bundles" contains a list of bundles into which the source was split. shards: DEPRECATED in favor of bundles. """ class OutcomeValueValuesEnum(_messages.Enum): """Indicates whether splitting happened and produced a list of bundles. If this is USE_CURRENT_SOURCE_AS_IS, the current source should be processed "as is" without splitting. "bundles" is ignored in this case. If this is SPLITTING_HAPPENED, then "bundles" contains a list of bundles into which the source was split. Values: SOURCE_SPLIT_OUTCOME_UNKNOWN: SOURCE_SPLIT_OUTCOME_USE_CURRENT: SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED: """ SOURCE_SPLIT_OUTCOME_UNKNOWN = 0 SOURCE_SPLIT_OUTCOME_USE_CURRENT = 1 SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED = 2 bundles = _messages.MessageField('DerivedSource', 1, repeated=True) outcome = _messages.EnumField('OutcomeValueValuesEnum', 2) shards = _messages.MessageField('SourceSplitShard', 3, repeated=True) class SourceSplitShard(_messages.Message): """DEPRECATED in favor of DerivedSource. Enums: DerivationModeValueValuesEnum: DEPRECATED Fields: derivationMode: DEPRECATED source: DEPRECATED """ class DerivationModeValueValuesEnum(_messages.Enum): """DEPRECATED Values: SOURCE_DERIVATION_MODE_UNKNOWN: SOURCE_DERIVATION_MODE_INDEPENDENT: SOURCE_DERIVATION_MODE_CHILD_OF_CURRENT: SOURCE_DERIVATION_MODE_SIBLING_OF_CURRENT: """ SOURCE_DERIVATION_MODE_UNKNOWN = 0 SOURCE_DERIVATION_MODE_INDEPENDENT = 1 SOURCE_DERIVATION_MODE_CHILD_OF_CURRENT = 2 SOURCE_DERIVATION_MODE_SIBLING_OF_CURRENT = 3 derivationMode = _messages.EnumField('DerivationModeValueValuesEnum', 1) source = _messages.MessageField('Source', 2) class StandardQueryParameters(_messages.Message): """Query parameters accepted by all methods. Fields: f__xgafv: V1 error format. access_token: OAuth access token. alt: Data format for response. bearer_token: OAuth bearer token. callback: JSONP fields: Selector specifying which fields to include in a partial response. key: API key. Your API key identifies your project and provides you with API access, quota, and reports. Required unless you provide an OAuth 2.0 token. oauth_token: OAuth 2.0 token for the current user. pp: Pretty-print response. prettyPrint: Returns response with indentations and line breaks. quotaUser: Available to use for quota purposes for server-side applications. Can be any arbitrary string assigned to a user, but should not exceed 40 characters. trace: A tracing token of the form "token:" to include in api requests. uploadType: Legacy upload protocol for media (e.g. "media", "multipart"). upload_protocol: Upload protocol for media (e.g. "raw", "multipart"). """ f__xgafv = _messages.StringField(1) access_token = _messages.StringField(2) alt = _messages.StringField(3, default=u'json') bearer_token = _messages.StringField(4) callback = _messages.StringField(5) fields = _messages.StringField(6) key = _messages.StringField(7) oauth_token = _messages.StringField(8) pp = _messages.BooleanField(9, default=True) prettyPrint = _messages.BooleanField(10, default=True) quotaUser = _messages.StringField(11) trace = _messages.StringField(12) uploadType = _messages.StringField(13) upload_protocol = _messages.StringField(14) class StateFamilyConfig(_messages.Message): """State family configuration. Fields: isRead: If true, this family corresponds to a read operation. stateFamily: The state family value. """ isRead = _messages.BooleanField(1) stateFamily = _messages.StringField(2) class Status(_messages.Message): """The `Status` type defines a logical error model that is suitable for different programming environments, including REST APIs and RPC APIs. It is used by [gRPC](https://github.com/grpc). The error model is designed to be: - Simple to use and understand for most users - Flexible enough to meet unexpected needs # Overview The `Status` message contains three pieces of data: error code, error message, and error details. The error code should be an enum value of google.rpc.Code, but it may accept additional error codes if needed. The error message should be a developer-facing English message that helps developers *understand* and *resolve* the error. If a localized user-facing error message is needed, put the localized message in the error details or localize it in the client. The optional error details may contain arbitrary information about the error. There is a predefined set of error detail types in the package `google.rpc` which can be used for common error conditions. # Language mapping The `Status` message is the logical representation of the error model, but it is not necessarily the actual wire format. When the `Status` message is exposed in different client libraries and different wire protocols, it can be mapped differently. For example, it will likely be mapped to some exceptions in Java, but more likely mapped to some error codes in C. # Other uses The error model and the `Status` message can be used in a variety of environments, either with or without APIs, to provide a consistent developer experience across different environments. Example uses of this error model include: - Partial errors. If a service needs to return partial errors to the client, it may embed the `Status` in the normal response to indicate the partial errors. - Workflow errors. A typical workflow has multiple steps. Each step may have a `Status` message for error reporting purpose. - Batch operations. If a client uses batch request and batch response, the `Status` message should be used directly inside batch response, one for each error sub-response. - Asynchronous operations. If an API call embeds asynchronous operation results in its response, the status of those operations should be represented directly using the `Status` message. - Logging. If some API errors are stored in logs, the message `Status` could be used directly after any stripping needed for security/privacy reasons. Messages: DetailsValueListEntry: A DetailsValueListEntry object. Fields: code: The status code, which should be an enum value of google.rpc.Code. details: A list of messages that carry the error details. There will be a common set of message types for APIs to use. message: A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the google.rpc.Status.details field, or localized by the client. """ @encoding.MapUnrecognizedFields('additionalProperties') class DetailsValueListEntry(_messages.Message): """A DetailsValueListEntry object. Messages: AdditionalProperty: An additional property for a DetailsValueListEntry object. Fields: additionalProperties: Properties of the object. Contains field @ype with type URL. """ class AdditionalProperty(_messages.Message): """An additional property for a DetailsValueListEntry object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) code = _messages.IntegerField(1, variant=_messages.Variant.INT32) details = _messages.MessageField('DetailsValueListEntry', 2, repeated=True) message = _messages.StringField(3) class Step(_messages.Message): """Defines a particular step within a Dataflow job. A job consists of multiple steps, each of which performs some specific operation as part of the overall job. Data is typically passed from one step to another as part of the job. Here's an example of a sequence of steps which together implement a Map-Reduce job: * Read a collection of data from some source, parsing the collection's elements. * Validate the elements. * Apply a user- defined function to map each element to some value and extract an element- specific key value. * Group elements with the same key into a single element with that key, transforming a multiply-keyed collection into a uniquely- keyed collection. * Write the elements out to some data sink. (Note that the Dataflow service may be used to run many different types of jobs, not just Map-Reduce). Messages: PropertiesValue: Named properties associated with the step. Each kind of predefined step has its own required set of properties. Fields: kind: The kind of step in the dataflow Job. name: Name identifying the step. This must be unique for each step with respect to all other steps in the dataflow Job. properties: Named properties associated with the step. Each kind of predefined step has its own required set of properties. """ @encoding.MapUnrecognizedFields('additionalProperties') class PropertiesValue(_messages.Message): """Named properties associated with the step. Each kind of predefined step has its own required set of properties. Messages: AdditionalProperty: An additional property for a PropertiesValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a PropertiesValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) kind = _messages.StringField(1) name = _messages.StringField(2) properties = _messages.MessageField('PropertiesValue', 3) class StreamLocation(_messages.Message): """Describes a stream of data, either as input to be processed or as output of a streaming Dataflow job. Fields: customSourceLocation: The stream is a custom source. pubsubLocation: The stream is a pubsub stream. sideInputLocation: The stream is a streaming side input. streamingStageLocation: The stream is part of another computation within the current streaming Dataflow job. """ customSourceLocation = _messages.MessageField('CustomSourceLocation', 1) pubsubLocation = _messages.MessageField('PubsubLocation', 2) sideInputLocation = _messages.MessageField('StreamingSideInputLocation', 3) streamingStageLocation = _messages.MessageField('StreamingStageLocation', 4) class StreamingComputationRanges(_messages.Message): """Describes full or partial data disk assignment information of the computation ranges. Fields: computationId: The ID of the computation. rangeAssignments: Data disk assignments for ranges from this computation. """ computationId = _messages.StringField(1) rangeAssignments = _messages.MessageField('KeyRangeDataDiskAssignment', 2, repeated=True) class StreamingComputationTask(_messages.Message): """A task which describes what action should be performed for the specified streaming computation ranges. Enums: TaskTypeValueValuesEnum: A type of streaming computation task. Fields: computationRanges: Contains ranges of a streaming computation this task should apply to. dataDisks: Describes the set of data disks this task should apply to. taskType: A type of streaming computation task. """ class TaskTypeValueValuesEnum(_messages.Enum): """A type of streaming computation task. Values: STREAMING_COMPUTATION_TASK_UNKNOWN: STREAMING_COMPUTATION_TASK_STOP: STREAMING_COMPUTATION_TASK_START: """ STREAMING_COMPUTATION_TASK_UNKNOWN = 0 STREAMING_COMPUTATION_TASK_STOP = 1 STREAMING_COMPUTATION_TASK_START = 2 computationRanges = _messages.MessageField('StreamingComputationRanges', 1, repeated=True) dataDisks = _messages.MessageField('MountedDataDisk', 2, repeated=True) taskType = _messages.EnumField('TaskTypeValueValuesEnum', 3) class StreamingSetupTask(_messages.Message): """A task which initializes part of a streaming Dataflow job. Fields: drain: The user has requested drain. receiveWorkPort: The TCP port on which the worker should listen for messages from other streaming computation workers. streamingComputationTopology: The global topology of the streaming Dataflow job. workerHarnessPort: The TCP port used by the worker to communicate with the Dataflow worker harness. """ drain = _messages.BooleanField(1) receiveWorkPort = _messages.IntegerField(2, variant=_messages.Variant.INT32) streamingComputationTopology = _messages.MessageField('TopologyConfig', 3) workerHarnessPort = _messages.IntegerField(4, variant=_messages.Variant.INT32) class StreamingSideInputLocation(_messages.Message): """Identifies the location of a streaming side input. Fields: stateFamily: Identifies the state family where this side input is stored. tag: Identifies the particular side input within the streaming Dataflow job. """ stateFamily = _messages.StringField(1) tag = _messages.StringField(2) class StreamingStageLocation(_messages.Message): """Identifies the location of a streaming computation stage, for stage-to- stage communication. Fields: streamId: Identifies the particular stream within the streaming Dataflow job. """ streamId = _messages.StringField(1) class TaskRunnerSettings(_messages.Message): """Taskrunner configuration settings. Fields: alsologtostderr: Also send taskrunner log info to stderr? baseTaskDir: Location on the worker for task-specific subdirectories. baseUrl: The base URL for the taskrunner to use when accessing Google Cloud APIs. When workers access Google Cloud APIs, they logically do so via relative URLs. If this field is specified, it supplies the base URL to use for resolving these relative URLs. The normative algorithm used is defined by RFC 1808, "Relative Uniform Resource Locators". If not specified, the default value is "http://www.googleapis.com/" commandlinesFileName: Store preprocessing commands in this file. continueOnException: Do we continue taskrunner if an exception is hit? dataflowApiVersion: API version of endpoint, e.g. "v1b3" harnessCommand: Command to launch the worker harness. languageHint: Suggested backend language. logDir: Directory on the VM to store logs. logToSerialconsole: Send taskrunner log into to Google Compute Engine VM serial console? logUploadLocation: Indicates where to put logs. If this is not specified, the logs will not be uploaded. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} oauthScopes: OAuth2 scopes to be requested by the taskrunner in order to access the dataflow API. parallelWorkerSettings: Settings to pass to the parallel worker harness. streamingWorkerMainClass: Streaming worker main class name. taskGroup: The UNIX group ID on the worker VM to use for tasks launched by taskrunner; e.g. "wheel". taskUser: The UNIX user ID on the worker VM to use for tasks launched by taskrunner; e.g. "root". tempStoragePrefix: The prefix of the resources the taskrunner should use for temporary storage. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} vmId: ID string of VM. workflowFileName: Store the workflow in this file. """ alsologtostderr = _messages.BooleanField(1) baseTaskDir = _messages.StringField(2) baseUrl = _messages.StringField(3) commandlinesFileName = _messages.StringField(4) continueOnException = _messages.BooleanField(5) dataflowApiVersion = _messages.StringField(6) harnessCommand = _messages.StringField(7) languageHint = _messages.StringField(8) logDir = _messages.StringField(9) logToSerialconsole = _messages.BooleanField(10) logUploadLocation = _messages.StringField(11) oauthScopes = _messages.StringField(12, repeated=True) parallelWorkerSettings = _messages.MessageField('WorkerSettings', 13) streamingWorkerMainClass = _messages.StringField(14) taskGroup = _messages.StringField(15) taskUser = _messages.StringField(16) tempStoragePrefix = _messages.StringField(17) vmId = _messages.StringField(18) workflowFileName = _messages.StringField(19) class TopologyConfig(_messages.Message): """Global topology of the streaming Dataflow job, including all computations and their sharded locations. Messages: UserStageToComputationNameMapValue: Maps user stage names to stable computation names. Fields: computations: The computations associated with a streaming Dataflow job. dataDiskAssignments: The disks assigned to a streaming Dataflow job. forwardingKeyBits: The size (in bits) of keys that will be assigned to source messages. persistentStateVersion: Version number for persistent state. userStageToComputationNameMap: Maps user stage names to stable computation names. """ @encoding.MapUnrecognizedFields('additionalProperties') class UserStageToComputationNameMapValue(_messages.Message): """Maps user stage names to stable computation names. Messages: AdditionalProperty: An additional property for a UserStageToComputationNameMapValue object. Fields: additionalProperties: Additional properties of type UserStageToComputationNameMapValue """ class AdditionalProperty(_messages.Message): """An additional property for a UserStageToComputationNameMapValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) computations = _messages.MessageField('ComputationTopology', 1, repeated=True) dataDiskAssignments = _messages.MessageField('DataDiskAssignment', 2, repeated=True) forwardingKeyBits = _messages.IntegerField(3, variant=_messages.Variant.INT32) persistentStateVersion = _messages.IntegerField(4, variant=_messages.Variant.INT32) userStageToComputationNameMap = _messages.MessageField('UserStageToComputationNameMapValue', 5) class WorkItem(_messages.Message): """WorkItem represents basic information about a WorkItem to be executed in the cloud. Fields: configuration: Work item-specific configuration as an opaque blob. id: Identifies this WorkItem. initialReportIndex: The initial index to use when reporting the status of the WorkItem. jobId: Identifies the workflow job this WorkItem belongs to. leaseExpireTime: Time when the lease on this Work will expire. mapTask: Additional information for MapTask WorkItems. packages: Any required packages that need to be fetched in order to execute this WorkItem. projectId: Identifies the cloud project this WorkItem belongs to. reportStatusInterval: Recommended reporting interval. seqMapTask: Additional information for SeqMapTask WorkItems. shellTask: Additional information for ShellTask WorkItems. sourceOperationTask: Additional information for source operation WorkItems. streamingComputationTask: Additional information for StreamingComputationTask WorkItems. streamingSetupTask: Additional information for StreamingSetupTask WorkItems. """ configuration = _messages.StringField(1) id = _messages.IntegerField(2) initialReportIndex = _messages.IntegerField(3) jobId = _messages.StringField(4) leaseExpireTime = _messages.StringField(5) mapTask = _messages.MessageField('MapTask', 6) packages = _messages.MessageField('Package', 7, repeated=True) projectId = _messages.StringField(8) reportStatusInterval = _messages.StringField(9) seqMapTask = _messages.MessageField('SeqMapTask', 10) shellTask = _messages.MessageField('ShellTask', 11) sourceOperationTask = _messages.MessageField('SourceOperationRequest', 12) streamingComputationTask = _messages.MessageField('StreamingComputationTask', 13) streamingSetupTask = _messages.MessageField('StreamingSetupTask', 14) class WorkItemServiceState(_messages.Message): """The Dataflow service's idea of the current state of a WorkItem being processed by a worker. Messages: HarnessDataValue: Other data returned by the service, specific to the particular worker harness. Fields: harnessData: Other data returned by the service, specific to the particular worker harness. leaseExpireTime: Time at which the current lease will expire. nextReportIndex: The index value to use for the next report sent by the worker. Note: If the report call fails for whatever reason, the worker should reuse this index for subsequent report attempts. reportStatusInterval: New recommended reporting interval. splitRequest: The progress point in the WorkItem where the Dataflow service suggests that the worker truncate the task. suggestedStopPoint: DEPRECATED in favor of split_request. suggestedStopPosition: Obsolete, always empty. """ @encoding.MapUnrecognizedFields('additionalProperties') class HarnessDataValue(_messages.Message): """Other data returned by the service, specific to the particular worker harness. Messages: AdditionalProperty: An additional property for a HarnessDataValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a HarnessDataValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) harnessData = _messages.MessageField('HarnessDataValue', 1) leaseExpireTime = _messages.StringField(2) nextReportIndex = _messages.IntegerField(3) reportStatusInterval = _messages.StringField(4) splitRequest = _messages.MessageField('ApproximateSplitRequest', 5) suggestedStopPoint = _messages.MessageField('ApproximateProgress', 6) suggestedStopPosition = _messages.MessageField('Position', 7) class WorkItemStatus(_messages.Message): """Conveys a worker's progress through the work described by a WorkItem. Fields: completed: True if the WorkItem was completed (successfully or unsuccessfully). dynamicSourceSplit: See documentation of stop_position. errors: Specifies errors which occurred during processing. If errors are provided, and completed = true, then the WorkItem is considered to have failed. metricUpdates: Worker output metrics (counters) for this WorkItem. progress: DEPRECATED in favor of reported_progress. reportIndex: The report index. When a WorkItem is leased, the lease will contain an initial report index. When a WorkItem's status is reported to the system, the report should be sent with that report index, and the response will contain the index the worker should use for the next report. Reports received with unexpected index values will be rejected by the service. In order to preserve idempotency, the worker should not alter the contents of a report, even if the worker must submit the same report multiple times before getting back a response. The worker should not submit a subsequent report until the response for the previous report had been received from the service. reportedProgress: The worker's progress through this WorkItem. requestedLeaseDuration: Amount of time the worker requests for its lease. sourceFork: DEPRECATED in favor of dynamic_source_split. sourceOperationResponse: If the work item represented a SourceOperationRequest, and the work is completed, contains the result of the operation. stopPosition: A worker may split an active map task in two parts, "primary" and "residual", continuing to process the primary part and returning the residual part into the pool of available work. This event is called a "dynamic split" and is critical to the dynamic work rebalancing feature. The two obtained sub-tasks are called "parts" of the split. The parts, if concatenated, must represent the same input as would be read by the current task if the split did not happen. The exact way in which the original task is decomposed into the two parts is specified either as a position demarcating them (stop_position), or explicitly as two DerivedSources, if this task consumes a user-defined source type (dynamic_source_split). The "current" task is adjusted as a result of the split: after a task with range [A, B) sends a stop_position update at C, its range is considered to be [A, C), e.g.: * Progress should be interpreted relative to the new range, e.g. "75% completed" means "75% of [A, C) completed" * The worker should interpret proposed_stop_position relative to the new range, e.g. "split at 68%" should be interpreted as "split at 68% of [A, C)". * If the worker chooses to split again using stop_position, only stop_positions in [A, C) will be accepted. * Etc. dynamic_source_split has similar semantics: e.g., if a task with source S splits using dynamic_source_split into {P, R} (where P and R must be together equivalent to S), then subsequent progress and proposed_stop_position should be interpreted relative to P, and in a potential subsequent dynamic_source_split into {P', R'}, P' and R' must be together equivalent to P, etc. workItemId: Identifies the WorkItem. """ completed = _messages.BooleanField(1) dynamicSourceSplit = _messages.MessageField('DynamicSourceSplit', 2) errors = _messages.MessageField('Status', 3, repeated=True) metricUpdates = _messages.MessageField('MetricUpdate', 4, repeated=True) progress = _messages.MessageField('ApproximateProgress', 5) reportIndex = _messages.IntegerField(6) reportedProgress = _messages.MessageField('ApproximateReportedProgress', 7) requestedLeaseDuration = _messages.StringField(8) sourceFork = _messages.MessageField('SourceFork', 9) sourceOperationResponse = _messages.MessageField('SourceOperationResponse', 10) stopPosition = _messages.MessageField('Position', 11) workItemId = _messages.StringField(12) class WorkerHealthReport(_messages.Message): """WorkerHealthReport contains information about the health of a worker. The VM should be identified by the labels attached to the WorkerMessage that this health ping belongs to. Messages: PodsValueListEntry: A PodsValueListEntry object. Fields: pods: The pods running on the worker. See: http://kubernetes.io/v1.1/docs /api-reference/v1/definitions.html#_v1_pod This field is used by the worker to send the status of the indvidual containers running on each worker. reportInterval: The interval at which the worker is sending health reports. The default value of 0 should be interpreted as the field is not being explicitly set by the worker. vmIsHealthy: Whether the VM is healthy. vmStartupTime: The time the VM was booted. """ @encoding.MapUnrecognizedFields('additionalProperties') class PodsValueListEntry(_messages.Message): """A PodsValueListEntry object. Messages: AdditionalProperty: An additional property for a PodsValueListEntry object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a PodsValueListEntry object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) pods = _messages.MessageField('PodsValueListEntry', 1, repeated=True) reportInterval = _messages.StringField(2) vmIsHealthy = _messages.BooleanField(3) vmStartupTime = _messages.StringField(4) class WorkerHealthReportResponse(_messages.Message): """WorkerHealthReportResponse contains information returned to the worker in response to a health ping. Fields: reportInterval: A positive value indicates the worker should change its reporting interval to the specified value. The default value of zero means no change in report rate is requested by the server. """ reportInterval = _messages.StringField(1) class WorkerMessage(_messages.Message): """WorkerMessage provides information to the backend about a worker. Messages: LabelsValue: Labels are used to group WorkerMessages. For example, a worker_message about a particular container might have the labels: { "JOB_ID": "2015-04-22", "WORKER_ID": "wordcount-vm-2015\u2026" "CONTAINER_TYPE": "worker", "CONTAINER_ID": "ac1234def"} Label tags typically correspond to Label enum values. However, for ease of development other strings can be used as tags. LABEL_UNSPECIFIED should not be used here. Fields: labels: Labels are used to group WorkerMessages. For example, a worker_message about a particular container might have the labels: { "JOB_ID": "2015-04-22", "WORKER_ID": "wordcount-vm-2015\u2026" "CONTAINER_TYPE": "worker", "CONTAINER_ID": "ac1234def"} Label tags typically correspond to Label enum values. However, for ease of development other strings can be used as tags. LABEL_UNSPECIFIED should not be used here. time: The timestamp of the worker_message. workerHealthReport: The health of a worker. workerMessageCode: A worker message code. """ @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): """Labels are used to group WorkerMessages. For example, a worker_message about a particular container might have the labels: { "JOB_ID": "2015-04-22", "WORKER_ID": "wordcount-vm-2015\u2026" "CONTAINER_TYPE": "worker", "CONTAINER_ID": "ac1234def"} Label tags typically correspond to Label enum values. However, for ease of development other strings can be used as tags. LABEL_UNSPECIFIED should not be used here. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): """An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) labels = _messages.MessageField('LabelsValue', 1) time = _messages.StringField(2) workerHealthReport = _messages.MessageField('WorkerHealthReport', 3) workerMessageCode = _messages.MessageField('WorkerMessageCode', 4) class WorkerMessageCode(_messages.Message): """A message code is used to report status and error messages to the service. The message codes are intended to be machine readable. The service will take care of translating these into user understandable messages if necessary. Example use cases: 1. Worker processes reporting successful startup. 2. Worker processes reporting specific errors (e.g. package staging failure). Messages: ParametersValue: Parameters contains specific information about the code. This is a struct to allow parameters of different types. Examples: 1. For a "HARNESS_STARTED" message parameters might provide the name of the worker and additional data like timing information. 2. For a "GCS_DOWNLOAD_ERROR" parameters might contain fields listing the GCS objects being downloaded and fields containing errors. In general complex data structures should be avoided. If a worker needs to send a specific and complicated data structure then please consider defining a new proto and adding it to the data oneof in WorkerMessageResponse. Conventions: Parameters should only be used for information that isn't typically passed as a label. hostname and other worker identifiers should almost always be passed as labels since they will be included on most messages. Fields: code: The code is a string intended for consumption by a machine that identifies the type of message being sent. Examples: 1. "HARNESS_STARTED" might be used to indicate the worker harness has started. 2. "GCS_DOWNLOAD_ERROR" might be used to indicate an error downloading a GCS file as part of the boot process of one of the worker containers. This is a string and not an enum to make it easy to add new codes without waiting for an API change. parameters: Parameters contains specific information about the code. This is a struct to allow parameters of different types. Examples: 1. For a "HARNESS_STARTED" message parameters might provide the name of the worker and additional data like timing information. 2. For a "GCS_DOWNLOAD_ERROR" parameters might contain fields listing the GCS objects being downloaded and fields containing errors. In general complex data structures should be avoided. If a worker needs to send a specific and complicated data structure then please consider defining a new proto and adding it to the data oneof in WorkerMessageResponse. Conventions: Parameters should only be used for information that isn't typically passed as a label. hostname and other worker identifiers should almost always be passed as labels since they will be included on most messages. """ @encoding.MapUnrecognizedFields('additionalProperties') class ParametersValue(_messages.Message): """Parameters contains specific information about the code. This is a struct to allow parameters of different types. Examples: 1. For a "HARNESS_STARTED" message parameters might provide the name of the worker and additional data like timing information. 2. For a "GCS_DOWNLOAD_ERROR" parameters might contain fields listing the GCS objects being downloaded and fields containing errors. In general complex data structures should be avoided. If a worker needs to send a specific and complicated data structure then please consider defining a new proto and adding it to the data oneof in WorkerMessageResponse. Conventions: Parameters should only be used for information that isn't typically passed as a label. hostname and other worker identifiers should almost always be passed as labels since they will be included on most messages. Messages: AdditionalProperty: An additional property for a ParametersValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): """An additional property for a ParametersValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) code = _messages.StringField(1) parameters = _messages.MessageField('ParametersValue', 2) class WorkerMessageResponse(_messages.Message): """A worker_message response allows the server to pass information to the sender. Fields: workerHealthReportResponse: The service's response to a worker's health report. """ workerHealthReportResponse = _messages.MessageField('WorkerHealthReportResponse', 1) class WorkerPool(_messages.Message): """Describes one particular pool of Dataflow workers to be instantiated by the Dataflow service in order to perform the computations required by a job. Note that a workflow job may use multiple pools, in order to match the various computational requirements of the various stages of the job. Enums: DefaultPackageSetValueValuesEnum: The default package set to install. This allows the service to select a default set of packages which are useful to worker harnesses written in a particular language. TeardownPolicyValueValuesEnum: Sets the policy for determining when to turndown worker pool. Allowed values are: TEARDOWN_ALWAYS, TEARDOWN_ON_SUCCESS, and TEARDOWN_NEVER. TEARDOWN_ALWAYS means workers are always torn down regardless of whether the job succeeds. TEARDOWN_ON_SUCCESS means workers are torn down if the job succeeds. TEARDOWN_NEVER means the workers are never torn down. If the workers are not torn down by the service, they will continue to run and use Google Compute Engine VM resources in the user's project until they are explicitly terminated by the user. Because of this, Google recommends using the TEARDOWN_ALWAYS policy except for small, manually supervised test jobs. If unknown or unspecified, the service will attempt to choose a reasonable default. Messages: MetadataValue: Metadata to set on the Google Compute Engine VMs. PoolArgsValue: Extra arguments for this worker pool. Fields: autoscalingSettings: Settings for autoscaling of this WorkerPool. dataDisks: Data disks that are used by a VM in this workflow. defaultPackageSet: The default package set to install. This allows the service to select a default set of packages which are useful to worker harnesses written in a particular language. diskSizeGb: Size of root disk for VMs, in GB. If zero or unspecified, the service will attempt to choose a reasonable default. diskSourceImage: Fully qualified source image for disks. diskType: Type of root disk for VMs. If empty or unspecified, the service will attempt to choose a reasonable default. kind: The kind of the worker pool; currently only 'harness' and 'shuffle' are supported. machineType: Machine type (e.g. "n1-standard-1"). If empty or unspecified, the service will attempt to choose a reasonable default. metadata: Metadata to set on the Google Compute Engine VMs. network: Network to which VMs will be assigned. If empty or unspecified, the service will use the network "default". numThreadsPerWorker: The number of threads per worker harness. If empty or unspecified, the service will choose a number of threads (according to the number of cores on the selected machine type for batch, or 1 by convention for streaming). numWorkers: Number of Google Compute Engine workers in this pool needed to execute the job. If zero or unspecified, the service will attempt to choose a reasonable default. onHostMaintenance: The action to take on host maintenance, as defined by the Google Compute Engine API. packages: Packages to be installed on workers. poolArgs: Extra arguments for this worker pool. subnetwork: Subnetwork to which VMs will be assigned, if desired. Expected to be of the form "zones/ZONE/subnetworks/SUBNETWORK". taskrunnerSettings: Settings passed through to Google Compute Engine workers when using the standard Dataflow task runner. Users should ignore this field. teardownPolicy: Sets the policy for determining when to turndown worker pool. Allowed values are: TEARDOWN_ALWAYS, TEARDOWN_ON_SUCCESS, and TEARDOWN_NEVER. TEARDOWN_ALWAYS means workers are always torn down regardless of whether the job succeeds. TEARDOWN_ON_SUCCESS means workers are torn down if the job succeeds. TEARDOWN_NEVER means the workers are never torn down. If the workers are not torn down by the service, they will continue to run and use Google Compute Engine VM resources in the user's project until they are explicitly terminated by the user. Because of this, Google recommends using the TEARDOWN_ALWAYS policy except for small, manually supervised test jobs. If unknown or unspecified, the service will attempt to choose a reasonable default. workerHarnessContainerImage: Docker container image that executes Dataflow worker harness, residing in Google Container Registry. Required. zone: Zone to run the worker pools in. If empty or unspecified, the service will attempt to choose a reasonable default. """ class DefaultPackageSetValueValuesEnum(_messages.Enum): """The default package set to install. This allows the service to select a default set of packages which are useful to worker harnesses written in a particular language. Values: DEFAULT_PACKAGE_SET_UNKNOWN: DEFAULT_PACKAGE_SET_NONE: DEFAULT_PACKAGE_SET_JAVA: DEFAULT_PACKAGE_SET_PYTHON: """ DEFAULT_PACKAGE_SET_UNKNOWN = 0 DEFAULT_PACKAGE_SET_NONE = 1 DEFAULT_PACKAGE_SET_JAVA = 2 DEFAULT_PACKAGE_SET_PYTHON = 3 class TeardownPolicyValueValuesEnum(_messages.Enum): """Sets the policy for determining when to turndown worker pool. Allowed values are: TEARDOWN_ALWAYS, TEARDOWN_ON_SUCCESS, and TEARDOWN_NEVER. TEARDOWN_ALWAYS means workers are always torn down regardless of whether the job succeeds. TEARDOWN_ON_SUCCESS means workers are torn down if the job succeeds. TEARDOWN_NEVER means the workers are never torn down. If the workers are not torn down by the service, they will continue to run and use Google Compute Engine VM resources in the user's project until they are explicitly terminated by the user. Because of this, Google recommends using the TEARDOWN_ALWAYS policy except for small, manually supervised test jobs. If unknown or unspecified, the service will attempt to choose a reasonable default. Values: TEARDOWN_POLICY_UNKNOWN: TEARDOWN_ALWAYS: TEARDOWN_ON_SUCCESS: TEARDOWN_NEVER: """ TEARDOWN_POLICY_UNKNOWN = 0 TEARDOWN_ALWAYS = 1 TEARDOWN_ON_SUCCESS = 2 TEARDOWN_NEVER = 3 @encoding.MapUnrecognizedFields('additionalProperties') class MetadataValue(_messages.Message): """Metadata to set on the Google Compute Engine VMs. Messages: AdditionalProperty: An additional property for a MetadataValue object. Fields: additionalProperties: Additional properties of type MetadataValue """ class AdditionalProperty(_messages.Message): """An additional property for a MetadataValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class PoolArgsValue(_messages.Message): """Extra arguments for this worker pool. Messages: AdditionalProperty: An additional property for a PoolArgsValue object. Fields: additionalProperties: Properties of the object. Contains field @ype with type URL. """ class AdditionalProperty(_messages.Message): """An additional property for a PoolArgsValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) autoscalingSettings = _messages.MessageField('AutoscalingSettings', 1) dataDisks = _messages.MessageField('Disk', 2, repeated=True) defaultPackageSet = _messages.EnumField('DefaultPackageSetValueValuesEnum', 3) diskSizeGb = _messages.IntegerField(4, variant=_messages.Variant.INT32) diskSourceImage = _messages.StringField(5) diskType = _messages.StringField(6) kind = _messages.StringField(7) machineType = _messages.StringField(8) metadata = _messages.MessageField('MetadataValue', 9) network = _messages.StringField(10) numThreadsPerWorker = _messages.IntegerField(11, variant=_messages.Variant.INT32) numWorkers = _messages.IntegerField(12, variant=_messages.Variant.INT32) onHostMaintenance = _messages.StringField(13) packages = _messages.MessageField('Package', 14, repeated=True) poolArgs = _messages.MessageField('PoolArgsValue', 15) subnetwork = _messages.StringField(16) taskrunnerSettings = _messages.MessageField('TaskRunnerSettings', 17) teardownPolicy = _messages.EnumField('TeardownPolicyValueValuesEnum', 18) workerHarnessContainerImage = _messages.StringField(19) zone = _messages.StringField(20) class WorkerSettings(_messages.Message): """Provides data to pass through to the worker harness. Fields: baseUrl: The base URL for accessing Google Cloud APIs. When workers access Google Cloud APIs, they logically do so via relative URLs. If this field is specified, it supplies the base URL to use for resolving these relative URLs. The normative algorithm used is defined by RFC 1808, "Relative Uniform Resource Locators". If not specified, the default value is "http://www.googleapis.com/" reportingEnabled: Send work progress updates to service. servicePath: The Dataflow service path relative to the root URL, for example, "dataflow/v1b3/projects". shuffleServicePath: The Shuffle service path relative to the root URL, for example, "shuffle/v1beta1". tempStoragePrefix: The prefix of the resources the system should use for temporary storage. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} workerId: ID of the worker running this pipeline. """ baseUrl = _messages.StringField(1) reportingEnabled = _messages.BooleanField(2) servicePath = _messages.StringField(3) shuffleServicePath = _messages.StringField(4) tempStoragePrefix = _messages.StringField(5) workerId = _messages.StringField(6) class WriteInstruction(_messages.Message): """An instruction that writes records. Takes one input, produces no outputs. Fields: input: The input. sink: The sink to write to. """ input = _messages.MessageField('InstructionInput', 1) sink = _messages.MessageField('Sink', 2) encoding.AddCustomJsonFieldMapping( StandardQueryParameters, 'f__xgafv', '$.xgafv', package=u'dataflow') DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/storage/000077500000000000000000000000001272340251000261235ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/storage/__init__.py000066400000000000000000000005401272340251000302330ustar00rootroot00000000000000"""Common imports for generated storage client library.""" # pylint:disable=wildcard-import import pkgutil from apitools.base.py import * from google.cloud.dataflow.internal.clients.storage.storage_v1_client import * from google.cloud.dataflow.internal.clients.storage.storage_v1_messages import * __path__ = pkgutil.extend_path(__path__, __name__) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/storage/storage_v1_client.py000066400000000000000000001250301272340251000321060ustar00rootroot00000000000000"""Generated client library for storage version v1.""" # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.py import base_api from google.cloud.dataflow.internal.clients.storage import storage_v1_messages as messages class StorageV1(base_api.BaseApiClient): """Generated client library for service storage version v1.""" MESSAGES_MODULE = messages _PACKAGE = u'storage' _SCOPES = [u'https://www.googleapis.com/auth/cloud-platform', u'https://www.googleapis.com/auth/cloud-platform.read-only', u'https://www.googleapis.com/auth/devstorage.full_control', u'https://www.googleapis.com/auth/devstorage.read_only', u'https://www.googleapis.com/auth/devstorage.read_write'] _VERSION = u'v1' _CLIENT_ID = '1042881264118.apps.googleusercontent.com' _CLIENT_SECRET = 'x_Tw5K8nnjoRAqULM9PFAC2b' _USER_AGENT = 'x_Tw5K8nnjoRAqULM9PFAC2b' _CLIENT_CLASS_NAME = u'StorageV1' _URL_VERSION = u'v1' _API_KEY = None def __init__(self, url='', credentials=None, get_credentials=True, http=None, model=None, log_request=False, log_response=False, credentials_args=None, default_global_params=None, additional_http_headers=None): """Create a new storage handle.""" url = url or u'https://www.googleapis.com/storage/v1/' super(StorageV1, self).__init__( url, credentials=credentials, get_credentials=get_credentials, http=http, model=model, log_request=log_request, log_response=log_response, credentials_args=credentials_args, default_global_params=default_global_params, additional_http_headers=additional_http_headers) self.bucketAccessControls = self.BucketAccessControlsService(self) self.buckets = self.BucketsService(self) self.channels = self.ChannelsService(self) self.defaultObjectAccessControls = self.DefaultObjectAccessControlsService(self) self.objectAccessControls = self.ObjectAccessControlsService(self) self.objects = self.ObjectsService(self) class BucketAccessControlsService(base_api.BaseApiService): """Service class for the bucketAccessControls resource.""" _NAME = u'bucketAccessControls' def __init__(self, client): super(StorageV1.BucketAccessControlsService, self).__init__(client) self._method_configs = { 'Delete': base_api.ApiMethodInfo( http_method=u'DELETE', method_id=u'storage.bucketAccessControls.delete', ordered_params=[u'bucket', u'entity'], path_params=[u'bucket', u'entity'], query_params=[], relative_path=u'b/{bucket}/acl/{entity}', request_field='', request_type_name=u'StorageBucketAccessControlsDeleteRequest', response_type_name=u'StorageBucketAccessControlsDeleteResponse', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.bucketAccessControls.get', ordered_params=[u'bucket', u'entity'], path_params=[u'bucket', u'entity'], query_params=[], relative_path=u'b/{bucket}/acl/{entity}', request_field='', request_type_name=u'StorageBucketAccessControlsGetRequest', response_type_name=u'BucketAccessControl', supports_download=False, ), 'Insert': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.bucketAccessControls.insert', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[], relative_path=u'b/{bucket}/acl', request_field='', request_type_name=u'BucketAccessControl', response_type_name=u'BucketAccessControl', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.bucketAccessControls.list', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[], relative_path=u'b/{bucket}/acl', request_field='', request_type_name=u'StorageBucketAccessControlsListRequest', response_type_name=u'BucketAccessControls', supports_download=False, ), 'Patch': base_api.ApiMethodInfo( http_method=u'PATCH', method_id=u'storage.bucketAccessControls.patch', ordered_params=[u'bucket', u'entity'], path_params=[u'bucket', u'entity'], query_params=[], relative_path=u'b/{bucket}/acl/{entity}', request_field='', request_type_name=u'BucketAccessControl', response_type_name=u'BucketAccessControl', supports_download=False, ), 'Update': base_api.ApiMethodInfo( http_method=u'PUT', method_id=u'storage.bucketAccessControls.update', ordered_params=[u'bucket', u'entity'], path_params=[u'bucket', u'entity'], query_params=[], relative_path=u'b/{bucket}/acl/{entity}', request_field='', request_type_name=u'BucketAccessControl', response_type_name=u'BucketAccessControl', supports_download=False, ), } self._upload_configs = { } def Delete(self, request, global_params=None): """Permanently deletes the ACL entry for the specified entity on the specified bucket. Args: request: (StorageBucketAccessControlsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (StorageBucketAccessControlsDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None): """Returns the ACL entry for the specified entity on the specified bucket. Args: request: (StorageBucketAccessControlsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BucketAccessControl) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) def Insert(self, request, global_params=None): """Creates a new ACL entry on the specified bucket. Args: request: (BucketAccessControl) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BucketAccessControl) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod( config, request, global_params=global_params) def List(self, request, global_params=None): """Retrieves ACL entries on the specified bucket. Args: request: (StorageBucketAccessControlsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BucketAccessControls) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Patch(self, request, global_params=None): """Updates an ACL entry on the specified bucket. This method supports patch semantics. Args: request: (BucketAccessControl) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BucketAccessControl) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod( config, request, global_params=global_params) def Update(self, request, global_params=None): """Updates an ACL entry on the specified bucket. Args: request: (BucketAccessControl) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BucketAccessControl) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) class BucketsService(base_api.BaseApiService): """Service class for the buckets resource.""" _NAME = u'buckets' def __init__(self, client): super(StorageV1.BucketsService, self).__init__(client) self._method_configs = { 'Delete': base_api.ApiMethodInfo( http_method=u'DELETE', method_id=u'storage.buckets.delete', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[u'ifMetagenerationMatch', u'ifMetagenerationNotMatch'], relative_path=u'b/{bucket}', request_field='', request_type_name=u'StorageBucketsDeleteRequest', response_type_name=u'StorageBucketsDeleteResponse', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.buckets.get', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'projection'], relative_path=u'b/{bucket}', request_field='', request_type_name=u'StorageBucketsGetRequest', response_type_name=u'Bucket', supports_download=False, ), 'Insert': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.buckets.insert', ordered_params=[u'project'], path_params=[], query_params=[u'predefinedAcl', u'predefinedDefaultObjectAcl', u'project', u'projection'], relative_path=u'b', request_field=u'bucket', request_type_name=u'StorageBucketsInsertRequest', response_type_name=u'Bucket', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.buckets.list', ordered_params=[u'project'], path_params=[], query_params=[u'maxResults', u'pageToken', u'prefix', u'project', u'projection'], relative_path=u'b', request_field='', request_type_name=u'StorageBucketsListRequest', response_type_name=u'Buckets', supports_download=False, ), 'Patch': base_api.ApiMethodInfo( http_method=u'PATCH', method_id=u'storage.buckets.patch', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'predefinedAcl', u'predefinedDefaultObjectAcl', u'projection'], relative_path=u'b/{bucket}', request_field=u'bucketResource', request_type_name=u'StorageBucketsPatchRequest', response_type_name=u'Bucket', supports_download=False, ), 'Update': base_api.ApiMethodInfo( http_method=u'PUT', method_id=u'storage.buckets.update', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'predefinedAcl', u'predefinedDefaultObjectAcl', u'projection'], relative_path=u'b/{bucket}', request_field=u'bucketResource', request_type_name=u'StorageBucketsUpdateRequest', response_type_name=u'Bucket', supports_download=False, ), } self._upload_configs = { } def Delete(self, request, global_params=None): """Permanently deletes an empty bucket. Args: request: (StorageBucketsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (StorageBucketsDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None): """Returns metadata for the specified bucket. Args: request: (StorageBucketsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Bucket) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) def Insert(self, request, global_params=None): """Creates a new bucket. Args: request: (StorageBucketsInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Bucket) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod( config, request, global_params=global_params) def List(self, request, global_params=None): """Retrieves a list of buckets for a given project. Args: request: (StorageBucketsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Buckets) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Patch(self, request, global_params=None): """Updates a bucket. This method supports patch semantics. Args: request: (StorageBucketsPatchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Bucket) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod( config, request, global_params=global_params) def Update(self, request, global_params=None): """Updates a bucket. Args: request: (StorageBucketsUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Bucket) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) class ChannelsService(base_api.BaseApiService): """Service class for the channels resource.""" _NAME = u'channels' def __init__(self, client): super(StorageV1.ChannelsService, self).__init__(client) self._method_configs = { 'Stop': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.channels.stop', ordered_params=[], path_params=[], query_params=[], relative_path=u'channels/stop', request_field='', request_type_name=u'Channel', response_type_name=u'StorageChannelsStopResponse', supports_download=False, ), } self._upload_configs = { } def Stop(self, request, global_params=None): """Stop watching resources through this channel. Args: request: (Channel) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (StorageChannelsStopResponse) The response message. """ config = self.GetMethodConfig('Stop') return self._RunMethod( config, request, global_params=global_params) class DefaultObjectAccessControlsService(base_api.BaseApiService): """Service class for the defaultObjectAccessControls resource.""" _NAME = u'defaultObjectAccessControls' def __init__(self, client): super(StorageV1.DefaultObjectAccessControlsService, self).__init__(client) self._method_configs = { 'Delete': base_api.ApiMethodInfo( http_method=u'DELETE', method_id=u'storage.defaultObjectAccessControls.delete', ordered_params=[u'bucket', u'entity'], path_params=[u'bucket', u'entity'], query_params=[], relative_path=u'b/{bucket}/defaultObjectAcl/{entity}', request_field='', request_type_name=u'StorageDefaultObjectAccessControlsDeleteRequest', response_type_name=u'StorageDefaultObjectAccessControlsDeleteResponse', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.defaultObjectAccessControls.get', ordered_params=[u'bucket', u'entity'], path_params=[u'bucket', u'entity'], query_params=[], relative_path=u'b/{bucket}/defaultObjectAcl/{entity}', request_field='', request_type_name=u'StorageDefaultObjectAccessControlsGetRequest', response_type_name=u'ObjectAccessControl', supports_download=False, ), 'Insert': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.defaultObjectAccessControls.insert', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[], relative_path=u'b/{bucket}/defaultObjectAcl', request_field='', request_type_name=u'ObjectAccessControl', response_type_name=u'ObjectAccessControl', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.defaultObjectAccessControls.list', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[u'ifMetagenerationMatch', u'ifMetagenerationNotMatch'], relative_path=u'b/{bucket}/defaultObjectAcl', request_field='', request_type_name=u'StorageDefaultObjectAccessControlsListRequest', response_type_name=u'ObjectAccessControls', supports_download=False, ), 'Patch': base_api.ApiMethodInfo( http_method=u'PATCH', method_id=u'storage.defaultObjectAccessControls.patch', ordered_params=[u'bucket', u'entity'], path_params=[u'bucket', u'entity'], query_params=[], relative_path=u'b/{bucket}/defaultObjectAcl/{entity}', request_field='', request_type_name=u'ObjectAccessControl', response_type_name=u'ObjectAccessControl', supports_download=False, ), 'Update': base_api.ApiMethodInfo( http_method=u'PUT', method_id=u'storage.defaultObjectAccessControls.update', ordered_params=[u'bucket', u'entity'], path_params=[u'bucket', u'entity'], query_params=[], relative_path=u'b/{bucket}/defaultObjectAcl/{entity}', request_field='', request_type_name=u'ObjectAccessControl', response_type_name=u'ObjectAccessControl', supports_download=False, ), } self._upload_configs = { } def Delete(self, request, global_params=None): """Permanently deletes the default object ACL entry for the specified entity on the specified bucket. Args: request: (StorageDefaultObjectAccessControlsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (StorageDefaultObjectAccessControlsDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None): """Returns the default object ACL entry for the specified entity on the specified bucket. Args: request: (StorageDefaultObjectAccessControlsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControl) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) def Insert(self, request, global_params=None): """Creates a new default object ACL entry on the specified bucket. Args: request: (ObjectAccessControl) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControl) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod( config, request, global_params=global_params) def List(self, request, global_params=None): """Retrieves default object ACL entries on the specified bucket. Args: request: (StorageDefaultObjectAccessControlsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControls) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Patch(self, request, global_params=None): """Updates a default object ACL entry on the specified bucket. This method supports patch semantics. Args: request: (ObjectAccessControl) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControl) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod( config, request, global_params=global_params) def Update(self, request, global_params=None): """Updates a default object ACL entry on the specified bucket. Args: request: (ObjectAccessControl) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControl) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) class ObjectAccessControlsService(base_api.BaseApiService): """Service class for the objectAccessControls resource.""" _NAME = u'objectAccessControls' def __init__(self, client): super(StorageV1.ObjectAccessControlsService, self).__init__(client) self._method_configs = { 'Delete': base_api.ApiMethodInfo( http_method=u'DELETE', method_id=u'storage.objectAccessControls.delete', ordered_params=[u'bucket', u'object', u'entity'], path_params=[u'bucket', u'entity', u'object'], query_params=[u'generation'], relative_path=u'b/{bucket}/o/{object}/acl/{entity}', request_field='', request_type_name=u'StorageObjectAccessControlsDeleteRequest', response_type_name=u'StorageObjectAccessControlsDeleteResponse', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.objectAccessControls.get', ordered_params=[u'bucket', u'object', u'entity'], path_params=[u'bucket', u'entity', u'object'], query_params=[u'generation'], relative_path=u'b/{bucket}/o/{object}/acl/{entity}', request_field='', request_type_name=u'StorageObjectAccessControlsGetRequest', response_type_name=u'ObjectAccessControl', supports_download=False, ), 'Insert': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.objectAccessControls.insert', ordered_params=[u'bucket', u'object'], path_params=[u'bucket', u'object'], query_params=[u'generation'], relative_path=u'b/{bucket}/o/{object}/acl', request_field=u'objectAccessControl', request_type_name=u'StorageObjectAccessControlsInsertRequest', response_type_name=u'ObjectAccessControl', supports_download=False, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.objectAccessControls.list', ordered_params=[u'bucket', u'object'], path_params=[u'bucket', u'object'], query_params=[u'generation'], relative_path=u'b/{bucket}/o/{object}/acl', request_field='', request_type_name=u'StorageObjectAccessControlsListRequest', response_type_name=u'ObjectAccessControls', supports_download=False, ), 'Patch': base_api.ApiMethodInfo( http_method=u'PATCH', method_id=u'storage.objectAccessControls.patch', ordered_params=[u'bucket', u'object', u'entity'], path_params=[u'bucket', u'entity', u'object'], query_params=[u'generation'], relative_path=u'b/{bucket}/o/{object}/acl/{entity}', request_field=u'objectAccessControl', request_type_name=u'StorageObjectAccessControlsPatchRequest', response_type_name=u'ObjectAccessControl', supports_download=False, ), 'Update': base_api.ApiMethodInfo( http_method=u'PUT', method_id=u'storage.objectAccessControls.update', ordered_params=[u'bucket', u'object', u'entity'], path_params=[u'bucket', u'entity', u'object'], query_params=[u'generation'], relative_path=u'b/{bucket}/o/{object}/acl/{entity}', request_field=u'objectAccessControl', request_type_name=u'StorageObjectAccessControlsUpdateRequest', response_type_name=u'ObjectAccessControl', supports_download=False, ), } self._upload_configs = { } def Delete(self, request, global_params=None): """Permanently deletes the ACL entry for the specified entity on the specified object. Args: request: (StorageObjectAccessControlsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (StorageObjectAccessControlsDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None): """Returns the ACL entry for the specified entity on the specified object. Args: request: (StorageObjectAccessControlsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControl) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) def Insert(self, request, global_params=None): """Creates a new ACL entry on the specified object. Args: request: (StorageObjectAccessControlsInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControl) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod( config, request, global_params=global_params) def List(self, request, global_params=None): """Retrieves ACL entries on the specified object. Args: request: (StorageObjectAccessControlsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControls) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Patch(self, request, global_params=None): """Updates an ACL entry on the specified object. This method supports patch semantics. Args: request: (StorageObjectAccessControlsPatchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControl) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod( config, request, global_params=global_params) def Update(self, request, global_params=None): """Updates an ACL entry on the specified object. Args: request: (StorageObjectAccessControlsUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ObjectAccessControl) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) class ObjectsService(base_api.BaseApiService): """Service class for the objects resource.""" _NAME = u'objects' def __init__(self, client): super(StorageV1.ObjectsService, self).__init__(client) self._method_configs = { 'Compose': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.objects.compose', ordered_params=[u'destinationBucket', u'destinationObject'], path_params=[u'destinationBucket', u'destinationObject'], query_params=[u'destinationPredefinedAcl', u'ifGenerationMatch', u'ifMetagenerationMatch'], relative_path=u'b/{destinationBucket}/o/{destinationObject}/compose', request_field=u'composeRequest', request_type_name=u'StorageObjectsComposeRequest', response_type_name=u'Object', supports_download=True, ), 'Copy': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.objects.copy', ordered_params=[u'sourceBucket', u'sourceObject', u'destinationBucket', u'destinationObject'], path_params=[u'destinationBucket', u'destinationObject', u'sourceBucket', u'sourceObject'], query_params=[u'destinationPredefinedAcl', u'ifGenerationMatch', u'ifGenerationNotMatch', u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'ifSourceGenerationMatch', u'ifSourceGenerationNotMatch', u'ifSourceMetagenerationMatch', u'ifSourceMetagenerationNotMatch', u'projection', u'sourceGeneration'], relative_path=u'b/{sourceBucket}/o/{sourceObject}/copyTo/b/{destinationBucket}/o/{destinationObject}', request_field=u'object', request_type_name=u'StorageObjectsCopyRequest', response_type_name=u'Object', supports_download=True, ), 'Delete': base_api.ApiMethodInfo( http_method=u'DELETE', method_id=u'storage.objects.delete', ordered_params=[u'bucket', u'object'], path_params=[u'bucket', u'object'], query_params=[u'generation', u'ifGenerationMatch', u'ifGenerationNotMatch', u'ifMetagenerationMatch', u'ifMetagenerationNotMatch'], relative_path=u'b/{bucket}/o/{object}', request_field='', request_type_name=u'StorageObjectsDeleteRequest', response_type_name=u'StorageObjectsDeleteResponse', supports_download=False, ), 'Get': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.objects.get', ordered_params=[u'bucket', u'object'], path_params=[u'bucket', u'object'], query_params=[u'generation', u'ifGenerationMatch', u'ifGenerationNotMatch', u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'projection'], relative_path=u'b/{bucket}/o/{object}', request_field='', request_type_name=u'StorageObjectsGetRequest', response_type_name=u'Object', supports_download=True, ), 'Insert': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.objects.insert', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[u'contentEncoding', u'ifGenerationMatch', u'ifGenerationNotMatch', u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'name', u'predefinedAcl', u'projection'], relative_path=u'b/{bucket}/o', request_field=u'object', request_type_name=u'StorageObjectsInsertRequest', response_type_name=u'Object', supports_download=True, ), 'List': base_api.ApiMethodInfo( http_method=u'GET', method_id=u'storage.objects.list', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[u'delimiter', u'maxResults', u'pageToken', u'prefix', u'projection', u'versions'], relative_path=u'b/{bucket}/o', request_field='', request_type_name=u'StorageObjectsListRequest', response_type_name=u'Objects', supports_download=False, ), 'Patch': base_api.ApiMethodInfo( http_method=u'PATCH', method_id=u'storage.objects.patch', ordered_params=[u'bucket', u'object'], path_params=[u'bucket', u'object'], query_params=[u'generation', u'ifGenerationMatch', u'ifGenerationNotMatch', u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'predefinedAcl', u'projection'], relative_path=u'b/{bucket}/o/{object}', request_field=u'objectResource', request_type_name=u'StorageObjectsPatchRequest', response_type_name=u'Object', supports_download=False, ), 'Rewrite': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.objects.rewrite', ordered_params=[u'sourceBucket', u'sourceObject', u'destinationBucket', u'destinationObject'], path_params=[u'destinationBucket', u'destinationObject', u'sourceBucket', u'sourceObject'], query_params=[u'destinationPredefinedAcl', u'ifGenerationMatch', u'ifGenerationNotMatch', u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'ifSourceGenerationMatch', u'ifSourceGenerationNotMatch', u'ifSourceMetagenerationMatch', u'ifSourceMetagenerationNotMatch', u'maxBytesRewrittenPerCall', u'projection', u'rewriteToken', u'sourceGeneration'], relative_path=u'b/{sourceBucket}/o/{sourceObject}/rewriteTo/b/{destinationBucket}/o/{destinationObject}', request_field=u'object', request_type_name=u'StorageObjectsRewriteRequest', response_type_name=u'RewriteResponse', supports_download=False, ), 'Update': base_api.ApiMethodInfo( http_method=u'PUT', method_id=u'storage.objects.update', ordered_params=[u'bucket', u'object'], path_params=[u'bucket', u'object'], query_params=[u'generation', u'ifGenerationMatch', u'ifGenerationNotMatch', u'ifMetagenerationMatch', u'ifMetagenerationNotMatch', u'predefinedAcl', u'projection'], relative_path=u'b/{bucket}/o/{object}', request_field=u'objectResource', request_type_name=u'StorageObjectsUpdateRequest', response_type_name=u'Object', supports_download=True, ), 'WatchAll': base_api.ApiMethodInfo( http_method=u'POST', method_id=u'storage.objects.watchAll', ordered_params=[u'bucket'], path_params=[u'bucket'], query_params=[u'delimiter', u'maxResults', u'pageToken', u'prefix', u'projection', u'versions'], relative_path=u'b/{bucket}/o/watch', request_field=u'channel', request_type_name=u'StorageObjectsWatchAllRequest', response_type_name=u'Channel', supports_download=False, ), } self._upload_configs = { 'Insert': base_api.ApiUploadInfo( accept=['*/*'], max_size=None, resumable_multipart=True, resumable_path=u'/resumable/upload/storage/v1/b/{bucket}/o', simple_multipart=True, simple_path=u'/upload/storage/v1/b/{bucket}/o', ), } def Compose(self, request, global_params=None, download=None): """Concatenates a list of existing objects into a new object in the same bucket. Args: request: (StorageObjectsComposeRequest) input message global_params: (StandardQueryParameters, default: None) global arguments download: (Download, default: None) If present, download data from the request via this stream. Returns: (Object) The response message. """ config = self.GetMethodConfig('Compose') return self._RunMethod( config, request, global_params=global_params, download=download) def Copy(self, request, global_params=None, download=None): """Copies a source object to a destination object. Optionally overrides metadata. Args: request: (StorageObjectsCopyRequest) input message global_params: (StandardQueryParameters, default: None) global arguments download: (Download, default: None) If present, download data from the request via this stream. Returns: (Object) The response message. """ config = self.GetMethodConfig('Copy') return self._RunMethod( config, request, global_params=global_params, download=download) def Delete(self, request, global_params=None): """Deletes an object and its metadata. Deletions are permanent if versioning is not enabled for the bucket, or if the generation parameter is used. Args: request: (StorageObjectsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (StorageObjectsDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod( config, request, global_params=global_params) def Get(self, request, global_params=None, download=None): """Retrieves an object or its metadata. Args: request: (StorageObjectsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments download: (Download, default: None) If present, download data from the request via this stream. Returns: (Object) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params, download=download) def Insert(self, request, global_params=None, upload=None, download=None): """Stores a new object and metadata. Args: request: (StorageObjectsInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments upload: (Upload, default: None) If present, upload this stream with the request. download: (Download, default: None) If present, download data from the request via this stream. Returns: (Object) The response message. """ config = self.GetMethodConfig('Insert') upload_config = self.GetUploadConfig('Insert') return self._RunMethod( config, request, global_params=global_params, upload=upload, upload_config=upload_config, download=download) def List(self, request, global_params=None): """Retrieves a list of objects matching the criteria. Args: request: (StorageObjectsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Objects) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) def Patch(self, request, global_params=None): """Updates an object's metadata. This method supports patch semantics. Args: request: (StorageObjectsPatchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Object) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod( config, request, global_params=global_params) def Rewrite(self, request, global_params=None): """Rewrites a source object to a destination object. Optionally overrides metadata. Args: request: (StorageObjectsRewriteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (RewriteResponse) The response message. """ config = self.GetMethodConfig('Rewrite') return self._RunMethod( config, request, global_params=global_params) def Update(self, request, global_params=None, download=None): """Updates an object's metadata. Args: request: (StorageObjectsUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments download: (Download, default: None) If present, download data from the request via this stream. Returns: (Object) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params, download=download) def WatchAll(self, request, global_params=None): """Watch for changes on all objects in a bucket. Args: request: (StorageObjectsWatchAllRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Channel) The response message. """ config = self.GetMethodConfig('WatchAll') return self._RunMethod( config, request, global_params=global_params) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/clients/storage/storage_v1_messages.py000066400000000000000000002273421272340251000324500ustar00rootroot00000000000000"""Generated message classes for storage version v1. Stores and retrieves potentially large, immutable data objects. """ # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.protorpclite import message_types as _message_types from apitools.base.protorpclite import messages as _messages from apitools.base.py import encoding from apitools.base.py import extra_types package = 'storage' class Bucket(_messages.Message): """A bucket. Messages: CorsValueListEntry: A CorsValueListEntry object. LifecycleValue: The bucket's lifecycle configuration. See lifecycle management for more information. LoggingValue: The bucket's logging configuration, which defines the destination bucket and optional name prefix for the current bucket's logs. OwnerValue: The owner of the bucket. This is always the project team's owner group. VersioningValue: The bucket's versioning configuration. WebsiteValue: The bucket's website configuration. Fields: acl: Access controls on the bucket. cors: The bucket's Cross-Origin Resource Sharing (CORS) configuration. defaultObjectAcl: Default access controls to apply to new objects when no ACL is provided. etag: HTTP 1.1 Entity tag for the bucket. id: The ID of the bucket. kind: The kind of item this is. For buckets, this is always storage#bucket. lifecycle: The bucket's lifecycle configuration. See lifecycle management for more information. location: The location of the bucket. Object data for objects in the bucket resides in physical storage within this region. Defaults to US. See the developer's guide for the authoritative list. logging: The bucket's logging configuration, which defines the destination bucket and optional name prefix for the current bucket's logs. metageneration: The metadata generation of this bucket. name: The name of the bucket. owner: The owner of the bucket. This is always the project team's owner group. projectNumber: The project number of the project the bucket belongs to. selfLink: The URI of this bucket. storageClass: The bucket's storage class. This defines how objects in the bucket are stored and determines the SLA and the cost of storage. Values include STANDARD, NEARLINE and DURABLE_REDUCED_AVAILABILITY. Defaults to STANDARD. For more information, see storage classes. timeCreated: The creation time of the bucket in RFC 3339 format. updated: The modification time of the bucket in RFC 3339 format. versioning: The bucket's versioning configuration. website: The bucket's website configuration. """ class CorsValueListEntry(_messages.Message): """A CorsValueListEntry object. Fields: maxAgeSeconds: The value, in seconds, to return in the Access-Control- Max-Age header used in preflight responses. method: The list of HTTP methods on which to include CORS response headers, (GET, OPTIONS, POST, etc) Note: "*" is permitted in the list of methods, and means "any method". origin: The list of Origins eligible to receive CORS response headers. Note: "*" is permitted in the list of origins, and means "any Origin". responseHeader: The list of HTTP headers other than the simple response headers to give permission for the user-agent to share across domains. """ maxAgeSeconds = _messages.IntegerField(1, variant=_messages.Variant.INT32) method = _messages.StringField(2, repeated=True) origin = _messages.StringField(3, repeated=True) responseHeader = _messages.StringField(4, repeated=True) class LifecycleValue(_messages.Message): """The bucket's lifecycle configuration. See lifecycle management for more information. Messages: RuleValueListEntry: A RuleValueListEntry object. Fields: rule: A lifecycle management rule, which is made of an action to take and the condition(s) under which the action will be taken. """ class RuleValueListEntry(_messages.Message): """A RuleValueListEntry object. Messages: ActionValue: The action to take. ConditionValue: The condition(s) under which the action will be taken. Fields: action: The action to take. condition: The condition(s) under which the action will be taken. """ class ActionValue(_messages.Message): """The action to take. Fields: type: Type of the action. Currently, only Delete is supported. """ type = _messages.StringField(1) class ConditionValue(_messages.Message): """The condition(s) under which the action will be taken. Fields: age: Age of an object (in days). This condition is satisfied when an object reaches the specified age. createdBefore: A date in RFC 3339 format with only the date part (for instance, "2013-01-15"). This condition is satisfied when an object is created before midnight of the specified date in UTC. isLive: Relevant only for versioned objects. If the value is true, this condition matches live objects; if the value is false, it matches archived objects. numNewerVersions: Relevant only for versioned objects. If the value is N, this condition is satisfied when there are at least N versions (including the live version) newer than this version of the object. """ age = _messages.IntegerField(1, variant=_messages.Variant.INT32) createdBefore = extra_types.DateField(2) isLive = _messages.BooleanField(3) numNewerVersions = _messages.IntegerField(4, variant=_messages.Variant.INT32) action = _messages.MessageField('ActionValue', 1) condition = _messages.MessageField('ConditionValue', 2) rule = _messages.MessageField('RuleValueListEntry', 1, repeated=True) class LoggingValue(_messages.Message): """The bucket's logging configuration, which defines the destination bucket and optional name prefix for the current bucket's logs. Fields: logBucket: The destination bucket where the current bucket's logs should be placed. logObjectPrefix: A prefix for log object names. """ logBucket = _messages.StringField(1) logObjectPrefix = _messages.StringField(2) class OwnerValue(_messages.Message): """The owner of the bucket. This is always the project team's owner group. Fields: entity: The entity, in the form project-owner-projectId. entityId: The ID for the entity. """ entity = _messages.StringField(1) entityId = _messages.StringField(2) class VersioningValue(_messages.Message): """The bucket's versioning configuration. Fields: enabled: While set to true, versioning is fully enabled for this bucket. """ enabled = _messages.BooleanField(1) class WebsiteValue(_messages.Message): """The bucket's website configuration. Fields: mainPageSuffix: Behaves as the bucket's directory index where missing objects are treated as potential directories. notFoundPage: The custom object to return when a requested resource is not found. """ mainPageSuffix = _messages.StringField(1) notFoundPage = _messages.StringField(2) acl = _messages.MessageField('BucketAccessControl', 1, repeated=True) cors = _messages.MessageField('CorsValueListEntry', 2, repeated=True) defaultObjectAcl = _messages.MessageField('ObjectAccessControl', 3, repeated=True) etag = _messages.StringField(4) id = _messages.StringField(5) kind = _messages.StringField(6, default=u'storage#bucket') lifecycle = _messages.MessageField('LifecycleValue', 7) location = _messages.StringField(8) logging = _messages.MessageField('LoggingValue', 9) metageneration = _messages.IntegerField(10) name = _messages.StringField(11) owner = _messages.MessageField('OwnerValue', 12) projectNumber = _messages.IntegerField(13, variant=_messages.Variant.UINT64) selfLink = _messages.StringField(14) storageClass = _messages.StringField(15) timeCreated = _message_types.DateTimeField(16) updated = _message_types.DateTimeField(17) versioning = _messages.MessageField('VersioningValue', 18) website = _messages.MessageField('WebsiteValue', 19) class BucketAccessControl(_messages.Message): """An access-control entry. Messages: ProjectTeamValue: The project team associated with the entity, if any. Fields: bucket: The name of the bucket. domain: The domain associated with the entity, if any. email: The email address associated with the entity, if any. entity: The entity holding the permission, in one of the following forms: - user-userId - user-email - group-groupId - group-email - domain- domain - project-team-projectId - allUsers - allAuthenticatedUsers Examples: - The user liz@example.com would be user-liz@example.com. - The group example@googlegroups.com would be group- example@googlegroups.com. - To refer to all members of the Google Apps for Business domain example.com, the entity would be domain-example.com. entityId: The ID for the entity, if any. etag: HTTP 1.1 Entity tag for the access-control entry. id: The ID of the access-control entry. kind: The kind of item this is. For bucket access control entries, this is always storage#bucketAccessControl. projectTeam: The project team associated with the entity, if any. role: The access permission for the entity. Can be READER, WRITER, or OWNER. selfLink: The link to this access-control entry. """ class ProjectTeamValue(_messages.Message): """The project team associated with the entity, if any. Fields: projectNumber: The project number. team: The team. Can be owners, editors, or viewers. """ projectNumber = _messages.StringField(1) team = _messages.StringField(2) bucket = _messages.StringField(1) domain = _messages.StringField(2) email = _messages.StringField(3) entity = _messages.StringField(4) entityId = _messages.StringField(5) etag = _messages.StringField(6) id = _messages.StringField(7) kind = _messages.StringField(8, default=u'storage#bucketAccessControl') projectTeam = _messages.MessageField('ProjectTeamValue', 9) role = _messages.StringField(10) selfLink = _messages.StringField(11) class BucketAccessControls(_messages.Message): """An access-control list. Fields: items: The list of items. kind: The kind of item this is. For lists of bucket access control entries, this is always storage#bucketAccessControls. """ items = _messages.MessageField('BucketAccessControl', 1, repeated=True) kind = _messages.StringField(2, default=u'storage#bucketAccessControls') class Buckets(_messages.Message): """A list of buckets. Fields: items: The list of items. kind: The kind of item this is. For lists of buckets, this is always storage#buckets. nextPageToken: The continuation token, used to page through large result sets. Provide this value in a subsequent request to return the next page of results. """ items = _messages.MessageField('Bucket', 1, repeated=True) kind = _messages.StringField(2, default=u'storage#buckets') nextPageToken = _messages.StringField(3) class Channel(_messages.Message): """An notification channel used to watch for resource changes. Messages: ParamsValue: Additional parameters controlling delivery channel behavior. Optional. Fields: address: The address where notifications are delivered for this channel. expiration: Date and time of notification channel expiration, expressed as a Unix timestamp, in milliseconds. Optional. id: A UUID or similar unique string that identifies this channel. kind: Identifies this as a notification channel used to watch for changes to a resource. Value: the fixed string "api#channel". params: Additional parameters controlling delivery channel behavior. Optional. payload: A Boolean value to indicate whether payload is wanted. Optional. resourceId: An opaque ID that identifies the resource being watched on this channel. Stable across different API versions. resourceUri: A version-specific identifier for the watched resource. token: An arbitrary string delivered to the target address with each notification delivered over this channel. Optional. type: The type of delivery mechanism used for this channel. """ @encoding.MapUnrecognizedFields('additionalProperties') class ParamsValue(_messages.Message): """Additional parameters controlling delivery channel behavior. Optional. Messages: AdditionalProperty: An additional property for a ParamsValue object. Fields: additionalProperties: Declares a new parameter by name. """ class AdditionalProperty(_messages.Message): """An additional property for a ParamsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) address = _messages.StringField(1) expiration = _messages.IntegerField(2) id = _messages.StringField(3) kind = _messages.StringField(4, default=u'api#channel') params = _messages.MessageField('ParamsValue', 5) payload = _messages.BooleanField(6) resourceId = _messages.StringField(7) resourceUri = _messages.StringField(8) token = _messages.StringField(9) type = _messages.StringField(10) class ComposeRequest(_messages.Message): """A Compose request. Messages: SourceObjectsValueListEntry: A SourceObjectsValueListEntry object. Fields: destination: Properties of the resulting object. kind: The kind of item this is. sourceObjects: The list of source objects that will be concatenated into a single object. """ class SourceObjectsValueListEntry(_messages.Message): """A SourceObjectsValueListEntry object. Messages: ObjectPreconditionsValue: Conditions that must be met for this operation to execute. Fields: generation: The generation of this object to use as the source. name: The source object's name. The source object's bucket is implicitly the destination bucket. objectPreconditions: Conditions that must be met for this operation to execute. """ class ObjectPreconditionsValue(_messages.Message): """Conditions that must be met for this operation to execute. Fields: ifGenerationMatch: Only perform the composition if the generation of the source object that would be used matches this value. If this value and a generation are both specified, they must be the same value or the call will fail. """ ifGenerationMatch = _messages.IntegerField(1) generation = _messages.IntegerField(1) name = _messages.StringField(2) objectPreconditions = _messages.MessageField('ObjectPreconditionsValue', 3) destination = _messages.MessageField('Object', 1) kind = _messages.StringField(2, default=u'storage#composeRequest') sourceObjects = _messages.MessageField('SourceObjectsValueListEntry', 3, repeated=True) class Object(_messages.Message): """An object. Messages: CustomerEncryptionValue: Metadata of customer-supplied encryption key, if the object is encrypted by such a key. MetadataValue: User-provided metadata, in key/value pairs. OwnerValue: The owner of the object. This will always be the uploader of the object. Fields: acl: Access controls on the object. bucket: The name of the bucket containing this object. cacheControl: Cache-Control directive for the object data. componentCount: Number of underlying components that make up this object. Components are accumulated by compose operations. contentDisposition: Content-Disposition of the object data. contentEncoding: Content-Encoding of the object data. contentLanguage: Content-Language of the object data. contentType: Content-Type of the object data. crc32c: CRC32c checksum, as described in RFC 4960, Appendix B; encoded using base64 in big-endian byte order. For more information about using the CRC32c checksum, see Hashes and ETags: Best Practices. customerEncryption: Metadata of customer-supplied encryption key, if the object is encrypted by such a key. etag: HTTP 1.1 Entity tag for the object. generation: The content generation of this object. Used for object versioning. id: The ID of the object. kind: The kind of item this is. For objects, this is always storage#object. md5Hash: MD5 hash of the data; encoded using base64. For more information about using the MD5 hash, see Hashes and ETags: Best Practices. mediaLink: Media download link. metadata: User-provided metadata, in key/value pairs. metageneration: The version of the metadata for this object at this generation. Used for preconditions and for detecting changes in metadata. A metageneration number is only meaningful in the context of a particular generation of a particular object. name: The name of this object. Required if not specified by URL parameter. owner: The owner of the object. This will always be the uploader of the object. selfLink: The link to this object. size: Content-Length of the data in bytes. storageClass: Storage class of the object. timeCreated: The creation time of the object in RFC 3339 format. timeDeleted: The deletion time of the object in RFC 3339 format. Will be returned if and only if this version of the object has been deleted. updated: The modification time of the object metadata in RFC 3339 format. """ class CustomerEncryptionValue(_messages.Message): """Metadata of customer-supplied encryption key, if the object is encrypted by such a key. Fields: encryptionAlgorithm: The encryption algorithm. keySha256: SHA256 hash value of the encryption key. """ encryptionAlgorithm = _messages.StringField(1) keySha256 = _messages.StringField(2) @encoding.MapUnrecognizedFields('additionalProperties') class MetadataValue(_messages.Message): """User-provided metadata, in key/value pairs. Messages: AdditionalProperty: An additional property for a MetadataValue object. Fields: additionalProperties: An individual metadata entry. """ class AdditionalProperty(_messages.Message): """An additional property for a MetadataValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) class OwnerValue(_messages.Message): """The owner of the object. This will always be the uploader of the object. Fields: entity: The entity, in the form user-userId. entityId: The ID for the entity. """ entity = _messages.StringField(1) entityId = _messages.StringField(2) acl = _messages.MessageField('ObjectAccessControl', 1, repeated=True) bucket = _messages.StringField(2) cacheControl = _messages.StringField(3) componentCount = _messages.IntegerField(4, variant=_messages.Variant.INT32) contentDisposition = _messages.StringField(5) contentEncoding = _messages.StringField(6) contentLanguage = _messages.StringField(7) contentType = _messages.StringField(8) crc32c = _messages.StringField(9) customerEncryption = _messages.MessageField('CustomerEncryptionValue', 10) etag = _messages.StringField(11) generation = _messages.IntegerField(12) id = _messages.StringField(13) kind = _messages.StringField(14, default=u'storage#object') md5Hash = _messages.StringField(15) mediaLink = _messages.StringField(16) metadata = _messages.MessageField('MetadataValue', 17) metageneration = _messages.IntegerField(18) name = _messages.StringField(19) owner = _messages.MessageField('OwnerValue', 20) selfLink = _messages.StringField(21) size = _messages.IntegerField(22, variant=_messages.Variant.UINT64) storageClass = _messages.StringField(23) timeCreated = _message_types.DateTimeField(24) timeDeleted = _message_types.DateTimeField(25) updated = _message_types.DateTimeField(26) class ObjectAccessControl(_messages.Message): """An access-control entry. Messages: ProjectTeamValue: The project team associated with the entity, if any. Fields: bucket: The name of the bucket. domain: The domain associated with the entity, if any. email: The email address associated with the entity, if any. entity: The entity holding the permission, in one of the following forms: - user-userId - user-email - group-groupId - group-email - domain- domain - project-team-projectId - allUsers - allAuthenticatedUsers Examples: - The user liz@example.com would be user-liz@example.com. - The group example@googlegroups.com would be group- example@googlegroups.com. - To refer to all members of the Google Apps for Business domain example.com, the entity would be domain-example.com. entityId: The ID for the entity, if any. etag: HTTP 1.1 Entity tag for the access-control entry. generation: The content generation of the object. id: The ID of the access-control entry. kind: The kind of item this is. For object access control entries, this is always storage#objectAccessControl. object: The name of the object. projectTeam: The project team associated with the entity, if any. role: The access permission for the entity. Can be READER or OWNER. selfLink: The link to this access-control entry. """ class ProjectTeamValue(_messages.Message): """The project team associated with the entity, if any. Fields: projectNumber: The project number. team: The team. Can be owners, editors, or viewers. """ projectNumber = _messages.StringField(1) team = _messages.StringField(2) bucket = _messages.StringField(1) domain = _messages.StringField(2) email = _messages.StringField(3) entity = _messages.StringField(4) entityId = _messages.StringField(5) etag = _messages.StringField(6) generation = _messages.IntegerField(7) id = _messages.StringField(8) kind = _messages.StringField(9, default=u'storage#objectAccessControl') object = _messages.StringField(10) projectTeam = _messages.MessageField('ProjectTeamValue', 11) role = _messages.StringField(12) selfLink = _messages.StringField(13) class ObjectAccessControls(_messages.Message): """An access-control list. Fields: items: The list of items. kind: The kind of item this is. For lists of object access control entries, this is always storage#objectAccessControls. """ items = _messages.MessageField('extra_types.JsonValue', 1, repeated=True) kind = _messages.StringField(2, default=u'storage#objectAccessControls') class Objects(_messages.Message): """A list of objects. Fields: items: The list of items. kind: The kind of item this is. For lists of objects, this is always storage#objects. nextPageToken: The continuation token, used to page through large result sets. Provide this value in a subsequent request to return the next page of results. prefixes: The list of prefixes of objects matching-but-not-listed up to and including the requested delimiter. """ items = _messages.MessageField('Object', 1, repeated=True) kind = _messages.StringField(2, default=u'storage#objects') nextPageToken = _messages.StringField(3) prefixes = _messages.StringField(4, repeated=True) class RewriteResponse(_messages.Message): """A rewrite response. Fields: done: true if the copy is finished; otherwise, false if the copy is in progress. This property is always present in the response. kind: The kind of item this is. objectSize: The total size of the object being copied in bytes. This property is always present in the response. resource: A resource containing the metadata for the copied-to object. This property is present in the response only when copying completes. rewriteToken: A token to use in subsequent requests to continue copying data. This token is present in the response only when there is more data to copy. totalBytesRewritten: The total bytes written so far, which can be used to provide a waiting user with a progress indicator. This property is always present in the response. """ done = _messages.BooleanField(1) kind = _messages.StringField(2, default=u'storage#rewriteResponse') objectSize = _messages.IntegerField(3, variant=_messages.Variant.UINT64) resource = _messages.MessageField('Object', 4) rewriteToken = _messages.StringField(5) totalBytesRewritten = _messages.IntegerField(6, variant=_messages.Variant.UINT64) class StandardQueryParameters(_messages.Message): """Query parameters accepted by all methods. Enums: AltValueValuesEnum: Data format for the response. Fields: alt: Data format for the response. fields: Selector specifying which fields to include in a partial response. key: API key. Your API key identifies your project and provides you with API access, quota, and reports. Required unless you provide an OAuth 2.0 token. oauth_token: OAuth 2.0 token for the current user. prettyPrint: Returns response with indentations and line breaks. quotaUser: Available to use for quota purposes for server-side applications. Can be any arbitrary string assigned to a user, but should not exceed 40 characters. Overrides userIp if both are provided. trace: A tracing token of the form "token:" to include in api requests. userIp: IP address of the site where the request originates. Use this if you want to enforce per-user limits. """ class AltValueValuesEnum(_messages.Enum): """Data format for the response. Values: json: Responses with Content-Type of application/json """ json = 0 alt = _messages.EnumField('AltValueValuesEnum', 1, default=u'json') fields = _messages.StringField(2) key = _messages.StringField(3) oauth_token = _messages.StringField(4) prettyPrint = _messages.BooleanField(5, default=True) quotaUser = _messages.StringField(6) trace = _messages.StringField(7) userIp = _messages.StringField(8) class StorageBucketAccessControlsDeleteRequest(_messages.Message): """A StorageBucketAccessControlsDeleteRequest object. Fields: bucket: Name of a bucket. entity: The entity holding the permission. Can be user-userId, user- emailAddress, group-groupId, group-emailAddress, allUsers, or allAuthenticatedUsers. """ bucket = _messages.StringField(1, required=True) entity = _messages.StringField(2, required=True) class StorageBucketAccessControlsDeleteResponse(_messages.Message): """An empty StorageBucketAccessControlsDelete response.""" class StorageBucketAccessControlsGetRequest(_messages.Message): """A StorageBucketAccessControlsGetRequest object. Fields: bucket: Name of a bucket. entity: The entity holding the permission. Can be user-userId, user- emailAddress, group-groupId, group-emailAddress, allUsers, or allAuthenticatedUsers. """ bucket = _messages.StringField(1, required=True) entity = _messages.StringField(2, required=True) class StorageBucketAccessControlsListRequest(_messages.Message): """A StorageBucketAccessControlsListRequest object. Fields: bucket: Name of a bucket. """ bucket = _messages.StringField(1, required=True) class StorageBucketsDeleteRequest(_messages.Message): """A StorageBucketsDeleteRequest object. Fields: bucket: Name of a bucket. ifMetagenerationMatch: If set, only deletes the bucket if its metageneration matches this value. ifMetagenerationNotMatch: If set, only deletes the bucket if its metageneration does not match this value. """ bucket = _messages.StringField(1, required=True) ifMetagenerationMatch = _messages.IntegerField(2) ifMetagenerationNotMatch = _messages.IntegerField(3) class StorageBucketsDeleteResponse(_messages.Message): """An empty StorageBucketsDelete response.""" class StorageBucketsGetRequest(_messages.Message): """A StorageBucketsGetRequest object. Enums: ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl. Fields: bucket: Name of a bucket. ifMetagenerationMatch: Makes the return of the bucket metadata conditional on whether the bucket's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the return of the bucket metadata conditional on whether the bucket's current metageneration does not match the given value. projection: Set of properties to return. Defaults to noAcl. """ class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl. Values: full: Include all properties. noAcl: Omit acl and defaultObjectAcl properties. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) ifMetagenerationMatch = _messages.IntegerField(2) ifMetagenerationNotMatch = _messages.IntegerField(3) projection = _messages.EnumField('ProjectionValueValuesEnum', 4) class StorageBucketsInsertRequest(_messages.Message): """A StorageBucketsInsertRequest object. Enums: PredefinedAclValueValuesEnum: Apply a predefined set of access controls to this bucket. PredefinedDefaultObjectAclValueValuesEnum: Apply a predefined set of default object access controls to this bucket. ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl, unless the bucket resource specifies acl or defaultObjectAcl properties, when it defaults to full. Fields: bucket: A Bucket resource to be passed as the request body. predefinedAcl: Apply a predefined set of access controls to this bucket. predefinedDefaultObjectAcl: Apply a predefined set of default object access controls to this bucket. project: A valid API project identifier. projection: Set of properties to return. Defaults to noAcl, unless the bucket resource specifies acl or defaultObjectAcl properties, when it defaults to full. """ class PredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to this bucket. Values: authenticatedRead: Project team owners get OWNER access, and allAuthenticatedUsers get READER access. private: Project team owners get OWNER access. projectPrivate: Project team members get access according to their roles. publicRead: Project team owners get OWNER access, and allUsers get READER access. publicReadWrite: Project team owners get OWNER access, and allUsers get WRITER access. """ authenticatedRead = 0 private = 1 projectPrivate = 2 publicRead = 3 publicReadWrite = 4 class PredefinedDefaultObjectAclValueValuesEnum(_messages.Enum): """Apply a predefined set of default object access controls to this bucket. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl, unless the bucket resource specifies acl or defaultObjectAcl properties, when it defaults to full. Values: full: Include all properties. noAcl: Omit acl and defaultObjectAcl properties. """ full = 0 noAcl = 1 bucket = _messages.MessageField('Bucket', 1) predefinedAcl = _messages.EnumField('PredefinedAclValueValuesEnum', 2) predefinedDefaultObjectAcl = _messages.EnumField('PredefinedDefaultObjectAclValueValuesEnum', 3) project = _messages.StringField(4, required=True) projection = _messages.EnumField('ProjectionValueValuesEnum', 5) class StorageBucketsListRequest(_messages.Message): """A StorageBucketsListRequest object. Enums: ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl. Fields: maxResults: Maximum number of buckets to return. pageToken: A previously-returned page token representing part of the larger set of results to view. prefix: Filter results to buckets whose names begin with this prefix. project: A valid API project identifier. projection: Set of properties to return. Defaults to noAcl. """ class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl. Values: full: Include all properties. noAcl: Omit acl and defaultObjectAcl properties. """ full = 0 noAcl = 1 maxResults = _messages.IntegerField(1, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(2) prefix = _messages.StringField(3) project = _messages.StringField(4, required=True) projection = _messages.EnumField('ProjectionValueValuesEnum', 5) class StorageBucketsPatchRequest(_messages.Message): """A StorageBucketsPatchRequest object. Enums: PredefinedAclValueValuesEnum: Apply a predefined set of access controls to this bucket. PredefinedDefaultObjectAclValueValuesEnum: Apply a predefined set of default object access controls to this bucket. ProjectionValueValuesEnum: Set of properties to return. Defaults to full. Fields: bucket: Name of a bucket. bucketResource: A Bucket resource to be passed as the request body. ifMetagenerationMatch: Makes the return of the bucket metadata conditional on whether the bucket's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the return of the bucket metadata conditional on whether the bucket's current metageneration does not match the given value. predefinedAcl: Apply a predefined set of access controls to this bucket. predefinedDefaultObjectAcl: Apply a predefined set of default object access controls to this bucket. projection: Set of properties to return. Defaults to full. """ class PredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to this bucket. Values: authenticatedRead: Project team owners get OWNER access, and allAuthenticatedUsers get READER access. private: Project team owners get OWNER access. projectPrivate: Project team members get access according to their roles. publicRead: Project team owners get OWNER access, and allUsers get READER access. publicReadWrite: Project team owners get OWNER access, and allUsers get WRITER access. """ authenticatedRead = 0 private = 1 projectPrivate = 2 publicRead = 3 publicReadWrite = 4 class PredefinedDefaultObjectAclValueValuesEnum(_messages.Enum): """Apply a predefined set of default object access controls to this bucket. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to full. Values: full: Include all properties. noAcl: Omit acl and defaultObjectAcl properties. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) bucketResource = _messages.MessageField('Bucket', 2) ifMetagenerationMatch = _messages.IntegerField(3) ifMetagenerationNotMatch = _messages.IntegerField(4) predefinedAcl = _messages.EnumField('PredefinedAclValueValuesEnum', 5) predefinedDefaultObjectAcl = _messages.EnumField('PredefinedDefaultObjectAclValueValuesEnum', 6) projection = _messages.EnumField('ProjectionValueValuesEnum', 7) class StorageBucketsUpdateRequest(_messages.Message): """A StorageBucketsUpdateRequest object. Enums: PredefinedAclValueValuesEnum: Apply a predefined set of access controls to this bucket. PredefinedDefaultObjectAclValueValuesEnum: Apply a predefined set of default object access controls to this bucket. ProjectionValueValuesEnum: Set of properties to return. Defaults to full. Fields: bucket: Name of a bucket. bucketResource: A Bucket resource to be passed as the request body. ifMetagenerationMatch: Makes the return of the bucket metadata conditional on whether the bucket's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the return of the bucket metadata conditional on whether the bucket's current metageneration does not match the given value. predefinedAcl: Apply a predefined set of access controls to this bucket. predefinedDefaultObjectAcl: Apply a predefined set of default object access controls to this bucket. projection: Set of properties to return. Defaults to full. """ class PredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to this bucket. Values: authenticatedRead: Project team owners get OWNER access, and allAuthenticatedUsers get READER access. private: Project team owners get OWNER access. projectPrivate: Project team members get access according to their roles. publicRead: Project team owners get OWNER access, and allUsers get READER access. publicReadWrite: Project team owners get OWNER access, and allUsers get WRITER access. """ authenticatedRead = 0 private = 1 projectPrivate = 2 publicRead = 3 publicReadWrite = 4 class PredefinedDefaultObjectAclValueValuesEnum(_messages.Enum): """Apply a predefined set of default object access controls to this bucket. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to full. Values: full: Include all properties. noAcl: Omit acl and defaultObjectAcl properties. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) bucketResource = _messages.MessageField('Bucket', 2) ifMetagenerationMatch = _messages.IntegerField(3) ifMetagenerationNotMatch = _messages.IntegerField(4) predefinedAcl = _messages.EnumField('PredefinedAclValueValuesEnum', 5) predefinedDefaultObjectAcl = _messages.EnumField('PredefinedDefaultObjectAclValueValuesEnum', 6) projection = _messages.EnumField('ProjectionValueValuesEnum', 7) class StorageChannelsStopResponse(_messages.Message): """An empty StorageChannelsStop response.""" class StorageDefaultObjectAccessControlsDeleteRequest(_messages.Message): """A StorageDefaultObjectAccessControlsDeleteRequest object. Fields: bucket: Name of a bucket. entity: The entity holding the permission. Can be user-userId, user- emailAddress, group-groupId, group-emailAddress, allUsers, or allAuthenticatedUsers. """ bucket = _messages.StringField(1, required=True) entity = _messages.StringField(2, required=True) class StorageDefaultObjectAccessControlsDeleteResponse(_messages.Message): """An empty StorageDefaultObjectAccessControlsDelete response.""" class StorageDefaultObjectAccessControlsGetRequest(_messages.Message): """A StorageDefaultObjectAccessControlsGetRequest object. Fields: bucket: Name of a bucket. entity: The entity holding the permission. Can be user-userId, user- emailAddress, group-groupId, group-emailAddress, allUsers, or allAuthenticatedUsers. """ bucket = _messages.StringField(1, required=True) entity = _messages.StringField(2, required=True) class StorageDefaultObjectAccessControlsListRequest(_messages.Message): """A StorageDefaultObjectAccessControlsListRequest object. Fields: bucket: Name of a bucket. ifMetagenerationMatch: If present, only return default ACL listing if the bucket's current metageneration matches this value. ifMetagenerationNotMatch: If present, only return default ACL listing if the bucket's current metageneration does not match the given value. """ bucket = _messages.StringField(1, required=True) ifMetagenerationMatch = _messages.IntegerField(2) ifMetagenerationNotMatch = _messages.IntegerField(3) class StorageObjectAccessControlsDeleteRequest(_messages.Message): """A StorageObjectAccessControlsDeleteRequest object. Fields: bucket: Name of a bucket. entity: The entity holding the permission. Can be user-userId, user- emailAddress, group-groupId, group-emailAddress, allUsers, or allAuthenticatedUsers. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. """ bucket = _messages.StringField(1, required=True) entity = _messages.StringField(2, required=True) generation = _messages.IntegerField(3) object = _messages.StringField(4, required=True) class StorageObjectAccessControlsDeleteResponse(_messages.Message): """An empty StorageObjectAccessControlsDelete response.""" class StorageObjectAccessControlsGetRequest(_messages.Message): """A StorageObjectAccessControlsGetRequest object. Fields: bucket: Name of a bucket. entity: The entity holding the permission. Can be user-userId, user- emailAddress, group-groupId, group-emailAddress, allUsers, or allAuthenticatedUsers. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. """ bucket = _messages.StringField(1, required=True) entity = _messages.StringField(2, required=True) generation = _messages.IntegerField(3) object = _messages.StringField(4, required=True) class StorageObjectAccessControlsInsertRequest(_messages.Message): """A StorageObjectAccessControlsInsertRequest object. Fields: bucket: Name of a bucket. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. objectAccessControl: A ObjectAccessControl resource to be passed as the request body. """ bucket = _messages.StringField(1, required=True) generation = _messages.IntegerField(2) object = _messages.StringField(3, required=True) objectAccessControl = _messages.MessageField('ObjectAccessControl', 4) class StorageObjectAccessControlsListRequest(_messages.Message): """A StorageObjectAccessControlsListRequest object. Fields: bucket: Name of a bucket. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. """ bucket = _messages.StringField(1, required=True) generation = _messages.IntegerField(2) object = _messages.StringField(3, required=True) class StorageObjectAccessControlsPatchRequest(_messages.Message): """A StorageObjectAccessControlsPatchRequest object. Fields: bucket: Name of a bucket. entity: The entity holding the permission. Can be user-userId, user- emailAddress, group-groupId, group-emailAddress, allUsers, or allAuthenticatedUsers. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. objectAccessControl: A ObjectAccessControl resource to be passed as the request body. """ bucket = _messages.StringField(1, required=True) entity = _messages.StringField(2, required=True) generation = _messages.IntegerField(3) object = _messages.StringField(4, required=True) objectAccessControl = _messages.MessageField('ObjectAccessControl', 5) class StorageObjectAccessControlsUpdateRequest(_messages.Message): """A StorageObjectAccessControlsUpdateRequest object. Fields: bucket: Name of a bucket. entity: The entity holding the permission. Can be user-userId, user- emailAddress, group-groupId, group-emailAddress, allUsers, or allAuthenticatedUsers. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. objectAccessControl: A ObjectAccessControl resource to be passed as the request body. """ bucket = _messages.StringField(1, required=True) entity = _messages.StringField(2, required=True) generation = _messages.IntegerField(3) object = _messages.StringField(4, required=True) objectAccessControl = _messages.MessageField('ObjectAccessControl', 5) class StorageObjectsComposeRequest(_messages.Message): """A StorageObjectsComposeRequest object. Enums: DestinationPredefinedAclValueValuesEnum: Apply a predefined set of access controls to the destination object. Fields: composeRequest: A ComposeRequest resource to be passed as the request body. destinationBucket: Name of the bucket in which to store the new object. destinationObject: Name of the new object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. destinationPredefinedAcl: Apply a predefined set of access controls to the destination object. ifGenerationMatch: Makes the operation conditional on whether the object's current generation matches the given value. ifMetagenerationMatch: Makes the operation conditional on whether the object's current metageneration matches the given value. """ class DestinationPredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to the destination object. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 composeRequest = _messages.MessageField('ComposeRequest', 1) destinationBucket = _messages.StringField(2, required=True) destinationObject = _messages.StringField(3, required=True) destinationPredefinedAcl = _messages.EnumField('DestinationPredefinedAclValueValuesEnum', 4) ifGenerationMatch = _messages.IntegerField(5) ifMetagenerationMatch = _messages.IntegerField(6) class StorageObjectsCopyRequest(_messages.Message): """A StorageObjectsCopyRequest object. Enums: DestinationPredefinedAclValueValuesEnum: Apply a predefined set of access controls to the destination object. ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. Fields: destinationBucket: Name of the bucket in which to store the new object. Overrides the provided object metadata's bucket value, if any.For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. destinationObject: Name of the new object. Required when the object metadata is not otherwise provided. Overrides the object metadata's name value, if any. destinationPredefinedAcl: Apply a predefined set of access controls to the destination object. ifGenerationMatch: Makes the operation conditional on whether the destination object's current generation matches the given value. ifGenerationNotMatch: Makes the operation conditional on whether the destination object's current generation does not match the given value. ifMetagenerationMatch: Makes the operation conditional on whether the destination object's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the operation conditional on whether the destination object's current metageneration does not match the given value. ifSourceGenerationMatch: Makes the operation conditional on whether the source object's generation matches the given value. ifSourceGenerationNotMatch: Makes the operation conditional on whether the source object's generation does not match the given value. ifSourceMetagenerationMatch: Makes the operation conditional on whether the source object's current metageneration matches the given value. ifSourceMetagenerationNotMatch: Makes the operation conditional on whether the source object's current metageneration does not match the given value. object: A Object resource to be passed as the request body. projection: Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. sourceBucket: Name of the bucket in which to find the source object. sourceGeneration: If present, selects a specific revision of the source object (as opposed to the latest version, the default). sourceObject: Name of the source object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. """ class DestinationPredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to the destination object. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. Values: full: Include all properties. noAcl: Omit the acl property. """ full = 0 noAcl = 1 destinationBucket = _messages.StringField(1, required=True) destinationObject = _messages.StringField(2, required=True) destinationPredefinedAcl = _messages.EnumField('DestinationPredefinedAclValueValuesEnum', 3) ifGenerationMatch = _messages.IntegerField(4) ifGenerationNotMatch = _messages.IntegerField(5) ifMetagenerationMatch = _messages.IntegerField(6) ifMetagenerationNotMatch = _messages.IntegerField(7) ifSourceGenerationMatch = _messages.IntegerField(8) ifSourceGenerationNotMatch = _messages.IntegerField(9) ifSourceMetagenerationMatch = _messages.IntegerField(10) ifSourceMetagenerationNotMatch = _messages.IntegerField(11) object = _messages.MessageField('Object', 12) projection = _messages.EnumField('ProjectionValueValuesEnum', 13) sourceBucket = _messages.StringField(14, required=True) sourceGeneration = _messages.IntegerField(15) sourceObject = _messages.StringField(16, required=True) class StorageObjectsDeleteRequest(_messages.Message): """A StorageObjectsDeleteRequest object. Fields: bucket: Name of the bucket in which the object resides. generation: If present, permanently deletes a specific revision of this object (as opposed to the latest version, the default). ifGenerationMatch: Makes the operation conditional on whether the object's current generation matches the given value. ifGenerationNotMatch: Makes the operation conditional on whether the object's current generation does not match the given value. ifMetagenerationMatch: Makes the operation conditional on whether the object's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the operation conditional on whether the object's current metageneration does not match the given value. object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. """ bucket = _messages.StringField(1, required=True) generation = _messages.IntegerField(2) ifGenerationMatch = _messages.IntegerField(3) ifGenerationNotMatch = _messages.IntegerField(4) ifMetagenerationMatch = _messages.IntegerField(5) ifMetagenerationNotMatch = _messages.IntegerField(6) object = _messages.StringField(7, required=True) class StorageObjectsDeleteResponse(_messages.Message): """An empty StorageObjectsDelete response.""" class StorageObjectsGetRequest(_messages.Message): """A StorageObjectsGetRequest object. Enums: ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl. Fields: bucket: Name of the bucket in which the object resides. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). ifGenerationMatch: Makes the operation conditional on whether the object's generation matches the given value. ifGenerationNotMatch: Makes the operation conditional on whether the object's generation does not match the given value. ifMetagenerationMatch: Makes the operation conditional on whether the object's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the operation conditional on whether the object's current metageneration does not match the given value. object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. projection: Set of properties to return. Defaults to noAcl. """ class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl. Values: full: Include all properties. noAcl: Omit the acl property. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) generation = _messages.IntegerField(2) ifGenerationMatch = _messages.IntegerField(3) ifGenerationNotMatch = _messages.IntegerField(4) ifMetagenerationMatch = _messages.IntegerField(5) ifMetagenerationNotMatch = _messages.IntegerField(6) object = _messages.StringField(7, required=True) projection = _messages.EnumField('ProjectionValueValuesEnum', 8) class StorageObjectsInsertRequest(_messages.Message): """A StorageObjectsInsertRequest object. Enums: PredefinedAclValueValuesEnum: Apply a predefined set of access controls to this object. ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. Fields: bucket: Name of the bucket in which to store the new object. Overrides the provided object metadata's bucket value, if any. contentEncoding: If set, sets the contentEncoding property of the final object to this value. Setting this parameter is equivalent to setting the contentEncoding metadata property. This can be useful when uploading an object with uploadType=media to indicate the encoding of the content being uploaded. ifGenerationMatch: Makes the operation conditional on whether the object's current generation matches the given value. ifGenerationNotMatch: Makes the operation conditional on whether the object's current generation does not match the given value. ifMetagenerationMatch: Makes the operation conditional on whether the object's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the operation conditional on whether the object's current metageneration does not match the given value. name: Name of the object. Required when the object metadata is not otherwise provided. Overrides the object metadata's name value, if any. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. object: A Object resource to be passed as the request body. predefinedAcl: Apply a predefined set of access controls to this object. projection: Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. """ class PredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to this object. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. Values: full: Include all properties. noAcl: Omit the acl property. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) contentEncoding = _messages.StringField(2) ifGenerationMatch = _messages.IntegerField(3) ifGenerationNotMatch = _messages.IntegerField(4) ifMetagenerationMatch = _messages.IntegerField(5) ifMetagenerationNotMatch = _messages.IntegerField(6) name = _messages.StringField(7) object = _messages.MessageField('Object', 8) predefinedAcl = _messages.EnumField('PredefinedAclValueValuesEnum', 9) projection = _messages.EnumField('ProjectionValueValuesEnum', 10) class StorageObjectsListRequest(_messages.Message): """A StorageObjectsListRequest object. Enums: ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl. Fields: bucket: Name of the bucket in which to look for objects. delimiter: Returns results in a directory-like mode. items will contain only objects whose names, aside from the prefix, do not contain delimiter. Objects whose names, aside from the prefix, contain delimiter will have their name, truncated after the delimiter, returned in prefixes. Duplicate prefixes are omitted. maxResults: Maximum number of items plus prefixes to return. As duplicate prefixes are omitted, fewer total results may be returned than requested. The default value of this parameter is 1,000 items. pageToken: A previously-returned page token representing part of the larger set of results to view. prefix: Filter results to objects whose names begin with this prefix. projection: Set of properties to return. Defaults to noAcl. versions: If true, lists all versions of an object as distinct results. The default is false. For more information, see Object Versioning. """ class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl. Values: full: Include all properties. noAcl: Omit the acl property. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) delimiter = _messages.StringField(2) maxResults = _messages.IntegerField(3, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(4) prefix = _messages.StringField(5) projection = _messages.EnumField('ProjectionValueValuesEnum', 6) versions = _messages.BooleanField(7) class StorageObjectsPatchRequest(_messages.Message): """A StorageObjectsPatchRequest object. Enums: PredefinedAclValueValuesEnum: Apply a predefined set of access controls to this object. ProjectionValueValuesEnum: Set of properties to return. Defaults to full. Fields: bucket: Name of the bucket in which the object resides. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). ifGenerationMatch: Makes the operation conditional on whether the object's current generation matches the given value. ifGenerationNotMatch: Makes the operation conditional on whether the object's current generation does not match the given value. ifMetagenerationMatch: Makes the operation conditional on whether the object's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the operation conditional on whether the object's current metageneration does not match the given value. object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. objectResource: A Object resource to be passed as the request body. predefinedAcl: Apply a predefined set of access controls to this object. projection: Set of properties to return. Defaults to full. """ class PredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to this object. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to full. Values: full: Include all properties. noAcl: Omit the acl property. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) generation = _messages.IntegerField(2) ifGenerationMatch = _messages.IntegerField(3) ifGenerationNotMatch = _messages.IntegerField(4) ifMetagenerationMatch = _messages.IntegerField(5) ifMetagenerationNotMatch = _messages.IntegerField(6) object = _messages.StringField(7, required=True) objectResource = _messages.MessageField('Object', 8) predefinedAcl = _messages.EnumField('PredefinedAclValueValuesEnum', 9) projection = _messages.EnumField('ProjectionValueValuesEnum', 10) class StorageObjectsRewriteRequest(_messages.Message): """A StorageObjectsRewriteRequest object. Enums: DestinationPredefinedAclValueValuesEnum: Apply a predefined set of access controls to the destination object. ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. Fields: destinationBucket: Name of the bucket in which to store the new object. Overrides the provided object metadata's bucket value, if any. destinationObject: Name of the new object. Required when the object metadata is not otherwise provided. Overrides the object metadata's name value, if any. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. destinationPredefinedAcl: Apply a predefined set of access controls to the destination object. ifGenerationMatch: Makes the operation conditional on whether the destination object's current generation matches the given value. ifGenerationNotMatch: Makes the operation conditional on whether the destination object's current generation does not match the given value. ifMetagenerationMatch: Makes the operation conditional on whether the destination object's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the operation conditional on whether the destination object's current metageneration does not match the given value. ifSourceGenerationMatch: Makes the operation conditional on whether the source object's generation matches the given value. ifSourceGenerationNotMatch: Makes the operation conditional on whether the source object's generation does not match the given value. ifSourceMetagenerationMatch: Makes the operation conditional on whether the source object's current metageneration matches the given value. ifSourceMetagenerationNotMatch: Makes the operation conditional on whether the source object's current metageneration does not match the given value. maxBytesRewrittenPerCall: The maximum number of bytes that will be rewritten per rewrite request. Most callers shouldn't need to specify this parameter - it is primarily in place to support testing. If specified the value must be an integral multiple of 1 MiB (1048576). Also, this only applies to requests where the source and destination span locations and/or storage classes. Finally, this value must not change across rewrite calls else you'll get an error that the rewriteToken is invalid. object: A Object resource to be passed as the request body. projection: Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. rewriteToken: Include this field (from the previous rewrite response) on each rewrite request after the first one, until the rewrite response 'done' flag is true. Calls that provide a rewriteToken can omit all other request fields, but if included those fields must match the values provided in the first rewrite request. sourceBucket: Name of the bucket in which to find the source object. sourceGeneration: If present, selects a specific revision of the source object (as opposed to the latest version, the default). sourceObject: Name of the source object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. """ class DestinationPredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to the destination object. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl, unless the object resource specifies the acl property, when it defaults to full. Values: full: Include all properties. noAcl: Omit the acl property. """ full = 0 noAcl = 1 destinationBucket = _messages.StringField(1, required=True) destinationObject = _messages.StringField(2, required=True) destinationPredefinedAcl = _messages.EnumField('DestinationPredefinedAclValueValuesEnum', 3) ifGenerationMatch = _messages.IntegerField(4) ifGenerationNotMatch = _messages.IntegerField(5) ifMetagenerationMatch = _messages.IntegerField(6) ifMetagenerationNotMatch = _messages.IntegerField(7) ifSourceGenerationMatch = _messages.IntegerField(8) ifSourceGenerationNotMatch = _messages.IntegerField(9) ifSourceMetagenerationMatch = _messages.IntegerField(10) ifSourceMetagenerationNotMatch = _messages.IntegerField(11) maxBytesRewrittenPerCall = _messages.IntegerField(12) object = _messages.MessageField('Object', 13) projection = _messages.EnumField('ProjectionValueValuesEnum', 14) rewriteToken = _messages.StringField(15) sourceBucket = _messages.StringField(16, required=True) sourceGeneration = _messages.IntegerField(17) sourceObject = _messages.StringField(18, required=True) class StorageObjectsUpdateRequest(_messages.Message): """A StorageObjectsUpdateRequest object. Enums: PredefinedAclValueValuesEnum: Apply a predefined set of access controls to this object. ProjectionValueValuesEnum: Set of properties to return. Defaults to full. Fields: bucket: Name of the bucket in which the object resides. generation: If present, selects a specific revision of this object (as opposed to the latest version, the default). ifGenerationMatch: Makes the operation conditional on whether the object's current generation matches the given value. ifGenerationNotMatch: Makes the operation conditional on whether the object's current generation does not match the given value. ifMetagenerationMatch: Makes the operation conditional on whether the object's current metageneration matches the given value. ifMetagenerationNotMatch: Makes the operation conditional on whether the object's current metageneration does not match the given value. object: Name of the object. For information about how to URL encode object names to be path safe, see Encoding URI Path Parts. objectResource: A Object resource to be passed as the request body. predefinedAcl: Apply a predefined set of access controls to this object. projection: Set of properties to return. Defaults to full. """ class PredefinedAclValueValuesEnum(_messages.Enum): """Apply a predefined set of access controls to this object. Values: authenticatedRead: Object owner gets OWNER access, and allAuthenticatedUsers get READER access. bucketOwnerFullControl: Object owner gets OWNER access, and project team owners get OWNER access. bucketOwnerRead: Object owner gets OWNER access, and project team owners get READER access. private: Object owner gets OWNER access. projectPrivate: Object owner gets OWNER access, and project team members get access according to their roles. publicRead: Object owner gets OWNER access, and allUsers get READER access. """ authenticatedRead = 0 bucketOwnerFullControl = 1 bucketOwnerRead = 2 private = 3 projectPrivate = 4 publicRead = 5 class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to full. Values: full: Include all properties. noAcl: Omit the acl property. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) generation = _messages.IntegerField(2) ifGenerationMatch = _messages.IntegerField(3) ifGenerationNotMatch = _messages.IntegerField(4) ifMetagenerationMatch = _messages.IntegerField(5) ifMetagenerationNotMatch = _messages.IntegerField(6) object = _messages.StringField(7, required=True) objectResource = _messages.MessageField('Object', 8) predefinedAcl = _messages.EnumField('PredefinedAclValueValuesEnum', 9) projection = _messages.EnumField('ProjectionValueValuesEnum', 10) class StorageObjectsWatchAllRequest(_messages.Message): """A StorageObjectsWatchAllRequest object. Enums: ProjectionValueValuesEnum: Set of properties to return. Defaults to noAcl. Fields: bucket: Name of the bucket in which to look for objects. channel: A Channel resource to be passed as the request body. delimiter: Returns results in a directory-like mode. items will contain only objects whose names, aside from the prefix, do not contain delimiter. Objects whose names, aside from the prefix, contain delimiter will have their name, truncated after the delimiter, returned in prefixes. Duplicate prefixes are omitted. maxResults: Maximum number of items plus prefixes to return. As duplicate prefixes are omitted, fewer total results may be returned than requested. The default value of this parameter is 1,000 items. pageToken: A previously-returned page token representing part of the larger set of results to view. prefix: Filter results to objects whose names begin with this prefix. projection: Set of properties to return. Defaults to noAcl. versions: If true, lists all versions of an object as distinct results. The default is false. For more information, see Object Versioning. """ class ProjectionValueValuesEnum(_messages.Enum): """Set of properties to return. Defaults to noAcl. Values: full: Include all properties. noAcl: Omit the acl property. """ full = 0 noAcl = 1 bucket = _messages.StringField(1, required=True) channel = _messages.MessageField('Channel', 2) delimiter = _messages.StringField(3) maxResults = _messages.IntegerField(4, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(5) prefix = _messages.StringField(6) projection = _messages.EnumField('ProjectionValueValuesEnum', 7) versions = _messages.BooleanField(8) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/json_value.py000066400000000000000000000107651272340251000255460ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """JSON conversion utility functions.""" from apitools.base.py import extra_types def _get_typed_value_descriptor(obj): """Converts a basic type into a @type/value dictionary. Args: obj: A basestring, bool, int, or float to be converted. Returns: A dictionary containing the keys '@type' and 'value' with the value for the @type of appropriate type. Raises: TypeError: if the Python object has a type that is not supported. """ if isinstance(obj, basestring): type_name = 'Text' elif isinstance(obj, bool): type_name = 'Boolean' elif isinstance(obj, int): type_name = 'Integer' elif isinstance(obj, float): type_name = 'Float' else: raise TypeError('Cannot get a type descriptor for %s.' % repr(obj)) return {'@type': 'http://schema.org/%s' % type_name, 'value': obj} def to_json_value(obj, with_type=False): """Converts Python objects into extra_types.JsonValue objects. Args: obj: Python object to be converted. with_type: If true then the basic types (string, int, float, bool) will be wrapped in @type/value dictionaries. Otherwise the straight value is encoded into a JsonValue. Returns: A JsonValue object using JsonValue, JsonArray and JsonObject types for the corresponding values, lists, or dictionaries. Raises: TypeError: if the Python object contains a type that is not supported. The types supported are str, bool, list, tuple, dict. The Dataflow API requires JsonValue(s) in many places, and it is quite convenient to be able to specify these hierarchical objects using Python syntax. """ if isinstance(obj, (list, tuple)): return extra_types.JsonValue( array_value=extra_types.JsonArray( entries=[to_json_value(o, with_type=with_type) for o in obj])) elif isinstance(obj, dict): json_object = extra_types.JsonObject() for k, v in obj.iteritems(): json_object.properties.append( extra_types.JsonObject.Property( key=k, value=to_json_value(v, with_type=with_type))) return extra_types.JsonValue(object_value=json_object) elif with_type: return to_json_value(_get_typed_value_descriptor(obj), with_type=False) elif isinstance(obj, basestring): return extra_types.JsonValue(string_value=obj) elif isinstance(obj, bool): return extra_types.JsonValue(boolean_value=obj) elif isinstance(obj, int): return extra_types.JsonValue(integer_value=obj) elif isinstance(obj, float): return extra_types.JsonValue(double_value=obj) else: raise TypeError('Cannot convert %s to a JSON value.' % repr(obj)) def from_json_value(v): """Converts extra_types.JsonValue objects into Python objects. Args: v: JsonValue object to be converted. Returns: A Python object structured as values, lists, and dictionaries corresponding to JsonValue, JsonArray and JsonObject types. Raises: TypeError: if the JsonValue object contains a type that is not supported. The types supported are str, bool, list, dict. The Dataflow API returns JsonValue(s) in many places and it is quite convenient to be able to convert these hierarchical objects to much simpler Python objects. """ if isinstance(v, extra_types.JsonValue): if v.string_value is not None: return v.string_value elif v.boolean_value is not None: return v.boolean_value elif v.integer_value is not None: return v.integer_value elif v.double_value is not None: return v.double_value elif v.array_value is not None: return from_json_value(v.array_value) elif v.object_value is not None: return from_json_value(v.object_value) elif v.is_null: return [] elif isinstance(v, extra_types.JsonArray): return [from_json_value(e) for e in v.entries] elif isinstance(v, extra_types.JsonObject): return {p.key: from_json_value(p.value) for p in v.properties} raise TypeError('Cannot convert %s from a JSON value.' % repr(v)) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/json_value_test.py000066400000000000000000000041741272340251000266020ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the json_value module.""" import unittest from apitools.base.py.extra_types import JsonValue from apitools.base.py.extra_types import JsonObject from google.cloud.dataflow.internal.json_value import from_json_value from google.cloud.dataflow.internal.json_value import to_json_value class JsonValueTest(unittest.TestCase): def test_string_to(self): self.assertEquals(JsonValue(string_value='abc'), to_json_value('abc')) def test_true_to(self): self.assertEquals(JsonValue(boolean_value=True), to_json_value(True)) def test_false_to(self): self.assertEquals(JsonValue(boolean_value=False), to_json_value(False)) def test_int_to(self): self.assertEquals(JsonValue(integer_value=14), to_json_value(14)) def test_float_to(self): self.assertEquals(JsonValue(double_value=2.75), to_json_value(2.75)) def test_string_from(self): self.assertEquals('WXYZ', from_json_value(to_json_value('WXYZ'))) def test_true_from(self): self.assertEquals(True, from_json_value(to_json_value(True))) def test_false_from(self): self.assertEquals(False, from_json_value(to_json_value(False))) def test_int_from(self): self.assertEquals(-27, from_json_value(to_json_value(-27))) def test_float_from(self): self.assertEquals(4.5, from_json_value(to_json_value(4.5))) def test_with_type(self): rt = from_json_value(to_json_value('abcd', with_type=True)) self.assertEquals('http://schema.org/Text', rt['@type']) self.assertEquals('abcd', rt['value']) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/module_test.py000066400000000000000000000025461272340251000257230ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Module used to define functions and classes used by the coder unit tests.""" import re class TopClass(object): class NestedClass(object): def __init__(self, datum): self.datum = 'X:%s' % datum class MiddleClass(object): class NestedClass(object): def __init__(self, datum): self.datum = 'Y:%s' % datum def get_lambda_with_globals(): return lambda s: re.findall(r'\w+', s) def get_lambda_with_closure(message): return lambda: 'closure: %s' % message class Xyz(object): """A class to be pickled.""" def foo(self, s): return re.findall(r'\w+', s) def create_class(datum): """Creates an unnamable class to be pickled.""" class Z(object): def get(self): return 'Z:%s' % datum return Z() XYZ_OBJECT = Xyz() DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/pickler.py000066400000000000000000000160301272340251000250210ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Pickler for values, functions, and classes. Pickles created by the pickling library contain non-ASCII characters, so we base64-encode the results so that we can put them in a JSON objects. The pickler is used to embed FlatMap callable objects into the workflow JSON description. The pickler module should be used to pickle functions and modules; for values, the coders.*PickleCoder classes should be used instead. """ import base64 import logging import sys import traceback import types import dill def is_nested_class(cls): """Returns true if argument is a class object that appears to be nested.""" return (isinstance(cls, type) and cls.__module__ != '__builtin__' and cls.__name__ not in sys.modules[cls.__module__].__dict__) def find_containing_class(nested_class): """Finds containing class of a nestec class passed as argument.""" def find_containing_class_inner(outer): for k, v in outer.__dict__.items(): if v is nested_class: return outer, k elif isinstance(v, (type, types.ClassType)) and hasattr(v, '__dict__'): res = find_containing_class_inner(v) if res: return res return find_containing_class_inner(sys.modules[nested_class.__module__]) def _nested_type_wrapper(fun): """A wrapper for the standard pickler handler for class objects. Args: fun: Original pickler handler for type objects. Returns: A wrapper for type objects that handles nested classes. The wrapper detects if an object being pickled is a nested class object. For nested class object only it will save the containing class object so the nested structure is recreated during unpickle. """ def wrapper(pickler, obj): # When the nested class is defined in the __main__ module we do not have to # do anything special because the pickler itself will save the constituent # parts of the type (i.e., name, base classes, dictionary) and then # recreate it during unpickling. if is_nested_class(obj) and obj.__module__ != '__main__': containing_class_and_name = find_containing_class(obj) if containing_class_and_name is not None: return pickler.save_reduce( getattr, containing_class_and_name, obj=obj) try: return fun(pickler, obj) except dill.dill.PicklingError: # pylint: disable=protected-access return pickler.save_reduce( dill.dill._create_type, (type(obj), obj.__name__, obj.__bases__, dill.dill._dict_from_dictproxy(obj.__dict__)), obj=obj) # pylint: enable=protected-access return wrapper # Monkey patch the standard pickler dispatch table entry for type objects. # Dill, for certain types, defers to the standard pickler (including type # objects). We wrap the standard handler using type_wrapper() because # for nested class we want to pickle the actual enclosing class object so we # can recreate it during unpickling. # TODO(silviuc): Make sure we submit the fix upstream to GitHub dill project. dill.dill.Pickler.dispatch[type] = _nested_type_wrapper( dill.dill.Pickler.dispatch[type]) # Dill pickles generators objects without complaint, but unpickling produces # TypeError: object.__new__(generator) is not safe, use generator.__new__() # on some versions of Python. def reject_generators(unused_pickler, unused_obj): raise TypeError("can't (safely) pickle generator objects") dill.dill.Pickler.dispatch[types.GeneratorType] = reject_generators # This if guards against dill not being full initialized when generating docs. if 'save_module' in dir(dill.dill): # Always pickle non-main modules by name. old_save_module = dill.dill.save_module @dill.dill.register(dill.dill.ModuleType) def save_module(pickler, obj): if dill.dill.is_dill(pickler) and obj is pickler._main: return old_save_module(pickler, obj) else: dill.dill.log.info('M2: %s' % obj) # pylint: disable=protected-access pickler.save_reduce(dill.dill._import_module, (obj.__name__,), obj=obj) # pylint: enable=protected-access dill.dill.log.info('# M2') # Pickle module dictionaries (commonly found in lambda's globals) # by referencing their module. old_save_module_dict = dill.dill.save_module_dict known_module_dicts = {} @dill.dill.register(dict) def new_save_module_dict(pickler, obj): obj_id = id(obj) if not known_module_dicts or '__file__' in obj or '__package__' in obj: if obj_id not in known_module_dicts: for m in sys.modules.values(): try: if m and m.__name__ != '__main__': known_module_dicts[id(m.__dict__)] = m except AttributeError: # Skip modules that do not have the __name__ attribute. pass # TODO(silviuc): Must investigate the disabled if branch below. if obj_id in known_module_dicts and dill.dill.is_dill(pickler) and False: return pickler.save_reduce( getattr, (known_module_dicts[obj_id], '__dict__'), obj=obj) else: return old_save_module_dict(pickler, obj) dill.dill.save_module_dict = new_save_module_dict def _nest_dill_logging(): """Prefix all dill logging with its depth in the callstack. Useful for debugging pickling of deeply nested structures. """ old_log_info = dill.dill.log.info def new_log_info(msg, *args, **kwargs): old_log_info( ('1 2 3 4 5 6 7 8 9 0 ' * 10)[:len(traceback.extract_stack())] + msg, *args, **kwargs) dill.dill.log.info = new_log_info # Turn off verbose logging from the dill pickler. logging.getLogger('dill').setLevel(logging.WARN) # TODO(ccy): Currently, there are still instances of pickler.dumps() and # pickler.loads() being used for data, which results in an unnecessary base64 # encoding. This should be cleaned up. def dumps(o): try: return base64.b64encode(dill.dumps(o)) except Exception: # pylint: disable=broad-except dill.dill._trace(True) # pylint: disable=protected-access return base64.b64encode(dill.dumps(o)) finally: dill.dill._trace(False) # pylint: disable=protected-access def loads(s): try: return dill.loads(base64.b64decode(s)) except Exception: # pylint: disable=broad-except dill.dill._trace(True) # pylint: disable=protected-access return dill.loads(base64.b64decode(s)) finally: dill.dill._trace(False) # pylint: disable=protected-access def dump_session(file_path): return dill.dump_session(file_path) def load_session(file_path): return dill.load_session(file_path) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/pickler_test.py000066400000000000000000000051671272340251000260710ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the pickler module.""" import unittest from google.cloud.dataflow.internal import module_test from google.cloud.dataflow.internal.pickler import dumps from google.cloud.dataflow.internal.pickler import loads class PicklerTest(unittest.TestCase): def test_basics(self): self.assertEquals([1, 'a', (u'z',)], loads(dumps([1, 'a', (u'z',)]))) fun = lambda x: 'xyz-%s' % x self.assertEquals('xyz-abc', loads(dumps(fun))('abc')) def test_lambda_with_globals(self): """Tests that the globals of a function are preserved.""" # The point of the test is that the lambda being called after unpickling # relies on having the re module being loaded. self.assertEquals( ['abc', 'def'], loads(dumps(module_test.get_lambda_with_globals()))('abc def')) def test_lambda_with_closure(self): """Tests that the closure of a function is preserved.""" self.assertEquals( 'closure: abc', loads(dumps(module_test.get_lambda_with_closure('abc')))()) def test_class(self): """Tests that a class object is pickled correctly.""" self.assertEquals( ['abc', 'def'], loads(dumps(module_test.Xyz))().foo('abc def')) def test_object(self): """Tests that a class instance is pickled correctly.""" self.assertEquals( ['abc', 'def'], loads(dumps(module_test.XYZ_OBJECT)).foo('abc def')) def test_nested_class(self): """Tests that a nested class object is pickled correctly.""" self.assertEquals( 'X:abc', loads(dumps(module_test.TopClass.NestedClass('abc'))).datum) self.assertEquals( 'Y:abc', loads(dumps(module_test.TopClass.MiddleClass.NestedClass('abc'))).datum) def test_dynamic_class(self): """Tests that a nested class object is pickled correctly.""" self.assertEquals( 'Z:abc', loads(dumps(module_test.create_class('abc'))).get()) def test_generators(self): with self.assertRaises(TypeError): dumps((_ for _ in range(10))) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/util.py000066400000000000000000000065051272340251000243530ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Utility functions used throughout the dataflow package.""" class ArgumentPlaceholder(object): """A place holder object replacing PValues in argument lists. A Fn object can take any number of "side inputs", which are PValues that will be evaluated during pipeline execution and will be provided to the function at the moment of its execution as positional or keyword arguments. This is used only internally and should never be used by user code. A custom Fn object by the time it executes will have such values replaced with real computed values. """ def __eq__(self, other): """Tests for equality of two placeholder objects. Args: other: Another placeholder object to compare to. This method is used only for test code. All placeholder objects are equal to each other. """ return isinstance(other, ArgumentPlaceholder) def remove_objects_from_args(args, kwargs, pvalue_classes): """Replaces all objects of a given type in args/kwargs with a placeholder. Args: args: A list of positional arguments. kwargs: A dictionary of keyword arguments. pvalue_classes: A tuple of class objects representing the types of the arguments that must be replaced with a placeholder value (instance of ArgumentPlaceholder) Returns: A 3-tuple containing a modified list of positional arguments, a modified dictionary of keyword arguments, and a list of all objects replaced with a placeholder value. """ pvals = [] def swapper(value): pvals.append(value) return ArgumentPlaceholder() new_args = [swapper(v) if isinstance(v, pvalue_classes) else v for v in args] # Make sure the order in which we process the dictionary keys is predictable # by sorting the entries first. This will be important when putting back # PValues. new_kwargs = dict((k, swapper(v)) if isinstance(v, pvalue_classes) else (k, v) for k, v in sorted(kwargs.iteritems())) return (new_args, new_kwargs, pvals) def insert_values_in_args(args, kwargs, values): """Replaces all placeholders in args/kwargs with actual values. Args: args: A list of positional arguments. kwargs: A dictionary of keyword arguments. values: A list of values that will be used to replace placeholder values. Returns: A 2-tuple containing a modified list of positional arguments, and a modified dictionary of keyword arguments. """ # Use a local iterator so that we don't modify values. v_iter = iter(values) new_args = [ v_iter.next() if isinstance(arg, ArgumentPlaceholder) else arg for arg in args] new_kwargs = dict( (k, v_iter.next()) if isinstance(v, ArgumentPlaceholder) else (k, v) for k, v in sorted(kwargs.iteritems())) return (new_args, new_kwargs) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/util_test.py000066400000000000000000000041111272340251000254010ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the util module.""" import unittest from google.cloud.dataflow.internal.util import ArgumentPlaceholder from google.cloud.dataflow.internal.util import insert_values_in_args from google.cloud.dataflow.internal.util import remove_objects_from_args class UtilTest(unittest.TestCase): def test_remove_objects_from_args(self): args, kwargs, objs = remove_objects_from_args( [1, 'a'], {'x': 1, 'y': 3.14}, (str, float)) self.assertEquals([1, ArgumentPlaceholder()], args) self.assertEquals({'x': 1, 'y': ArgumentPlaceholder()}, kwargs) self.assertEquals(['a', 3.14], objs) def test_remove_objects_from_args_nothing_to_remove(self): args, kwargs, objs = remove_objects_from_args( [1, 2], {'x': 1, 'y': 2}, (str, float)) self.assertEquals([1, 2], args) self.assertEquals({'x': 1, 'y': 2}, kwargs) self.assertEquals([], objs) def test_insert_values_in_args(self): values = ['a', 'b'] args = [1, ArgumentPlaceholder()] kwargs = {'x': 1, 'y': ArgumentPlaceholder()} args, kwargs = insert_values_in_args(args, kwargs, values) self.assertEquals([1, 'a'], args) self.assertEquals({'x': 1, 'y': 'b'}, kwargs) def test_insert_values_in_args_nothing_to_insert(self): values = [] args = [1, 'a'] kwargs = {'x': 1, 'y': 'b'} args, kwargs = insert_values_in_args(args, kwargs, values) self.assertEquals([1, 'a'], args) self.assertEquals({'x': 1, 'y': 'b'}, kwargs) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/windmill_pb2.py000066400000000000000000002611071272340251000257610ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. # Generated by the protocol buffer compiler. DO NOT EDIT! # source: windmill.proto import sys _b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) from google.protobuf import descriptor as _descriptor from google.protobuf import message as _message from google.protobuf import reflection as _reflection from google.protobuf import symbol_database as _symbol_database from google.protobuf import descriptor_pb2 # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() DESCRIPTOR = _descriptor.FileDescriptor( name='windmill.proto', package='windmill', syntax='proto2', serialized_pb=_b('\n\x0ewindmill.proto\x12\x08windmill\"R\n\x07Message\x12\'\n\ttimestamp\x18\x01 \x02(\x03:\x14-9223372036854775808\x12\x0c\n\x04\x64\x61ta\x18\x02 \x02(\x0c\x12\x10\n\x08metadata\x18\x03 \x01(\x0c\"\xbf\x01\n\x05Timer\x12\x0b\n\x03tag\x18\x01 \x02(\x0c\x12\'\n\ttimestamp\x18\x02 \x01(\x03:\x14-9223372036854775808\x12-\n\x04type\x18\x03 \x01(\x0e\x32\x14.windmill.Timer.Type:\tWATERMARK\x12\x14\n\x0cstate_family\x18\x04 \x01(\t\";\n\x04Type\x12\r\n\tWATERMARK\x10\x00\x12\x0c\n\x08REALTIME\x10\x01\x12\x16\n\x12\x44\x45PENDENT_REALTIME\x10\x02\"X\n\x12InputMessageBundle\x12\x1d\n\x15source_computation_id\x18\x01 \x02(\t\x12#\n\x08messages\x18\x02 \x03(\x0b\x32\x11.windmill.Message\"r\n\x12KeyedMessageBundle\x12\x0b\n\x03key\x18\x01 \x02(\x0c\x12\x14\n\x0csharding_key\x18\x04 \x01(\x06\x12#\n\x08messages\x18\x02 \x03(\x0b\x32\x11.windmill.Message\x12\x14\n\x0cmessages_ids\x18\x03 \x03(\x0c\"\x87\x01\n\x13OutputMessageBundle\x12\"\n\x1a\x64\x65stination_computation_id\x18\x01 \x01(\t\x12\x1d\n\x15\x64\x65stination_stream_id\x18\x03 \x01(\t\x12-\n\x07\x62undles\x18\x02 \x03(\x0b\x32\x1c.windmill.KeyedMessageBundle\"t\n\x13PubSubMessageBundle\x12\r\n\x05topic\x18\x01 \x02(\t\x12#\n\x08messages\x18\x02 \x03(\x0b\x32\x11.windmill.Message\x12\x17\n\x0ftimestamp_label\x18\x03 \x01(\t\x12\x10\n\x08id_label\x18\x04 \x01(\t\".\n\x0bTimerBundle\x12\x1f\n\x06timers\x18\x01 \x03(\x0b\x32\x0f.windmill.Timer\">\n\x05Value\x12\'\n\ttimestamp\x18\x01 \x02(\x03:\x14-9223372036854775808\x12\x0c\n\x04\x64\x61ta\x18\x02 \x02(\x0c\"M\n\x08TagValue\x12\x0b\n\x03tag\x18\x01 \x02(\x0c\x12\x1e\n\x05value\x18\x02 \x01(\x0b\x32\x0f.windmill.Value\x12\x14\n\x0cstate_family\x18\x03 \x01(\t\"\xdb\x01\n\x07TagList\x12\x0b\n\x03tag\x18\x01 \x02(\x0c\x12+\n\rend_timestamp\x18\x02 \x01(\x03:\x14-9223372036854775808\x12\x1f\n\x06values\x18\x03 \x03(\x0b\x32\x0f.windmill.Value\x12\x14\n\x0cstate_family\x18\x04 \x01(\t\x12\x15\n\rrequest_token\x18\x07 \x01(\x0c\x12\x1a\n\x12\x63ontinuation_token\x18\x05 \x01(\x0c\x12,\n\x0f\x66\x65tch_max_bytes\x18\x06 \x01(\x03:\x13\x39\x32\x32\x33\x33\x37\x32\x30\x33\x36\x38\x35\x34\x37\x37\x35\x38\x30\x37\",\n\x0cGlobalDataId\x12\x0b\n\x03tag\x18\x01 \x02(\t\x12\x0f\n\x07version\x18\x02 \x02(\x0c\"k\n\nGlobalData\x12\'\n\x07\x64\x61ta_id\x18\x01 \x02(\x0b\x32\x16.windmill.GlobalDataId\x12\x10\n\x08is_ready\x18\x02 \x01(\x08\x12\x0c\n\x04\x64\x61ta\x18\x03 \x01(\x0c\x12\x14\n\x0cstate_family\x18\x04 \x01(\t\"2\n\x0bSourceState\x12\r\n\x05state\x18\x01 \x01(\x0c\x12\x14\n\x0c\x66inalize_ids\x18\x02 \x03(\x06\"Y\n\rWatermarkHold\x12\x0b\n\x03tag\x18\x01 \x02(\x0c\x12\x16\n\ntimestamps\x18\x02 \x03(\x03\x42\x02\x10\x01\x12\r\n\x05reset\x18\x03 \x01(\x08\x12\x14\n\x0cstate_family\x18\x04 \x01(\t\"\xd4\x02\n\x08WorkItem\x12\x0b\n\x03key\x18\x01 \x02(\x0c\x12\x12\n\nwork_token\x18\x02 \x02(\x06\x12\x14\n\x0csharding_key\x18\t \x01(\x06\x12\x13\n\x0b\x63\x61\x63he_token\x18\x07 \x01(\x06\x12\x35\n\x0fmessage_bundles\x18\x03 \x03(\x0b\x32\x1c.windmill.InputMessageBundle\x12%\n\x06timers\x18\x04 \x01(\x0b\x32\x15.windmill.TimerBundle\x12<\n\x1cglobal_data_id_notifications\x18\x05 \x03(\x0b\x32\x16.windmill.GlobalDataId\x12+\n\x0csource_state\x18\x06 \x01(\x0b\x32\x15.windmill.SourceState\x12\x33\n\x15output_data_watermark\x18\x08 \x01(\x03:\x14-9223372036854775808\"\xc6\x01\n\x14\x43omputationWorkItems\x12\x16\n\x0e\x63omputation_id\x18\x01 \x02(\t\x12 \n\x04work\x18\x02 \x03(\x0b\x32\x12.windmill.WorkItem\x12\x32\n\x14input_data_watermark\x18\x03 \x01(\x03:\x14-9223372036854775808\x12@\n\"dependent_realtime_input_watermark\x18\x04 \x01(\x03:\x14-9223372036854775808\"\x8d\x01\n\x0eGetWorkRequest\x12\x11\n\tclient_id\x18\x01 \x02(\x06\x12\x11\n\tworker_id\x18\x04 \x01(\t\x12\x0e\n\x06job_id\x18\x05 \x01(\t\x12\x1d\n\tmax_items\x18\x02 \x01(\x03:\n4294967295\x12&\n\tmax_bytes\x18\x03 \x01(\x03:\x13\x39\x32\x32\x33\x33\x37\x32\x30\x33\x36\x38\x35\x34\x37\x37\x35\x38\x30\x37\"?\n\x0fGetWorkResponse\x12,\n\x04work\x18\x01 \x03(\x0b\x32\x1e.windmill.ComputationWorkItems\"\xdf\x01\n\x13KeyedGetDataRequest\x12\x0b\n\x03key\x18\x01 \x02(\x0c\x12\x12\n\nwork_token\x18\x02 \x02(\x06\x12\x14\n\x0csharding_key\x18\x06 \x01(\x06\x12+\n\x0fvalues_to_fetch\x18\x03 \x03(\x0b\x32\x12.windmill.TagValue\x12)\n\x0elists_to_fetch\x18\x04 \x03(\x0b\x32\x11.windmill.TagList\x12\x39\n\x18watermark_holds_to_fetch\x18\x05 \x03(\x0b\x32\x17.windmill.WatermarkHold\"d\n\x19\x43omputationGetDataRequest\x12\x16\n\x0e\x63omputation_id\x18\x01 \x02(\t\x12/\n\x08requests\x18\x02 \x03(\x0b\x32\x1d.windmill.KeyedGetDataRequest\"\xce\x01\n\x0eGetDataRequest\x12\x0e\n\x06job_id\x18\x04 \x01(\t\x12\x35\n\x08requests\x18\x01 \x03(\x0b\x32#.windmill.ComputationGetDataRequest\x12?\n\x1aglobal_data_fetch_requests\x18\x03 \x03(\x0b\x32\x1b.windmill.GlobalDataRequest\x12\x34\n\x14global_data_to_fetch\x18\x02 \x03(\x0b\x32\x16.windmill.GlobalDataId\"\xab\x01\n\x14KeyedGetDataResponse\x12\x0b\n\x03key\x18\x01 \x02(\x0c\x12\x0e\n\x06\x66\x61iled\x18\x02 \x01(\x08\x12\"\n\x06values\x18\x03 \x03(\x0b\x32\x12.windmill.TagValue\x12 \n\x05lists\x18\x04 \x03(\x0b\x32\x11.windmill.TagList\x12\x30\n\x0fwatermark_holds\x18\x05 \x03(\x0b\x32\x17.windmill.WatermarkHold\"b\n\x1a\x43omputationGetDataResponse\x12\x16\n\x0e\x63omputation_id\x18\x01 \x02(\t\x12,\n\x04\x64\x61ta\x18\x02 \x03(\x0b\x32\x1e.windmill.KeyedGetDataResponse\"p\n\x0fGetDataResponse\x12\x32\n\x04\x64\x61ta\x18\x01 \x03(\x0b\x32$.windmill.ComputationGetDataResponse\x12)\n\x0bglobal_data\x18\x02 \x03(\x0b\x32\x14.windmill.GlobalData\"\xbd\x01\n\x07\x43ounter\x12\x0c\n\x04name\x18\x01 \x01(\t\x12$\n\x04kind\x18\x02 \x01(\x0e\x32\x16.windmill.Counter.Kind\x12\x15\n\rdouble_scalar\x18\x03 \x01(\x01\x12\x12\n\nint_scalar\x18\x04 \x01(\x03\x12\x12\n\nmean_count\x18\x06 \x01(\x03\x12\x12\n\ncumulative\x18\x07 \x01(\x08\"+\n\x04Kind\x12\x07\n\x03SUM\x10\x00\x12\x07\n\x03MAX\x10\x01\x12\x07\n\x03MIN\x10\x02\x12\x08\n\x04MEAN\x10\x03\"\x8d\x01\n\x11GlobalDataRequest\x12\'\n\x07\x64\x61ta_id\x18\x01 \x02(\x0b\x32\x16.windmill.GlobalDataId\x12\x39\n\x1c\x65xistence_watermark_deadline\x18\x02 \x01(\x03:\x13\x39\x32\x32\x33\x33\x37\x32\x30\x33\x36\x38\x35\x34\x37\x37\x35\x38\x30\x37\x12\x14\n\x0cstate_family\x18\x03 \x01(\t\"\xa4\x05\n\x15WorkItemCommitRequest\x12\x0b\n\x03key\x18\x01 \x02(\x0c\x12\x12\n\nwork_token\x18\x02 \x02(\x06\x12\x14\n\x0csharding_key\x18\x0f \x01(\x06\x12\x36\n\x0foutput_messages\x18\x03 \x03(\x0b\x32\x1d.windmill.OutputMessageBundle\x12\x36\n\x0fpubsub_messages\x18\x07 \x03(\x0b\x32\x1d.windmill.PubSubMessageBundle\x12&\n\routput_timers\x18\x04 \x03(\x0b\x32\x0f.windmill.Timer\x12)\n\rvalue_updates\x18\x05 \x03(\x0b\x32\x12.windmill.TagValue\x12\'\n\x0clist_updates\x18\x06 \x03(\x0b\x32\x11.windmill.TagList\x12*\n\x0f\x63ounter_updates\x18\x08 \x03(\x0b\x32\x11.windmill.Counter\x12\x39\n\x14global_data_requests\x18\x0b \x03(\x0b\x32\x1b.windmill.GlobalDataRequest\x12\x31\n\x13global_data_updates\x18\n \x03(\x0b\x32\x14.windmill.GlobalData\x12\x33\n\x14source_state_updates\x18\x0c \x01(\x0b\x32\x15.windmill.SourceState\x12.\n\x10source_watermark\x18\r \x01(\x03:\x14-9223372036854775808\x12\x30\n\x0fwatermark_holds\x18\x0e \x03(\x0b\x32\x17.windmill.WatermarkHold\x12\x37\n\x17global_data_id_requests\x18\t \x03(\x0b\x32\x16.windmill.GlobalDataId\"i\n\x1c\x43omputationCommitWorkRequest\x12\x16\n\x0e\x63omputation_id\x18\x01 \x02(\t\x12\x31\n\x08requests\x18\x02 \x03(\x0b\x32\x1f.windmill.WorkItemCommitRequest\"]\n\x11\x43ommitWorkRequest\x12\x0e\n\x06job_id\x18\x02 \x01(\t\x12\x38\n\x08requests\x18\x01 \x03(\x0b\x32&.windmill.ComputationCommitWorkRequest\"\x14\n\x12\x43ommitWorkResponse\"8\n\x10GetConfigRequest\x12\x0e\n\x06job_id\x18\x02 \x01(\t\x12\x14\n\x0c\x63omputations\x18\x01 \x03(\t\"\xd8\x02\n\x11GetConfigResponse\x12\x13\n\x0b\x63loud_works\x18\x01 \x03(\t\x12:\n\x08name_map\x18\x02 \x03(\x0b\x32(.windmill.GetConfigResponse.NameMapEntry\x12h\n!system_name_to_computation_id_map\x18\x03 \x03(\x0b\x32=.windmill.GetConfigResponse.SystemNameToComputationIdMapEntry\x1a\x36\n\x0cNameMapEntry\x12\x11\n\tuser_name\x18\x01 \x01(\t\x12\x13\n\x0bsystem_name\x18\x02 \x01(\t\x1aP\n!SystemNameToComputationIdMapEntry\x12\x13\n\x0bsystem_name\x18\x01 \x01(\t\x12\x16\n\x0e\x63omputation_id\x18\x02 \x01(\t\"E\n\tException\x12\x14\n\x0cstack_frames\x18\x01 \x03(\t\x12\"\n\x05\x63\x61use\x18\x02 \x01(\x0b\x32\x13.windmill.Exception\"\xc8\x01\n\x12ReportStatsRequest\x12\x0e\n\x06job_id\x18\x06 \x01(\t\x12\x16\n\x0e\x63omputation_id\x18\x01 \x01(\t\x12\x0b\n\x03key\x18\x02 \x01(\x0c\x12\x12\n\nwork_token\x18\x03 \x01(\x06\x12\x14\n\x0csharding_key\x18\x07 \x01(\x06\x12\'\n\nexceptions\x18\x04 \x03(\x0b\x32\x13.windmill.Exception\x12*\n\x0f\x63ounter_updates\x18\x05 \x03(\x0b\x32\x11.windmill.Counter\"%\n\x13ReportStatsResponse\x12\x0e\n\x06\x66\x61iled\x18\x01 \x01(\x08\x42\x41\n5com.google.cloud.dataflow.sdk.runners.worker.windmillB\x08Windmill') ) _sym_db.RegisterFileDescriptor(DESCRIPTOR) _TIMER_TYPE = _descriptor.EnumDescriptor( name='Type', full_name='windmill.Timer.Type', filename=None, file=DESCRIPTOR, values=[ _descriptor.EnumValueDescriptor( name='WATERMARK', index=0, number=0, options=None, type=None), _descriptor.EnumValueDescriptor( name='REALTIME', index=1, number=1, options=None, type=None), _descriptor.EnumValueDescriptor( name='DEPENDENT_REALTIME', index=2, number=2, options=None, type=None), ], containing_type=None, options=None, serialized_start=245, serialized_end=304, ) _sym_db.RegisterEnumDescriptor(_TIMER_TYPE) _COUNTER_KIND = _descriptor.EnumDescriptor( name='Kind', full_name='windmill.Counter.Kind', filename=None, file=DESCRIPTOR, values=[ _descriptor.EnumValueDescriptor( name='SUM', index=0, number=0, options=None, type=None), _descriptor.EnumValueDescriptor( name='MAX', index=1, number=1, options=None, type=None), _descriptor.EnumValueDescriptor( name='MIN', index=2, number=2, options=None, type=None), _descriptor.EnumValueDescriptor( name='MEAN', index=3, number=3, options=None, type=None), ], containing_type=None, options=None, serialized_start=3304, serialized_end=3347, ) _sym_db.RegisterEnumDescriptor(_COUNTER_KIND) _MESSAGE = _descriptor.Descriptor( name='Message', full_name='windmill.Message', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='timestamp', full_name='windmill.Message.timestamp', index=0, number=1, type=3, cpp_type=2, label=2, has_default_value=True, default_value=-9223372036854775808, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='data', full_name='windmill.Message.data', index=1, number=2, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='metadata', full_name='windmill.Message.metadata', index=2, number=3, type=12, cpp_type=9, label=1, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=28, serialized_end=110, ) _TIMER = _descriptor.Descriptor( name='Timer', full_name='windmill.Timer', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='tag', full_name='windmill.Timer.tag', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='timestamp', full_name='windmill.Timer.timestamp', index=1, number=2, type=3, cpp_type=2, label=1, has_default_value=True, default_value=-9223372036854775808, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='type', full_name='windmill.Timer.type', index=2, number=3, type=14, cpp_type=8, label=1, has_default_value=True, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='state_family', full_name='windmill.Timer.state_family', index=3, number=4, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ _TIMER_TYPE, ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=113, serialized_end=304, ) _INPUTMESSAGEBUNDLE = _descriptor.Descriptor( name='InputMessageBundle', full_name='windmill.InputMessageBundle', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='source_computation_id', full_name='windmill.InputMessageBundle.source_computation_id', index=0, number=1, type=9, cpp_type=9, label=2, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='messages', full_name='windmill.InputMessageBundle.messages', index=1, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=306, serialized_end=394, ) _KEYEDMESSAGEBUNDLE = _descriptor.Descriptor( name='KeyedMessageBundle', full_name='windmill.KeyedMessageBundle', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='key', full_name='windmill.KeyedMessageBundle.key', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='sharding_key', full_name='windmill.KeyedMessageBundle.sharding_key', index=1, number=4, type=6, cpp_type=4, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='messages', full_name='windmill.KeyedMessageBundle.messages', index=2, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='messages_ids', full_name='windmill.KeyedMessageBundle.messages_ids', index=3, number=3, type=12, cpp_type=9, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=396, serialized_end=510, ) _OUTPUTMESSAGEBUNDLE = _descriptor.Descriptor( name='OutputMessageBundle', full_name='windmill.OutputMessageBundle', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='destination_computation_id', full_name='windmill.OutputMessageBundle.destination_computation_id', index=0, number=1, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='destination_stream_id', full_name='windmill.OutputMessageBundle.destination_stream_id', index=1, number=3, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='bundles', full_name='windmill.OutputMessageBundle.bundles', index=2, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=513, serialized_end=648, ) _PUBSUBMESSAGEBUNDLE = _descriptor.Descriptor( name='PubSubMessageBundle', full_name='windmill.PubSubMessageBundle', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='topic', full_name='windmill.PubSubMessageBundle.topic', index=0, number=1, type=9, cpp_type=9, label=2, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='messages', full_name='windmill.PubSubMessageBundle.messages', index=1, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='timestamp_label', full_name='windmill.PubSubMessageBundle.timestamp_label', index=2, number=3, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='id_label', full_name='windmill.PubSubMessageBundle.id_label', index=3, number=4, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=650, serialized_end=766, ) _TIMERBUNDLE = _descriptor.Descriptor( name='TimerBundle', full_name='windmill.TimerBundle', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='timers', full_name='windmill.TimerBundle.timers', index=0, number=1, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=768, serialized_end=814, ) _VALUE = _descriptor.Descriptor( name='Value', full_name='windmill.Value', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='timestamp', full_name='windmill.Value.timestamp', index=0, number=1, type=3, cpp_type=2, label=2, has_default_value=True, default_value=-9223372036854775808, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='data', full_name='windmill.Value.data', index=1, number=2, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=816, serialized_end=878, ) _TAGVALUE = _descriptor.Descriptor( name='TagValue', full_name='windmill.TagValue', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='tag', full_name='windmill.TagValue.tag', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='value', full_name='windmill.TagValue.value', index=1, number=2, type=11, cpp_type=10, label=1, has_default_value=False, default_value=None, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='state_family', full_name='windmill.TagValue.state_family', index=2, number=3, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=880, serialized_end=957, ) _TAGLIST = _descriptor.Descriptor( name='TagList', full_name='windmill.TagList', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='tag', full_name='windmill.TagList.tag', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='end_timestamp', full_name='windmill.TagList.end_timestamp', index=1, number=2, type=3, cpp_type=2, label=1, has_default_value=True, default_value=-9223372036854775808, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='values', full_name='windmill.TagList.values', index=2, number=3, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='state_family', full_name='windmill.TagList.state_family', index=3, number=4, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='request_token', full_name='windmill.TagList.request_token', index=4, number=7, type=12, cpp_type=9, label=1, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='continuation_token', full_name='windmill.TagList.continuation_token', index=5, number=5, type=12, cpp_type=9, label=1, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='fetch_max_bytes', full_name='windmill.TagList.fetch_max_bytes', index=6, number=6, type=3, cpp_type=2, label=1, has_default_value=True, default_value=9223372036854775807, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=960, serialized_end=1179, ) _GLOBALDATAID = _descriptor.Descriptor( name='GlobalDataId', full_name='windmill.GlobalDataId', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='tag', full_name='windmill.GlobalDataId.tag', index=0, number=1, type=9, cpp_type=9, label=2, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='version', full_name='windmill.GlobalDataId.version', index=1, number=2, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=1181, serialized_end=1225, ) _GLOBALDATA = _descriptor.Descriptor( name='GlobalData', full_name='windmill.GlobalData', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='data_id', full_name='windmill.GlobalData.data_id', index=0, number=1, type=11, cpp_type=10, label=2, has_default_value=False, default_value=None, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='is_ready', full_name='windmill.GlobalData.is_ready', index=1, number=2, type=8, cpp_type=7, label=1, has_default_value=False, default_value=False, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='data', full_name='windmill.GlobalData.data', index=2, number=3, type=12, cpp_type=9, label=1, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='state_family', full_name='windmill.GlobalData.state_family', index=3, number=4, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=1227, serialized_end=1334, ) _SOURCESTATE = _descriptor.Descriptor( name='SourceState', full_name='windmill.SourceState', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='state', full_name='windmill.SourceState.state', index=0, number=1, type=12, cpp_type=9, label=1, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='finalize_ids', full_name='windmill.SourceState.finalize_ids', index=1, number=2, type=6, cpp_type=4, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=1336, serialized_end=1386, ) _WATERMARKHOLD = _descriptor.Descriptor( name='WatermarkHold', full_name='windmill.WatermarkHold', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='tag', full_name='windmill.WatermarkHold.tag', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='timestamps', full_name='windmill.WatermarkHold.timestamps', index=1, number=2, type=3, cpp_type=2, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))), _descriptor.FieldDescriptor( name='reset', full_name='windmill.WatermarkHold.reset', index=2, number=3, type=8, cpp_type=7, label=1, has_default_value=False, default_value=False, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='state_family', full_name='windmill.WatermarkHold.state_family', index=3, number=4, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=1388, serialized_end=1477, ) _WORKITEM = _descriptor.Descriptor( name='WorkItem', full_name='windmill.WorkItem', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='key', full_name='windmill.WorkItem.key', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='work_token', full_name='windmill.WorkItem.work_token', index=1, number=2, type=6, cpp_type=4, label=2, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='sharding_key', full_name='windmill.WorkItem.sharding_key', index=2, number=9, type=6, cpp_type=4, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='cache_token', full_name='windmill.WorkItem.cache_token', index=3, number=7, type=6, cpp_type=4, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='message_bundles', full_name='windmill.WorkItem.message_bundles', index=4, number=3, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='timers', full_name='windmill.WorkItem.timers', index=5, number=4, type=11, cpp_type=10, label=1, has_default_value=False, default_value=None, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='global_data_id_notifications', full_name='windmill.WorkItem.global_data_id_notifications', index=6, number=5, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='source_state', full_name='windmill.WorkItem.source_state', index=7, number=6, type=11, cpp_type=10, label=1, has_default_value=False, default_value=None, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='output_data_watermark', full_name='windmill.WorkItem.output_data_watermark', index=8, number=8, type=3, cpp_type=2, label=1, has_default_value=True, default_value=-9223372036854775808, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=1480, serialized_end=1820, ) _COMPUTATIONWORKITEMS = _descriptor.Descriptor( name='ComputationWorkItems', full_name='windmill.ComputationWorkItems', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='computation_id', full_name='windmill.ComputationWorkItems.computation_id', index=0, number=1, type=9, cpp_type=9, label=2, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='work', full_name='windmill.ComputationWorkItems.work', index=1, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='input_data_watermark', full_name='windmill.ComputationWorkItems.input_data_watermark', index=2, number=3, type=3, cpp_type=2, label=1, has_default_value=True, default_value=-9223372036854775808, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='dependent_realtime_input_watermark', full_name='windmill.ComputationWorkItems.dependent_realtime_input_watermark', index=3, number=4, type=3, cpp_type=2, label=1, has_default_value=True, default_value=-9223372036854775808, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=1823, serialized_end=2021, ) _GETWORKREQUEST = _descriptor.Descriptor( name='GetWorkRequest', full_name='windmill.GetWorkRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='client_id', full_name='windmill.GetWorkRequest.client_id', index=0, number=1, type=6, cpp_type=4, label=2, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='worker_id', full_name='windmill.GetWorkRequest.worker_id', index=1, number=4, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='job_id', full_name='windmill.GetWorkRequest.job_id', index=2, number=5, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='max_items', full_name='windmill.GetWorkRequest.max_items', index=3, number=2, type=3, cpp_type=2, label=1, has_default_value=True, default_value=4294967295, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='max_bytes', full_name='windmill.GetWorkRequest.max_bytes', index=4, number=3, type=3, cpp_type=2, label=1, has_default_value=True, default_value=9223372036854775807, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=2024, serialized_end=2165, ) _GETWORKRESPONSE = _descriptor.Descriptor( name='GetWorkResponse', full_name='windmill.GetWorkResponse', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='work', full_name='windmill.GetWorkResponse.work', index=0, number=1, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=2167, serialized_end=2230, ) _KEYEDGETDATAREQUEST = _descriptor.Descriptor( name='KeyedGetDataRequest', full_name='windmill.KeyedGetDataRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='key', full_name='windmill.KeyedGetDataRequest.key', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='work_token', full_name='windmill.KeyedGetDataRequest.work_token', index=1, number=2, type=6, cpp_type=4, label=2, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='sharding_key', full_name='windmill.KeyedGetDataRequest.sharding_key', index=2, number=6, type=6, cpp_type=4, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='values_to_fetch', full_name='windmill.KeyedGetDataRequest.values_to_fetch', index=3, number=3, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='lists_to_fetch', full_name='windmill.KeyedGetDataRequest.lists_to_fetch', index=4, number=4, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='watermark_holds_to_fetch', full_name='windmill.KeyedGetDataRequest.watermark_holds_to_fetch', index=5, number=5, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=2233, serialized_end=2456, ) _COMPUTATIONGETDATAREQUEST = _descriptor.Descriptor( name='ComputationGetDataRequest', full_name='windmill.ComputationGetDataRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='computation_id', full_name='windmill.ComputationGetDataRequest.computation_id', index=0, number=1, type=9, cpp_type=9, label=2, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='requests', full_name='windmill.ComputationGetDataRequest.requests', index=1, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=2458, serialized_end=2558, ) _GETDATAREQUEST = _descriptor.Descriptor( name='GetDataRequest', full_name='windmill.GetDataRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='job_id', full_name='windmill.GetDataRequest.job_id', index=0, number=4, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='requests', full_name='windmill.GetDataRequest.requests', index=1, number=1, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='global_data_fetch_requests', full_name='windmill.GetDataRequest.global_data_fetch_requests', index=2, number=3, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='global_data_to_fetch', full_name='windmill.GetDataRequest.global_data_to_fetch', index=3, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=2561, serialized_end=2767, ) _KEYEDGETDATARESPONSE = _descriptor.Descriptor( name='KeyedGetDataResponse', full_name='windmill.KeyedGetDataResponse', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='key', full_name='windmill.KeyedGetDataResponse.key', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='failed', full_name='windmill.KeyedGetDataResponse.failed', index=1, number=2, type=8, cpp_type=7, label=1, has_default_value=False, default_value=False, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='values', full_name='windmill.KeyedGetDataResponse.values', index=2, number=3, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='lists', full_name='windmill.KeyedGetDataResponse.lists', index=3, number=4, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='watermark_holds', full_name='windmill.KeyedGetDataResponse.watermark_holds', index=4, number=5, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=2770, serialized_end=2941, ) _COMPUTATIONGETDATARESPONSE = _descriptor.Descriptor( name='ComputationGetDataResponse', full_name='windmill.ComputationGetDataResponse', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='computation_id', full_name='windmill.ComputationGetDataResponse.computation_id', index=0, number=1, type=9, cpp_type=9, label=2, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='data', full_name='windmill.ComputationGetDataResponse.data', index=1, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=2943, serialized_end=3041, ) _GETDATARESPONSE = _descriptor.Descriptor( name='GetDataResponse', full_name='windmill.GetDataResponse', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='data', full_name='windmill.GetDataResponse.data', index=0, number=1, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='global_data', full_name='windmill.GetDataResponse.global_data', index=1, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=3043, serialized_end=3155, ) _COUNTER = _descriptor.Descriptor( name='Counter', full_name='windmill.Counter', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='name', full_name='windmill.Counter.name', index=0, number=1, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='kind', full_name='windmill.Counter.kind', index=1, number=2, type=14, cpp_type=8, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='double_scalar', full_name='windmill.Counter.double_scalar', index=2, number=3, type=1, cpp_type=5, label=1, has_default_value=False, default_value=float(0), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='int_scalar', full_name='windmill.Counter.int_scalar', index=3, number=4, type=3, cpp_type=2, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='mean_count', full_name='windmill.Counter.mean_count', index=4, number=6, type=3, cpp_type=2, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='cumulative', full_name='windmill.Counter.cumulative', index=5, number=7, type=8, cpp_type=7, label=1, has_default_value=False, default_value=False, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ _COUNTER_KIND, ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=3158, serialized_end=3347, ) _GLOBALDATAREQUEST = _descriptor.Descriptor( name='GlobalDataRequest', full_name='windmill.GlobalDataRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='data_id', full_name='windmill.GlobalDataRequest.data_id', index=0, number=1, type=11, cpp_type=10, label=2, has_default_value=False, default_value=None, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='existence_watermark_deadline', full_name='windmill.GlobalDataRequest.existence_watermark_deadline', index=1, number=2, type=3, cpp_type=2, label=1, has_default_value=True, default_value=9223372036854775807, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='state_family', full_name='windmill.GlobalDataRequest.state_family', index=2, number=3, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=3350, serialized_end=3491, ) _WORKITEMCOMMITREQUEST = _descriptor.Descriptor( name='WorkItemCommitRequest', full_name='windmill.WorkItemCommitRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='key', full_name='windmill.WorkItemCommitRequest.key', index=0, number=1, type=12, cpp_type=9, label=2, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='work_token', full_name='windmill.WorkItemCommitRequest.work_token', index=1, number=2, type=6, cpp_type=4, label=2, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='sharding_key', full_name='windmill.WorkItemCommitRequest.sharding_key', index=2, number=15, type=6, cpp_type=4, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='output_messages', full_name='windmill.WorkItemCommitRequest.output_messages', index=3, number=3, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='pubsub_messages', full_name='windmill.WorkItemCommitRequest.pubsub_messages', index=4, number=7, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='output_timers', full_name='windmill.WorkItemCommitRequest.output_timers', index=5, number=4, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='value_updates', full_name='windmill.WorkItemCommitRequest.value_updates', index=6, number=5, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='list_updates', full_name='windmill.WorkItemCommitRequest.list_updates', index=7, number=6, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='counter_updates', full_name='windmill.WorkItemCommitRequest.counter_updates', index=8, number=8, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='global_data_requests', full_name='windmill.WorkItemCommitRequest.global_data_requests', index=9, number=11, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='global_data_updates', full_name='windmill.WorkItemCommitRequest.global_data_updates', index=10, number=10, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='source_state_updates', full_name='windmill.WorkItemCommitRequest.source_state_updates', index=11, number=12, type=11, cpp_type=10, label=1, has_default_value=False, default_value=None, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='source_watermark', full_name='windmill.WorkItemCommitRequest.source_watermark', index=12, number=13, type=3, cpp_type=2, label=1, has_default_value=True, default_value=-9223372036854775808, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='watermark_holds', full_name='windmill.WorkItemCommitRequest.watermark_holds', index=13, number=14, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='global_data_id_requests', full_name='windmill.WorkItemCommitRequest.global_data_id_requests', index=14, number=9, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=3494, serialized_end=4170, ) _COMPUTATIONCOMMITWORKREQUEST = _descriptor.Descriptor( name='ComputationCommitWorkRequest', full_name='windmill.ComputationCommitWorkRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='computation_id', full_name='windmill.ComputationCommitWorkRequest.computation_id', index=0, number=1, type=9, cpp_type=9, label=2, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='requests', full_name='windmill.ComputationCommitWorkRequest.requests', index=1, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4172, serialized_end=4277, ) _COMMITWORKREQUEST = _descriptor.Descriptor( name='CommitWorkRequest', full_name='windmill.CommitWorkRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='job_id', full_name='windmill.CommitWorkRequest.job_id', index=0, number=2, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='requests', full_name='windmill.CommitWorkRequest.requests', index=1, number=1, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4279, serialized_end=4372, ) _COMMITWORKRESPONSE = _descriptor.Descriptor( name='CommitWorkResponse', full_name='windmill.CommitWorkResponse', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4374, serialized_end=4394, ) _GETCONFIGREQUEST = _descriptor.Descriptor( name='GetConfigRequest', full_name='windmill.GetConfigRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='job_id', full_name='windmill.GetConfigRequest.job_id', index=0, number=2, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='computations', full_name='windmill.GetConfigRequest.computations', index=1, number=1, type=9, cpp_type=9, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4396, serialized_end=4452, ) _GETCONFIGRESPONSE_NAMEMAPENTRY = _descriptor.Descriptor( name='NameMapEntry', full_name='windmill.GetConfigResponse.NameMapEntry', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='user_name', full_name='windmill.GetConfigResponse.NameMapEntry.user_name', index=0, number=1, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='system_name', full_name='windmill.GetConfigResponse.NameMapEntry.system_name', index=1, number=2, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4663, serialized_end=4717, ) _GETCONFIGRESPONSE_SYSTEMNAMETOCOMPUTATIONIDMAPENTRY = _descriptor.Descriptor( name='SystemNameToComputationIdMapEntry', full_name='windmill.GetConfigResponse.SystemNameToComputationIdMapEntry', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='system_name', full_name='windmill.GetConfigResponse.SystemNameToComputationIdMapEntry.system_name', index=0, number=1, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='computation_id', full_name='windmill.GetConfigResponse.SystemNameToComputationIdMapEntry.computation_id', index=1, number=2, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4719, serialized_end=4799, ) _GETCONFIGRESPONSE = _descriptor.Descriptor( name='GetConfigResponse', full_name='windmill.GetConfigResponse', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='cloud_works', full_name='windmill.GetConfigResponse.cloud_works', index=0, number=1, type=9, cpp_type=9, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='name_map', full_name='windmill.GetConfigResponse.name_map', index=1, number=2, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='system_name_to_computation_id_map', full_name='windmill.GetConfigResponse.system_name_to_computation_id_map', index=2, number=3, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[_GETCONFIGRESPONSE_NAMEMAPENTRY, _GETCONFIGRESPONSE_SYSTEMNAMETOCOMPUTATIONIDMAPENTRY, ], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4455, serialized_end=4799, ) _EXCEPTION = _descriptor.Descriptor( name='Exception', full_name='windmill.Exception', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='stack_frames', full_name='windmill.Exception.stack_frames', index=0, number=1, type=9, cpp_type=9, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='cause', full_name='windmill.Exception.cause', index=1, number=2, type=11, cpp_type=10, label=1, has_default_value=False, default_value=None, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4801, serialized_end=4870, ) _REPORTSTATSREQUEST = _descriptor.Descriptor( name='ReportStatsRequest', full_name='windmill.ReportStatsRequest', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='job_id', full_name='windmill.ReportStatsRequest.job_id', index=0, number=6, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='computation_id', full_name='windmill.ReportStatsRequest.computation_id', index=1, number=1, type=9, cpp_type=9, label=1, has_default_value=False, default_value=_b("").decode('utf-8'), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='key', full_name='windmill.ReportStatsRequest.key', index=2, number=2, type=12, cpp_type=9, label=1, has_default_value=False, default_value=_b(""), message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='work_token', full_name='windmill.ReportStatsRequest.work_token', index=3, number=3, type=6, cpp_type=4, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='sharding_key', full_name='windmill.ReportStatsRequest.sharding_key', index=4, number=7, type=6, cpp_type=4, label=1, has_default_value=False, default_value=0, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='exceptions', full_name='windmill.ReportStatsRequest.exceptions', index=5, number=4, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), _descriptor.FieldDescriptor( name='counter_updates', full_name='windmill.ReportStatsRequest.counter_updates', index=6, number=5, type=11, cpp_type=10, label=3, has_default_value=False, default_value=[], message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=4873, serialized_end=5073, ) _REPORTSTATSRESPONSE = _descriptor.Descriptor( name='ReportStatsResponse', full_name='windmill.ReportStatsResponse', filename=None, file=DESCRIPTOR, containing_type=None, fields=[ _descriptor.FieldDescriptor( name='failed', full_name='windmill.ReportStatsResponse.failed', index=0, number=1, type=8, cpp_type=7, label=1, has_default_value=False, default_value=False, message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), ], extensions=[ ], nested_types=[], enum_types=[ ], options=None, is_extendable=False, syntax='proto2', extension_ranges=[], oneofs=[ ], serialized_start=5075, serialized_end=5112, ) _TIMER.fields_by_name['type'].enum_type = _TIMER_TYPE _TIMER_TYPE.containing_type = _TIMER _INPUTMESSAGEBUNDLE.fields_by_name['messages'].message_type = _MESSAGE _KEYEDMESSAGEBUNDLE.fields_by_name['messages'].message_type = _MESSAGE _OUTPUTMESSAGEBUNDLE.fields_by_name['bundles'].message_type = _KEYEDMESSAGEBUNDLE _PUBSUBMESSAGEBUNDLE.fields_by_name['messages'].message_type = _MESSAGE _TIMERBUNDLE.fields_by_name['timers'].message_type = _TIMER _TAGVALUE.fields_by_name['value'].message_type = _VALUE _TAGLIST.fields_by_name['values'].message_type = _VALUE _GLOBALDATA.fields_by_name['data_id'].message_type = _GLOBALDATAID _WORKITEM.fields_by_name['message_bundles'].message_type = _INPUTMESSAGEBUNDLE _WORKITEM.fields_by_name['timers'].message_type = _TIMERBUNDLE _WORKITEM.fields_by_name['global_data_id_notifications'].message_type = _GLOBALDATAID _WORKITEM.fields_by_name['source_state'].message_type = _SOURCESTATE _COMPUTATIONWORKITEMS.fields_by_name['work'].message_type = _WORKITEM _GETWORKRESPONSE.fields_by_name['work'].message_type = _COMPUTATIONWORKITEMS _KEYEDGETDATAREQUEST.fields_by_name['values_to_fetch'].message_type = _TAGVALUE _KEYEDGETDATAREQUEST.fields_by_name['lists_to_fetch'].message_type = _TAGLIST _KEYEDGETDATAREQUEST.fields_by_name['watermark_holds_to_fetch'].message_type = _WATERMARKHOLD _COMPUTATIONGETDATAREQUEST.fields_by_name['requests'].message_type = _KEYEDGETDATAREQUEST _GETDATAREQUEST.fields_by_name['requests'].message_type = _COMPUTATIONGETDATAREQUEST _GETDATAREQUEST.fields_by_name['global_data_fetch_requests'].message_type = _GLOBALDATAREQUEST _GETDATAREQUEST.fields_by_name['global_data_to_fetch'].message_type = _GLOBALDATAID _KEYEDGETDATARESPONSE.fields_by_name['values'].message_type = _TAGVALUE _KEYEDGETDATARESPONSE.fields_by_name['lists'].message_type = _TAGLIST _KEYEDGETDATARESPONSE.fields_by_name['watermark_holds'].message_type = _WATERMARKHOLD _COMPUTATIONGETDATARESPONSE.fields_by_name['data'].message_type = _KEYEDGETDATARESPONSE _GETDATARESPONSE.fields_by_name['data'].message_type = _COMPUTATIONGETDATARESPONSE _GETDATARESPONSE.fields_by_name['global_data'].message_type = _GLOBALDATA _COUNTER.fields_by_name['kind'].enum_type = _COUNTER_KIND _COUNTER_KIND.containing_type = _COUNTER _GLOBALDATAREQUEST.fields_by_name['data_id'].message_type = _GLOBALDATAID _WORKITEMCOMMITREQUEST.fields_by_name['output_messages'].message_type = _OUTPUTMESSAGEBUNDLE _WORKITEMCOMMITREQUEST.fields_by_name['pubsub_messages'].message_type = _PUBSUBMESSAGEBUNDLE _WORKITEMCOMMITREQUEST.fields_by_name['output_timers'].message_type = _TIMER _WORKITEMCOMMITREQUEST.fields_by_name['value_updates'].message_type = _TAGVALUE _WORKITEMCOMMITREQUEST.fields_by_name['list_updates'].message_type = _TAGLIST _WORKITEMCOMMITREQUEST.fields_by_name['counter_updates'].message_type = _COUNTER _WORKITEMCOMMITREQUEST.fields_by_name['global_data_requests'].message_type = _GLOBALDATAREQUEST _WORKITEMCOMMITREQUEST.fields_by_name['global_data_updates'].message_type = _GLOBALDATA _WORKITEMCOMMITREQUEST.fields_by_name['source_state_updates'].message_type = _SOURCESTATE _WORKITEMCOMMITREQUEST.fields_by_name['watermark_holds'].message_type = _WATERMARKHOLD _WORKITEMCOMMITREQUEST.fields_by_name['global_data_id_requests'].message_type = _GLOBALDATAID _COMPUTATIONCOMMITWORKREQUEST.fields_by_name['requests'].message_type = _WORKITEMCOMMITREQUEST _COMMITWORKREQUEST.fields_by_name['requests'].message_type = _COMPUTATIONCOMMITWORKREQUEST _GETCONFIGRESPONSE_NAMEMAPENTRY.containing_type = _GETCONFIGRESPONSE _GETCONFIGRESPONSE_SYSTEMNAMETOCOMPUTATIONIDMAPENTRY.containing_type = _GETCONFIGRESPONSE _GETCONFIGRESPONSE.fields_by_name['name_map'].message_type = _GETCONFIGRESPONSE_NAMEMAPENTRY _GETCONFIGRESPONSE.fields_by_name['system_name_to_computation_id_map'].message_type = _GETCONFIGRESPONSE_SYSTEMNAMETOCOMPUTATIONIDMAPENTRY _EXCEPTION.fields_by_name['cause'].message_type = _EXCEPTION _REPORTSTATSREQUEST.fields_by_name['exceptions'].message_type = _EXCEPTION _REPORTSTATSREQUEST.fields_by_name['counter_updates'].message_type = _COUNTER DESCRIPTOR.message_types_by_name['Message'] = _MESSAGE DESCRIPTOR.message_types_by_name['Timer'] = _TIMER DESCRIPTOR.message_types_by_name['InputMessageBundle'] = _INPUTMESSAGEBUNDLE DESCRIPTOR.message_types_by_name['KeyedMessageBundle'] = _KEYEDMESSAGEBUNDLE DESCRIPTOR.message_types_by_name['OutputMessageBundle'] = _OUTPUTMESSAGEBUNDLE DESCRIPTOR.message_types_by_name['PubSubMessageBundle'] = _PUBSUBMESSAGEBUNDLE DESCRIPTOR.message_types_by_name['TimerBundle'] = _TIMERBUNDLE DESCRIPTOR.message_types_by_name['Value'] = _VALUE DESCRIPTOR.message_types_by_name['TagValue'] = _TAGVALUE DESCRIPTOR.message_types_by_name['TagList'] = _TAGLIST DESCRIPTOR.message_types_by_name['GlobalDataId'] = _GLOBALDATAID DESCRIPTOR.message_types_by_name['GlobalData'] = _GLOBALDATA DESCRIPTOR.message_types_by_name['SourceState'] = _SOURCESTATE DESCRIPTOR.message_types_by_name['WatermarkHold'] = _WATERMARKHOLD DESCRIPTOR.message_types_by_name['WorkItem'] = _WORKITEM DESCRIPTOR.message_types_by_name['ComputationWorkItems'] = _COMPUTATIONWORKITEMS DESCRIPTOR.message_types_by_name['GetWorkRequest'] = _GETWORKREQUEST DESCRIPTOR.message_types_by_name['GetWorkResponse'] = _GETWORKRESPONSE DESCRIPTOR.message_types_by_name['KeyedGetDataRequest'] = _KEYEDGETDATAREQUEST DESCRIPTOR.message_types_by_name['ComputationGetDataRequest'] = _COMPUTATIONGETDATAREQUEST DESCRIPTOR.message_types_by_name['GetDataRequest'] = _GETDATAREQUEST DESCRIPTOR.message_types_by_name['KeyedGetDataResponse'] = _KEYEDGETDATARESPONSE DESCRIPTOR.message_types_by_name['ComputationGetDataResponse'] = _COMPUTATIONGETDATARESPONSE DESCRIPTOR.message_types_by_name['GetDataResponse'] = _GETDATARESPONSE DESCRIPTOR.message_types_by_name['Counter'] = _COUNTER DESCRIPTOR.message_types_by_name['GlobalDataRequest'] = _GLOBALDATAREQUEST DESCRIPTOR.message_types_by_name['WorkItemCommitRequest'] = _WORKITEMCOMMITREQUEST DESCRIPTOR.message_types_by_name['ComputationCommitWorkRequest'] = _COMPUTATIONCOMMITWORKREQUEST DESCRIPTOR.message_types_by_name['CommitWorkRequest'] = _COMMITWORKREQUEST DESCRIPTOR.message_types_by_name['CommitWorkResponse'] = _COMMITWORKRESPONSE DESCRIPTOR.message_types_by_name['GetConfigRequest'] = _GETCONFIGREQUEST DESCRIPTOR.message_types_by_name['GetConfigResponse'] = _GETCONFIGRESPONSE DESCRIPTOR.message_types_by_name['Exception'] = _EXCEPTION DESCRIPTOR.message_types_by_name['ReportStatsRequest'] = _REPORTSTATSREQUEST DESCRIPTOR.message_types_by_name['ReportStatsResponse'] = _REPORTSTATSRESPONSE Message = _reflection.GeneratedProtocolMessageType('Message', (_message.Message,), dict( DESCRIPTOR = _MESSAGE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.Message) )) _sym_db.RegisterMessage(Message) Timer = _reflection.GeneratedProtocolMessageType('Timer', (_message.Message,), dict( DESCRIPTOR = _TIMER, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.Timer) )) _sym_db.RegisterMessage(Timer) InputMessageBundle = _reflection.GeneratedProtocolMessageType('InputMessageBundle', (_message.Message,), dict( DESCRIPTOR = _INPUTMESSAGEBUNDLE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.InputMessageBundle) )) _sym_db.RegisterMessage(InputMessageBundle) KeyedMessageBundle = _reflection.GeneratedProtocolMessageType('KeyedMessageBundle', (_message.Message,), dict( DESCRIPTOR = _KEYEDMESSAGEBUNDLE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.KeyedMessageBundle) )) _sym_db.RegisterMessage(KeyedMessageBundle) OutputMessageBundle = _reflection.GeneratedProtocolMessageType('OutputMessageBundle', (_message.Message,), dict( DESCRIPTOR = _OUTPUTMESSAGEBUNDLE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.OutputMessageBundle) )) _sym_db.RegisterMessage(OutputMessageBundle) PubSubMessageBundle = _reflection.GeneratedProtocolMessageType('PubSubMessageBundle', (_message.Message,), dict( DESCRIPTOR = _PUBSUBMESSAGEBUNDLE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.PubSubMessageBundle) )) _sym_db.RegisterMessage(PubSubMessageBundle) TimerBundle = _reflection.GeneratedProtocolMessageType('TimerBundle', (_message.Message,), dict( DESCRIPTOR = _TIMERBUNDLE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.TimerBundle) )) _sym_db.RegisterMessage(TimerBundle) Value = _reflection.GeneratedProtocolMessageType('Value', (_message.Message,), dict( DESCRIPTOR = _VALUE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.Value) )) _sym_db.RegisterMessage(Value) TagValue = _reflection.GeneratedProtocolMessageType('TagValue', (_message.Message,), dict( DESCRIPTOR = _TAGVALUE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.TagValue) )) _sym_db.RegisterMessage(TagValue) TagList = _reflection.GeneratedProtocolMessageType('TagList', (_message.Message,), dict( DESCRIPTOR = _TAGLIST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.TagList) )) _sym_db.RegisterMessage(TagList) GlobalDataId = _reflection.GeneratedProtocolMessageType('GlobalDataId', (_message.Message,), dict( DESCRIPTOR = _GLOBALDATAID, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GlobalDataId) )) _sym_db.RegisterMessage(GlobalDataId) GlobalData = _reflection.GeneratedProtocolMessageType('GlobalData', (_message.Message,), dict( DESCRIPTOR = _GLOBALDATA, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GlobalData) )) _sym_db.RegisterMessage(GlobalData) SourceState = _reflection.GeneratedProtocolMessageType('SourceState', (_message.Message,), dict( DESCRIPTOR = _SOURCESTATE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.SourceState) )) _sym_db.RegisterMessage(SourceState) WatermarkHold = _reflection.GeneratedProtocolMessageType('WatermarkHold', (_message.Message,), dict( DESCRIPTOR = _WATERMARKHOLD, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.WatermarkHold) )) _sym_db.RegisterMessage(WatermarkHold) WorkItem = _reflection.GeneratedProtocolMessageType('WorkItem', (_message.Message,), dict( DESCRIPTOR = _WORKITEM, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.WorkItem) )) _sym_db.RegisterMessage(WorkItem) ComputationWorkItems = _reflection.GeneratedProtocolMessageType('ComputationWorkItems', (_message.Message,), dict( DESCRIPTOR = _COMPUTATIONWORKITEMS, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.ComputationWorkItems) )) _sym_db.RegisterMessage(ComputationWorkItems) GetWorkRequest = _reflection.GeneratedProtocolMessageType('GetWorkRequest', (_message.Message,), dict( DESCRIPTOR = _GETWORKREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GetWorkRequest) )) _sym_db.RegisterMessage(GetWorkRequest) GetWorkResponse = _reflection.GeneratedProtocolMessageType('GetWorkResponse', (_message.Message,), dict( DESCRIPTOR = _GETWORKRESPONSE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GetWorkResponse) )) _sym_db.RegisterMessage(GetWorkResponse) KeyedGetDataRequest = _reflection.GeneratedProtocolMessageType('KeyedGetDataRequest', (_message.Message,), dict( DESCRIPTOR = _KEYEDGETDATAREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.KeyedGetDataRequest) )) _sym_db.RegisterMessage(KeyedGetDataRequest) ComputationGetDataRequest = _reflection.GeneratedProtocolMessageType('ComputationGetDataRequest', (_message.Message,), dict( DESCRIPTOR = _COMPUTATIONGETDATAREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.ComputationGetDataRequest) )) _sym_db.RegisterMessage(ComputationGetDataRequest) GetDataRequest = _reflection.GeneratedProtocolMessageType('GetDataRequest', (_message.Message,), dict( DESCRIPTOR = _GETDATAREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GetDataRequest) )) _sym_db.RegisterMessage(GetDataRequest) KeyedGetDataResponse = _reflection.GeneratedProtocolMessageType('KeyedGetDataResponse', (_message.Message,), dict( DESCRIPTOR = _KEYEDGETDATARESPONSE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.KeyedGetDataResponse) )) _sym_db.RegisterMessage(KeyedGetDataResponse) ComputationGetDataResponse = _reflection.GeneratedProtocolMessageType('ComputationGetDataResponse', (_message.Message,), dict( DESCRIPTOR = _COMPUTATIONGETDATARESPONSE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.ComputationGetDataResponse) )) _sym_db.RegisterMessage(ComputationGetDataResponse) GetDataResponse = _reflection.GeneratedProtocolMessageType('GetDataResponse', (_message.Message,), dict( DESCRIPTOR = _GETDATARESPONSE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GetDataResponse) )) _sym_db.RegisterMessage(GetDataResponse) Counter = _reflection.GeneratedProtocolMessageType('Counter', (_message.Message,), dict( DESCRIPTOR = _COUNTER, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.Counter) )) _sym_db.RegisterMessage(Counter) GlobalDataRequest = _reflection.GeneratedProtocolMessageType('GlobalDataRequest', (_message.Message,), dict( DESCRIPTOR = _GLOBALDATAREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GlobalDataRequest) )) _sym_db.RegisterMessage(GlobalDataRequest) WorkItemCommitRequest = _reflection.GeneratedProtocolMessageType('WorkItemCommitRequest', (_message.Message,), dict( DESCRIPTOR = _WORKITEMCOMMITREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.WorkItemCommitRequest) )) _sym_db.RegisterMessage(WorkItemCommitRequest) ComputationCommitWorkRequest = _reflection.GeneratedProtocolMessageType('ComputationCommitWorkRequest', (_message.Message,), dict( DESCRIPTOR = _COMPUTATIONCOMMITWORKREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.ComputationCommitWorkRequest) )) _sym_db.RegisterMessage(ComputationCommitWorkRequest) CommitWorkRequest = _reflection.GeneratedProtocolMessageType('CommitWorkRequest', (_message.Message,), dict( DESCRIPTOR = _COMMITWORKREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.CommitWorkRequest) )) _sym_db.RegisterMessage(CommitWorkRequest) CommitWorkResponse = _reflection.GeneratedProtocolMessageType('CommitWorkResponse', (_message.Message,), dict( DESCRIPTOR = _COMMITWORKRESPONSE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.CommitWorkResponse) )) _sym_db.RegisterMessage(CommitWorkResponse) GetConfigRequest = _reflection.GeneratedProtocolMessageType('GetConfigRequest', (_message.Message,), dict( DESCRIPTOR = _GETCONFIGREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GetConfigRequest) )) _sym_db.RegisterMessage(GetConfigRequest) GetConfigResponse = _reflection.GeneratedProtocolMessageType('GetConfigResponse', (_message.Message,), dict( NameMapEntry = _reflection.GeneratedProtocolMessageType('NameMapEntry', (_message.Message,), dict( DESCRIPTOR = _GETCONFIGRESPONSE_NAMEMAPENTRY, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GetConfigResponse.NameMapEntry) )) , SystemNameToComputationIdMapEntry = _reflection.GeneratedProtocolMessageType('SystemNameToComputationIdMapEntry', (_message.Message,), dict( DESCRIPTOR = _GETCONFIGRESPONSE_SYSTEMNAMETOCOMPUTATIONIDMAPENTRY, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GetConfigResponse.SystemNameToComputationIdMapEntry) )) , DESCRIPTOR = _GETCONFIGRESPONSE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.GetConfigResponse) )) _sym_db.RegisterMessage(GetConfigResponse) _sym_db.RegisterMessage(GetConfigResponse.NameMapEntry) _sym_db.RegisterMessage(GetConfigResponse.SystemNameToComputationIdMapEntry) Exception = _reflection.GeneratedProtocolMessageType('Exception', (_message.Message,), dict( DESCRIPTOR = _EXCEPTION, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.Exception) )) _sym_db.RegisterMessage(Exception) ReportStatsRequest = _reflection.GeneratedProtocolMessageType('ReportStatsRequest', (_message.Message,), dict( DESCRIPTOR = _REPORTSTATSREQUEST, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.ReportStatsRequest) )) _sym_db.RegisterMessage(ReportStatsRequest) ReportStatsResponse = _reflection.GeneratedProtocolMessageType('ReportStatsResponse', (_message.Message,), dict( DESCRIPTOR = _REPORTSTATSRESPONSE, __module__ = 'windmill_pb2' # @@protoc_insertion_point(class_scope:windmill.ReportStatsResponse) )) _sym_db.RegisterMessage(ReportStatsResponse) DESCRIPTOR.has_options = True DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n5com.google.cloud.dataflow.sdk.runners.worker.windmillB\010Windmill')) _WATERMARKHOLD.fields_by_name['timestamps'].has_options = True _WATERMARKHOLD.fields_by_name['timestamps']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001')) from grpc.beta import implementations as beta_implementations from grpc.beta import interfaces as beta_interfaces from grpc.framework.common import cardinality from grpc.framework.interfaces.face import utilities as face_utilities # @@protoc_insertion_point(module_scope) DataflowPythonSDK-0.2.5/google/cloud/dataflow/internal/windmill_service_pb2.py000066400000000000000000000224631272340251000275010ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. # Generated by the protocol buffer compiler. DO NOT EDIT! # source: windmill_service.proto import sys _b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) from google.protobuf import descriptor as _descriptor from google.protobuf import message as _message from google.protobuf import reflection as _reflection from google.protobuf import symbol_database as _symbol_database from google.protobuf import descriptor_pb2 # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() import windmill_pb2 as windmill__pb2 DESCRIPTOR = _descriptor.FileDescriptor( name='windmill_service.proto', package='google.dataflow.windmillservice.v1alpha1', syntax='proto2', serialized_pb=_b('\n\x16windmill_service.proto\x12(google.dataflow.windmillservice.v1alpha1\x1a\x0ewindmill.proto2\xf9\x02\n\x1c\x43loudWindmillServiceV1Alpha1\x12>\n\x07GetWork\x12\x18.windmill.GetWorkRequest\x1a\x19.windmill.GetWorkResponse\x12>\n\x07GetData\x12\x18.windmill.GetDataRequest\x1a\x19.windmill.GetDataResponse\x12G\n\nCommitWork\x12\x1b.windmill.CommitWorkRequest\x1a\x1c.windmill.CommitWorkResponse\x12\x44\n\tGetConfig\x12\x1a.windmill.GetConfigRequest\x1a\x1b.windmill.GetConfigResponse\x12J\n\x0bReportStats\x12\x1c.windmill.ReportStatsRequest\x1a\x1d.windmill.ReportStatsResponseB7\n5com.google.cloud.dataflow.sdk.runners.worker.windmill') , dependencies=[windmill__pb2.DESCRIPTOR,]) _sym_db.RegisterFileDescriptor(DESCRIPTOR) DESCRIPTOR.has_options = True DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n5com.google.cloud.dataflow.sdk.runners.worker.windmill')) from grpc.beta import implementations as beta_implementations from grpc.beta import interfaces as beta_interfaces from grpc.framework.common import cardinality from grpc.framework.interfaces.face import utilities as face_utilities class BetaCloudWindmillServiceV1Alpha1Servicer(object): """The Cloud Windmill Service API used by GCE to acquire and process streaming Dataflow work. """ def GetWork(self, request, context): """Gets streaming Dataflow work. """ context.code(beta_interfaces.StatusCode.UNIMPLEMENTED) def GetData(self, request, context): """Gets data from Windmill. """ context.code(beta_interfaces.StatusCode.UNIMPLEMENTED) def CommitWork(self, request, context): """Commits previously acquired work. """ context.code(beta_interfaces.StatusCode.UNIMPLEMENTED) def GetConfig(self, request, context): """Gets dependant configuration from windmill. """ context.code(beta_interfaces.StatusCode.UNIMPLEMENTED) def ReportStats(self, request, context): """Reports stats to Windmill. """ context.code(beta_interfaces.StatusCode.UNIMPLEMENTED) class BetaCloudWindmillServiceV1Alpha1Stub(object): """The Cloud Windmill Service API used by GCE to acquire and process streaming Dataflow work. """ def GetWork(self, request, timeout, metadata=None, with_call=False, protocol_options=None): """Gets streaming Dataflow work. """ raise NotImplementedError() GetWork.future = None def GetData(self, request, timeout, metadata=None, with_call=False, protocol_options=None): """Gets data from Windmill. """ raise NotImplementedError() GetData.future = None def CommitWork(self, request, timeout, metadata=None, with_call=False, protocol_options=None): """Commits previously acquired work. """ raise NotImplementedError() CommitWork.future = None def GetConfig(self, request, timeout, metadata=None, with_call=False, protocol_options=None): """Gets dependant configuration from windmill. """ raise NotImplementedError() GetConfig.future = None def ReportStats(self, request, timeout, metadata=None, with_call=False, protocol_options=None): """Reports stats to Windmill. """ raise NotImplementedError() ReportStats.future = None def beta_create_CloudWindmillServiceV1Alpha1_server(servicer, pool=None, pool_size=None, default_timeout=None, maximum_timeout=None): request_deserializers = { ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'CommitWork'): windmill__pb2.CommitWorkRequest.FromString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetConfig'): windmill__pb2.GetConfigRequest.FromString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetData'): windmill__pb2.GetDataRequest.FromString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetWork'): windmill__pb2.GetWorkRequest.FromString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'ReportStats'): windmill__pb2.ReportStatsRequest.FromString, } response_serializers = { ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'CommitWork'): windmill__pb2.CommitWorkResponse.SerializeToString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetConfig'): windmill__pb2.GetConfigResponse.SerializeToString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetData'): windmill__pb2.GetDataResponse.SerializeToString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetWork'): windmill__pb2.GetWorkResponse.SerializeToString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'ReportStats'): windmill__pb2.ReportStatsResponse.SerializeToString, } method_implementations = { ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'CommitWork'): face_utilities.unary_unary_inline(servicer.CommitWork), ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetConfig'): face_utilities.unary_unary_inline(servicer.GetConfig), ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetData'): face_utilities.unary_unary_inline(servicer.GetData), ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetWork'): face_utilities.unary_unary_inline(servicer.GetWork), ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'ReportStats'): face_utilities.unary_unary_inline(servicer.ReportStats), } server_options = beta_implementations.server_options(request_deserializers=request_deserializers, response_serializers=response_serializers, thread_pool=pool, thread_pool_size=pool_size, default_timeout=default_timeout, maximum_timeout=maximum_timeout) return beta_implementations.server(method_implementations, options=server_options) def beta_create_CloudWindmillServiceV1Alpha1_stub(channel, host=None, metadata_transformer=None, pool=None, pool_size=None): request_serializers = { ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'CommitWork'): windmill__pb2.CommitWorkRequest.SerializeToString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetConfig'): windmill__pb2.GetConfigRequest.SerializeToString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetData'): windmill__pb2.GetDataRequest.SerializeToString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetWork'): windmill__pb2.GetWorkRequest.SerializeToString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'ReportStats'): windmill__pb2.ReportStatsRequest.SerializeToString, } response_deserializers = { ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'CommitWork'): windmill__pb2.CommitWorkResponse.FromString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetConfig'): windmill__pb2.GetConfigResponse.FromString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetData'): windmill__pb2.GetDataResponse.FromString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'GetWork'): windmill__pb2.GetWorkResponse.FromString, ('google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', 'ReportStats'): windmill__pb2.ReportStatsResponse.FromString, } cardinalities = { 'CommitWork': cardinality.Cardinality.UNARY_UNARY, 'GetConfig': cardinality.Cardinality.UNARY_UNARY, 'GetData': cardinality.Cardinality.UNARY_UNARY, 'GetWork': cardinality.Cardinality.UNARY_UNARY, 'ReportStats': cardinality.Cardinality.UNARY_UNARY, } stub_options = beta_implementations.stub_options(host=host, metadata_transformer=metadata_transformer, request_serializers=request_serializers, response_deserializers=response_deserializers, thread_pool=pool, thread_pool_size=pool_size) return beta_implementations.dynamic_stub(channel, 'google.dataflow.windmillservice.v1alpha1.CloudWindmillServiceV1Alpha1', cardinalities, options=stub_options) # @@protoc_insertion_point(module_scope) DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/000077500000000000000000000000001272340251000216115ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/__init__.py000066400000000000000000000021031272340251000237160ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A package defining several input sources and output sinks.""" # pylint: disable=wildcard-import from google.cloud.dataflow.io.bigquery import * from google.cloud.dataflow.io.fileio import * from google.cloud.dataflow.io.iobase import Read from google.cloud.dataflow.io.iobase import Sink from google.cloud.dataflow.io.iobase import Write from google.cloud.dataflow.io.iobase import Writer from google.cloud.dataflow.io.pubsub import * from google.cloud.dataflow.io.range_trackers import * DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/bigquery.py000066400000000000000000001050201272340251000240100ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """BigQuery sources and sinks. This module implements reading from and writing to BigQuery tables. It relies on several classes exposed by the BigQuery API: TableSchema, TableFieldSchema, TableRow, and TableCell. The default mode is to return table rows read from a BigQuery source as dictionaries. Similarly a Write transform to a BigQuerySink accepts PCollections of dictionaries. This is done for more convenient programming. If desired, the native TableRow objects can be used throughout to represent rows (use an instance of TableRowJsonCoder as a coder argument when creating the sources or sinks respectively). Also, for programming convenience, instances of TableReference and TableSchema have a string representation that can be used for the corresponding arguments: - TableReference can be a PROJECT:DATASET.TABLE or DATASET.TABLE string. - TableSchema can be a NAME:TYPE{,NAME:TYPE}* string (e.g. 'month:STRING,event_count:INTEGER'). The syntax supported is described here: https://cloud.google.com/bigquery/bq-command-line-tool-quickstart BigQuery sources can be used as main inputs or side inputs. A main input (common case) is expected to be massive and the Dataflow service will make sure it is split into manageable chunks and processed in parallel. Side inputs are expected to be small and will be read completely every time a ParDo DoFn gets executed. In the example below the lambda function implementing the DoFn for the Map transform will get on each call *one* row of the main table and *all* rows of the side table. The execution framework may use some caching techniques to share the side inputs between calls in order to avoid excessive reading:: main_table = pipeline | df.io.Read(df.io.BigQuerySource('very_big_table') side_table = pipeline | df.io.Read(df.io.BigQuerySource('not_so_big_table') results = ( main_table | df.Map('process data', lambda element, side_input: ..., AsList(side_table))) There is no difference in how main and side inputs are read. What makes the side_table a 'side input' is the AsList wrapper used when passing the table as a parameter to the Map transform. AsList signals to the execution framework that its input should be made available whole. The main and side inputs are implemented differently. Reading a BigQuery table as main input entails exporting the table to a set of GCS files (currently in JSON format) and then processing those files. Reading the same table as a side input entails querying the table for all its rows. The coder argument on BigQuerySource controls the reading of the lines in the export files (i.e., transform a JSON object into a PCollection element). The coder is not involved when the same table is read as a side input since there is no intermediate format involved. We get the table rows directly from the BigQuery service with a query. Users may provide a query to read from rather than reading all of a BigQuery table. If specified, the result obtained by executing the specified query will be used as the data of the input transform. query_results = pipeline | df.io.Read(df.io.BigQuerySource( query='SELECT year, mean_temp FROM samples.weather_stations')) When creating a BigQuery input transform, users should provide either a query or a table. Pipeline construction will fail with a validation error if neither or both are specified. *** Short introduction to BigQuery concepts *** Tables have rows (TableRow) and each row has cells (TableCell). A table has a schema (TableSchema), which in turn describes the schema of each cell (TableFieldSchema). The terms field and cell are used interchangeably. TableSchema: Describes the schema (types and order) for values in each row. Has one attribute, 'field', which is list of TableFieldSchema objects. TableFieldSchema: Describes the schema (type, name) for one field. Has several attributes, including 'name' and 'type'. Common values for the type attribute are: 'STRING', 'INTEGER', 'FLOAT', 'BOOLEAN'. All possible values are described at: https://cloud.google.com/bigquery/preparing-data-for-bigquery#datatypes TableRow: Holds all values in a table row. Has one attribute, 'f', which is a list of TableCell instances. TableCell: Holds the value for one cell (or field). Has one attribute, 'v', which is a JsonValue instance. This class is defined in apitools.base.py.extra_types.py module. """ from __future__ import absolute_import import collections import json import logging import re import time import uuid from google.cloud.dataflow import coders from google.cloud.dataflow.internal import auth from google.cloud.dataflow.internal.json_value import from_json_value from google.cloud.dataflow.internal.json_value import to_json_value from google.cloud.dataflow.io import iobase from google.cloud.dataflow.utils import retry from google.cloud.dataflow.utils.options import GoogleCloudOptions from apitools.base.py.exceptions import HttpError # Protect against environments where bigquery library is not available. # pylint: disable=g-import-not-at-top try: from google.cloud.dataflow.internal.clients import bigquery except ImportError: pass # pylint: enable=g-import-not-at-top __all__ = [ 'TableRowJsonCoder', 'BigQueryDisposition', 'BigQuerySource', 'BigQuerySink', ] class RowAsDictJsonCoder(coders.Coder): """A coder for a table row (represented as a dict) to/from a JSON string. This is the default coder for sources and sinks if the coder argument is not specified. """ def encode(self, table_row): return json.dumps(table_row) def decode(self, encoded_table_row): return json.loads(encoded_table_row) class TableRowJsonCoder(coders.Coder): """A coder for a TableRow instance to/from a JSON string. Note that the encoding operation (used when writing to sinks) requires the table schema in order to obtain the ordered list of field names. Reading from sources on the other hand does not need the table schema. """ def __init__(self, table_schema=None): # The table schema is needed for encoding TableRows as JSON (writing to # sinks) because the ordered list of field names is used in the JSON # representation. self.table_schema = table_schema # Precompute field names since we need them for row encoding. if self.table_schema: self.field_names = tuple(fs.name for fs in self.table_schema.fields) def encode(self, table_row): if self.table_schema is None: raise AttributeError( 'The TableRowJsonCoder requires a table schema for ' 'encoding operations. Please specify a table_schema argument.') return json.dumps( collections.OrderedDict( zip(self.field_names, [from_json_value(f.v) for f in table_row.f]))) def decode(self, encoded_table_row): od = json.loads( encoded_table_row, object_pairs_hook=collections.OrderedDict) return bigquery.TableRow( f=[bigquery.TableCell(v=to_json_value(e)) for e in od.itervalues()]) class BigQueryDisposition(object): """Class holding standard strings used for create and write dispositions.""" CREATE_NEVER = 'CREATE_NEVER' CREATE_IF_NEEDED = 'CREATE_IF_NEEDED' WRITE_TRUNCATE = 'WRITE_TRUNCATE' WRITE_APPEND = 'WRITE_APPEND' WRITE_EMPTY = 'WRITE_EMPTY' @staticmethod def validate_create(disposition): values = (BigQueryDisposition.CREATE_NEVER, BigQueryDisposition.CREATE_IF_NEEDED) if disposition not in values: raise ValueError( 'Invalid create disposition %s. Expecting %s' % (disposition, values)) return disposition @staticmethod def validate_write(disposition): values = (BigQueryDisposition.WRITE_TRUNCATE, BigQueryDisposition.WRITE_APPEND, BigQueryDisposition.WRITE_EMPTY) if disposition not in values: raise ValueError( 'Invalid write disposition %s. Expecting %s' % (disposition, values)) return disposition def _parse_table_reference(table, dataset=None, project=None): """Parses a table reference into a (project, dataset, table) tuple. Args: table: The ID of the table. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). If dataset argument is None then the table argument must contain the entire table reference: 'DATASET.TABLE' or 'PROJECT:DATASET.TABLE'. This argument can be a bigquery.TableReference instance in which case dataset and project are ignored and the reference is returned as a result. dataset: The ID of the dataset containing this table or null if the table reference is specified entirely by the table argument. project: The ID of the project containing this table or null if the table reference is specified entirely by the table (and possibly dataset) argument. Returns: A bigquery.TableReference object. The object has the following attributes: projectId, datasetId, and tableId. Raises: ValueError: if the table reference as a string does not match the expected format. """ if isinstance(table, bigquery.TableReference): return table table_reference = bigquery.TableReference() # If dataset argument is not specified, the expectation is that the # table argument will contain a full table reference instead of just a # table name. if dataset is None: match = re.match( r'^((?P.+):)?(?P\w+)\.(?P\w+)$', table) if not match: raise ValueError( 'Expected a table reference (PROJECT:DATASET.TABLE or ' 'DATASET.TABLE) instead of %s.' % table) table_reference.projectId = match.group('project') table_reference.datasetId = match.group('dataset') table_reference.tableId = match.group('table') else: table_reference.projectId = project table_reference.datasetId = dataset table_reference.tableId = table return table_reference # ----------------------------------------------------------------------------- # BigQuerySource, BigQuerySink. class BigQuerySource(iobase.NativeSource): """A source based on a BigQuery table.""" def __init__(self, table=None, dataset=None, project=None, query=None, validate=False, coder=None): """Initialize a BigQuerySource. Args: table: The ID of a BigQuery table. If specified all data of the table will be used as input of the current source. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). If dataset and query arguments are None then the table argument must contain the entire table reference specified as: 'DATASET.TABLE' or 'PROJECT:DATASET.TABLE'. dataset: The ID of the dataset containing this table or null if the table reference is specified entirely by the table argument or a query is specified. project: The ID of the project containing this table or null if the table reference is specified entirely by the table argument or a query is specified. query: A query to be used instead of arguments table, dataset, and project. validate: If true, various checks will be done when source gets initialized (e.g., is table present?). This should be True for most scenarios in order to catch errors as early as possible (pipeline construction instead of pipeline execution). It should be False if the table is created during pipeline execution by a previous step. coder: The coder for the table rows if serialized to disk. If None, then the default coder is RowAsDictJsonCoder, which will interpret every line in a file as a JSON serialized dictionary. This argument needs a value only in special cases when returning table rows as dictionaries is not desirable. Raises: ValueError: if any of the following is true (1) the table reference as a string does not match the expected format (2) neither a table nor a query is specified (3) both a table and a query is specified. """ if table is not None and query is not None: raise ValueError('Both a BigQuery table and a query were specified.' ' Please specify only one of these.') elif table is None and query is None: raise ValueError('A BigQuery table or a query must be specified') elif table is not None: self.table_reference = _parse_table_reference(table, dataset, project) self.query = None else: self.query = query self.table_reference = None self.validate = validate self.coder = coder or RowAsDictJsonCoder() @property def format(self): """Source format name required for remote execution.""" return 'bigquery' def reader(self, test_bigquery_client=None): return BigQueryReader( source=self, test_bigquery_client=test_bigquery_client) class BigQuerySink(iobase.NativeSink): """A sink based on a BigQuery table.""" def __init__(self, table, dataset=None, project=None, schema=None, create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=BigQueryDisposition.WRITE_EMPTY, validate=False, coder=None): """Initialize a BigQuerySink. Args: table: The ID of the table. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). If dataset argument is None then the table argument must contain the entire table reference specified as: 'DATASET.TABLE' or 'PROJECT:DATASET.TABLE'. dataset: The ID of the dataset containing this table or null if the table reference is specified entirely by the table argument. project: The ID of the project containing this table or null if the table reference is specified entirely by the table argument. schema: The schema to be used if the BigQuery table to write has to be created. This can be either specified as a 'bigquery.TableSchema' object or a single string of the form 'field1:type1,field2:type2,field3:type3' that defines a comma separated list of fields. Here 'type' should specify the BigQuery type of the field. Single string based schemas do not support nested fields, repeated fields, or specifying a BigQuery mode for fields (mode will always be set to 'NULLABLE'). create_disposition: A string describing what happens if the table does not exist. Possible values are: - BigQueryDisposition.CREATE_IF_NEEDED: create if does not exist. - BigQueryDisposition.CREATE_NEVER: fail the write if does not exist. write_disposition: A string describing what happens if the table has already some data. Possible values are: - BigQueryDisposition.WRITE_TRUNCATE: delete existing rows. - BigQueryDisposition.WRITE_APPEND: add to existing rows. - BigQueryDisposition.WRITE_EMPTY: fail the write if table not empty. validate: If true, various checks will be done when sink gets initialized (e.g., is table present given the disposition arguments?). This should be True for most scenarios in order to catch errors as early as possible (pipeline construction instead of pipeline execution). It should be False if the table is created during pipeline execution by a previous step. coder: The coder for the table rows if serialized to disk. If None, then the default coder is RowAsDictJsonCoder, which will interpret every element written to the sink as a dictionary that will be JSON serialized as a line in a file. This argument needs a value only in special cases when writing table rows as dictionaries is not desirable. Raises: TypeError: if the schema argument is not a string or a TableSchema object. ValueError: if the table reference as a string does not match the expected format. """ self.table_reference = _parse_table_reference(table, dataset, project) # Transform the table schema into a bigquery.TableSchema instance. if isinstance(schema, basestring): # TODO(silviuc): Should add a regex-based validation of the format. table_schema = bigquery.TableSchema() schema_list = [s.strip(' ') for s in schema.split(',')] for field_and_type in schema_list: field_name, field_type = field_and_type.split(':') field_schema = bigquery.TableFieldSchema() field_schema.name = field_name field_schema.type = field_type field_schema.mode = 'NULLABLE' table_schema.fields.append(field_schema) self.table_schema = table_schema elif schema is None: # TODO(silviuc): Should check that table exists if no schema specified. self.table_schema = schema elif isinstance(schema, bigquery.TableSchema): self.table_schema = schema else: raise TypeError('Unexpected schema argument: %s.' % schema) self.create_disposition = BigQueryDisposition.validate_create( create_disposition) self.write_disposition = BigQueryDisposition.validate_write( write_disposition) self.validate = validate self.coder = coder or RowAsDictJsonCoder() def schema_as_json(self): """Returns the TableSchema associated with the sink as a JSON string.""" def schema_list_as_object(schema_list): """Returns a list of TableFieldSchema objects as a list of dicts.""" fields = [] for f in schema_list: fs = {'name': f.name, 'type': f.type} if f.description is not None: fs['description'] = f.description if f.mode is not None: fs['mode'] = f.mode if f.type.lower() == 'record': fs['fields'] = schema_list_as_object(f.fields) fields.append(fs) return fields return json.dumps( {'fields': schema_list_as_object(self.table_schema.fields)}) @property def format(self): """Sink format name required for remote execution.""" return 'bigquery' def writer(self, test_bigquery_client=None, buffer_size=None): return BigQueryWriter( sink=self, test_bigquery_client=test_bigquery_client, buffer_size=buffer_size) # ----------------------------------------------------------------------------- # BigQueryReader, BigQueryWriter. class BigQueryReader(iobase.NativeSourceReader): """A reader for a BigQuery source.""" def __init__(self, source, test_bigquery_client=None): self.source = source self.test_bigquery_client = test_bigquery_client if auth.is_running_in_gce: self.executing_project = auth.executing_project elif hasattr(source, 'pipeline_options'): self.executing_project = ( source.pipeline_options.view_as(GoogleCloudOptions).project) else: self.executing_project = None # TODO(silviuc): Try to automatically get it from gcloud config info. if not self.executing_project and test_bigquery_client is None: raise RuntimeError( 'Missing executing project information. Please use the --project ' 'command line option to specify it.') self.row_as_dict = isinstance(self.source.coder, RowAsDictJsonCoder) # Schema for the rows being read by the reader. It is initialized the # first time something gets read from the table. It is not required # for reading the field values in each row but could be useful for # getting additional details. self.schema = None if self.source.query is None: # If table schema did not define a project we default to executing # project. project_id = self.source.table_reference.projectId if not project_id: project_id = self.executing_project self.query = 'SELECT * FROM [%s:%s.%s];' % ( project_id, self.source.table_reference.datasetId, self.source.table_reference.tableId) else: self.query = self.source.query def __enter__(self): self.client = BigQueryWrapper(client=self.test_bigquery_client) return self def __exit__(self, exception_type, exception_value, traceback): pass def __iter__(self): for rows, schema in self.client.run_query( project_id=self.executing_project, query=self.query): if self.schema is None: self.schema = schema for row in rows: if self.row_as_dict: yield self.client.convert_row_to_dict(row, schema) else: yield row class BigQueryWriter(iobase.NativeSinkWriter): """The sink writer for a BigQuerySink.""" def __init__(self, sink, test_bigquery_client=None, buffer_size=None): self.sink = sink self.test_bigquery_client = test_bigquery_client self.row_as_dict = isinstance(self.sink.coder, RowAsDictJsonCoder) # Buffer used to batch written rows so we reduce communication with the # BigQuery service. self.rows_buffer = [] self.rows_buffer_flush_threshold = buffer_size or 1000 # Figure out the project, dataset, and table used for the sink. self.project_id = self.sink.table_reference.projectId # If table schema did not define a project we default to executing project. if self.project_id is None and hasattr(sink, 'pipeline_options'): self.project_id = ( sink.pipeline_options.view_as(GoogleCloudOptions).project) assert self.project_id is not None self.dataset_id = self.sink.table_reference.datasetId self.table_id = self.sink.table_reference.tableId def _flush_rows_buffer(self): if self.rows_buffer: logging.info('Writing %d rows to %s:%s.%s table.', len(self.rows_buffer), self.project_id, self.dataset_id, self.table_id) passed, errors = self.client.insert_rows( project_id=self.project_id, dataset_id=self.dataset_id, table_id=self.table_id, rows=self.rows_buffer) self.rows_buffer = [] if not passed: raise RuntimeError('Could not successfully insert rows to BigQuery' ' table [%s:%s.%s]. Errors: %s'% (self.project_id, self.dataset_id, self.table_id, errors)) def __enter__(self): self.client = BigQueryWrapper(client=self.test_bigquery_client) self.client.get_or_create_table( self.project_id, self.dataset_id, self.table_id, self.sink.table_schema, self.sink.create_disposition, self.sink.write_disposition) return self def __exit__(self, exception_type, exception_value, traceback): self._flush_rows_buffer() def Write(self, row): self.rows_buffer.append(row) if len(self.rows_buffer) > self.rows_buffer_flush_threshold: self._flush_rows_buffer() # ----------------------------------------------------------------------------- # BigQueryWrapper. class BigQueryWrapper(object): """BigQuery client wrapper with utilities for querying. The wrapper is used to organize all the BigQuery integration points and offer a common place where retry logic for failures can be controlled. In addition it offers various functions used both in sources and sinks (e.g., find and create tables, query a table, etc.). """ def __init__(self, client=None): self.client = client or bigquery.BigqueryV2( credentials=auth.get_service_credentials()) self._unique_row_id = 0 # For testing scenarios where we pass in a client we do not want a # randomized prefix for row IDs. self._row_id_prefix = '' if client else uuid.uuid4() @property def unique_row_id(self): """Returns a unique row ID (str) used to avoid multiple insertions. If the row ID is provided, BigQuery will make a best effort to not insert the same row multiple times for fail and retry scenarios in which the insert request may be issued several times. This comes into play for sinks executed in a local runner. Returns: a unique row ID string """ self._unique_row_id += 1 return '%s_%d' % (self._row_id_prefix, self._unique_row_id) @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def _start_query_job(self, project_id, query, dry_run=False): request = bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=bigquery.Job( configuration=bigquery.JobConfiguration( dryRun=dry_run, query=bigquery.JobConfigurationQuery( query=query)))) response = self.client.jobs.Insert(request) return response.jobReference.jobId @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def _get_query_results(self, project_id, job_id, page_token=None, max_results=10000): request = bigquery.BigqueryJobsGetQueryResultsRequest( jobId=job_id, pageToken=page_token, projectId=project_id, maxResults=max_results) response = self.client.jobs.GetQueryResults(request) return response @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def _insert_all_rows(self, project_id, dataset_id, table_id, rows): # The rows argument is a list of # bigquery.TableDataInsertAllRequest.RowsValueListEntry instances as # required bu the InsertAll() method. request = bigquery.BigqueryTabledataInsertAllRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id, tableDataInsertAllRequest=bigquery.TableDataInsertAllRequest( # TODO(silviuc): Should have an option for skipInvalidRows? # TODO(silviuc): Should have an option for ignoreUnknownValues? rows=rows)) response = self.client.tabledata.InsertAll(request) # response.insertErrors is not [] if errors encountered. return not response.insertErrors, response.insertErrors @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def _get_table(self, project_id, dataset_id, table_id): request = bigquery.BigqueryTablesGetRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id) response = self.client.tables.Get(request) # The response is a bigquery.Table instance. return response @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def _create_table(self, project_id, dataset_id, table_id, schema): table = bigquery.Table( tableReference=bigquery.TableReference( projectId=project_id, datasetId=dataset_id, tableId=table_id), schema=schema) request = bigquery.BigqueryTablesInsertRequest( projectId=project_id, datasetId=dataset_id, table=table) response = self.client.tables.Insert(request) # The response is a bigquery.Table instance. return response @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def _is_table_empty(self, project_id, dataset_id, table_id): request = bigquery.BigqueryTabledataListRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id, maxResults=1) response = self.client.tabledata.List(request) # The response is a bigquery.TableDataList instance. return response.totalRows == 0 @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def _delete_table(self, project_id, dataset_id, table_id): request = bigquery.BigqueryTablesDeleteRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id) self.client.tables.Delete(request) def get_or_create_table( self, project_id, dataset_id, table_id, schema, create_disposition, write_disposition): """Gets or creates a table based on create and write dispositions. The function mimics the behavior of BigQuery import jobs when using the same create and write dispositions. Args: project_id: The project id owning the table. dataset_id: The dataset id owning the table. table_id: The table id. schema: A bigquery.TableSchema instance or None. create_disposition: CREATE_NEVER or CREATE_IF_NEEDED. write_disposition: WRITE_APPEND, WRITE_EMPTY or WRITE_TRUNCATE. Returns: A bigquery.Table instance if table was found or created. Raises: RuntimeError: For various mismatches between the state of the table and the create/write dispositions passed in. For example if the table is not empty and WRITE_EMPTY was specified then an error will be raised since the table was expected to be empty. """ found_table = None try: found_table = self._get_table(project_id, dataset_id, table_id) except HttpError as exn: if exn.status_code == 404: if create_disposition == BigQueryDisposition.CREATE_NEVER: raise RuntimeError( 'Table %s:%s.%s not found but create disposition is CREATE_NEVER.' % (project_id, dataset_id, table_id)) else: raise # If table exists already then handle the semantics for WRITE_EMPTY and # WRITE_TRUNCATE write dispositions. if found_table: table_empty = self._is_table_empty(project_id, dataset_id, table_id) if (not table_empty and write_disposition == BigQueryDisposition.WRITE_EMPTY): raise RuntimeError( 'Table %s:%s.%s is not empty but write disposition is WRITE_EMPTY.' % (project_id, dataset_id, table_id)) # Delete the table and recreate it (later) if WRITE_TRUNCATE was # specified. if write_disposition == BigQueryDisposition.WRITE_TRUNCATE: self._delete_table(project_id, dataset_id, table_id) # Create a new table potentially reusing the schema from a previously # found table in case the schema was not specified. if schema is None and found_table is None: raise RuntimeError( 'Table %s:%s.%s requires a schema. None can be inferred because the ' 'table does not exist.' % (project_id, dataset_id, table_id)) if found_table and write_disposition != BigQueryDisposition.WRITE_TRUNCATE: return found_table else: # if write_disposition == BigQueryDisposition.WRITE_TRUNCATE we delete # the table before this point. return self._create_table(project_id=project_id, dataset_id=dataset_id, table_id=table_id, schema=schema or found_table.schema) def run_query(self, project_id, query, dry_run=False): job_id = self._start_query_job(project_id, query, dry_run) if dry_run: # If this was a dry run then the fact that we get here means the # query has no errors. The start_query_job would raise an error otherwise. return page_token = None while True: response = self._get_query_results(project_id, job_id, page_token) if not response.jobComplete: # The jobComplete field can be False if the query request times out # (default is 10 seconds). Note that this is a timeout for the query # request not for the actual execution of the query in the service. If # the request times out we keep trying. This situation is quite possible # if the query will return a large number of rows. logging.info('Waiting on response from query: %s ...', query) time.sleep(1.0) continue # We got some results. The last page is signalled by a missing pageToken. yield response.rows, response.schema if not response.pageToken: break page_token = response.pageToken def insert_rows(self, project_id, dataset_id, table_id, rows): """Inserts rows into the specified table. Args: project_id: The project id owning the table. dataset_id: The dataset id owning the table. table_id: The table id. rows: A list of plain Python dictionaries. Each dictionary is a row and each key in it is the name of a field. Returns: A tuple (bool, errors). If first element is False then the second element will be a bigquery.InserttErrorsValueListEntry instance containing specific errors. """ # Prepare rows for insertion. Of special note is the row ID that we add to # each row in order to help BigQuery avoid inserting a row multiple times. # BigQuery will do a best-effort if unique IDs are provided. This situation # can happen during retries on failures. # TODO(silviuc): Must add support to writing TableRow's instead of dicts. final_rows = [] for row in rows: json_object = bigquery.JsonObject() for k, v in row.iteritems(): json_object.additionalProperties.append( bigquery.JsonObject.AdditionalProperty( key=k, value=to_json_value(v))) final_rows.append( bigquery.TableDataInsertAllRequest.RowsValueListEntry( insertId=str(self.unique_row_id), json=json_object)) result, errors = self._insert_all_rows( project_id, dataset_id, table_id, final_rows) return result, errors def convert_row_to_dict(self, row, schema): """Converts a TableRow instance using the schema to a Python dict.""" result = {} for index, field in enumerate(schema.fields): cell = row.f[index] if cell.v is None: continue # Field not present in the row. # The JSON values returned by BigQuery for table fields in a row have # always set the string_value attribute, which means the value below will # be a string. Converting to the appropriate type is not tricky except # for boolean values. For such values the string values are 'true' or # 'false', which cannot be converted by simply calling bool() (it will # return True for both!). value = from_json_value(cell.v) if field.type == 'STRING': value = value elif field.type == 'BOOLEAN': value = value == 'true' elif field.type == 'INTEGER': value = int(value) elif field.type == 'FLOAT': value = float(value) elif field.type == 'TIMESTAMP': value = float(value) elif field.type == 'BYTES': value = value else: # Note that a schema field object supports also a RECORD type. However # when querying, the repeated and/or record fields always come # flattened. For more details please read: # https://cloud.google.com/bigquery/docs/data raise RuntimeError('Unexpected field type: %s' % field.type) result[field.name] = value return result DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/bigquery_test.py000066400000000000000000000445531272340251000250640ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for BigQuery sources and sinks.""" import json import logging import time import unittest import mock import google.cloud.dataflow as df from google.cloud.dataflow.internal.json_value import to_json_value from google.cloud.dataflow.io.bigquery import RowAsDictJsonCoder from google.cloud.dataflow.io.bigquery import TableRowJsonCoder from google.cloud.dataflow.utils.options import PipelineOptions from apitools.base.py.exceptions import HttpError from google.cloud.dataflow.internal.clients import bigquery class TestRowAsDictJsonCoder(unittest.TestCase): def test_row_as_dict(self): coder = RowAsDictJsonCoder() test_value = {'s': 'abc', 'i': 123, 'f': 123.456, 'b': True} self.assertEqual(test_value, coder.decode(coder.encode(test_value))) class TestTableRowJsonCoder(unittest.TestCase): def test_row_as_table_row(self): schema_definition = [ ('s', 'STRING'), ('i', 'INTEGER'), ('f', 'FLOAT'), ('b', 'BOOLEAN')] schema = bigquery.TableSchema( fields=[bigquery.TableFieldSchema(name=k, type=v) for k, v in schema_definition]) coder = TableRowJsonCoder(table_schema=schema) test_row = bigquery.TableRow( f=[bigquery.TableCell(v=to_json_value(e)) for e in ['abc', 123, 123.456, True]]) self.assertEqual('{"s": "abc", "i": 123, "f": 123.456, "b": true}', coder.encode(test_row)) self.assertEqual(test_row, coder.decode(coder.encode(test_row))) # A coder without schema can still decode. self.assertEqual( test_row, TableRowJsonCoder().decode(coder.encode(test_row))) def test_row_and_no_schema(self): coder = TableRowJsonCoder() test_row = bigquery.TableRow( f=[bigquery.TableCell(v=to_json_value(e)) for e in ['abc', 123, 123.456, True]]) with self.assertRaises(AttributeError) as ctx: coder.encode(test_row) self.assertTrue( ctx.exception.message.startswith('The TableRowJsonCoder requires')) class TestBigQuerySource(unittest.TestCase): def test_parse_table_reference(self): source = df.io.BigQuerySource('dataset.table') self.assertEqual(source.table_reference.datasetId, 'dataset') self.assertEqual(source.table_reference.tableId, 'table') source = df.io.BigQuerySource('project:dataset.table') self.assertEqual(source.table_reference.projectId, 'project') self.assertEqual(source.table_reference.datasetId, 'dataset') self.assertEqual(source.table_reference.tableId, 'table') source = df.io.BigQuerySource('xyz.com:project:dataset.table') self.assertEqual(source.table_reference.projectId, 'xyz.com:project') self.assertEqual(source.table_reference.datasetId, 'dataset') self.assertEqual(source.table_reference.tableId, 'table') def test_specify_query_without_table(self): source = df.io.BigQuerySource(query='my_query') self.assertEqual(source.query, 'my_query') self.assertIsNone(source.table_reference) class TestBigQuerySink(unittest.TestCase): def test_parse_schema_descriptor(self): sink = df.io.BigQuerySink( 'dataset.table', schema='s:STRING, n:INTEGER') self.assertEqual(sink.table_reference.datasetId, 'dataset') self.assertEqual(sink.table_reference.tableId, 'table') result_schema = { field.name: field.type for field in sink.table_schema.fields} self.assertEqual({'n': 'INTEGER', 's': 'STRING'}, result_schema) def test_simple_schema_as_json(self): sink = df.io.BigQuerySink( 'dataset.table', schema='s:STRING, n:INTEGER') self.assertEqual( json.dumps({'fields': [ {'name': 's', 'type': 'STRING', 'mode': 'NULLABLE'}, {'name': 'n', 'type': 'INTEGER', 'mode': 'NULLABLE'}]}), sink.schema_as_json()) def test_nested_schema_as_json(self): string_field = bigquery.TableFieldSchema( name='s', type='STRING', mode='NULLABLE', description='s description') number_field = bigquery.TableFieldSchema( name='n', type='INTEGER', mode='REQUIRED', description='n description') record_field = bigquery.TableFieldSchema( name='r', type='RECORD', mode='REQUIRED', description='r description', fields=[string_field, number_field]) schema = bigquery.TableSchema(fields=[record_field]) sink = df.io.BigQuerySink('dataset.table', schema=schema) self.assertEqual( {'fields': [ {'name': 'r', 'type': 'RECORD', 'mode': 'REQUIRED', 'description': 'r description', 'fields': [ {'name': 's', 'type': 'STRING', 'mode': 'NULLABLE', 'description': 's description'}, {'name': 'n', 'type': 'INTEGER', 'mode': 'REQUIRED', 'description': 'n description'}]}]}, json.loads(sink.schema_as_json())) class TestBigQueryReader(unittest.TestCase): def get_test_rows(self): now = time.time() expected_rows = [ {'i': 1, 's': 'abc', 'f': 2.3, 'b': True, 't': now}, {'i': 10, 's': 'xyz', 'f': -3.14, 'b': False}] schema = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name='b', type='BOOLEAN', mode='REQUIRED'), bigquery.TableFieldSchema( name='f', type='FLOAT', mode='REQUIRED'), bigquery.TableFieldSchema( name='i', type='INTEGER', mode='REQUIRED'), bigquery.TableFieldSchema( name='s', type='STRING', mode='REQUIRED'), bigquery.TableFieldSchema( name='t', type='TIMESTAMP', mode='NULLABLE')]) table_rows = [ bigquery.TableRow(f=[ bigquery.TableCell(v=to_json_value('true')), bigquery.TableCell(v=to_json_value(str(2.3))), bigquery.TableCell(v=to_json_value(str(1))), bigquery.TableCell(v=to_json_value('abc')), # For timestamps cannot use str() because it will truncate the # number representing the timestamp. bigquery.TableCell(v=to_json_value('%f' % now))]), bigquery.TableRow(f=[ bigquery.TableCell(v=to_json_value('false')), bigquery.TableCell(v=to_json_value(str(-3.14))), bigquery.TableCell(v=to_json_value(str(10))), bigquery.TableCell(v=to_json_value('xyz')), bigquery.TableCell(v=None)])] return table_rows, schema, expected_rows def test_read_from_table(self): client = mock.Mock() client.jobs.Insert.return_value = bigquery.Job( jobReference=bigquery.JobReference( jobId='somejob')) table_rows, schema, expected_rows = self.get_test_rows() client.jobs.GetQueryResults.return_value = bigquery.GetQueryResultsResponse( jobComplete=True, rows=table_rows, schema=schema) actual_rows = [] with df.io.BigQuerySource('dataset.table').reader(client) as reader: for row in reader: actual_rows.append(row) self.assertEqual(actual_rows, expected_rows) self.assertEqual(schema, reader.schema) def test_read_from_query(self): client = mock.Mock() client.jobs.Insert.return_value = bigquery.Job( jobReference=bigquery.JobReference( jobId='somejob')) table_rows, schema, expected_rows = self.get_test_rows() client.jobs.GetQueryResults.return_value = bigquery.GetQueryResultsResponse( jobComplete=True, rows=table_rows, schema=schema) actual_rows = [] with df.io.BigQuerySource(query='query').reader(client) as reader: for row in reader: actual_rows.append(row) self.assertEqual(actual_rows, expected_rows) self.assertEqual(schema, reader.schema) def test_using_both_query_and_table_fails(self): with self.assertRaises(ValueError) as exn: df.io.BigQuerySource(table='dataset.table', query='query') self.assertEqual(exn.exception.message, 'Both a BigQuery table and a' ' query were specified. Please specify only one of ' 'these.') def test_using_neither_query_nor_table_fails(self): with self.assertRaises(ValueError) as exn: df.io.BigQuerySource() self.assertEqual(exn.exception.message, 'A BigQuery table or a query' ' must be specified') def test_read_from_table_as_tablerows(self): client = mock.Mock() client.jobs.Insert.return_value = bigquery.Job( jobReference=bigquery.JobReference( jobId='somejob')) table_rows, schema, _ = self.get_test_rows() client.jobs.GetQueryResults.return_value = bigquery.GetQueryResultsResponse( jobComplete=True, rows=table_rows, schema=schema) actual_rows = [] # We set the coder to TableRowJsonCoder, which is a signal that # the caller wants to see the rows as TableRows. with df.io.BigQuerySource( 'dataset.table', coder=TableRowJsonCoder).reader(client) as reader: for row in reader: actual_rows.append(row) self.assertEqual(actual_rows, table_rows) self.assertEqual(schema, reader.schema) def test_read_from_table_and_job_complete_retry(self): client = mock.Mock() client.jobs.Insert.return_value = bigquery.Job( jobReference=bigquery.JobReference( jobId='somejob')) table_rows, schema, expected_rows = self.get_test_rows() # Return jobComplete=False on first call to trigger the code path where # query needs to handle waiting a bit. client.jobs.GetQueryResults.side_effect = [ bigquery.GetQueryResultsResponse( jobComplete=False), bigquery.GetQueryResultsResponse( jobComplete=True, rows=table_rows, schema=schema)] actual_rows = [] with df.io.BigQuerySource('dataset.table').reader(client) as reader: for row in reader: actual_rows.append(row) self.assertEqual(actual_rows, expected_rows) def test_read_from_table_and_multiple_pages(self): client = mock.Mock() client.jobs.Insert.return_value = bigquery.Job( jobReference=bigquery.JobReference( jobId='somejob')) table_rows, schema, expected_rows = self.get_test_rows() # Return a pageToken on first call to trigger the code path where # query needs to handle multiple pages of results. client.jobs.GetQueryResults.side_effect = [ bigquery.GetQueryResultsResponse( jobComplete=True, rows=table_rows, schema=schema, pageToken='token'), bigquery.GetQueryResultsResponse( jobComplete=True, rows=table_rows, schema=schema)] actual_rows = [] with df.io.BigQuerySource('dataset.table').reader(client) as reader: for row in reader: actual_rows.append(row) # We return expected rows for each of the two pages of results so we # adjust our expectation below accordingly. self.assertEqual(actual_rows, expected_rows * 2) def test_table_schema_without_project(self): # Reader should pick executing project by default. source = df.io.BigQuerySource(table='mydataset.mytable') options = PipelineOptions(flags=['--project', 'myproject']) source.pipeline_options = options reader = source.reader() self.assertEquals('SELECT * FROM [myproject:mydataset.mytable];', reader.query) class TestBigQueryWriter(unittest.TestCase): def test_no_table_and_create_never(self): client = mock.Mock() client.tables.Get.side_effect = HttpError( response={'status': '404'}, url='', content='') create_disposition = df.io.BigQueryDisposition.CREATE_NEVER with self.assertRaises(RuntimeError) as exn: with df.io.BigQuerySink( 'project:dataset.table', create_disposition=create_disposition).writer(client): pass self.assertEqual( exn.exception.message, 'Table project:dataset.table not found but create disposition is ' 'CREATE_NEVER.') def test_no_table_and_create_if_needed(self): client = mock.Mock() table = bigquery.Table( tableReference=bigquery.TableReference( projectId='project', datasetId='dataset', tableId='table'), schema=bigquery.TableSchema()) client.tables.Get.side_effect = HttpError( response={'status': '404'}, url='', content='') client.tables.Insert.return_value = table create_disposition = df.io.BigQueryDisposition.CREATE_IF_NEEDED with df.io.BigQuerySink( 'project:dataset.table', schema='somefield:INTEGER', create_disposition=create_disposition).writer(client): pass self.assertTrue(client.tables.Get.called) self.assertTrue(client.tables.Insert.called) def test_no_table_and_create_if_needed_and_no_schema(self): client = mock.Mock() client.tables.Get.side_effect = HttpError( response={'status': '404'}, url='', content='') create_disposition = df.io.BigQueryDisposition.CREATE_IF_NEEDED with self.assertRaises(RuntimeError) as exn: with df.io.BigQuerySink( 'project:dataset.table', create_disposition=create_disposition).writer(client): pass self.assertEqual( exn.exception.message, 'Table project:dataset.table requires a schema. None can be inferred ' 'because the table does not exist.') def test_table_not_empty_and_write_disposition_empty(self): client = mock.Mock() client.tables.Get.return_value = bigquery.Table( tableReference=bigquery.TableReference( projectId='project', datasetId='dataset', tableId='table'), schema=bigquery.TableSchema()) client.tabledata.List.return_value = bigquery.TableDataList(totalRows=1) write_disposition = df.io.BigQueryDisposition.WRITE_EMPTY with self.assertRaises(RuntimeError) as exn: with df.io.BigQuerySink( 'project:dataset.table', write_disposition=write_disposition).writer(client): pass self.assertEqual( exn.exception.message, 'Table project:dataset.table is not empty but write disposition is ' 'WRITE_EMPTY.') def test_table_empty_and_write_disposition_empty(self): client = mock.Mock() table = bigquery.Table( tableReference=bigquery.TableReference( projectId='project', datasetId='dataset', tableId='table'), schema=bigquery.TableSchema()) client.tables.Get.return_value = table client.tabledata.List.return_value = bigquery.TableDataList(totalRows=0) client.tables.Insert.return_value = table write_disposition = df.io.BigQueryDisposition.WRITE_EMPTY with df.io.BigQuerySink( 'project:dataset.table', write_disposition=write_disposition).writer(client): pass self.assertTrue(client.tables.Get.called) self.assertTrue(client.tabledata.List.called) self.assertFalse(client.tables.Delete.called) self.assertFalse(client.tables.Insert.called) def test_table_with_write_disposition_truncate(self): client = mock.Mock() table = bigquery.Table( tableReference=bigquery.TableReference( projectId='project', datasetId='dataset', tableId='table'), schema=bigquery.TableSchema()) client.tables.Get.return_value = table client.tables.Insert.return_value = table write_disposition = df.io.BigQueryDisposition.WRITE_TRUNCATE with df.io.BigQuerySink( 'project:dataset.table', write_disposition=write_disposition).writer(client): pass self.assertTrue(client.tables.Get.called) self.assertTrue(client.tables.Delete.called) self.assertTrue(client.tables.Insert.called) def test_table_with_write_disposition_append(self): client = mock.Mock() table = bigquery.Table( tableReference=bigquery.TableReference( projectId='project', datasetId='dataset', tableId='table'), schema=bigquery.TableSchema()) client.tables.Get.return_value = table client.tables.Insert.return_value = table write_disposition = df.io.BigQueryDisposition.WRITE_APPEND with df.io.BigQuerySink( 'project:dataset.table', write_disposition=write_disposition).writer(client): pass self.assertTrue(client.tables.Get.called) self.assertFalse(client.tables.Delete.called) self.assertFalse(client.tables.Insert.called) def test_rows_are_written(self): client = mock.Mock() table = bigquery.Table( tableReference=bigquery.TableReference( projectId='project', datasetId='dataset', tableId='table'), schema=bigquery.TableSchema()) client.tables.Get.return_value = table write_disposition = df.io.BigQueryDisposition.WRITE_APPEND insert_response = mock.Mock() insert_response.insertErrors = [] client.tabledata.InsertAll.return_value = insert_response with df.io.BigQuerySink( 'project:dataset.table', write_disposition=write_disposition).writer(client) as writer: writer.Write({'i': 1, 'b': True, 's': 'abc', 'f': 3.14}) sample_row = {'i': 1, 'b': True, 's': 'abc', 'f': 3.14} expected_rows = [] json_object = bigquery.JsonObject() for k, v in sample_row.iteritems(): json_object.additionalProperties.append( bigquery.JsonObject.AdditionalProperty( key=k, value=to_json_value(v))) expected_rows.append( bigquery.TableDataInsertAllRequest.RowsValueListEntry( insertId='_1', # First row ID generated with prefix '' json=json_object)) client.tabledata.InsertAll.assert_called_with( bigquery.BigqueryTabledataInsertAllRequest( projectId='project', datasetId='dataset', tableId='table', tableDataInsertAllRequest=bigquery.TableDataInsertAllRequest( rows=expected_rows))) def test_table_schema_without_project(self): # Writer should pick executing project by default. sink = df.io.BigQuerySink(table='mydataset.mytable') options = PipelineOptions(flags=['--project', 'myproject']) sink.pipeline_options = options writer = sink.writer() self.assertEquals('myproject', writer.project_id) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/fileio.py000066400000000000000000000614151272340251000234410ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """File-based sources and sinks.""" from __future__ import absolute_import import glob import gzip import logging import os import re import shutil import tempfile import time from google.cloud.dataflow import coders from google.cloud.dataflow.io import iobase from google.cloud.dataflow.io import range_trackers from google.cloud.dataflow.utils import processes from google.cloud.dataflow.utils import retry __all__ = ['TextFileSource', 'TextFileSink'] # Retrying is needed because there are transient errors that can happen. @retry.with_exponential_backoff(num_retries=4, retry_filter=lambda _: True) def _gcs_file_copy(from_path, to_path, encoding=''): """Copy a local file to a GCS location with retries for transient errors.""" if not encoding: command_args = ['gsutil', '-m', '-q', 'cp', from_path, to_path] else: encoding = 'Content-Type:' + encoding command_args = ['gsutil', '-m', '-q', '-h', encoding, 'cp', from_path, to_path] logging.info('Executing command: %s', command_args) popen = processes.Popen(command_args, stdout=processes.PIPE, stderr=processes.PIPE) stdoutdata, stderrdata = popen.communicate() if popen.returncode != 0: raise ValueError( 'Failed to copy GCS file from %s to %s (stdout=%s, stderr=%s).' % ( from_path, to_path, stdoutdata, stderrdata)) # ----------------------------------------------------------------------------- # TextFileSource, TextFileSink. class TextFileSource(iobase.NativeSource): """A source for a GCS or local text file. Parses a text file as newline-delimited elements, by default assuming UTF-8 encoding. """ def __init__(self, file_path, start_offset=None, end_offset=None, compression_type='AUTO', strip_trailing_newlines=True, coder=coders.StrUtf8Coder()): """Initialize a TextSource. Args: file_path: The file path to read from as a local file path or a GCS gs:// path. The path can contain glob characters (*, ?, and [...] sets). start_offset: The byte offset in the source text file that the reader should start reading. By default is 0 (beginning of file). end_offset: The byte offset in the file that the reader should stop reading. By default it is the end of the file. compression_type: Used to handle compressed input files. Typical value is 'AUTO'. strip_trailing_newlines: Indicates whether this source should remove the newline char in each line it reads before decoding that line. coder: Coder used to decode each line. Raises: TypeError: if file_path is not a string. If the file_path contains glob characters then the start_offset and end_offset must not be specified. The 'start_offset' and 'end_offset' pair provide a mechanism to divide the text file into multiple pieces for individual sources. Because the offset is measured by bytes, some complication arises when the offset splits in the middle of a text line. To avoid the scenario where two adjacent sources each get a fraction of a line we adopt the following rules: If start_offset falls inside a line (any character except the firt one) then the source will skip the line and start with the next one. If end_offset falls inside a line (any character except the first one) then the source will contain that entire line. """ if not isinstance(file_path, basestring): raise TypeError( '%s: file_path must be a string; got %r instead' % (self.__class__.__name__, file_path)) self.file_path = file_path self.start_offset = start_offset self.end_offset = end_offset self.compression_type = compression_type self.strip_trailing_newlines = strip_trailing_newlines self.coder = coder self.is_gcs_source = file_path.startswith('gs://') @property def format(self): """Source format name required for remote execution.""" return 'text' def __eq__(self, other): return (self.file_path == other.file_path and self.start_offset == other.start_offset and self.end_offset == other.end_offset and self.strip_trailing_newlines == other.strip_trailing_newlines and self.coder == other.coder) @property def path(self): return self.file_path def reader(self): # If a multi-file pattern was specified as a source then make sure the # start/end offsets use the default values for reading the entire file. if re.search(r'[*?\[\]]', self.file_path) is not None: if self.start_offset is not None: raise ValueError( 'start offset cannot be specified for a multi-file source: ' '%s' % self.file_path) if self.end_offset is not None: raise ValueError( 'End offset cannot be specified for a multi-file source: ' '%s' % self.file_path) return TextMultiFileReader(self) else: return TextFileReader(self) def TextFileSink(file_path_prefix, # pylint: disable=invalid-name append_trailing_newlines=True, file_name_suffix='', num_shards=0, shard_name_template=None, validate=True, coder=coders.ToStringCoder()): """Initialize a TextSink. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. append_trailing_newlines: indicate whether this sink should write an additional newline char after writing each element. file_name_suffix: Suffix for the files written. num_shards: The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. Constraining the number of shards is likely to reduce the performance of a pipeline. Setting this value is not recommended unless you require a specific number of output files. shard_name_template: A template string containing placeholders for the shard number and shard count. Currently only '' and '-SSSSS-of-NNNNN' are patterns accepted by the service. When constructing a filename for a particular shard number, the upper-case letters 'S' and 'N' are replaced with the 0-padded shard number and shard count respectively. This argument can be '' in which case it behaves as if num_shards was set to 1 and only one file will be generated. The default pattern used is '-SSSSS-of-NNNNN'. validate: Enable path validation on pipeline creation. coder: Coder used to encode each line. Raises: TypeError: if file_path is not a string. ValueError: if shard_name_template is not of expected format. Returns: A TextFileSink object usable for writing. """ if not isinstance(file_path_prefix, basestring): raise TypeError( 'TextFileSink: file_path_prefix must be a string; got %r instead' % file_path_prefix) if not isinstance(file_name_suffix, basestring): raise TypeError( 'TextFileSink: file_name_suffix must be a string; got %r instead' % file_name_suffix) if shard_name_template not in (None, '', '-SSSSS-of-NNNNN'): raise ValueError( 'The shard_name_template argument must be an empty string or the ' 'pattern -SSSSS-of-NNNNN instead of %s' % shard_name_template) if shard_name_template == '': # pylint: disable=g-explicit-bool-comparison num_shards = 1 if num_shards: return NativeTextFileSink(file_path_prefix, append_trailing_newlines=append_trailing_newlines, file_name_suffix=file_name_suffix, num_shards=num_shards, shard_name_template=shard_name_template, validate=validate, coder=coder) else: return PureTextFileSink(file_path_prefix, append_trailing_newlines=append_trailing_newlines, file_name_suffix=file_name_suffix, coder=coder) class ChannelFactory(object): # TODO(robertwb): Generalize into extensible framework. @staticmethod def mkdir(path): if path.startswith('gs://'): return else: try: os.makedirs(path) except OSError as err: raise IOError(err) @staticmethod def open(path, mode, mime_type): if path.startswith('gs://'): # pylint: disable=g-import-not-at-top from google.cloud.dataflow.io import gcsio return gcsio.GcsIO().open(path, mode, mime_type) else: return open(path, mode) @staticmethod def rename(src, dst): if src.startswith('gs://'): assert dst.startswith('gs://'), dst # pylint: disable=g-import-not-at-top from google.cloud.dataflow.io import gcsio gcsio.GcsIO().rename(src, dst) else: try: os.rename(src, dst) except OSError as err: raise IOError(err) @staticmethod def exists(path): if path.startswith('gs://'): # pylint: disable=g-import-not-at-top from google.cloud.dataflow.io import gcsio return gcsio.GcsIO().exists() else: return os.path.exists(path) @staticmethod def rmdir(path): if path.startswith('gs://'): # pylint: disable=g-import-not-at-top from google.cloud.dataflow.io import gcsio gcs = gcsio.GcsIO() if not path.endswith('/'): path += '/' # TODO(robertwb): Threadpool? for entry in gcs.glob(path + '*'): gcs.delete(entry) else: try: shutil.rmtree(path) except OSError as err: raise IOError(err) @staticmethod def rm(path): if path.startswith('gs://'): # pylint: disable=g-import-not-at-top from google.cloud.dataflow.io import gcsio gcsio.GcsIO().delete(path) else: try: os.remove(path) except OSError as err: raise IOError(err) @staticmethod def glob(path): if path.startswith('gs://'): # pylint: disable=g-import-not-at-top from google.cloud.dataflow.io import gcsio return gcsio.GcsIO().glob(path) else: return glob.glob(path) class _CompressionType(object): """Object representing single compression type.""" def __init__(self, identifier): self.identifier = identifier def __eq__(self, other): return self.identifier == other.identifier class CompressionTypes(object): """Enum-like class representing known compression types.""" NO_COMPRESSION = _CompressionType(1) # No compression. DEFLATE = _CompressionType(2) # 'Deflate' ie gzip compression. @staticmethod def valid_compression_type(compression_type): """Returns true for valid compression types, false otherwise.""" return isinstance(compression_type, _CompressionType) class FileSink(iobase.Sink): """A sink to a GCS or local files. To implement a file-based sink, extend this class and override either ``write_record()`` or ``write_encoded_record()``. If needed, also overwrite ``open()`` and/or ``close()`` to customize the file handling or write headers and footers. The output of this write is a PCollection of all written shards. """ def __init__(self, file_path_prefix, coder, file_name_suffix='', mime_type='application/octet-stream'): self.file_path_prefix = file_path_prefix self.file_name_suffix = file_name_suffix self.coder = coder self.mime_type = mime_type def open(self, temp_path): """Opens ``temp_path``, returning an opaque file handle object. The returned file handle is passed to ``write_[encoded_]record`` and ``close``. """ return ChannelFactory.open(temp_path, 'wb', self.mime_type) def write_record(self, file_handle, value): """Writes a single record go the file handle returned by ``open()``. By default, calls ``write_encoded_record`` after encoding the record with this sink's Coder. """ self.write_encoded_record(file_handle, self.coder.encode(value)) def write_encoded_record(self, file_handle, encoded_value): """Writes a single encoded record to the file handle returned by ``open()``. """ raise NotImplementedError def close(self, file_handle): """Finalize and close the file handle returned from ``open()``. Called after all records are written. By default, calls ``file_handle.close()`` iff it is not None. """ if file_handle is not None: file_handle.close() def initialize_write(self): tmp_dir = self.file_path_prefix + self.file_name_suffix + time.strftime( '-temp-%Y-%m-%d_%H-%M-%S') ChannelFactory().mkdir(tmp_dir) return tmp_dir def open_writer(self, init_result, uid): return FileSinkWriter(self, os.path.join(init_result, uid)) def finalize_write(self, init_result, writer_results): writer_results = sorted(writer_results) num_shards = len(writer_results) # TODO(robertwb): Threadpool? channel_factory = ChannelFactory() for shard_num, shard in enumerate(writer_results): final_name = '%s-%05d-of-%05d%s' % (self.file_path_prefix, shard_num, num_shards, self.file_name_suffix) try: channel_factory.rename(shard, final_name) except IOError: # May have already been copied. print shard, final_name, os.path.exists(final_name) if not channel_factory.exists(final_name): raise yield final_name try: channel_factory.rmdir(init_result) except IOError: # May have already been removed. pass def __eq__(self, other): # TODO(robertwb): Clean up workitem_test which uses this. # pylint: disable=unidiomatic-typecheck return type(self) == type(other) and self.__dict__ == other.__dict__ class FileSinkWriter(iobase.Writer): """The writer for FileSink. """ def __init__(self, sink, temp_shard_path): self.sink = sink self.temp_shard_path = temp_shard_path self.temp_handle = self.sink.open(temp_shard_path) def write(self, value): self.sink.write_record(self.temp_handle, value) def close(self): self.sink.close(self.temp_handle) return self.temp_shard_path class PureTextFileSink(FileSink): """A sink to a GCS or local text file or files.""" def __init__(self, file_path_prefix, file_name_suffix='', coder=coders.ToStringCoder(), compression_type=CompressionTypes.NO_COMPRESSION, append_trailing_newlines=True): """Initialize a PureTextFileSink. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. file_name_suffix: Suffix for the files written. coder: Coder used to encode each line. compression_type: Type of compression to use for this sink. append_trailing_newlines: indicate whether this sink should write an additional newline char after writing each element. Raises: TypeError: if file_path is not a string or if compression_type is not member of CompressionTypes. Returns: A PureTextFileSink object usable for writing. """ if not CompressionTypes.valid_compression_type(compression_type): raise TypeError('compression_type must be CompressionType object but ' 'was %s' % type(compression_type)) if compression_type == CompressionTypes.DEFLATE: mime_type = 'application/x-gzip' else: mime_type = 'text/plain' self.compression_type = compression_type super(PureTextFileSink, self).__init__(file_path_prefix, file_name_suffix=file_name_suffix, coder=coder, mime_type=mime_type) self.append_trailing_newlines = append_trailing_newlines def open(self, temp_path): """Opens ''temp_path'', returning a writeable file object.""" fobj = ChannelFactory.open(temp_path, 'wb', self.mime_type) if self.compression_type == CompressionTypes.DEFLATE: return gzip.GzipFile(fileobj=fobj) return fobj def write_encoded_record(self, file_handle, encoded_value): file_handle.write(encoded_value) if self.append_trailing_newlines: file_handle.write('\n') class NativeTextFileSink(iobase.NativeSink): """A sink to a GCS or local text file or files.""" def __init__(self, file_path_prefix, append_trailing_newlines=True, file_name_suffix='', num_shards=0, shard_name_template=None, validate=True, coder=coders.ToStringCoder()): # We initialize a file_path attribute containing just the prefix part for # local runner environment. For now, sharding is not supported in the local # runner and sharding options (template, num, suffix) are ignored. # The attribute is also used in the worker environment when we just write # to a specific file. # TODO(silviuc): Add support for file sharding in the local runner. self.file_path = file_path_prefix self.append_trailing_newlines = append_trailing_newlines self.coder = coder self.is_gcs_sink = self.file_path.startswith('gs://') self.file_name_prefix = file_path_prefix self.file_name_suffix = file_name_suffix self.num_shards = num_shards # TODO(silviuc): Update this when the service supports more patterns. self.shard_name_template = ('-SSSSS-of-NNNNN' if shard_name_template is None else shard_name_template) # TODO(silviuc): Implement sink validation. self.validate = validate @property def format(self): """Sink format name required for remote execution.""" return 'text' @property def path(self): return self.file_path def writer(self): return TextFileWriter(self) def __eq__(self, other): return (self.file_path == other.file_path and self.append_trailing_newlines == other.append_trailing_newlines and self.coder == other.coder and self.file_name_prefix == other.file_name_prefix and self.file_name_suffix == other.file_name_suffix and self.num_shards == other.num_shards and self.shard_name_template == other.shard_name_template and self.validate == other.validate) # ----------------------------------------------------------------------------- # TextFileReader, TextMultiFileReader. class TextFileReader(iobase.NativeSourceReader): """A reader for a text file source.""" def __init__(self, source): self.source = source self.start_offset = self.source.start_offset or 0 self.end_offset = self.source.end_offset self.current_offset = self.start_offset def __enter__(self): if self.source.is_gcs_source: # pylint: disable=g-import-not-at-top from google.cloud.dataflow.io import gcsio self._file = gcsio.GcsIO().open(self.source.file_path, 'rb') else: self._file = open(self.source.file_path, 'rb') # Determine the real end_offset. # If not specified it will be the length of the file. if self.end_offset is None: self._file.seek(0, os.SEEK_END) self.end_offset = self._file.tell() if self.start_offset is None: self.start_offset = 0 self.current_offset = self.start_offset if self.start_offset > 0: # Read one byte before. This operation will either consume a previous # newline if start_offset was at the beginning of a line or consume the # line if we were in the middle of it. Either way we get the read position # exactly where we wanted: at the begining of the first full line. self._file.seek(self.start_offset - 1) self.current_offset -= 1 line = self._file.readline() self.current_offset += len(line) else: self._file.seek(self.start_offset) # Initializing range tracker after start and end offsets are finalized. self.range_tracker = range_trackers.OffsetRangeTracker(self.start_offset, self.end_offset) return self def __exit__(self, exception_type, exception_value, traceback): self._file.close() def __iter__(self): while True: if not self.range_tracker.try_claim( record_start=self.current_offset): # Reader has completed reading the set of records in its range. Note # that the end offset of the range may be smaller than the original # end offset defined when creating the reader due to reader accepting # a dynamic split request from the service. return line = self._file.readline() self.current_offset += len(line) if self.source.strip_trailing_newlines: line = line.rstrip('\n') yield self.source.coder.decode(line) def get_progress(self): return iobase.ReaderProgress(position=iobase.ReaderPosition( byte_offset=self.range_tracker.last_record_start)) def request_dynamic_split(self, dynamic_split_request): assert dynamic_split_request is not None progress = dynamic_split_request.progress split_position = progress.position if split_position is None: percent_complete = progress.percent_complete if percent_complete is not None: if percent_complete <= 0 or percent_complete >= 1: logging.warning( 'FileBasedReader cannot be split since the provided percentage ' 'of work to be completed is out of the valid range (0, ' '1). Requested: %r', dynamic_split_request) return split_position = iobase.ReaderPosition() split_position.byte_offset = ( self.range_tracker.position_at_fraction(percent_complete)) else: logging.warning( 'TextReader requires either a position or a percentage of work to ' 'be complete to perform a dynamic split request. Requested: %r', dynamic_split_request) return if self.range_tracker.try_split(split_position.byte_offset): return iobase.DynamicSplitResultWithPosition(split_position) else: return class TextMultiFileReader(iobase.NativeSourceReader): """A reader for a multi-file text source.""" def __init__(self, source): self.source = source self.file_paths = ChannelFactory.glob(self.source.file_path) if not self.file_paths: raise RuntimeError( 'No files found for path: %s' % self.source.file_path) def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): pass def __iter__(self): index = 0 for path in self.file_paths: index += 1 logging.info('Reading from %s (%d/%d)', path, index, len(self.file_paths)) with TextFileSource( path, strip_trailing_newlines=self.source.strip_trailing_newlines, coder=self.source.coder).reader() as reader: for line in reader: yield line # ----------------------------------------------------------------------------- # TextFileWriter. class TextFileWriter(iobase.NativeSinkWriter): """The sink writer for a TextFileSink.""" def __init__(self, sink): self.sink = sink def __enter__(self): if self.sink.is_gcs_sink: # TODO(silviuc): Use the storage library instead of gsutil for writes. self.temp_path = os.path.join(tempfile.mkdtemp(), 'gcsfile') self._file = open(self.temp_path, 'wb') else: self._file = open(self.sink.file_path, 'wb') return self def __exit__(self, exception_type, exception_value, traceback): self._file.close() if hasattr(self, 'temp_path'): _gcs_file_copy(self.temp_path, self.sink.file_path, 'text/plain') def Write(self, line): self._file.write(self.sink.coder.encode(line)) if self.sink.append_trailing_newlines: self._file.write('\n') DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/fileio_test.py000066400000000000000000000363631272340251000245040ustar00rootroot00000000000000# -*- coding: utf-8 -*- # Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for local and GCS sources and sinks.""" import glob import gzip import logging import tempfile import unittest from google.cloud.dataflow import coders from google.cloud.dataflow.io import fileio from google.cloud.dataflow.io import iobase class TestTextFileSource(unittest.TestCase): def create_temp_file(self, text): temp = tempfile.NamedTemporaryFile(delete=False) with temp.file as tmp: tmp.write(text) return temp.name def read_with_offsets(self, input_lines, output_lines, start_offset=None, end_offset=None): source = fileio.TextFileSource( file_path=self.create_temp_file('\n'.join(input_lines)), start_offset=start_offset, end_offset=end_offset) read_lines = [] with source.reader() as reader: for line in reader: read_lines.append(line) self.assertEqual(read_lines, output_lines) def progress_with_offsets(self, input_lines, start_offset=None, end_offset=None): source = fileio.TextFileSource( file_path=self.create_temp_file('\n'.join(input_lines)), start_offset=start_offset, end_offset=end_offset) progress_record = [] with source.reader() as reader: self.assertEqual(reader.get_progress().position.byte_offset, -1) for line in reader: self.assertIsNotNone(line) progress_record.append(reader.get_progress().position.byte_offset) previous = 0 for current in progress_record: self.assertGreater(current, previous) previous = current def test_read_entire_file(self): lines = ['First', 'Second', 'Third'] source = fileio.TextFileSource( file_path=self.create_temp_file('\n'.join(lines))) read_lines = [] with source.reader() as reader: for line in reader: read_lines.append(line) self.assertEqual(read_lines, lines) def test_progress_entire_file(self): lines = ['First', 'Second', 'Third'] source = fileio.TextFileSource( file_path=self.create_temp_file('\n'.join(lines))) progress_record = [] with source.reader() as reader: self.assertEqual(-1, reader.get_progress().position.byte_offset) for line in reader: self.assertIsNotNone(line) progress_record.append(reader.get_progress().position.byte_offset) self.assertEqual(13, reader.get_progress().position.byte_offset) self.assertEqual(len(progress_record), 3) self.assertEqual(progress_record, [0, 6, 13]) def try_splitting_reader_at(self, reader, split_request, expected_response): actual_response = reader.request_dynamic_split(split_request) if expected_response is None: self.assertIsNone(actual_response) else: self.assertIsNotNone(actual_response.stop_position) self.assertIsInstance(actual_response.stop_position, iobase.ReaderPosition) self.assertIsNotNone(actual_response.stop_position.byte_offset) self.assertEqual(expected_response.stop_position.byte_offset, actual_response.stop_position.byte_offset) return actual_response def test_update_stop_position_for_percent_complete(self): lines = ['aaaa', 'bbbb', 'cccc', 'dddd', 'eeee'] source = fileio.TextFileSource( file_path=self.create_temp_file('\n'.join(lines))) with source.reader() as reader: # Reading two lines reader_iter = iter(reader) next(reader_iter) next(reader_iter) next(reader_iter) # Splitting at end of the range should be unsuccessful self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress(percent_complete=0)), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress(percent_complete=1)), None) # Splitting at positions on or before start offset of the last record self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress(percent_complete= 0.2)), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress(percent_complete= 0.4)), None) # Splitting at a position after the start offset of the last record should # be successful self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress(percent_complete= 0.6)), iobase.DynamicSplitResultWithPosition(iobase.ReaderPosition( byte_offset=15))) def test_update_stop_position_percent_complete_for_position(self): lines = ['aaaa', 'bbbb', 'cccc', 'dddd', 'eeee'] source = fileio.TextFileSource( file_path=self.create_temp_file('\n'.join(lines))) with source.reader() as reader: # Reading two lines reader_iter = iter(reader) next(reader_iter) next(reader_iter) next(reader_iter) # Splitting at end of the range should be unsuccessful self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition(byte_offset=0))), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition(byte_offset=25))), None) # Splitting at positions on or before start offset of the last record self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition(byte_offset=5))), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition(byte_offset=10))), None) # Splitting at a position after the start offset of the last record should # be successful self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition(byte_offset=15))), iobase.DynamicSplitResultWithPosition(iobase.ReaderPosition( byte_offset=15))) def run_update_stop_position_exhaustive(self, lines, newline): """An exhaustive test for dynamic splitting. For the given set of data items, try to perform a split at all possible combinations of following. * start position * original stop position * updated stop position * number of items read Args: lines: set of data items to be used to create the file newline: separater to be used when writing give set of lines to a text file. """ file_path = self.create_temp_file(newline.join(lines)) total_records = len(lines) total_bytes = 0 for line in lines: total_bytes += len(line) total_bytes += len(newline) * (total_records - 1) for start in xrange(0, total_bytes - 1): for end in xrange(start + 1, total_bytes): for stop in xrange(start, end): for records_to_read in range(0, total_records): self.run_update_stop_position(start, end, stop, records_to_read, file_path) def test_update_stop_position_exhaustive(self): self.run_update_stop_position_exhaustive( ['aaaa', 'bbbb', 'cccc', 'dddd', 'eeee'], '\n') def test_update_stop_position_exhaustive_with_empty_lines(self): self.run_update_stop_position_exhaustive( ['', 'aaaa', '', 'bbbb', 'cccc', '', 'dddd', 'eeee', ''], '\n') def test_update_stop_position_exhaustive_windows_newline(self): self.run_update_stop_position_exhaustive( ['aaaa', 'bbbb', 'cccc', 'dddd', 'eeee'], '\r\n') def test_update_stop_position_exhaustive_multi_byte(self): self.run_update_stop_position_exhaustive( [u'අඅඅඅ'.encode('utf-8'), u'බබබබ'.encode('utf-8'), u'කකකක'.encode('utf-8')], '\n') def run_update_stop_position(self, start_offset, end_offset, stop_offset, records_to_read, file_path): source = fileio.TextFileSource(file_path, start_offset, end_offset) records_of_first_split = '' with source.reader() as reader: reader_iter = iter(reader) i = 0 try: while i < records_to_read: records_of_first_split += next(reader_iter) i += 1 except StopIteration: # Invalid case, given source does not contain this many records. return last_record_start_after_reading = reader.range_tracker.last_record_start if stop_offset <= last_record_start_after_reading: expected_split_response = None elif stop_offset == start_offset or stop_offset == end_offset: expected_split_response = None elif records_to_read == 0: expected_split_response = None # unstarted else: expected_split_response = iobase.DynamicSplitResultWithPosition( stop_position=iobase.ReaderPosition(byte_offset=stop_offset)) split_response = self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(progress=iobase.ReaderProgress( iobase.ReaderPosition(byte_offset=stop_offset))), expected_split_response) # Reading remaining records from the updated reader. for line in reader: records_of_first_split += line if split_response is not None: # Total contents received by reading the two splits should be equal to the # result obtained by reading the original source. records_of_original = '' records_of_second_split = '' with source.reader() as original_reader: for line in original_reader: records_of_original += line new_source = fileio.TextFileSource( file_path, split_response.stop_position.byte_offset, end_offset) with new_source.reader() as reader: for line in reader: records_of_second_split += line self.assertEqual(records_of_original, records_of_first_split + records_of_second_split) def test_various_offset_combination_with_local_file_for_read(self): lines = ['01234', '6789012', '456789012'] self.read_with_offsets(lines, lines[1:], start_offset=5) self.read_with_offsets(lines, lines[1:], start_offset=6) self.read_with_offsets(lines, lines[2:], start_offset=7) self.read_with_offsets(lines, lines[1:2], start_offset=5, end_offset=13) self.read_with_offsets(lines, lines[1:2], start_offset=5, end_offset=14) self.read_with_offsets(lines, lines[1:], start_offset=5, end_offset=16) self.read_with_offsets(lines, lines[2:], start_offset=14, end_offset=20) self.read_with_offsets(lines, lines[2:], start_offset=14) self.read_with_offsets(lines, [], start_offset=20, end_offset=20) def test_various_offset_combination_with_local_file_for_progress(self): lines = ['01234', '6789012', '456789012'] self.progress_with_offsets(lines, start_offset=5) self.progress_with_offsets(lines, start_offset=6) self.progress_with_offsets(lines, start_offset=7) self.progress_with_offsets(lines, start_offset=5, end_offset=13) self.progress_with_offsets(lines, start_offset=5, end_offset=14) self.progress_with_offsets(lines, start_offset=5, end_offset=16) self.progress_with_offsets(lines, start_offset=14, end_offset=20) self.progress_with_offsets(lines, start_offset=14) self.progress_with_offsets(lines, start_offset=20, end_offset=20) class NativeTestTextFileSink(unittest.TestCase): def create_temp_file(self): temp = tempfile.NamedTemporaryFile(delete=False) return temp.name def test_write_entire_file(self): lines = ['First', 'Second', 'Third'] file_path = self.create_temp_file() sink = fileio.NativeTextFileSink(file_path) with sink.writer() as writer: for line in lines: writer.Write(line) with open(file_path, 'r') as f: self.assertEqual(f.read().splitlines(), lines) class TestPureTextFileSink(unittest.TestCase): def setUp(self): self.lines = ['Line %d' % d for d in range(100)] self.path = tempfile.NamedTemporaryFile().name def _write_lines(self, sink, lines): f = sink.open(self.path) for line in lines: sink.write_record(f, line) sink.close(f) def test_write_text_file(self): sink = fileio.PureTextFileSink(self.path) self._write_lines(sink, self.lines) with open(self.path, 'r') as f: self.assertEqual(f.read().splitlines(), self.lines) def test_write_gzip_file(self): sink = fileio.PureTextFileSink( self.path, compression_type=fileio.CompressionTypes.DEFLATE) self._write_lines(sink, self.lines) with gzip.GzipFile(self.path, 'r') as f: self.assertEqual(f.read().splitlines(), self.lines) class MyFileSink(fileio.FileSink): def open(self, temp_path): file_handle = super(MyFileSink, self).open(temp_path) file_handle.write('[start]') return file_handle def write_encoded_record(self, file_handle, encoded_value): file_handle.write('[') file_handle.write(encoded_value) file_handle.write(']') def close(self, file_handle): file_handle.write('[end]') file_handle = super(MyFileSink, self).close(file_handle) class TestFileSink(unittest.TestCase): def test_file_sink_writing(self): temp_path = tempfile.NamedTemporaryFile().name sink = MyFileSink(temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder()) # Manually invoke the generic Sink API. init_token = sink.initialize_write() writer1 = sink.open_writer(init_token, '1') writer1.write('a') writer1.write('b') res1 = writer1.close() writer2 = sink.open_writer(init_token, '2') writer2.write('x') writer2.write('y') writer2.write('z') res2 = writer2.close() res = list(sink.finalize_write(init_token, [res1, res2])) # Retry the finalize operation (as if the first attempt was lost). res = list(sink.finalize_write(init_token, [res1, res2])) # Check the results. shard1 = temp_path + '-00000-of-00002.foo' shard2 = temp_path + '-00001-of-00002.foo' self.assertEqual(res, [shard1, shard2]) self.assertEqual(open(shard1).read(), '[start][a][b][end]') self.assertEqual(open(shard2).read(), '[start][x][y][z][end]') # Check that any temp files are deleted. self.assertEqual([shard1, shard2], sorted(glob.glob(temp_path + '*'))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/gcsio.py000066400000000000000000000431071272340251000232740ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Google Cloud Storage client. This library evolved from the Google App Engine GCS client available at https://github.com/GoogleCloudPlatform/appengine-gcs-client. """ import errno import fnmatch import logging import multiprocessing import os import re import StringIO import threading from google.cloud.dataflow.internal import auth from google.cloud.dataflow.utils import retry from apitools.base.py.exceptions import HttpError import apitools.base.py.transfer as transfer # Issue a friendlier error message if the storage library is not available. # TODO(silviuc): Remove this guard when storage is available everywhere. try: # pylint: disable=g-import-not-at-top from google.cloud.dataflow.internal.clients import storage except ImportError: raise RuntimeError( 'Google Cloud Storage I/O not supported for this execution environment ' '(could not import storage API client).') DEFAULT_READ_BUFFER_SIZE = 1024 * 1024 def parse_gcs_path(gcs_path): """Return the bucket and object names of the given gs:// path.""" match = re.match('^gs://([^/]+)/(.+)$', gcs_path) if match is None: raise ValueError('GCS path must be in the form gs:///.') return match.group(1), match.group(2) class GcsIO(object): """Google Cloud Storage I/O client.""" _instance = None def __new__(cls, storage_client=None): if storage_client: return super(GcsIO, cls).__new__(cls, storage_client) else: # Create a single storage client for each thread. We would like to avoid # creating more than one storage client for each thread, since each # initialization requires the relatively expensive step of initializing # credentaials. local_state = threading.local() if getattr(local_state, 'gcsio_instance', None) is None: credentials = auth.get_service_credentials() storage_client = storage.StorageV1(credentials=credentials) local_state.gcsio_instance = ( super(GcsIO, cls).__new__(cls, storage_client)) local_state.gcsio_instance.client = storage_client return local_state.gcsio_instance def __init__(self, storage_client=None): # We must do this check on storage_client because the client attribute may # have already been set in __new__ for the singleton case when # storage_client is None. if storage_client is not None: self.client = storage_client def open(self, filename, mode='r', read_buffer_size=DEFAULT_READ_BUFFER_SIZE, mime_type='application/octet-stream'): """Open a GCS file path for reading or writing. Args: filename: GCS file path in the form gs:///. mode: 'r' for reading or 'w' for writing. read_buffer_size: Buffer size to use during read operations. mime_type: Mime type to set for write operations. Returns: file object. Raises: ValueError: Invalid open file mode. """ if mode == 'r' or mode == 'rb': return GcsBufferedReader(self.client, filename, buffer_size=read_buffer_size) elif mode == 'w' or mode == 'wb': return GcsBufferedWriter(self.client, filename, mime_type=mime_type) else: raise ValueError('Invalid file open mode: %s.' % mode) @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def glob(self, pattern): """Return the GCS path names matching a given path name pattern. Path name patterns are those recognized by fnmatch.fnmatch(). The path can contain glob characters (*, ?, and [...] sets). Args: pattern: GCS file path pattern in the form gs:///. Returns: list of GCS file paths matching the given pattern. """ bucket, name_pattern = parse_gcs_path(pattern) # Get the prefix with which we can list objects in the given bucket. prefix = re.match('^[^[*?]*', name_pattern).group(0) request = storage.StorageObjectsListRequest(bucket=bucket, prefix=prefix) object_paths = [] while True: response = self.client.objects.List(request) for item in response.items: if fnmatch.fnmatch(item.name, name_pattern): object_paths.append('gs://%s/%s' % (item.bucket, item.name)) if response.nextPageToken: request.pageToken = response.nextPageToken else: break return object_paths @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def delete(self, path): """Deletes the given gcs object. Args: path: GCS file path pattern in the form gs:///. """ bucket, object_path = parse_gcs_path(path) request = storage.StorageObjectsDeleteRequest(bucket=bucket, object=object_path) self.client.objects.Delete(request) @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def rename(self, src, dst): """Renames the given gcs object from src to dst. Args: src: GCS file path pattern in the form gs:///. dst: GCS file path pattern in the form gs:///. """ src_bucket, src_path = parse_gcs_path(src) dst_bucket, dst_path = parse_gcs_path(dst) request = storage.StorageObjectsCopyRequest(sourceBucket=src_bucket, sourceObject=src_path, destinationBucket=dst_bucket, destinationObject=dst_path) self.client.objects.Copy(request) request = storage.StorageObjectsDeleteRequest(bucket=src_bucket, object=src_path) self.client.objects.Delete(request) @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def exists(self, path): """Returns whether the given gcs object exists. Args: path: GCS file path pattern in the form gs:///. """ bucket, object_path = parse_gcs_path(path) try: request = storage.StorageObjectsGetRequest(bucket=bucket, object=object_path) self.client.objects.Get(request) # metadata return True except IOError: return False class GcsBufferedReader(object): """A class for reading Google Cloud Storage files.""" def __init__(self, client, path, buffer_size=DEFAULT_READ_BUFFER_SIZE): self.client = client self.path = path self.bucket, self.name = parse_gcs_path(path) self.buffer_size = buffer_size # Get object state. get_request = ( storage.StorageObjectsGetRequest( bucket=self.bucket, object=self.name)) try: metadata = self._get_object_metadata(get_request) except HttpError as http_error: if http_error.status_code == 404: raise IOError(errno.ENOENT, 'Not found: %s' % self.path) else: logging.error( 'HTTP error while requesting file %s: %s', self.path, http_error) raise self.size = metadata.size # Ensure read is from file of the correct generation. get_request.generation = metadata.generation # Initialize read buffer state. self.download_stream = StringIO.StringIO() self.downloader = transfer.Download( self.download_stream, auto_transfer=False) self.client.objects.Get(get_request, download=self.downloader) self.position = 0 self.buffer = '' self.buffer_start_position = 0 self.closed = False @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def _get_object_metadata(self, get_request): return self.client.objects.Get(get_request) def read(self, size=-1): """Read data from a GCS file. Args: size: Number of bytes to read. Actual number of bytes read is always equal to size unless EOF is reached. If size is negative or unspecified, read the entire file. Returns: data read as str. Raises: IOError: When this buffer is closed. """ return self._read_inner(size=size, readline=False) def readline(self, size=-1): """Read one line delimited by '\\n' from the file. Mimics behavior of the readline() method on standard file objects. A trailing newline character is kept in the string. It may be absent when a file ends with an incomplete line. If the size argument is non-negative, it specifies the maximum string size (counting the newline) to return. A negative size is the same as unspecified. Empty string is returned only when EOF is encountered immediately. Args: size: Maximum number of bytes to read. If not specified, readline stops only on '\\n' or EOF. Returns: The data read as a string. Raises: IOError: When this buffer is closed. """ return self._read_inner(size=size, readline=True) def _read_inner(self, size=-1, readline=False): """Shared implementation of read() and readline().""" self._check_open() if not self._remaining(): return '' # Prepare to read. data_list = [] if size is None: size = -1 to_read = min(size, self._remaining()) if to_read < 0: to_read = self._remaining() break_after = False while to_read > 0: # If we have exhausted the buffer, get the next segment. # TODO(ccy): We should consider prefetching the next block in another # thread. self._fetch_next_if_buffer_exhausted() # Determine number of bytes to read from buffer. buffer_bytes_read = self.position - self.buffer_start_position bytes_to_read_from_buffer = min( len(self.buffer) - buffer_bytes_read, to_read) # If readline is set, we only want to read up to and including the next # newline character. if readline: next_newline_position = self.buffer.find( '\n', buffer_bytes_read, len(self.buffer)) if next_newline_position != -1: bytes_to_read_from_buffer = (1 + next_newline_position - buffer_bytes_read) break_after = True # Read bytes. data_list.append( self.buffer[buffer_bytes_read:buffer_bytes_read + bytes_to_read_from_buffer]) self.position += bytes_to_read_from_buffer to_read -= bytes_to_read_from_buffer if break_after: break return ''.join(data_list) def _fetch_next_if_buffer_exhausted(self): if not self.buffer or (self.buffer_start_position + len(self.buffer) <= self.position): bytes_to_request = min(self._remaining(), self.buffer_size) self.buffer_start_position = self.position self.buffer = self._get_segment(self.position, bytes_to_request) def _remaining(self): return self.size - self.position def close(self): """Close the current GCS file.""" self.closed = True self.download_stream = None self.downloader = None self.buffer = None def _get_segment(self, start, size): """Get the given segment of the current GCS file.""" if size == 0: return '' end = start + size - 1 self.downloader.GetRange(start, end) value = self.download_stream.getvalue() # Clear the StringIO object after we've read its contents. self.download_stream.truncate(0) assert len(value) == size return value def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): self.close() def seek(self, offset, whence=os.SEEK_SET): """Set the file's current offset. Note if the new offset is out of bound, it is adjusted to either 0 or EOF. Args: offset: seek offset as number. whence: seek mode. Supported modes are os.SEEK_SET (absolute seek), os.SEEK_CUR (seek relative to the current position), and os.SEEK_END (seek relative to the end, offset should be negative). Raises: IOError: When this buffer is closed. ValueError: When whence is invalid. """ self._check_open() self.buffer = '' self.buffer_start_position = -1 if whence == os.SEEK_SET: self.position = offset elif whence == os.SEEK_CUR: self.position += offset elif whence == os.SEEK_END: self.position = self.size + offset else: raise ValueError('Whence mode %r is invalid.' % whence) self.position = min(self.position, self.size) self.position = max(self.position, 0) def tell(self): """Tell the file's current offset. Returns: current offset in reading this file. Raises: IOError: When this buffer is closed. """ self._check_open() return self.position def _check_open(self): if self.closed: raise IOError('Buffer is closed.') def seekable(self): return True def readable(self): return True def writable(self): return False class GcsBufferedWriter(object): """A class for writing Google Cloud Storage files.""" class PipeStream(object): """A class that presents a pipe connection as a readable stream.""" def __init__(self, recv_pipe): self.conn = recv_pipe self.closed = False self.position = 0 self.remaining = '' def read(self, size): """Read data from the wrapped pipe connection. Args: size: Number of bytes to read. Actual number of bytes read is always equal to size unless EOF is reached. Returns: data read as str. """ data_list = [] bytes_read = 0 while bytes_read < size: bytes_from_remaining = min(size - bytes_read, len(self.remaining)) data_list.append(self.remaining[0:bytes_from_remaining]) self.remaining = self.remaining[bytes_from_remaining:] self.position += bytes_from_remaining bytes_read += bytes_from_remaining if not self.remaining: try: self.remaining = self.conn.recv_bytes() except EOFError: break return ''.join(data_list) def tell(self): """Tell the file's current offset. Returns: current offset in reading this file. Raises: IOError: When this stream is closed. """ self._check_open() return self.position def seek(self, offset, whence=os.SEEK_SET): # The apitools.base.py.transfer.Upload class insists on seeking to the end # of a stream to do a check before completing an upload, so we must have # this no-op method here in that case. if whence == os.SEEK_END and offset == 0: return elif whence == os.SEEK_SET and offset == self.position: return raise NotImplementedError def _check_open(self): if self.closed: raise IOError('Stream is closed.') def __init__(self, client, path, mime_type='application/octet-stream'): self.client = client self.path = path self.bucket, self.name = parse_gcs_path(path) self.closed = False self.position = 0 # Set up communication with uploading thread. parent_conn, child_conn = multiprocessing.Pipe() self.conn = parent_conn # Set up uploader. self.insert_request = ( storage.StorageObjectsInsertRequest( bucket=self.bucket, name=self.name)) self.upload = transfer.Upload(GcsBufferedWriter.PipeStream(child_conn), mime_type) self.upload.strategy = transfer.RESUMABLE_UPLOAD # Start uploading thread. self.upload_thread = threading.Thread(target=self._start_upload) self.upload_thread.daemon = True self.upload_thread.start() # TODO(silviuc): Refactor so that retry logic can be applied. # There is retry logic in the underlying transfer library but we should make # it more explicit so we can control the retry parameters. @retry.no_retries # Using no_retries marks this as an integration point. def _start_upload(self): # This starts the uploader thread. We are forced to run the uploader in # another thread because the apitools uploader insists on taking a stream # as input. Happily, this also means we get asynchronous I/O to GCS. # # The uploader by default transfers data in chunks of 1024 * 1024 bytes at # a time, buffering writes until that size is reached. self.client.objects.Insert(self.insert_request, upload=self.upload) def write(self, data): """Write data to a GCS file. Args: data: data to write as str. Raises: IOError: When this buffer is closed. """ self._check_open() if not data: return self.conn.send_bytes(data) self.position += len(data) def tell(self): """Return the total number of bytes passed to write() so far.""" return self.position def close(self): """Close the current GCS file.""" self.conn.close() self.upload_thread.join() def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): self.close() def _check_open(self): if self.closed: raise IOError('Buffer is closed.') def seekable(self): return False def readable(self): return False def writable(self): return True DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/gcsio_test.py000066400000000000000000000307521272340251000243350ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for Google Cloud Storage client.""" import logging import multiprocessing import os import random import threading import unittest from google.cloud.dataflow.io import gcsio from google.cloud.dataflow.internal.clients import storage class FakeGcsClient(object): # Fake storage client. Usage in gcsio.py is client.objects.Get(...) and # client.objects.Insert(...). def __init__(self): self.objects = FakeGcsObjects() class FakeFile(object): def __init__(self, bucket, obj, contents, generation): self.bucket = bucket self.object = obj self.contents = contents self.generation = generation def get_metadata(self): return storage.Object(bucket=self.bucket, name=self.object, generation=self.generation, size=len(self.contents)) class FakeGcsObjects(object): def __init__(self): self.files = {} self.list_page_tokens = {} def add_file(self, f): self.files[(f.bucket, f.object)] = f def get_file(self, bucket, obj): return self.files.get((bucket, obj), None) def Get(self, get_request, download=None): # pylint: disable=invalid-name f = self.get_file(get_request.bucket, get_request.object) if f is None: raise ValueError('Specified object does not exist.') if download is None: return f.get_metadata() else: stream = download.stream def get_range_callback(start, end): assert start >= 0 and end >= start and end < len(f.contents) stream.write(f.contents[start:end + 1]) download.GetRange = get_range_callback def Insert(self, insert_request, upload=None): # pylint: disable=invalid-name assert upload is not None generation = 1 f = self.get_file(insert_request.bucket, insert_request.name) if f is not None: generation = f.generation + 1 f = FakeFile(insert_request.bucket, insert_request.name, '', generation) # Stream data into file. stream = upload.stream data_list = [] while True: data = stream.read(1024 * 1024) if not data: break data_list.append(data) f.contents = ''.join(data_list) self.add_file(f) def List(self, list_request): # pylint: disable=invalid-name bucket = list_request.bucket prefix = list_request.prefix or '' matching_files = [] for file_bucket, file_name in sorted(iter(self.files)): if bucket == file_bucket and file_name.startswith(prefix): file_object = self.files[(file_bucket, file_name)].get_metadata() matching_files.append(file_object) # Handle pagination. items_per_page = 5 if not list_request.pageToken: range_start = 0 else: if list_request.pageToken not in self.list_page_tokens: raise ValueError('Invalid page token.') range_start = self.list_page_tokens[list_request.pageToken] del self.list_page_tokens[list_request.pageToken] result = storage.Objects( items=matching_files[range_start:range_start + items_per_page]) if range_start + items_per_page < len(matching_files): next_range_start = range_start + items_per_page next_page_token = '_page_token_%s_%s_%d' % (bucket, prefix, next_range_start) self.list_page_tokens[next_page_token] = next_range_start result.nextPageToken = next_page_token return result class TestGCSPathParser(unittest.TestCase): def test_gcs_path(self): self.assertEqual( gcsio.parse_gcs_path('gs://bucket/name'), ('bucket', 'name')) self.assertEqual( gcsio.parse_gcs_path('gs://bucket/name/sub'), ('bucket', 'name/sub')) def test_bad_gcs_path(self): self.assertRaises(ValueError, gcsio.parse_gcs_path, 'gs://') self.assertRaises(ValueError, gcsio.parse_gcs_path, 'gs://bucket') self.assertRaises(ValueError, gcsio.parse_gcs_path, 'gs://bucket/') self.assertRaises(ValueError, gcsio.parse_gcs_path, 'gs:///name') self.assertRaises(ValueError, gcsio.parse_gcs_path, 'gs:///') self.assertRaises(ValueError, gcsio.parse_gcs_path, 'gs:/blah/bucket/name') class TestGCSIO(unittest.TestCase): def _insert_random_file(self, client, path, size, generation=1): bucket, name = gcsio.parse_gcs_path(path) f = FakeFile(bucket, name, os.urandom(size), generation) client.objects.add_file(f) return f def setUp(self): self.client = FakeGcsClient() self.gcs = gcsio.GcsIO(self.client) def test_full_file_read(self): file_name = 'gs://gcsio-test/full_file' file_size = 5 * 1024 * 1024 + 100 random_file = self._insert_random_file(self.client, file_name, file_size) f = self.gcs.open(file_name) f.seek(0, os.SEEK_END) self.assertEqual(f.tell(), file_size) self.assertEqual(f.read(), '') f.seek(0) self.assertEqual(f.read(), random_file.contents) def test_file_random_seek(self): file_name = 'gs://gcsio-test/seek_file' file_size = 5 * 1024 * 1024 - 100 random_file = self._insert_random_file(self.client, file_name, file_size) f = self.gcs.open(file_name) random.seed(0) for _ in range(0, 10): a = random.randint(0, file_size - 1) b = random.randint(0, file_size - 1) start, end = min(a, b), max(a, b) f.seek(start) self.assertEqual(f.tell(), start) self.assertEqual(f.read(end - start + 1), random_file.contents[start:end + 1]) self.assertEqual(f.tell(), end + 1) def test_file_read_line(self): file_name = 'gs://gcsio-test/read_line_file' lines = [] # Set a small buffer size to exercise refilling the buffer. # First line is carefully crafted so the newline falls as the last character # of the buffer to exercise this code path. read_buffer_size = 1024 lines.append('x' * 1023 + '\n') for _ in range(1, 1000): line_length = random.randint(100, 500) line = os.urandom(line_length).replace('\n', ' ') + '\n' lines.append(line) contents = ''.join(lines) file_size = len(contents) bucket, name = gcsio.parse_gcs_path(file_name) self.client.objects.add_file(FakeFile(bucket, name, contents, 1)) f = self.gcs.open(file_name, read_buffer_size=read_buffer_size) # Test read of first two lines. f.seek(0) self.assertEqual(f.readline(), lines[0]) self.assertEqual(f.tell(), len(lines[0])) self.assertEqual(f.readline(), lines[1]) # Test read at line boundary. f.seek(file_size - len(lines[-1]) - 1) self.assertEqual(f.readline(), '\n') # Test read at end of file. f.seek(file_size) self.assertEqual(f.readline(), '') # Test reads at random positions. random.seed(0) for _ in range(0, 10): start = random.randint(0, file_size - 1) line_index = 0 # Find line corresponding to start index. chars_left = start while True: next_line_length = len(lines[line_index]) if chars_left - next_line_length < 0: break chars_left -= next_line_length line_index += 1 f.seek(start) self.assertEqual(f.readline(), lines[line_index][chars_left:]) def test_file_write(self): file_name = 'gs://gcsio-test/write_file' file_size = 5 * 1024 * 1024 + 2000 contents = os.urandom(file_size) f = self.gcs.open(file_name, 'w') f.write(contents[0:1000]) f.write(contents[1000:1024 * 1024]) f.write(contents[1024 * 1024:]) f.close() bucket, name = gcsio.parse_gcs_path(file_name) self.assertEqual( self.client.objects.get_file(bucket, name).contents, contents) def test_context_manager(self): # Test writing with a context manager. file_name = 'gs://gcsio-test/context_manager_file' file_size = 1024 contents = os.urandom(file_size) with self.gcs.open(file_name, 'w') as f: f.write(contents) bucket, name = gcsio.parse_gcs_path(file_name) self.assertEqual( self.client.objects.get_file(bucket, name).contents, contents) # Test reading with a context manager. with self.gcs.open(file_name) as f: self.assertEqual(f.read(), contents) # Test that exceptions are not swallowed by the context manager. with self.assertRaises(ZeroDivisionError): with self.gcs.open(file_name) as f: f.read(0 / 0) def test_glob(self): bucket_name = 'gcsio-test' object_names = [ 'cow/cat/fish', 'cow/cat/blubber', 'cow/dog/blubber', 'apple/dog/blubber', 'apple/fish/blubber', 'apple/fish/blowfish', 'apple/fish/bambi', 'apple/fish/balloon', 'apple/fish/cat', 'apple/fish/cart', 'apple/fish/carl', 'apple/dish/bat', 'apple/dish/cat', 'apple/dish/carl', ] for object_name in object_names: file_name = 'gs://%s/%s' % (bucket_name, object_name) self._insert_random_file(self.client, file_name, 0) test_cases = [ ('gs://gcsio-test/*', [ 'cow/cat/fish', 'cow/cat/blubber', 'cow/dog/blubber', 'apple/dog/blubber', 'apple/fish/blubber', 'apple/fish/blowfish', 'apple/fish/bambi', 'apple/fish/balloon', 'apple/fish/cat', 'apple/fish/cart', 'apple/fish/carl', 'apple/dish/bat', 'apple/dish/cat', 'apple/dish/carl', ]), ('gs://gcsio-test/cow/*', [ 'cow/cat/fish', 'cow/cat/blubber', 'cow/dog/blubber', ]), ('gs://gcsio-test/cow/ca*', [ 'cow/cat/fish', 'cow/cat/blubber', ]), ('gs://gcsio-test/apple/[df]ish/ca*', [ 'apple/fish/cat', 'apple/fish/cart', 'apple/fish/carl', 'apple/dish/cat', 'apple/dish/carl', ]), ('gs://gcsio-test/apple/fish/car?', [ 'apple/fish/cart', 'apple/fish/carl', ]), ('gs://gcsio-test/apple/fish/b*', [ 'apple/fish/blubber', 'apple/fish/blowfish', 'apple/fish/bambi', 'apple/fish/balloon', ]), ('gs://gcsio-test/apple/dish/[cb]at', [ 'apple/dish/bat', 'apple/dish/cat', ]), ] for file_pattern, expected_object_names in test_cases: expected_file_names = ['gs://%s/%s' % (bucket_name, o) for o in expected_object_names] self.assertEqual(set(self.gcs.glob(file_pattern)), set(expected_file_names)) class TestPipeStream(unittest.TestCase): def _read_and_verify(self, stream, expected, buffer_size): data_list = [] bytes_read = 0 seen_last_block = False while True: data = stream.read(buffer_size) self.assertLessEqual(len(data), buffer_size) if len(data) < buffer_size: # Test the constraint that the pipe stream returns less than the buffer # size only when at the end of the stream. if data: self.assertFalse(seen_last_block) seen_last_block = True if not data: break data_list.append(data) bytes_read += len(data) self.assertEqual(stream.tell(), bytes_read) self.assertEqual(''.join(data_list), expected) def test_pipe_stream(self): block_sizes = list(4 ** i for i in range(0, 12)) data_blocks = list(os.urandom(size) for size in block_sizes) expected = ''.join(data_blocks) buffer_sizes = [100001, 512 * 1024, 1024 * 1024] for buffer_size in buffer_sizes: parent_conn, child_conn = multiprocessing.Pipe() stream = gcsio.GcsBufferedWriter.PipeStream(child_conn) child_thread = threading.Thread(target=self._read_and_verify, args=(stream, expected, buffer_size)) child_thread.start() for data in data_blocks: parent_conn.send_bytes(data) parent_conn.close() child_thread.join() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/iobase.py000066400000000000000000001155441272340251000234370ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Sources and sinks. A Source manages record-oriented data input from a particular kind of source (e.g. a set of files, a database table, etc.). The reader() method of a source returns a reader object supporting the iterator protocol; iteration yields raw records of unprocessed, serialized data. A Sink manages record-oriented data output to a particular kind of sink (e.g. a set of files, a database table, etc.). The writer() method of a sink returns a writer object supporting writing records of serialized data to the sink. """ from collections import namedtuple import logging import uuid from google.cloud.dataflow import pvalue from google.cloud.dataflow.coders import PickleCoder from google.cloud.dataflow.pvalue import AsIter from google.cloud.dataflow.pvalue import AsSingleton from google.cloud.dataflow.transforms import core from google.cloud.dataflow.transforms import ptransform from google.cloud.dataflow.transforms import window def _dict_printable_fields(dict_object, skip_fields): """Returns a list of strings for the interesting fields of a dict.""" return ['%s=%r' % (name, value) for name, value in dict_object.iteritems() # want to output value 0 but not None nor [] if (value or value == 0) and name not in skip_fields] _minor_fields = ['coder', 'key_coder', 'value_coder', 'config_bytes', 'elements', 'append_trailing_newlines', 'strip_trailing_newlines', 'compression_type'] class NativeSource(object): """A source implemented by Dataflow service. This class is to be only inherited by sources natively implemented by Cloud Dataflow service, hence should not be sub-classed by users. This class is deprecated and should not be used to define new sources. """ def reader(self): """Returns a NativeSourceReader instance associated with this source.""" raise NotImplementedError def __repr__(self): return '<{name} {vals}>'.format( name=self.__class__.__name__, vals=', '.join(_dict_printable_fields(self.__dict__, _minor_fields))) class NativeSourceReader(object): """A reader for a source implemented by Dataflow service.""" def __enter__(self): """Opens everything necessary for a reader to function properly.""" raise NotImplementedError def __exit__(self, exception_type, exception_value, traceback): """Cleans up after a reader executed.""" raise NotImplementedError def __iter__(self): """Returns an iterator over all the records of the source.""" raise NotImplementedError @property def returns_windowed_values(self): """Returns whether this reader returns windowed values.""" return False def get_progress(self): """Returns a representation of how far the reader has read. Returns: A SourceReaderProgress object that gives the current progress of the reader. """ return def request_dynamic_split(self, dynamic_split_request): """Attempts to split the input in two parts. The two parts are named the "primary" part and the "residual" part. The current 'NativeSourceReader' keeps processing the primary part, while the residual part will be processed elsewhere (e.g. perhaps on a different worker). The primary and residual parts, if concatenated, must represent the same input as the current input of this 'NativeSourceReader' before this call. The boundary between the primary part and the residual part is specified in a framework-specific way using 'DynamicSplitRequest' e.g., if the framework supports the notion of positions, it might be a position at which the input is asked to split itself (which is not necessarily the same position at which it *will* split itself); it might be an approximate fraction of input, or something else. This function returns a 'DynamicSplitResult', which encodes, in a framework-specific way, the information sufficient to construct a description of the resulting primary and residual inputs. For example, it might, again, be a position demarcating these parts, or it might be a pair of fully-specified input descriptions, or something else. After a successful call to 'request_dynamic_split()', subsequent calls should be interpreted relative to the new primary. Args: dynamic_split_request: A 'DynamicSplitRequest' describing the split request. Returns: 'None' if the 'DynamicSplitRequest' cannot be honored (in that case the input represented by this 'NativeSourceReader' stays the same), or a 'DynamicSplitResult' describing how the input was split into a primary and residual part. """ logging.debug( 'SourceReader %r does not support dynamic splitting. Ignoring dynamic ' 'split request: %r', self, dynamic_split_request) return class ReaderProgress(object): """A representation of how far a NativeSourceReader has read.""" def __init__(self, position=None, percent_complete=None, remaining_time=None): self._position = position if percent_complete is not None: percent_complete = float(percent_complete) if percent_complete < 0 or percent_complete > 1: raise ValueError( 'The percent_complete argument was %f. Must be in range [0, 1].' % percent_complete) self._percent_complete = percent_complete self._remaining_time = remaining_time @property def position(self): """Returns progress, represented as a ReaderPosition object.""" return self._position @property def percent_complete(self): """Returns progress, represented as a percentage of total work. Progress range from 0.0 (beginning, nothing complete) to 1.0 (end of the work range, entire WorkItem complete). Returns: Progress represented as a percentage of total work. """ return self._percent_complete @property def remaining_time(self): """Returns progress, represented as an estimated time remaining.""" return self._remaining_time class ReaderPosition(object): """A representation of position in an iteration of a 'NativeSourceReader'.""" def __init__(self, end=None, key=None, byte_offset=None, record_index=None, shuffle_position=None, concat_position=None): """Initializes ReaderPosition. A ReaderPosition may get instantiated for one of these position types. Only one of these should be specified. Args: end: position is past all other positions. For example, this may be used to represent the end position of an unbounded range. key: position is a string key. byte_offset: position is a byte offset. record_index: position is a record index shuffle_position: position is a base64 encoded shuffle position. concat_position: position is a 'ConcatPosition'. """ self.end = end self.key = key self.byte_offset = byte_offset self.record_index = record_index self.shuffle_position = shuffle_position if concat_position is not None: assert isinstance(concat_position, ConcatPosition) self.concat_position = concat_position class ConcatPosition(object): """A position that encapsulate an inner position and an index. This is used to represent the position of a source that encapsulate several other sources. """ def __init__(self, index, position): """Initializes ConcatPosition. Args: index: index of the source currently being read. position: inner position within the source currently being read. """ if position is not None: assert isinstance(position, ReaderPosition) self.index = index self.position = position class DynamicSplitRequest(object): """Specifies how 'NativeSourceReader.request_dynamic_split' should split. """ def __init__(self, progress): assert isinstance(progress, ReaderProgress) self.progress = progress class DynamicSplitResult(object): pass class DynamicSplitResultWithPosition(DynamicSplitResult): def __init__(self, stop_position): assert isinstance(stop_position, ReaderPosition) self.stop_position = stop_position class NativeSink(object): """A sink implemented by Dataflow service. This class is to be only inherited by sinks natively implemented by Cloud Dataflow service, hence should not be sub-classed by users. """ def writer(self): """Returns a SinkWriter for this source.""" raise NotImplementedError def __repr__(self): return '<{name} {vals}>'.format( name=self.__class__.__name__, vals=_dict_printable_fields(self.__dict__, _minor_fields)) class NativeSinkWriter(object): """A writer for a sink implemented by Dataflow service.""" def __enter__(self): """Opens everything necessary for a writer to function properly.""" raise NotImplementedError def __exit__(self, exception_type, exception_value, traceback): """Cleans up after a writer executed.""" raise NotImplementedError @property def takes_windowed_values(self): """Returns whether this writer takes windowed values.""" return False def Write(self, o): # pylint: disable=invalid-name """Writes a record to the sink associated with this writer.""" raise NotImplementedError # Encapsulates information about a bundle of a source generated when method # BoundedSource.split() is invoked. # This is a named 4-tuple that has following fields. # * weight - a number that represents the size of the bundle. This value will # be used to compare the relative sizes of bundles generated by the # current source. # The weight returned here could be specified using a unit of your # choice (for example, bundles of sizes 100MB, 200MB, and 700MB may # specify weights 100, 200, 700 or 1, 2, 7) but all bundles of a # source should specify the weight using the same unit. # * source - a BoundedSource object for the bundle. # * start_position - starting position of the bundle # * stop_position - ending position of the bundle. # # Type for start and stop positions are specific to the bounded source and must # be consistent throughout. SourceBundle = namedtuple( 'SourceBundle', 'weight source start_position stop_position') class BoundedSource(object): """A Dataflow source that reads a finite amount of input records. This class defines following operations which can be used to read the source efficiently. * Size estimation - method ``estimate_size()`` may return an accurate estimation in bytes for the size of the source. * Splitting into bundles of a given size - method ``split()`` can be used to split the source into a set of sub-sources (bundles) based on a desired bundle size. * Getting a RangeTracker - method ``get_range_tracker() should return a ``RangeTracker`` object for a given position range for the position type of the records returned by the source. * Reading the data - method ``read()`` can be used to read data from the source while respecting the boundaries defined by a given ``RangeTracker``. """ def estimate_size(self): """Estimates the size of source in bytes. An estimate of the total size (in bytes) of the data that would be read from this source. This estimate is in terms of external storage size, before performing decompression or other processing. Returns: estimated size of the source if the size can be determined, ``None`` otherwise. """ raise NotImplementedError def split(self, desired_bundle_size, start_position=None, stop_position=None): """Splits the source into a set of bundles. Bundles should be approximately of size ``desired_bundle_size`` bytes. Args: desired_bundle_size: the desired size (in bytes) of the bundles returned. start_position: if specified the given position must be used as the starting position of the first bundle. stop_position: if specified the given position must be used as the ending position of the last bundle. Returns: an iterator of objects of type 'SourceBundle' that gives information about the generated bundles. """ raise NotImplementedError def get_range_tracker(self, start_position, stop_position): """Returns a RangeTracker for a given position range. Framework may invoke ``read()`` method with the RangeTracker object returned here to read data from the source. Args: start_position: starting position of the range. stop_position: ending position of the range. Returns: a ``RangeTracker`` for the given position range. """ raise NotImplementedError def read(self, range_tracker): """Returns an iterator that reads data from the source. The returned set of data must respect the boundaries defined by the given ``RangeTracker`` object. For example: * Returned set of data must be for the range ``[range_tracker.start_position, range_tracker.stop_position)``. Note that a source may decide to return records that start after ``range_tracker.stop_position``. See documentation in class ``RangeTracker`` for more details. Also, note that framework might invoke ``range_tracker.try_split()`` to perform dynamic split operations. range_tracker.stop_position may be updated dynamically due to successful dynamic split operations. * Method ``range_tracker.try_split()`` must be invoked for every record that starts at a split point. * Method ``range_tracker.record_current_position()`` may be invoked for records that do not start at split points. Args: range_tracker: a ``RangeTracker`` whose boundaries must be respected when reading data from the source. If 'None' all records represented by the current source should be read. Returns: an iterator of data read by the source. """ raise NotImplementedError def default_output_coder(self): """Coder that should be used for the records returned by the source.""" return PickleCoder() class RangeTracker(object): """A thread safe object used by Dataflow source framework. A Dataflow source is defined using a ''BoundedSource'' and a ''RangeTracker'' pair. A ''RangeTracker'' is used by Dataflow source framework to perform dynamic work rebalancing of position-based sources. **Position-based sources** A position-based source is one where the source can be described by a range of positions of an ordered type and the records returned by the reader can be described by positions of the same type. In case a record occupies a range of positions in the source, the most important thing about the record is the position where it starts. Defining the semantics of positions for a source is entirely up to the source class, however the chosen definitions have to obey certain properties in order to make it possible to correctly split the source into parts, including dynamic splitting. Two main aspects need to be defined: 1. How to assign starting positions to records. 2. Which records should be read by a source with a range '[A, B)'. Moreover, reading a range must be *efficient*, i.e., the performance of reading a range should not significantly depend on the location of the range. For example, reading the range [A, B) should not require reading all data before 'A'. The sections below explain exactly what properties these definitions must satisfy, and how to use a ``RangeTracker`` with a properly defined source. **Properties of position-based sources** The main requirement for position-based sources is *associativity*: reading records from '[A, B)' and records from '[B, C)' should give the same records as reading from '[A, C)', where 'A <= B <= C'. This property ensures that no matter how a range of positions is split into arbitrarily many sub-ranges, the total set of records described by them stays the same. The other important property is how the source's range relates to positions of records in the source. In many sources each record can be identified by a unique starting position. In this case: * All records returned by a source '[A, B)' must have starting positions in this range. * All but the last record should end within this range. The last record may or may not extend past the end of the range. * Records should not overlap. Such sources should define "read '[A, B)'" as "read from the first record starting at or after 'A', up to but not including the first record starting at or after 'B'". Some examples of such sources include reading lines or CSV from a text file, reading keys and values from a BigTable, etc. The concept of *split points* allows to extend the definitions for dealing with sources where some records cannot be identified by a unique starting position. In all cases, all records returned by a source '[A, B)' must *start* at or after 'A'. **Split points** Some sources may have records that are not directly addressable. For example, imagine a file format consisting of a sequence of compressed blocks. Each block can be assigned an offset, but records within the block cannot be directly addressed without decompressing the block. Let us refer to this hypothetical format as CBF (Compressed Blocks Format). Many such formats can still satisfy the associativity property. For example, in CBF, reading '[A, B)' can mean "read all the records in all blocks whose starting offset is in '[A, B)'". To support such complex formats, we introduce the notion of *split points*. We say that a record is a split point if there exists a position 'A' such that the record is the first one to be returned when reading the range '[A, infinity)'. In CBF, the only split points would be the first records in each block. Split points allow us to define the meaning of a record's position and a source's range in all cases: * For a record that is at a split point, its position is defined to be the largest 'A' such that reading a source with the range '[A, infinity)' returns this record. * Positions of other records are only required to be non-decreasing. * Reading the source '[A, B)' must return records starting from the first split point at or after 'A', up to but not including the first split point at or after 'B'. In particular, this means that the first record returned by a source MUST always be a split point. * Positions of split points must be unique. As a result, for any decomposition of the full range of the source into position ranges, the total set of records will be the full set of records in the source, and each record will be read exactly once. **Consumed positions** As the source is being read, and records read from it are being passed to the downstream transforms in the pipeline, we say that positions in the source are being *consumed*. When a reader has read a record (or promised to a caller that a record will be returned), positions up to and including the record's start position are considered *consumed*. Dynamic splitting can happen only at *unconsumed* positions. If the reader just returned a record at offset 42 in a file, dynamic splitting can happen only at offset 43 or beyond, as otherwise that record could be read twice (by the current reader and by a reader of the task starting at 43). """ def start_position(self): """Returns the starting position of the current range, inclusive.""" raise NotImplementedError def stop_position(self): """Returns the ending position of the current range, exclusive.""" raise NotImplementedError def try_claim(self, position): # pylint: disable=unused-argument """Atomically determines if a record at a split point is within the range. This method should be called **if and only if** the record is at a split point. This method may modify the internal state of the ``RangeTracker`` by updating the last-consumed position to ``position``. ** Thread safety ** This method along with several other methods of this class may be invoked by multiple threads, hence must be made thread-safe, e.g. by using a single lock object. Args: position: starting position of a record being read by a source. Returns: ``True``, if the given position falls within the current range, returns ``False`` otherwise. """ raise NotImplementedError def set_current_position(self, position): """Updates the last-consumed position to the given position. A source may invoke this method for records that do not start at split points. This may modify the internal state of the ``RangeTracker``. If the record starts at a split point, method ``try_claim()`` **must** be invoked instead of this method. Args: position: starting position of a record being read by a source. """ raise NotImplementedError def position_at_fraction(self, fraction): """Returns the position at the given fraction. Given a fraction within the range [0.0, 1.0) this method will return the position at the given fraction compared the the position range [self.start_position, self.stop_position). ** Thread safety ** This method along with several other methods of this class may be invoked by multiple threads, hence must be made thread-safe, e.g. by using a single lock object. Args: fraction: a float value within the range [0.0, 1.0). Returns: a position within the range [self.start_position, self.stop_position). """ raise NotImplementedError def try_split(self, position): """Atomically splits the current range. Determines a position to split the current range, split_position, based on the given position. In most cases split_position and position will be the same. Splits the current range '[self.start_position, self.stop_position)' into a "primary" part '[self.start_position, split_position)' and a "residual" part '[split_position, self.stop_position)', assuming the current last-consumed position is within '[self.start_position, split_position)' (i.e., split_position has not been consumed yet). If successful, updates the current range to be the primary and returns a tuple (split_position, split_fraction). split_fraction should be the fraction of size of range '[self.start_position, split_position)' compared to the original (before split) range '[self.start_position, self.stop_position)'. If the split_position has already been consumed, returns ``None``. ** Thread safety ** This method along with several other methods of this class may be invoked by multiple threads, hence must be made thread-safe, e.g. by using a single lock object. Args: position: suggested position where the current range should try to be split at. Returns: a tuple containing the split position and split fraction. """ raise NotImplementedError def fraction_consumed(self): """Returns the approximate fraction of consumed positions in the source. ** Thread safety ** This method along with several other methods of this class may be invoked by multiple threads, hence must be made thread-safe, e.g. by using a single lock object. Returns: the approximate fraction of positions that have been consumed by successful 'try_split()' and 'report_current_position()' calls, or 0.0 if no such calls have happened. """ raise NotImplementedError class Sink(object): """A resource that can be written to using the ``df.io.Write`` transform. Here ``df`` stands for Dataflow Python code imported in following manner. ``import google.cloud.dataflow as df``. A parallel write to an ``iobase.Sink`` consists of three phases: 1. A sequential *initialization* phase (e.g., creating a temporary output directory, etc.) 2. A parallel write phase where workers write *bundles* of records 3. A sequential *finalization* phase (e.g., committing the writes, merging output files, etc.) For exact definition of a Dataflow bundle please see https://cloud.google.com/dataflow/faq. Implementing a new sink requires extending two classes. 1. iobase.Sink ``iobase.Sink`` is an immutable logical description of the location/resource to write to. Depending on the type of sink, it may contain fields such as the path to an output directory on a filesystem, a database table name, etc. ``iobase.Sink`` provides methods for performing a write operation to the sink described by it. To this end, implementors of an extension of ``iobase.Sink`` must implement three methods: ``initialize_write()``, ``open_writer()``, and ``finalize_write()``. 2. iobase.Writer ``iobase.Writer`` is used to write a single bundle of records. An ``iobase.Writer`` defines two methods: ``write()`` which writes a single record from the bundle and ``close()`` which is called once at the end of writing a bundle. See also ``df.io.fileio.FileSink`` which provides a simpler API for writing sinks that produce files. **Execution of the Write transform** ``initialize_write()`` and ``finalize_write()`` are conceptually called once: at the beginning and end of a ``Write`` transform. However, implementors must ensure that these methods are *idempotent*, as they may be called multiple times on different machines in the case of failure/retry or for redundancy. ``initialize_write()`` should perform any initialization that needs to be done prior to writing to the sink. ``initialize_write()`` may return a result (let's call this ``init_result``) that contains any parameters it wants to pass on to its writers about the sink. For example, a sink that writes to a file system may return an ``init_result`` that contains a dynamically generated unique directory to which data should be written. To perform writing of a bundle of elements, Dataflow execution engine will create an ``iobase.Writer`` using the implementation of ``iobase.Sink.open_writer()``. When invoking ``open_writer()`` execution engine will provide the ``init_result`` returned by ``initialize_write()`` invocation as well as a *bundle id* (let's call this ``bundle_id``) that is unique for each invocation of ``open_writer()``. Execution engine will then invoke ``iobase.Writer.write()`` implementation for each element that has to be written. Once all elements of a bundle are written, execution engine will invoke ``iobase.Writer.close()`` implementation which should return a result (let's call this ``write_result``) that contains information that encodes the result of the write and, in most cases, some encoding of the unique bundle id. For example, if each bundle is written to a unique temporary file, ``close()`` method may return an object that contains the temporary file name. After writing of all bundles is complete, execution engine will invoke ``finalize_write()`` implementation. As parameters to this invocation execution engine will provide ``init_result`` as well as an iterable of ``write_result``. The execution of a write transform can be illustrated using following pseudo code (assume that the outer for loop happens in parallel across many machines):: init_result = sink.initialize_write() write_results = [] for bundle in partition(pcoll): writer = sink.open_writer(init_result, generate_bundle_id()) for elem in bundle: writer.write(elem) write_results.append(writer.close()) sink.finalize_write(init_result, write_results) **init_result** Methods of 'iobase.Sink' should agree on the 'init_result' type that will be returned when initializing the sink. This type can be a client-defined object or an existing type. The returned type must be picklable using Dataflow coder ``coders.PickleCoder``. Returning an init_result is optional. **bundle_id** In order to ensure fault-tolerance, a bundle may be executed multiple times (e.g., in the event of failure/retry or for redundancy). However, exactly one of these executions will have its result passed to the ``iobase.Sink.finalize_write()`` method. Each call to ``iobase.Sink.open_writer()`` is passed a unique bundle id when it is called by the ``WriteImpl`` transform, so even redundant or retried bundles will have a unique way of identifying their output. The bundle id should be used to guarantee that a bundle's output is unique. This uniqueness guarantee is important; if a bundle is to be output to a file, for example, the name of the file must be unique to avoid conflicts with other writers. The bundle id should be encoded in the writer result returned by the writer and subsequently used by the ``finalize_write()`` method to identify the results of successful writes. For example, consider the scenario where a Writer writes files containing serialized records and the ``finalize_write()`` is to merge or rename these output files. In this case, a writer may use its unique id to name its output file (to avoid conflicts) and return the name of the file it wrote as its writer result. The ``finalize_write()`` will then receive an ``Iterable`` of output file names that it can then merge or rename using some bundle naming scheme. **write_result** ``iobase.Writer.close()`` and ``finalize_write()`` implementations must agree on type of the ``write_result`` object returned when invoking ``iobase.Writer.close()``. This type can be a client-defined object or an existing type. The returned type must be picklable using Dataflow coder ``coders.PickleCoder``. Returning a ``write_result`` when ``iobase.Writer.close()`` is invoked is optional but if unique ``write_result`` objects are not returned, sink should, guarantee idempotency when same bundle is written multiple times due to failure/retry or redundancy. **More information** For more information on creating new sinks please refer to the official documentation at ``https://cloud.google.com/dataflow/model/custom-io#creating-sinks``. """ def initialize_write(self): """Initializes the sink before writing begins. Invoked before any data is written to the sink. Please see documentation in ``iobase.Sink`` for an example. Returns: An object that contains any sink specific state generated by initialization. This object will be passed to open_writer() and finalize_write() methods. """ raise NotImplementedError def open_writer(self, init_result, uid): """Opens a writer for writing a bundle of elements to the sink. Args: init_result: the result of initialize_write() invocation. uid: a unique identifier generated by the system. Returns: an ``iobase.Writer`` that can be used to write a bundle of records to the current sink. """ raise NotImplementedError def finalize_write(self, init_result, writer_results): """Finalizes the sink after all data is written to it. Given the result of initialization and an iterable of results from bundle writes, performs finalization after writing and closes the sink. Called after all bundle writes are complete. The bundle write results that are passed to finalize are those returned by bundles that completed successfully. Although bundles may have been run multiple times (for fault-tolerance), only one writer result will be passed to finalize for each bundle. An implementation of finalize should perform clean up of any failed and successfully retried bundles. Note that these failed bundles will not have their writer result passed to finalize, so finalize should be capable of locating any temporary/partial output written by failed bundles. If all retries of a bundle fails, the whole pipeline will fail *without* finalize_write() being invoked. A best practice is to make finalize atomic. If this is impossible given the semantics of the sink, finalize should be idempotent, as it may be called multiple times in the case of failure/retry or for redundancy. Note that the iteration order of the writer results is not guaranteed to be consistent if finalize is called multiple times. Args: init_result: the result of ``initialize_write()`` invocation. writer_results: an iterable containing results of ``Writer.close()`` invocations. This will only contain results of successful writes, and will only contain the result of a single successful write for a given bundle. """ raise NotImplementedError class Writer(object): """Writes a bundle of elements from a ``PCollection`` to a sink. A Writer ``iobase.Writer.write()`` writes and elements to the sink while ``iobase.Writer.close()`` is called after all elements in the bundle have been written. See ``iobase.Sink`` for more detailed documentation about the process of writing to a sink. """ def write(self, value): """Writes a value to the sink using the current writer.""" raise NotImplementedError def close(self): """Closes the current writer. Please see documentation in ``iobase.Sink`` for an example. Returns: An object representing the writes that were performed by the current writer. """ raise NotImplementedError class _NativeWrite(ptransform.PTransform): """A PTransform for writing to a Dataflow native sink. These are sinks that are implemented natively by the Dataflow service and hence should not be updated by users. These sinks are processed using a Dataflow native write transform. Applying this transform results in a ``pvalue.PDone``. """ def __init__(self, *args, **kwargs): """Initializes a Write transform. Args: *args: A tuple of position arguments. **kwargs: A dictionary of keyword arguments. The *args, **kwargs are expected to be (label, sink) or (sink). """ label, sink = self.parse_label_and_arg(args, kwargs, 'sink') super(_NativeWrite, self).__init__(label) self.sink = sink def apply(self, pcoll): self._check_pcollection(pcoll) return pvalue.PDone(pcoll.pipeline) class Read(ptransform.PTransform): """A transform that reads a PCollection.""" def __init__(self, *args, **kwargs): """Initializes a Read transform. Args: *args: A tuple of position arguments. **kwargs: A dictionary of keyword arguments. The *args, **kwargs are expected to be (label, source) or (source). """ label, source = self.parse_label_and_arg(args, kwargs, 'source') super(Read, self).__init__(label) self.source = source def apply(self, pbegin): assert isinstance(pbegin, pvalue.PBegin) self.pipeline = pbegin.pipeline return pvalue.PCollection(self.pipeline) def get_windowing(self, unused_inputs): return core.Windowing(window.GlobalWindows()) class Write(ptransform.PTransform): """A ``PTransform`` that writes to a sink. A sink should inherit ``iobase.Sink``. Such implementations are handled using a composite transform that consists of three ``ParDo``s - (1) a ``ParDo`` performing a global initialization (2) a ``ParDo`` performing a parallel write and (3) a ``ParDo`` performing a global finalization. In the case of an empty ``PCollection``, only the global initialization and finalization will be performed. Currently only batch workflows support custom sinks. Example usage:: pcollection | df.io.Write(MySink()) This returns a ``pvalue.PValue`` object that represents the end of the Pipeline. The sink argument may also be a full PTransform, in which case it will be applied directly. This allows composite sink-like transforms (e.g. a sink with some pre-processing DoFns) to be used the same as all other sinks. This transform also supports sinks that inherit ``iobase.NativeSink``. These are sinks that are implemented natively by the Dataflow service and hence should not be updated by users. These sinks are processed using a Dataflow native write transform. """ def __init__(self, *args, **kwargs): """Initializes a Write transform. Args: *args: A tuple of position arguments. **kwargs: A dictionary of keyword arguments. The *args, **kwargs are expected to be (label, sink) or (sink). """ label, sink = self.parse_label_and_arg(args, kwargs, 'sink') super(Write, self).__init__(label) self.sink = sink def apply(self, pcoll): from google.cloud.dataflow.io import iobase if isinstance(self.sink, iobase.NativeSink): # A native sink return pcoll | _NativeWrite('native_write', self.sink) elif isinstance(self.sink, iobase.Sink): # A custom sink return pcoll | WriteImpl(self.sink) elif isinstance(self.sink, ptransform.PTransform): # This allows "composite" sinks to be used like non-composite ones. return pcoll | self.sink else: raise ValueError('A sink must inherit iobase.Sink, iobase.NativeSink, ' 'or be a PTransform. Received : %r', self.sink) class WriteImpl(ptransform.PTransform): """Implements the writing of custom sinks.""" def __init__(self, sink): super(WriteImpl, self).__init__() self.sink = sink def apply(self, pcoll): sink_coll = pcoll.pipeline | core.Create('create_sink_collection', [self.sink]) init_result_coll = sink_coll | core.Map( 'initialize_write', lambda sink: sink.initialize_write()) write_result_coll = pcoll | core.ParDo( 'write_bundles', _WriteBundleDoFn(), AsSingleton(sink_coll), AsSingleton(init_result_coll)) return sink_coll | core.FlatMap( 'finalize_write', lambda sink, init_result, write_results: (window.TimestampedValue(v, window.MAX_TIMESTAMP) for v in sink.finalize_write(init_result, write_results) or ()), AsSingleton(init_result_coll), AsIter(write_result_coll)) class _WriteBundleDoFn(core.DoFn): """A DoFn for writing elements to an iobase.Writer. Opens a writer at the first element and closes the writer at finish_bundle(). """ def __init__(self): self.writer = None def process(self, context, sink, init_result): if self.writer is None: self.writer = sink.open_writer(init_result, str(uuid.uuid4())) self.writer.write(context.element) def finish_bundle(self, context, *args, **kwargs): if self.writer is not None: yield window.TimestampedValue(self.writer.close(), window.MAX_TIMESTAMP) DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/pubsub.py000066400000000000000000000050561272340251000234710ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Google Cloud PubSub sources and sinks. Cloud Pub/Sub sources and sinks are currently supported only in streaming pipelines, during remote execution. """ from __future__ import absolute_import from google.cloud.dataflow import coders from google.cloud.dataflow.io import iobase class PubSubSource(iobase.NativeSource): """Source for reading from a given Cloud Pub/Sub topic. Attributes: topic: Cloud Pub/Sub topic in the form "/topics//". subscription: Optional existing Cloud Pub/Sub subscription to use in the form "projects//subscriptions/". id_label: The attribute on incoming Pub/Sub messages to use as a unique record identifier. When specified, the value of this attribute (which can be any string that uniquely identifies the record) will be used for deduplication of messages. If not provided, Dataflow cannot guarantee that no duplicate data will be delivered on the Pub/Sub stream. In this case, deduplication of the stream will be strictly best effort. coder: The Coder to use for decoding incoming Pub/Sub messages. """ def __init__(self, topic, subscription=None, id_label=None, coder=coders.StrUtf8Coder()): self.topic = topic self.subscription = subscription self.id_label = id_label self.coder = coder @property def format(self): """Source format name required for remote execution.""" return 'pubsub' def reader(self): raise NotImplementedError( 'PubSubSource is not supported in local execution.') class PubSubSink(iobase.NativeSink): """Sink for writing to a given Cloud Pub/Sub topic.""" def __init__(self, topic, coder=coders.StrUtf8Coder()): self.topic = topic self.coder = coder @property def format(self): """Sink format name required for remote execution.""" return 'pubsub' def writer(self): raise NotImplementedError( 'PubSubSink is not supported in local execution.') DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/range_trackers.py000066400000000000000000000255121272340251000251620ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """iobase.RangeTracker implementations provided with Dataflow SDK. """ import logging import math import threading from google.cloud.dataflow.io import iobase class OffsetRangeTracker(iobase.RangeTracker): """A 'RangeTracker' for non-negative positions of type 'long'.""" # Offset corresponding to infinity. This can only be used as the upper-bound # of a range, and indicates reading all of the records until the end without # specifying exactly what the end is. # Infinite ranges cannot be split because it is impossible to estimate # progress within them. OFFSET_INFINITY = float('inf') def __init__(self, start, end): super(OffsetRangeTracker, self).__init__() self._start_offset = start self._stop_offset = end self._last_record_start = -1 self._offset_of_last_split_point = -1 self._lock = threading.Lock() def start_position(self): return self._start_offset def stop_position(self): return self._stop_offset @property def last_record_start(self): return self._last_record_start def _validate_record_start(self, record_start, split_point): # This function must only be called under the lock self.lock. if not self._lock.locked(): raise ValueError( 'This function must only be called under the lock self.lock.') if record_start < self._last_record_start: raise ValueError( 'Trying to return a record [starting at %d] which is before the ' 'last-returned record [starting at %d]' % (record_start, self._last_record_start)) if split_point: if (self._offset_of_last_split_point != -1 and record_start == self._offset_of_last_split_point): raise ValueError( 'Record at a split point has same offset as the previous split ' 'point: %d' % record_start) elif self._last_record_start == -1: raise ValueError( 'The first record [starting at %d] must be at a split point' % record_start) if (split_point and self._offset_of_last_split_point is not -1 and record_start is self._offset_of_last_split_point): raise ValueError( 'Record at a split point has same offset as the previous split ' 'point: %d' % record_start) if not split_point and self._last_record_start == -1: raise ValueError( 'The first record [starting at %d] must be at a split point' % record_start) def try_claim(self, record_start): with self._lock: self._validate_record_start(record_start, True) if record_start >= self.stop_position(): return False self._offset_of_last_split_point = record_start self._last_record_start = record_start return True def set_current_position(self, record_start): with self._lock: self._validate_record_start(record_start, False) self._last_record_start = record_start def try_split(self, split_offset): with self._lock: if self._stop_offset == OffsetRangeTracker.OFFSET_INFINITY: logging.debug('refusing to split %r at %d: stop position unspecified', self, split_offset) return if self._last_record_start == -1: logging.debug('Refusing to split %r at %d: unstarted', self, split_offset) return if split_offset <= self._last_record_start: logging.debug( 'Refusing to split %r at %d: already past proposed stop offset', self, split_offset) return if (split_offset < self.start_position() or split_offset >= self.stop_position()): logging.debug( 'Refusing to split %r at %d: proposed split position out of range', self, split_offset) return logging.debug('Agreeing to split %r at %d', self, split_offset) self._stop_offset = split_offset split_fraction = (float(split_offset - self._start_offset) / ( self._stop_offset - self._start_offset)) return self._stop_offset, split_fraction def fraction_consumed(self): with self._lock: fraction = ((1.0 * (self._last_record_start - self.start_position()) / (self.stop_position() - self.start_position())) if self.stop_position() != self.start_position() else 0.0) # self.last_record_start may become larger than self.end_offset when # reading the records since any record that starts before the first 'split # point' at or after the defined 'stop offset' is considered to be within # the range of the OffsetRangeTracker. Hence fraction could be > 1. # self.last_record_start is initialized to -1, hence fraction may be < 0. # Bounding the to range [0, 1]. return max(0.0, min(1.0, fraction)) def position_at_fraction(self, fraction): if self.stop_position() == OffsetRangeTracker.OFFSET_INFINITY: raise Exception( 'get_position_for_fraction_consumed is not applicable for an ' 'unbounded range') return (math.ceil(self.start_position() + fraction * ( self.stop_position() - self.start_position()))) class GroupedShuffleRangeTracker(iobase.RangeTracker): """A 'RangeTracker' for positions used by'GroupedShuffleReader'. These positions roughly correspond to hashes of keys. In case of hash collisions, multiple groups can have the same position. In that case, the first group at a particular position is considered a split point (because it is the first to be returned when reading a position range starting at this position), others are not. """ def __init__(self, decoded_start_pos, decoded_stop_pos): super(GroupedShuffleRangeTracker, self).__init__() self._decoded_start_pos = decoded_start_pos self._decoded_stop_pos = decoded_stop_pos self._decoded_last_group_start = None self._last_group_was_at_a_split_point = False self._lock = threading.Lock() def start_position(self): return self._decoded_start_pos def stop_position(self): return self._decoded_stop_pos def last_group_start(self): return self._decoded_last_group_start def _validate_decoded_group_start(self, decoded_group_start, split_point): if self.start_position() and decoded_group_start < self.start_position(): raise ValueError('Trying to return record at %r which is before the' ' starting position at %r' % (decoded_group_start, self.start_position())) if (self.last_group_start() and decoded_group_start < self.last_group_start()): raise ValueError('Trying to return group at %r which is before the' ' last-returned group at %r' % (decoded_group_start, self.last_group_start())) if (split_point and self.last_group_start() and self.last_group_start() == decoded_group_start): raise ValueError('Trying to return a group at a split point with ' 'same position as the previous group: both at %r, ' 'last group was %sat a split point.' % (decoded_group_start, ('' if self._last_group_was_at_a_split_point else 'not '))) if not split_point: if self.last_group_start() is None: raise ValueError('The first group [at %r] must be at a split point' % decoded_group_start) if self.last_group_start() != decoded_group_start: # This case is not a violation of general RangeTracker semantics, but it # is contrary to how GroupingShuffleReader in particular works. Hitting # it would mean it's behaving unexpectedly. raise ValueError('Trying to return a group not at a split point, but ' 'with a different position than the previous group: ' 'last group was %r at %r, current at a %s split' ' point.' % (self.last_group_start() , decoded_group_start , ('' if self._last_group_was_at_a_split_point else 'non-'))) def try_claim(self, decoded_group_start): with self._lock: self._validate_decoded_group_start(decoded_group_start, True) if (self.stop_position() and decoded_group_start >= self.stop_position()): return False self._decoded_last_group_start = decoded_group_start self._last_group_was_at_a_split_point = True return True def set_current_position(self, decoded_group_start): with self._lock: self._validate_decoded_group_start(decoded_group_start, False) self._decoded_last_group_start = decoded_group_start self._last_group_was_at_a_split_point = False def try_split(self, decoded_split_position): with self._lock: if self.last_group_start() is None: logging.info('Refusing to split %r at %r: unstarted' , self, decoded_split_position) return if decoded_split_position <= self.last_group_start(): logging.info('Refusing to split %r at %r: already past proposed split ' 'position' , self, decoded_split_position) return if ((self.stop_position() and decoded_split_position >= self.stop_position()) or (self.start_position() and decoded_split_position <= self.start_position())): logging.error('Refusing to split %r at %r: proposed split position out ' 'of range', self, decoded_split_position) return logging.debug('Agreeing to split %r at %r' , self, decoded_split_position) self._decoded_stop_pos = decoded_split_position # Since GroupedShuffleRangeTracker cannot determine relative sizes of the # two splits, returning 0.5 as the fraction below so that the framework # assumes the splits to be of the same size. return self._decoded_stop_pos, 0.5 def fraction_consumed(self): # GroupingShuffle sources have special support on the service and the # service will estimate progress from positions for us. raise RuntimeError('GroupedShuffleRangeTracker does not measure fraction' ' consumed due to positions being opaque strings' ' that are interpretted by the service') DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/range_trackers_test.py000066400000000000000000000320441272340251000262170ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the range_trackers module.""" import array import copy import logging import unittest from google.cloud.dataflow.io import range_trackers class OffsetRangeTrackerTest(unittest.TestCase): def test_try_return_record_simple_sparse(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertTrue(tracker.try_claim(110)) self.assertTrue(tracker.try_claim(140)) self.assertTrue(tracker.try_claim(183)) self.assertFalse(tracker.try_claim(210)) def test_try_return_record_simple_dense(self): tracker = range_trackers.OffsetRangeTracker(3, 6) self.assertTrue(tracker.try_claim(3)) self.assertTrue(tracker.try_claim(4)) self.assertTrue(tracker.try_claim(5)) self.assertFalse(tracker.try_claim(6)) def test_try_return_record_continuous_until_split_point(self): tracker = range_trackers.OffsetRangeTracker(9, 18) # Return records with gaps of 2; every 3rd record is a split point. self.assertTrue(tracker.try_claim(10)) tracker.set_current_position(12) tracker.set_current_position(14) self.assertTrue(tracker.try_claim(16)) # Out of range, but not a split point... tracker.set_current_position(18) tracker.set_current_position(20) # Out of range AND a split point. self.assertFalse(tracker.try_claim(22)) def test_split_at_offset_fails_if_unstarted(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertFalse(tracker.try_split(150)) def test_split_at_offset(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertTrue(tracker.try_claim(110)) # Example positions we shouldn't split at, when last record starts at 110: self.assertFalse(tracker.try_split(109)) self.assertFalse(tracker.try_split(110)) self.assertFalse(tracker.try_split(200)) self.assertFalse(tracker.try_split(210)) # Example positions we *should* split at: self.assertTrue(copy.copy(tracker).try_split(111)) self.assertTrue(copy.copy(tracker).try_split(129)) self.assertTrue(copy.copy(tracker).try_split(130)) self.assertTrue(copy.copy(tracker).try_split(131)) self.assertTrue(copy.copy(tracker).try_split(150)) self.assertTrue(copy.copy(tracker).try_split(199)) # If we split at 170 and then at 150: self.assertTrue(tracker.try_split(170)) self.assertTrue(tracker.try_split(150)) # Should be able to return a record starting before the new stop offset. # Returning records starting at the same offset is ok. self.assertTrue(copy.copy(tracker).try_claim(135)) self.assertTrue(copy.copy(tracker).try_claim(135)) # Should be able to return a record starting right before the new stop # offset. self.assertTrue(copy.copy(tracker).try_claim(149)) # Should not be able to return a record starting at or after the new stop # offset. self.assertFalse(tracker.try_claim(150)) self.assertFalse(tracker.try_claim(151)) # Should accept non-splitpoint records starting after stop offset. tracker.set_current_position(135) tracker.set_current_position(152) tracker.set_current_position(160) tracker.set_current_position(171) def test_get_position_for_fraction_dense(self): # Represents positions 3, 4, 5. tracker = range_trackers.OffsetRangeTracker(3, 6) # [3, 3) represents 0.0 of [3, 6) self.assertEqual(3, tracker.position_at_fraction(0.0)) # [3, 4) represents up to 1/3 of [3, 6) self.assertEqual(4, tracker.position_at_fraction(1.0 / 6)) self.assertEqual(4, tracker.position_at_fraction(0.333)) # [3, 5) represents up to 2/3 of [3, 6) self.assertEqual(5, tracker.position_at_fraction(0.334)) self.assertEqual(5, tracker.position_at_fraction(0.666)) # Any fraction consumed over 2/3 means the whole [3, 6) has been consumed. self.assertEqual(6, tracker.position_at_fraction(0.667)) def test_get_fraction_consumed_dense(self): tracker = range_trackers.OffsetRangeTracker(3, 6) self.assertEqual(0, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(3)) self.assertEqual(0.0, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(4)) self.assertEqual(1.0 / 3, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(5)) self.assertEqual(2.0 / 3, tracker.fraction_consumed()) tracker.set_current_position(6) self.assertEqual(1.0, tracker.fraction_consumed()) tracker.set_current_position(7) self.assertFalse(tracker.try_claim(7)) def test_get_fraction_consumed_sparse(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertEqual(0, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(110)) # Consumed positions through 110 = total 10 positions of 100 done. self.assertEqual(0.10, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(150)) self.assertEqual(0.50, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(195)) self.assertEqual(0.95, tracker.fraction_consumed()) def test_everything_with_unbounded_range(self): tracker = range_trackers.OffsetRangeTracker( 100, range_trackers.OffsetRangeTracker.OFFSET_INFINITY) self.assertTrue(tracker.try_claim(150)) self.assertTrue(tracker.try_claim(250)) # get_position_for_fraction_consumed should fail for an unbounded range with self.assertRaises(Exception): tracker.position_at_fraction(0.5) def test_try_return_first_record_not_split_point(self): with self.assertRaises(Exception): range_trackers.OffsetRangeTracker(100, 200).set_current_position(120) def test_try_return_record_non_monotonic(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertTrue(tracker.try_claim(120)) with self.assertRaises(Exception): tracker.try_claim(110) class GroupedShuffleRangeTrackerTest(unittest.TestCase): def bytes_to_position(self, bytes_array): return array.array('B', bytes_array).tostring() def test_try_return_record_in_infinite_range(self): tracker = range_trackers.GroupedShuffleRangeTracker('', '') self.assertTrue(tracker.try_claim( self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([1, 2, 5]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([3, 6, 8, 10]))) def test_try_return_record_finite_range(self): tracker = range_trackers.GroupedShuffleRangeTracker( self.bytes_to_position([1, 0, 0]), self.bytes_to_position([5, 0, 0])) self.assertTrue(tracker.try_claim( self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([1, 2, 5]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([3, 6, 8, 10]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([4, 255, 255, 255]))) # Should fail for positions that are lexicographically equal to or larger # than the defined stop position. self.assertFalse(copy.copy(tracker).try_claim( self.bytes_to_position([5, 0, 0]))) self.assertFalse(copy.copy(tracker).try_claim( self.bytes_to_position([5, 0, 1]))) self.assertFalse(copy.copy(tracker).try_claim( self.bytes_to_position([6, 0, 0]))) def test_try_return_record_with_non_split_point(self): tracker = range_trackers.GroupedShuffleRangeTracker( self.bytes_to_position([1, 0, 0]), self.bytes_to_position([5, 0, 0])) self.assertTrue(tracker.try_claim( self.bytes_to_position([1, 2, 3]))) tracker.set_current_position(self.bytes_to_position([1, 2, 3])) tracker.set_current_position(self.bytes_to_position([1, 2, 3])) self.assertTrue(tracker.try_claim( self.bytes_to_position([1, 2, 5]))) tracker.set_current_position(self.bytes_to_position([1, 2, 5])) self.assertTrue(tracker.try_claim( self.bytes_to_position([3, 6, 8, 10]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([4, 255, 255, 255]))) def test_first_record_non_split_point(self): tracker = range_trackers.GroupedShuffleRangeTracker( self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0])) with self.assertRaises(ValueError): tracker.set_current_position(self.bytes_to_position([3, 4, 5])) def test_non_split_point_record_with_different_position(self): tracker = range_trackers.GroupedShuffleRangeTracker( self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0])) self.assertTrue(tracker.try_claim(self.bytes_to_position([3, 4, 5]))) with self.assertRaises(ValueError): tracker.set_current_position(self.bytes_to_position([3, 4, 6])) def test_try_return_record_before_start(self): tracker = range_trackers.GroupedShuffleRangeTracker( self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0])) with self.assertRaises(ValueError): tracker.try_claim(self.bytes_to_position([1, 2, 3])) def test_try_return_non_monotonic(self): tracker = range_trackers.GroupedShuffleRangeTracker( self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0])) self.assertTrue(tracker.try_claim(self.bytes_to_position([3, 4, 5]))) self.assertTrue(tracker.try_claim(self.bytes_to_position([3, 4, 6]))) with self.assertRaises(ValueError): tracker.try_claim(self.bytes_to_position([3, 2, 1])) def test_try_return_identical_positions(self): tracker = range_trackers.GroupedShuffleRangeTracker( self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0])) self.assertTrue(tracker.try_claim( self.bytes_to_position([3, 4, 5]))) with self.assertRaises(ValueError): tracker.try_claim(self.bytes_to_position([3, 4, 5])) def test_try_split_at_position_infinite_range(self): tracker = range_trackers.GroupedShuffleRangeTracker('', '') # Should fail before first record is returned. self.assertFalse(tracker.try_split( self.bytes_to_position([3, 4, 5, 6]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([1, 2, 3]))) # Should now succeed. self.assertIsNotNone(tracker.try_split( self.bytes_to_position([3, 4, 5, 6]))) # Should not split at same or larger position. self.assertIsNone(tracker.try_split( self.bytes_to_position([3, 4, 5, 6]))) self.assertIsNone(tracker.try_split( self.bytes_to_position([3, 4, 5, 6, 7]))) self.assertIsNone(tracker.try_split( self.bytes_to_position([4, 5, 6, 7]))) # Should split at smaller position. self.assertIsNotNone(tracker.try_split( self.bytes_to_position([3, 2, 1]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([2, 3, 4]))) # Should not split at a position we're already past. self.assertIsNone(tracker.try_split( self.bytes_to_position([2, 3, 4]))) self.assertIsNone(tracker.try_split( self.bytes_to_position([2, 3, 3]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([3, 2, 0]))) self.assertFalse(tracker.try_claim( self.bytes_to_position([3, 2, 1]))) def test_try_test_split_at_position_finite_range(self): tracker = range_trackers.GroupedShuffleRangeTracker( self.bytes_to_position([0, 0, 0]), self.bytes_to_position([10, 20, 30])) # Should fail before first record is returned. self.assertFalse(tracker.try_split( self.bytes_to_position([0, 0, 0]))) self.assertFalse(tracker.try_split( self.bytes_to_position([3, 4, 5, 6]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([1, 2, 3]))) # Should now succeed. self.assertTrue(tracker.try_split( self.bytes_to_position([3, 4, 5, 6]))) # Should not split at same or larger position. self.assertFalse(tracker.try_split( self.bytes_to_position([3, 4, 5, 6]))) self.assertFalse(tracker.try_split( self.bytes_to_position([3, 4, 5, 6, 7]))) self.assertFalse(tracker.try_split( self.bytes_to_position([4, 5, 6, 7]))) # Should split at smaller position. self.assertTrue(tracker.try_split( self.bytes_to_position([3, 2, 1]))) # But not at a position at or before last returned record. self.assertFalse(tracker.try_split( self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([2, 3, 4]))) self.assertTrue(tracker.try_claim( self.bytes_to_position([3, 2, 0]))) self.assertFalse(tracker.try_claim( self.bytes_to_position([3, 2, 1]))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/io/sources_test.py000066400000000000000000000037051272340251000247120ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the sources framework.""" import logging import tempfile import unittest import google.cloud.dataflow as df from google.cloud.dataflow.io import iobase from google.cloud.dataflow.transforms.util import assert_that from google.cloud.dataflow.transforms.util import equal_to class LineSource(iobase.BoundedSource): """A simple source that reads lines from a given file.""" def __init__(self, file_name): self._file_name = file_name def read(self, _): with open(self._file_name) as f: for line in f: yield line.rstrip('\n') class SourcesTest(unittest.TestCase): def _create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def test_read_from_source(self): file_name = self._create_temp_file('aaaa\nbbbb\ncccc\ndddd') source = LineSource(file_name) result = [line for line in source.read(None)] self.assertItemsEqual(['aaaa', 'bbbb', 'cccc', 'dddd'], result) def test_run_direct(self): file_name = self._create_temp_file('aaaa\nbbbb\ncccc\ndddd') pipeline = df.Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Read(LineSource(file_name)) assert_that(pcoll, equal_to(['aaaa', 'bbbb', 'cccc', 'dddd'])) pipeline.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/pipeline.py000066400000000000000000000404661272340251000233730ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Pipeline, the top-level Dataflow object. A pipeline holds a DAG of data transforms. Conceptually the nodes of the DAG are transforms (PTransform objects) and the edges are values (mostly PCollection objects). The transforms take as inputs one or more PValues and output one or more PValues. The pipeline offers functionality to traverse the graph. The actual operation to be executed for each node visited is specified through a runner object. Typical usage: # Create a pipeline object using a local runner for execution. pipeline = Pipeline(runner=DirectPipelineRunner()) # Add to the pipeline a "Create" transform. When executed this # transform will produce a PCollection object with the specified values. pcoll = pipeline.create('label', [1, 2, 3]) # run() will execute the DAG stored in the pipeline. The execution of the # nodes visited is done using the specified local runner. pipeline.run() """ from __future__ import absolute_import import collections import logging import os import shutil import tempfile from google.cloud.dataflow import pvalue from google.cloud.dataflow import typehints from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.runners import create_runner from google.cloud.dataflow.runners import PipelineRunner from google.cloud.dataflow.transforms import format_full_label from google.cloud.dataflow.transforms import ptransform from google.cloud.dataflow.typehints import TypeCheckError from google.cloud.dataflow.utils.options import PipelineOptions from google.cloud.dataflow.utils.options import SetupOptions from google.cloud.dataflow.utils.options import StandardOptions from google.cloud.dataflow.utils.options import TypeOptions from google.cloud.dataflow.utils.pipeline_options_validator import PipelineOptionsValidator class Pipeline(object): """A pipeline object that manages a DAG of PValues and their PTransforms. Conceptually the PValues are the DAG's nodes and the PTransforms computing the PValues are the edges. All the transforms applied to the pipeline must have distinct full labels. If same transform instance needs to be applied then a clone should be created with a new label (e.g., transform.clone('new label')). """ def __init__(self, runner=None, options=None, argv=None): """Initialize a pipeline object. Args: runner: An object of type 'PipelineRunner' that will be used to execute the pipeline. For registered runners, the runner name can be specified, otherwise a runner object must be supplied. options: A configured 'PipelineOptions' object containing arguments that should be used for running the Dataflow job. argv: a list of arguments (such as sys.argv) to be used for building a 'PipelineOptions' object. This will only be used if argument 'options' is None. Raises: ValueError: if either the runner or options argument is not of the expected type. """ if options is not None: if isinstance(options, PipelineOptions): self.options = options else: raise ValueError( 'Parameter options, if specified, must be of type PipelineOptions. ' 'Received : %r', options) elif argv is not None: if isinstance(argv, list): self.options = PipelineOptions(argv) else: raise ValueError( 'Parameter argv, if specified, must be a list. Received : %r', argv) else: self.options = None if runner is None and self.options is not None: runner = self.options.view_as(StandardOptions).runner if runner is None: runner = StandardOptions.DEFAULT_RUNNER logging.info(('Missing pipeline option (runner). Executing pipeline ' 'using the default runner: %s.'), runner) if isinstance(runner, str): runner = create_runner(runner) elif not isinstance(runner, PipelineRunner): raise TypeError('Runner must be a PipelineRunner object or the ' 'name of a registered runner.') # Validate pipeline options if self.options is not None: errors = PipelineOptionsValidator(self.options, runner).validate() if errors: raise ValueError( 'Pipeline has validations errors: \n' + '\n'.join(errors)) # Default runner to be used. self.runner = runner # Stack of transforms generated by nested apply() calls. The stack will # contain a root node as an enclosing (parent) node for top transforms. self.transforms_stack = [AppliedPTransform(None, None, '', None)] # Set of transform labels (full labels) applied to the pipeline. # If a transform is applied and the full label is already in the set # then the transform will have to be cloned with a new label. self.applied_labels = set() # Store cache of views created from PCollections. For reference, see # pvalue._cache_view(). self._view_cache = {} def _current_transform(self): """Returns the transform currently on the top of the stack.""" return self.transforms_stack[-1] def _root_transform(self): """Returns the root transform of the transform stack.""" return self.transforms_stack[0] def run(self): """Runs the pipeline. Returns whatever our runner returns after running.""" if not self.options or self.options.view_as(SetupOptions).save_main_session: # If this option is chosen, verify we can pickle the main session early. tmpdir = tempfile.mkdtemp() try: pickler.dump_session(os.path.join(tmpdir, 'main_session.pickle')) finally: shutil.rmtree(tmpdir) return self.runner.run(self) def visit(self, visitor): """Visits depth-first every node of a pipeline's DAG. Args: visitor: PipelineVisitor object whose callbacks will be called for each node visited. See PipelineVisitor comments. Raises: TypeError: if node is specified and is not a PValue. pipeline.PipelineError: if node is specified and does not belong to this pipeline instance. """ visited = set() self._root_transform().visit(visitor, self, visited) def apply(self, transform, pvalueish=None): """Applies a custom transform using the pvalueish specified. Args: transform: the PTranform (or callable) to apply. pvalueish: the input for the PTransform (typically a PCollection). Raises: TypeError: if the transform object extracted from the argument list is not a callable type or a descendant from PTransform. RuntimeError: if the transform object was already applied to this pipeline and needs to be cloned in order to apply again. """ if not isinstance(transform, ptransform.PTransform): transform = _CallableWrapperPTransform(transform) full_label = format_full_label(self._current_transform(), transform) if full_label in self.applied_labels: raise RuntimeError( 'Transform "%s" does not have a stable unique label. ' 'This will prevent updating of pipelines. ' 'To clone a transform with a new label use: ' 'transform.clone("NEW LABEL").' % full_label) self.applied_labels.add(full_label) pvalueish, inputs = transform._extract_input_pvalues(pvalueish) try: inputs = tuple(inputs) for leaf_input in inputs: if not isinstance(leaf_input, pvalue.PValue): raise TypeError except TypeError: raise NotImplementedError( 'Unable to extract PValue inputs from %s; either %s does not accept ' 'inputs of this format, or it does not properly override ' '_extract_input_values' % (pvalueish, transform)) current = AppliedPTransform( self._current_transform(), transform, full_label, inputs) self._current_transform().add_part(current) self.transforms_stack.append(current) if self.options is not None: type_options = self.options.view_as(TypeOptions) else: type_options = None if type_options is not None and type_options.pipeline_type_check: transform.type_check_inputs(pvalueish) pvalueish_result = self.runner.apply(transform, pvalueish) if type_options is not None and type_options.pipeline_type_check: transform.type_check_outputs(pvalueish_result) for result in ptransform.GetPValues().visit(pvalueish_result): assert isinstance(result, (pvalue.PValue, pvalue.DoOutputsTuple)) # Make sure we set the producer only for a leaf node in the transform DAG. # This way we preserve the last transform of a composite transform as # being the real producer of the result. if result.producer is None: result.producer = current # TODO(robertwb): Multi-input, multi-output inference. # TODO(robertwb): Ideally we'd do intersection here. if (type_options is not None and type_options.pipeline_type_check and isinstance(result, (pvalue.PCollection, pvalue.PCollectionView)) and not result.element_type): input_element_type = ( inputs[0].element_type if len(inputs) == 1 else typehints.Any) type_hints = transform.get_type_hints() declared_output_type = type_hints.simple_output_type(transform.label) if declared_output_type: input_types = type_hints.input_types if input_types and input_types[0]: declared_input_type = input_types[0][0] result.element_type = typehints.bind_type_variables( declared_output_type, typehints.match_type_variables(declared_input_type, input_element_type)) else: result.element_type = declared_output_type else: result.element_type = transform.infer_output_type(input_element_type) assert isinstance(result.producer.inputs, tuple) current.add_output(result) if (type_options is not None and type_options.type_check_strictness == 'ALL_REQUIRED' and transform.get_type_hints().output_types is None): ptransform_name = '%s(%s)' % (transform.__class__.__name__, full_label) raise TypeCheckError('Pipeline type checking is enabled, however no ' 'output type-hint was found for the ' 'PTransform %s' % ptransform_name) current.update_input_refcounts() self.transforms_stack.pop() return pvalueish_result class _CallableWrapperPTransform(ptransform.PTransform): def __init__(self, callee): assert callable(callee) super(_CallableWrapperPTransform, self).__init__( label=getattr(callee, '__name__', 'Callable')) self._callee = callee def apply(self, *args, **kwargs): return self._callee(*args, **kwargs) class PipelineVisitor(object): """Visitor pattern class used to traverse a DAG of transforms. This is an internal class used for bookkeeping by a Pipeline. """ def visit_value(self, value, producer_node): """Callback for visiting a PValue in the pipeline DAG. Args: value: PValue visited (typically a PCollection instance). producer_node: AppliedPTransform object whose transform produced the pvalue. """ pass def visit_transform(self, transform_node): """Callback for visiting a transform node in the pipeline DAG.""" pass def enter_composite_transform(self, transform_node): """Callback for entering traversal of a composite transform node.""" pass def leave_composite_transform(self, transform_node): """Callback for leaving traversal of a composite transform node.""" pass class AppliedPTransform(object): """A transform node representing an instance of applying a PTransform. This is an internal class used for bookkeeping by a Pipeline. """ def __init__(self, parent, transform, full_label, inputs): self.parent = parent self.transform = transform # Note that we want the PipelineVisitor classes to use the full_label, # inputs, side_inputs, and outputs fields from this instance instead of the # ones of the PTransform instance associated with it. Doing this permits # reusing PTransform instances in different contexts (apply() calls) without # any interference. This is particularly useful for composite transforms. self.full_label = full_label self.inputs = inputs or () self.side_inputs = () if transform is None else tuple(transform.side_inputs) self.outputs = {} self.parts = [] # Per tag refcount dictionary for PValues for which this node is a # root producer. self.refcounts = collections.defaultdict(int) def update_input_refcounts(self): """Increment refcounts for all transforms providing inputs.""" def real_producer(pv): real = pv.producer while real.parts: real = real.parts[-1] return real if not self.is_composite(): for main_input in self.inputs: if not isinstance(main_input, pvalue.PBegin): real_producer(main_input).refcounts[main_input.tag] += 1 for side_input in self.side_inputs: real_producer(side_input).refcounts[side_input.tag] += 1 def add_output(self, output, tag=None): assert (isinstance(output, pvalue.PValue) or isinstance(output, pvalue.DoOutputsTuple)) if tag is None: tag = len(self.outputs) assert tag not in self.outputs self.outputs[tag] = output def add_part(self, part): assert isinstance(part, AppliedPTransform) self.parts.append(part) def is_composite(self): """Returns whether this is a composite transform. A composite transform has parts (inner transforms) or isn't the producer for any of its outputs. (An example of a transform that is not a producer is one that returns its inputs instead.) """ return bool(self.parts) or all( pval.producer is not self for pval in self.outputs.values()) def visit(self, visitor, pipeline, visited): """Visits all nodes reachable from the current node.""" for pval in self.inputs: if pval not in visited and not isinstance(pval, pvalue.PBegin): assert pval.producer is not None pval.producer.visit(visitor, pipeline, visited) # The value should be visited now since we visit outputs too. assert pval in visited, pval # Visit side inputs. for pval in self.side_inputs: if isinstance(pval, pvalue.PCollectionView) and pval not in visited: assert pval.producer is not None pval.producer.visit(visitor, pipeline, visited) # The value should be visited now since we visit outputs too. assert pval in visited # TODO(silviuc): Is there a way to signal that we are visiting a side # value? The issue is that the same PValue can be reachable through # multiple paths and therefore it is not guaranteed that the value # will be visited as a side value. # Visit a composite or primitive transform. if self.is_composite(): visitor.enter_composite_transform(self) for part in self.parts: part.visit(visitor, pipeline, visited) visitor.leave_composite_transform(self) else: visitor.visit_transform(self) # Visit the outputs (one or more). It is essential to mark as visited the # tagged PCollections of the DoOutputsTuple object. A tagged PCollection is # connected directly with its producer (a multi-output ParDo), but the # output of such a transform is the containing DoOutputsTuple, not the # PCollection inside it. Without the code below a tagged PCollection will # not be marked as visited while visiting its producer. for pval in self.outputs.values(): if isinstance(pval, pvalue.DoOutputsTuple): pvals = (v for v in pval) else: pvals = (pval,) for v in pvals: if v not in visited: visited.add(v) visitor.visit_value(v, self) DataflowPythonSDK-0.2.5/google/cloud/dataflow/pipeline_test.py000066400000000000000000000262501272340251000244250ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the Pipeline class.""" import gc import logging import unittest from google.cloud.dataflow.io.iobase import NativeSource from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.pipeline import PipelineOptions from google.cloud.dataflow.pipeline import PipelineVisitor from google.cloud.dataflow.pvalue import AsIter from google.cloud.dataflow.pvalue import SideOutputValue from google.cloud.dataflow.transforms import CombinePerKey from google.cloud.dataflow.transforms import Create from google.cloud.dataflow.transforms import FlatMap from google.cloud.dataflow.transforms import Flatten from google.cloud.dataflow.transforms import Map from google.cloud.dataflow.transforms import PTransform from google.cloud.dataflow.transforms import Read from google.cloud.dataflow.transforms.util import assert_that, equal_to class FakeSource(NativeSource): """Fake source returning a fixed list of values.""" class _Reader(object): def __init__(self, vals): self._vals = vals def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): pass def __iter__(self): for v in self._vals: yield v def __init__(self, vals): self._vals = vals def reader(self): return FakeSource._Reader(self._vals) class PipelineTest(unittest.TestCase): def setUp(self): self.runner_name = 'DirectPipelineRunner' @staticmethod def custom_callable(pcoll): return pcoll | FlatMap('+1', lambda x: [x + 1]) # Some of these tests designate a runner by name, others supply a runner. # This variation is just to verify that both means of runner specification # work and is not related to other aspects of the tests. class CustomTransform(PTransform): def apply(self, pcoll): return pcoll | FlatMap('+1', lambda x: [x + 1]) class Visitor(PipelineVisitor): def __init__(self, visited): self.visited = visited self.enter_composite = [] self.leave_composite = [] def visit_value(self, value, _): self.visited.append(value) def enter_composite_transform(self, transform_node): self.enter_composite.append(transform_node) def leave_composite_transform(self, transform_node): self.leave_composite.append(transform_node) def test_create(self): pipeline = Pipeline(self.runner_name) pcoll = pipeline | Create('label1', [1, 2, 3]) assert_that(pcoll, equal_to([1, 2, 3])) # Test if initial value is an iterator object. pcoll2 = pipeline | Create('label2', iter((4, 5, 6))) pcoll3 = pcoll2 | FlatMap('do', lambda x: [x + 10]) assert_that(pcoll3, equal_to([14, 15, 16]), label='pcoll3') pipeline.run() def test_create_singleton_pcollection(self): pipeline = Pipeline(self.runner_name) pcoll = pipeline | Create('label', [[1, 2, 3]]) assert_that(pcoll, equal_to([[1, 2, 3]])) pipeline.run() def test_read(self): pipeline = Pipeline(self.runner_name) pcoll = pipeline | Read('read', FakeSource([1, 2, 3])) assert_that(pcoll, equal_to([1, 2, 3])) pipeline.run() def test_visit_entire_graph(self): pipeline = Pipeline(self.runner_name) pcoll1 = pipeline | Create('pcoll', [1, 2, 3]) pcoll2 = pcoll1 | FlatMap('do1', lambda x: [x + 1]) pcoll3 = pcoll2 | FlatMap('do2', lambda x: [x + 1]) pcoll4 = pcoll2 | FlatMap('do3', lambda x: [x + 1]) transform = PipelineTest.CustomTransform() pcoll5 = pcoll4 | transform visitor = PipelineTest.Visitor(visited=[]) pipeline.visit(visitor) self.assertEqual(set([pcoll1, pcoll2, pcoll3, pcoll4, pcoll5]), set(visitor.visited)) self.assertEqual(set(visitor.enter_composite), set(visitor.leave_composite)) self.assertEqual(2, len(visitor.enter_composite)) self.assertEqual(visitor.enter_composite[1].transform, transform) self.assertEqual(visitor.leave_composite[0].transform, transform) def test_apply_custom_transform(self): pipeline = Pipeline(self.runner_name) pcoll = pipeline | Create('pcoll', [1, 2, 3]) result = pcoll | PipelineTest.CustomTransform() assert_that(result, equal_to([2, 3, 4])) pipeline.run() def test_reuse_custom_transform_instance(self): pipeline = Pipeline(self.runner_name) pcoll1 = pipeline | Create('pcoll1', [1, 2, 3]) pcoll2 = pipeline | Create('pcoll2', [4, 5, 6]) transform = PipelineTest.CustomTransform() pcoll1 | transform with self.assertRaises(RuntimeError) as cm: pipeline.apply(transform, pcoll2) self.assertEqual( cm.exception.message, 'Transform "CustomTransform" does not have a stable unique label. ' 'This will prevent updating of pipelines. ' 'To clone a transform with a new label use: ' 'transform.clone("NEW LABEL").') def test_reuse_cloned_custom_transform_instance(self): pipeline = Pipeline(self.runner_name) pcoll1 = pipeline | Create('pcoll1', [1, 2, 3]) pcoll2 = pipeline | Create('pcoll2', [4, 5, 6]) transform = PipelineTest.CustomTransform() result1 = pcoll1 | transform result2 = pcoll2 | transform.clone('new label') assert_that(result1, equal_to([2, 3, 4]), label='r1') assert_that(result2, equal_to([5, 6, 7]), label='r2') pipeline.run() def test_apply_custom_callable(self): pipeline = Pipeline(self.runner_name) pcoll = pipeline | Create('pcoll', [1, 2, 3]) result = pipeline.apply(PipelineTest.custom_callable, pcoll) assert_that(result, equal_to([2, 3, 4])) pipeline.run() def test_transform_no_super_init(self): class AddSuffix(PTransform): def __init__(self, suffix): # No call to super(...).__init__ self.suffix = suffix def apply(self, pcoll): return pcoll | Map(lambda x: x + self.suffix) self.assertEqual( ['a-x', 'b-x', 'c-x'], sorted(['a', 'b', 'c'] | AddSuffix('-x'))) def test_cached_pvalues_are_refcounted(self): """Test that cached PValues are refcounted and deleted. The intermediary PValues computed by the workflow below contain one million elements so if the refcounting does not work the number of objects tracked by the garbage collector will increase by a few millions by the time we execute the final Map checking the objects tracked. Anything that is much larger than what we started with will fail the test. """ def check_memory(value, count_threshold): gc.collect() objects_count = len(gc.get_objects()) if objects_count > count_threshold: raise RuntimeError( 'PValues are not refcounted: %s, %s' % ( objects_count, count_threshold)) return value def create_dupes(o, _): yield o yield SideOutputValue('side', o) pipeline = Pipeline('DirectPipelineRunner') gc.collect() count_threshold = len(gc.get_objects()) + 10000 biglist = pipeline | Create('oom:create', ['x'] * 1000000) dupes = ( biglist | Map('oom:addone', lambda x: (x, 1)) | FlatMap('oom:dupes', create_dupes, AsIter(biglist)).with_outputs('side', main='main')) result = ( (dupes.side, dupes.main, dupes.side) | Flatten('oom:flatten') | CombinePerKey('oom:combine', sum) | Map('oom:check', check_memory, count_threshold)) assert_that(result, equal_to([('x', 3000000)])) pipeline.run() self.assertEqual( pipeline.runner.debug_counters['element_counts'], { 'oom:flatten': 3000000, ('oom:combine/GroupByKey/reify_windows', None): 3000000, ('oom:dupes/oom:dupes', 'side'): 1000000, ('oom:dupes/oom:dupes', None): 1000000, 'oom:create': 1000000, ('oom:addone', None): 1000000, 'oom:combine/GroupByKey/group_by_key': 1, ('oom:check', None): 1, 'assert_that/singleton': 1, ('assert_that/Map(match)', None): 1, ('oom:combine/GroupByKey/group_by_window', None): 1, ('oom:combine/Combine/ParDo(CombineValuesDoFn)', None): 1}) def test_eager_pipeline(self): p = Pipeline('EagerPipelineRunner') self.assertEqual([1, 4, 9], p | Create([1, 2, 3]) | Map(lambda x: x*x)) class DiskCachedRunnerPipelineTest(PipelineTest): def setUp(self): self.runner_name = 'DiskCachedPipelineRunner' def test_cached_pvalues_are_refcounted(self): # Takes long with disk spilling. pass def test_eager_pipeline(self): # Tests eager runner only pass class Bacon(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--slices', type=int) class Eggs(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--style', default='scrambled') class Breakfast(Bacon, Eggs): pass class PipelineOptionsTest(unittest.TestCase): def test_flag_parsing(self): options = Breakfast(['--slices=3', '--style=sunny side up', '--ignored']) self.assertEquals(3, options.slices) self.assertEquals('sunny side up', options.style) def test_keyword_parsing(self): options = Breakfast( ['--slices=3', '--style=sunny side up', '--ignored'], slices=10) self.assertEquals(10, options.slices) self.assertEquals('sunny side up', options.style) def test_attribute_setting(self): options = Breakfast(slices=10) self.assertEquals(10, options.slices) options.slices = 20 self.assertEquals(20, options.slices) def test_view_as(self): generic_options = PipelineOptions(['--slices=3']) self.assertEquals(3, generic_options.view_as(Bacon).slices) self.assertEquals(3, generic_options.view_as(Breakfast).slices) generic_options.view_as(Breakfast).slices = 10 self.assertEquals(10, generic_options.view_as(Bacon).slices) with self.assertRaises(AttributeError): generic_options.slices # pylint: disable=pointless-statement with self.assertRaises(AttributeError): generic_options.view_as(Eggs).slices # pylint: disable=expression-not-assigned def test_defaults(self): options = Breakfast(['--slices=3']) self.assertEquals(3, options.slices) self.assertEquals('scrambled', options.style) def test_dir(self): options = Breakfast() self.assertEquals( ['from_dictionary', 'get_all_options', 'slices', 'style', 'view_as'], [attr for attr in dir(options) if not attr.startswith('_')]) self.assertEquals( ['from_dictionary', 'get_all_options', 'style', 'view_as'], [attr for attr in dir(options.view_as(Eggs)) if not attr.startswith('_')]) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/pvalue.py000066400000000000000000000372371272340251000230640ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """PValue, PCollection: one node of a dataflow graph. A node of a dataflow processing graph is a PValue. Currently, there is only one type: PCollection (a potentially very large set of arbitrary values). Once created, a PValue belongs to a pipeline and has an associated transform (of type PTransform), which describes how the value will be produced when the pipeline gets executed. """ from __future__ import absolute_import import collections class PValue(object): """Base class for PCollection. Dataflow users should not construct PValue objects directly in their pipelines. A PValue has the following main characteristics: (1) Belongs to a pipeline. Added during object initialization. (2) Has a transform that can compute the value if executed. (3) Has a value which is meaningful if the transform was executed. """ def __init__(self, pipeline, tag=None, element_type=None): """Initializes a PValue with all arguments hidden behind keyword arguments. Args: pipeline: Pipeline object for this PValue. tag: Tag of this PValue. element_type: The type of this PValue. """ self.pipeline = pipeline self.tag = tag self.element_type = element_type # The AppliedPTransform instance for the application of the PTransform # generating this PValue. The field gets initialized when a transform # gets applied. self.producer = None def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s transform=%s' % ( self.__class__.__name__, self.producer.transform if self.producer else 'n/a') def apply(self, *args, **kwargs): """Applies a transform or callable to a PValue. Args: *args: positional arguments. **kwargs: keyword arguments. The method will insert the pvalue as the next argument following an optional first label and a transform/callable object. It will call the pipeline.apply() method with this modified argument list. """ if isinstance(args[0], str): # TODO(robertwb): Make sure labels are properly passed during # ptransform construction and drop this argument. args = args[1:] arglist = list(args) arglist.insert(1, self) return self.pipeline.apply(*arglist, **kwargs) def __or__(self, ptransform): return self.pipeline.apply(ptransform, self) class PCollection(PValue): """A multiple values (potentially huge) container. Dataflow users should not construct PCollection objects directly in their pipelines. """ def __init__(self, pipeline, **kwargs): """Initializes a PCollection. Do not call directly.""" super(PCollection, self).__init__(pipeline, **kwargs) @property def windowing(self): if not hasattr(self, '_windowing'): self._windowing = self.producer.transform.get_windowing( self.producer.inputs) return self._windowing def __reduce_ex__(self, unused_version): # Pickling a PCollection is almost always the wrong thing to do, but we # can't prohibit it as it often gets implicitly picked up (e.g. as part # of a closure). return _InvalidUnpickledPCollection, () class _InvalidUnpickledPCollection(object): pass class PBegin(PValue): """A pipeline begin marker used as input to create/read transforms. The class is used internally to represent inputs to Create and Read transforms. This allows us to have transforms that uniformly take PValue(s) as inputs. """ pass class PDone(PValue): """PDone is the output of a transform that has a trivial result such as Write. """ pass class DoOutputsTuple(object): """An object grouping the multiple outputs of a ParDo or FlatMap transform.""" def __init__(self, pipeline, transform, tags, main_tag): self._pipeline = pipeline self._tags = tags self._main_tag = main_tag self._transform = transform # The ApplyPTransform instance for the application of the multi FlatMap # generating this value. The field gets initialized when a transform # gets applied. self.producer = None # Dictionary of PCollections already associated with tags. self._pcolls = {} def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s main_tag=%s tags=%s transform=%s' % ( self.__class__.__name__, self._main_tag, self._tags, self._transform) def __iter__(self): """Iterates over tags returning for each call a (tag, pvalue) pair.""" if self._main_tag is not None: yield self[self._main_tag] for tag in self._tags: yield self[tag] def __getattr__(self, tag): # Special methods which may be accessed before the object is # fully constructed (e.g. in unpickling). if tag[:2] == tag[-2:] == '__': return object.__getattr__(self, tag) return self[tag] def __getitem__(self, tag): # Accept int tags so that we can look at Partition tags with the # same ints that we used in the partition function. # TODO(gildea): Consider requiring string-based tags everywhere. # This will require a partition function that does not return ints. if isinstance(tag, int): tag = str(tag) if tag == self._main_tag: tag = None elif self._tags and tag not in self._tags: raise ValueError( 'Tag %s is neither the main tag %s nor any of the side tags %s' % ( tag, self._main_tag, self._tags)) # Check if we accessed this tag before. if tag in self._pcolls: return self._pcolls[tag] if tag is not None: self._transform.side_output_tags.add(tag) pcoll = PCollection(self._pipeline, tag=tag) # Transfer the producer from the DoOutputsTuple to the resulting # PCollection. pcoll.producer = self.producer self.producer.add_output(pcoll, tag) self._pcolls[tag] = pcoll return pcoll class SideOutputValue(object): """An object representing a tagged value. ParDo, Map, and FlatMap transforms can emit values on multiple outputs which are distinguished by string tags. The DoFn will return plain values if it wants to emit on the main output and SideOutputValue objects if it wants to emit a value on a specific tagged output. """ def __init__(self, tag, value): if not isinstance(tag, basestring): raise TypeError( 'Attempting to create a SideOutputValue with non-string tag %s' % tag) self.tag = tag self.value = value class PCollectionView(PValue): """An immutable view of a PCollection that can be used as a side input.""" def __init__(self, pipeline): """Initializes a PCollectionView. Do not call directly.""" super(PCollectionView, self).__init__(pipeline) @property def windowing(self): if not hasattr(self, '_windowing'): self._windowing = self.producer.transform.get_windowing( self.producer.inputs) return self._windowing def _view_options(self): """Internal options corresponding to specific view. Intended for internal use by runner implementations. Returns: Tuple of options for the given view. """ return () class SingletonPCollectionView(PCollectionView): """A PCollectionView that contains a single object.""" def __init__(self, pipeline, has_default, default_value): super(SingletonPCollectionView, self).__init__(pipeline) self.has_default = has_default self.default_value = default_value def _view_options(self): return (self.has_default, self.default_value) class IterablePCollectionView(PCollectionView): """A PCollectionView that can be treated as an iterable.""" pass class ListPCollectionView(PCollectionView): """A PCollectionView that can be treated as a list.""" pass class DictPCollectionView(PCollectionView): """A PCollectionView that can be treated as a dict.""" pass def _get_cached_view(pipeline, key): return pipeline._view_cache.get(key, None) # pylint: disable=protected-access def _cache_view(pipeline, key, view): pipeline._view_cache[key] = view # pylint: disable=protected-access def can_take_label_as_first_argument(callee): """Decorator to allow the "label" kwarg to be passed as the first argument. For example, since AsSingleton is annotated with this decorator, this allows the call "AsSingleton(pcoll, label='label1')" to be written more succinctly as "AsSingleton('label1', pcoll)". Args: callee: The callable to be called with an optional label argument. Returns: Callable that allows (but does not require) a string label as its first argument. """ def _inner(maybe_label, *args, **kwargs): if isinstance(maybe_label, basestring): return callee(*args, label=maybe_label, **kwargs) return callee(*((maybe_label,) + args), **kwargs) return _inner def _format_view_label(pcoll): # The monitoring UI doesn't like '/' character in transform labels. if not pcoll.producer: return str(pcoll.tag) return '%s.%s' % (pcoll.producer.full_label.replace('/', '|'), pcoll.tag) _SINGLETON_NO_DEFAULT = object() @can_take_label_as_first_argument def AsSingleton(pcoll, default_value=_SINGLETON_NO_DEFAULT, label=None): # pylint: disable=invalid-name """Create a SingletonPCollectionView from the contents of input PCollection. The input PCollection should contain at most one element (per window) and the resulting PCollectionView can then be used as a side input to PTransforms. If the PCollectionView is empty (for a given window), the side input value will be the default_value, if specified; otherwise, it will be an EmptySideInput object. Args: pcoll: Input pcollection. default_value: Default value for the singleton view. label: Label to be specified if several AsSingleton's with different defaults for the same PCollection. Returns: A singleton PCollectionView containing the element as above. """ label = label or _format_view_label(pcoll) has_default = default_value is not _SINGLETON_NO_DEFAULT if not has_default: default_value = None # Don't recreate the view if it was already created. hashable_default_value = ('val', default_value) if not isinstance(default_value, collections.Hashable): # Massage default value to treat as hash key. hashable_default_value = ('id', id(default_value)) cache_key = (pcoll, AsSingleton, has_default, hashable_default_value) cached_view = _get_cached_view(pcoll.pipeline, cache_key) if cached_view: return cached_view # Local import is required due to dependency loop; even though the # implementation of this function requires concepts defined in modules that # depend on pvalue, it lives in this module to reduce user workload. from google.cloud.dataflow.transforms import sideinputs # pylint: disable=g-import-not-at-top view = (pcoll | sideinputs.ViewAsSingleton(has_default, default_value, label=label)) _cache_view(pcoll.pipeline, cache_key, view) return view @can_take_label_as_first_argument def AsIter(pcoll, label=None): # pylint: disable=invalid-name """Create an IterablePCollectionView from the elements of input PCollection. The contents of the given PCollection will be available as an iterable in PTransforms that use the returned PCollectionView as a side input. Args: pcoll: Input pcollection. label: Label to be specified if several AsIter's for the same PCollection. Returns: An iterable PCollectionView containing the elements as above. """ label = label or _format_view_label(pcoll) # Don't recreate the view if it was already created. cache_key = (pcoll, AsIter) cached_view = _get_cached_view(pcoll.pipeline, cache_key) if cached_view: return cached_view # Local import is required due to dependency loop; even though the # implementation of this function requires concepts defined in modules that # depend on pvalue, it lives in this module to reduce user workload. from google.cloud.dataflow.transforms import sideinputs # pylint: disable=g-import-not-at-top view = (pcoll | sideinputs.ViewAsIterable(label=label)) _cache_view(pcoll.pipeline, cache_key, view) return view @can_take_label_as_first_argument def AsList(pcoll, label=None): # pylint: disable=invalid-name """Create a ListPCollectionView from the elements of input PCollection. The contents of the given PCollection will be available as a list-like object in PTransforms that use the returned PCollectionView as a side input. Args: pcoll: Input pcollection. label: Label to be specified if several AsList's for the same PCollection. Returns: A list PCollectionView containing the elements as above. """ label = label or _format_view_label(pcoll) # Don't recreate the view if it was already created. cache_key = (pcoll, AsList) cached_view = _get_cached_view(pcoll.pipeline, cache_key) if cached_view: return cached_view # Local import is required due to dependency loop; even though the # implementation of this function requires concepts defined in modules that # depend on pvalue, it lives in this module to reduce user workload. from google.cloud.dataflow.transforms import sideinputs # pylint: disable=g-import-not-at-top view = (pcoll | sideinputs.ViewAsList(label=label)) _cache_view(pcoll.pipeline, cache_key, view) return view @can_take_label_as_first_argument def AsDict(pcoll, label=None): # pylint: disable=invalid-name """Create a DictPCollectionView from the elements of input PCollection. The contents of the given PCollection whose elements are 2-tuples of key and value will be available as a dict-like object in PTransforms that use the returned PCollectionView as a side input. Args: pcoll: Input pcollection containing 2-tuples of key and value. label: Label to be specified if several AsDict's for the same PCollection. Returns: A dict PCollectionView containing the dict as above. """ label = label or _format_view_label(pcoll) # Don't recreate the view if it was already created. cache_key = (pcoll, AsDict) cached_view = _get_cached_view(pcoll.pipeline, cache_key) if cached_view: return cached_view # Local import is required due to dependency loop; even though the # implementation of this function requires concepts defined in modules that # depend on pvalue, it lives in this module to reduce user workload. from google.cloud.dataflow.transforms import sideinputs # pylint: disable=g-import-not-at-top view = (pcoll | sideinputs.ViewAsDict(label=label)) _cache_view(pcoll.pipeline, cache_key, view) return view class EmptySideInput(object): """Value indicating when a singleton side input was empty. If a PCollection was furnished as a singleton side input to a PTransform, and that PCollection was empty, then this value is supplied to the DoFn in the place where a value from a non-empty PCollection would have gone. This alerts the DoFn that the side input PCollection was empty. Users may want to check whether side input values are EmptySideInput, but they will very likely never want to create new instances of this class themselves. """ pass DataflowPythonSDK-0.2.5/google/cloud/dataflow/pvalue_test.py000066400000000000000000000045071272340251000241150ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the PValue and PCollection classes.""" import unittest from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.pvalue import AsDict from google.cloud.dataflow.pvalue import AsIter from google.cloud.dataflow.pvalue import AsList from google.cloud.dataflow.pvalue import AsSingleton from google.cloud.dataflow.pvalue import PValue from google.cloud.dataflow.transforms import Create class FakePipeline(Pipeline): """Fake pipeline object used to check if apply() receives correct args.""" def apply(self, *args, **kwargs): self.args = args self.kwargs = kwargs class PValueTest(unittest.TestCase): def test_pvalue_expected_arguments(self): pipeline = Pipeline('DirectPipelineRunner') value = PValue(pipeline) self.assertEqual(pipeline, value.pipeline) def test_pcollectionview_not_recreated(self): pipeline = Pipeline('DirectPipelineRunner') value = pipeline | Create('create1', [1, 2, 3]) value2 = pipeline | Create('create2', [(1, 1), (2, 2), (3, 3)]) self.assertEqual(AsSingleton(value), AsSingleton(value)) self.assertEqual(AsSingleton('new', value, default_value=1), AsSingleton('new', value, default_value=1)) self.assertNotEqual(AsSingleton(value), AsSingleton('new', value, default_value=1)) self.assertEqual(AsIter(value), AsIter(value)) self.assertEqual(AsList(value), AsList(value)) self.assertEqual(AsDict(value2), AsDict(value2)) self.assertNotEqual(AsSingleton(value), AsSingleton(value2)) self.assertNotEqual(AsIter(value), AsIter(value2)) self.assertNotEqual(AsList(value), AsList(value2)) self.assertNotEqual(AsDict(value), AsDict(value2)) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/python_sdk_releases.py000066400000000000000000000034311272340251000256220ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Descriptions of the versions of the SDK. This manages the features and tests supported by different versions of the Dataflow SDK for Python. To add feature 'foo' to a particular release, add a 'properties' value with 'feature_foo': True. To remove feature 'foo' from a particular release, add a 'properties' value with 'feature_foo': False. Features are cumulative and can be added and removed multiple times. By default, all tests are enabled. To remove test 'bar' from a particular release, add a 'properties' value with 'test_bar': False. To add it back to a subsequent release, add a 'properties' value with 'test_bar': True. Tests are cumulative and can be removed and added multiple times. See go/dataflow-testing for more information. """ OLDEST_SUPPORTED_PYTHON_SDK = 'python-0.1.4' RELEASES = [ {'name': 'python-0.2.5',}, {'name': 'python-0.2.4',}, {'name': 'python-0.2.3',}, {'name': 'python-0.2.2',}, {'name': 'python-0.2.1',}, {'name': 'python-0.2.0',}, {'name': 'python-0.1.5',}, {'name': 'python-0.1.4',}, {'name': 'python-0.1.3',}, {'name': 'python-0.1.2',}, {'name': 'python-0.1.1', 'properties': { 'feature_python': True, } }, ] DataflowPythonSDK-0.2.5/google/cloud/dataflow/runners/000077500000000000000000000000001272340251000226765ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/runners/__init__.py000066400000000000000000000020421272340251000250050ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Runner objects execute a Pipeline. This package defines runners, which are used to execute a pipeline. """ from google.cloud.dataflow.runners.dataflow_runner import DataflowPipelineRunner from google.cloud.dataflow.runners.direct_runner import DirectPipelineRunner from google.cloud.dataflow.runners.runner import create_runner from google.cloud.dataflow.runners.runner import PipelineRunner from google.cloud.dataflow.runners.runner import PipelineState DataflowPythonSDK-0.2.5/google/cloud/dataflow/runners/common.pxd000066400000000000000000000016171272340251000247100ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. cdef type SideOutputValue, TimestampedValue, WindowedValue cdef class DoFnRunner(object): cdef object dofn cdef object window_fn cdef object context cdef object tagged_receivers cdef object logger cdef object step_name cdef object main_receivers cpdef _process_outputs(self, element, results) DataflowPythonSDK-0.2.5/google/cloud/dataflow/runners/common.py000066400000000000000000000135031272340251000245420ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. # cython: profile=True """Worker operations executor.""" import sys from google.cloud.dataflow.internal import util from google.cloud.dataflow.pvalue import SideOutputValue from google.cloud.dataflow.transforms import core from google.cloud.dataflow.transforms.window import TimestampedValue from google.cloud.dataflow.transforms.window import WindowedValue from google.cloud.dataflow.transforms.window import WindowFn class FakeLogger(object): def PerThreadLoggingContext(self, *unused_args, **unused_kwargs): return self def __enter__(self): pass def __exit__(self, *unused_args): pass class DoFnRunner(object): """A helper class for executing ParDo operations. """ def __init__(self, fn, args, kwargs, side_inputs, windowing, context, tagged_receivers, logger=None, step_name=None): if not args and not kwargs: self.dofn = fn else: args, kwargs = util.insert_values_in_args(args, kwargs, side_inputs) class CurriedFn(core.DoFn): def start_bundle(self, context): return fn.start_bundle(context, *args, **kwargs) def process(self, context): return fn.process(context, *args, **kwargs) def finish_bundle(self, context): return fn.finish_bundle(context, *args, **kwargs) self.dofn = CurriedFn() self.window_fn = windowing.windowfn self.context = context self.tagged_receivers = tagged_receivers self.logger = logger or FakeLogger() self.step_name = step_name # Optimize for the common case. self.main_receivers = tagged_receivers[None] def start(self): self.context.set_element(None) try: self._process_outputs(None, self.dofn.start_bundle(self.context)) except BaseException as exn: self.reraise_augmented(exn) def finish(self): self.context.set_element(None) try: self._process_outputs(None, self.dofn.finish_bundle(self.context)) except BaseException as exn: self.reraise_augmented(exn) def process(self, element): try: with self.logger.PerThreadLoggingContext(step_name=self.step_name): self.context.set_element(element) self._process_outputs(element, self.dofn.process(self.context)) except BaseException as exn: self.reraise_augmented(exn) def reraise_augmented(self, exn): if getattr(exn, '_tagged_with_step', False) or not self.step_name: raise args = exn.args if args and isinstance(args[0], str): args = (args[0] + " [while running '%s']" % self.step_name,) + args[1:] # Poor man's exception chaining. raise type(exn), args, sys.exc_info()[2] else: raise def _process_outputs(self, element, results): """Dispatch the result of computation to the appropriate receivers. A value wrapped in a SideOutputValue object will be unwrapped and then dispatched to the appropriate indexed output. """ if results is None: return for result in results: tag = None if isinstance(result, SideOutputValue): tag = result.tag if not isinstance(tag, basestring): raise TypeError('In %s, tag %s is not a string' % (self, tag)) result = result.value if isinstance(result, WindowedValue): windowed_value = result elif element is None: # Start and finish have no element from which to grab context, # but may emit elements. if isinstance(result, TimestampedValue): value = result.value timestamp = result.timestamp assign_context = NoContext(value, timestamp) else: value = result timestamp = -1 assign_context = NoContext(value) windowed_value = WindowedValue( value, timestamp, self.window_fn.assign(assign_context)) elif isinstance(result, TimestampedValue): assign_context = WindowFn.AssignContext( result.timestamp, result.value, element.windows) windowed_value = WindowedValue( result.value, result.timestamp, self.window_fn.assign(assign_context)) else: windowed_value = element.with_value(result) if tag is None: self.main_receivers.output(windowed_value) else: self.tagged_receivers[tag].output(windowed_value) class NoContext(WindowFn.AssignContext): """An uninspectable WindowFn.AssignContext.""" NO_VALUE = object() def __init__(self, value, timestamp=NO_VALUE): self.value = value self._timestamp = timestamp @property def timestamp(self): if self._timestamp is self.NO_VALUE: raise ValueError('No timestamp in this context.') else: return self._timestamp @property def existing_windows(self): raise ValueError('No existing_windows in this context.') class DoFnState(object): """Keeps track of state that DoFns want, currently, user counters. """ def __init__(self, counter_factory): self.step_name = '' self._counter_factory = counter_factory def counter_for(self, aggregator): """Looks up the counter for this aggregator, creating one if necessary.""" return self._counter_factory.get_aggregator_counter( self.step_name, aggregator) DataflowPythonSDK-0.2.5/google/cloud/dataflow/runners/dataflow_runner.py000066400000000000000000000703701272340251000264510ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A runner implementation that submits a job for remote execution. The runner will create a JSON description of the job graph and then submit it to the Dataflow Service for remote execution by a worker. """ import base64 import logging import threading import time from google.cloud.dataflow import coders from google.cloud.dataflow import pvalue from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.io import iobase from google.cloud.dataflow.pvalue import PCollectionView from google.cloud.dataflow.runners.runner import PipelineResult from google.cloud.dataflow.runners.runner import PipelineRunner from google.cloud.dataflow.runners.runner import PipelineState from google.cloud.dataflow.runners.runner import PValueCache from google.cloud.dataflow.typehints import typehints from google.cloud.dataflow.utils import names from google.cloud.dataflow.utils.names import PropertyNames from google.cloud.dataflow.utils.names import TransformNames from google.cloud.dataflow.utils.options import StandardOptions from google.cloud.dataflow.internal.clients import dataflow as dataflow_api class DataflowPipelineRunner(PipelineRunner): """A runner that creates job graphs and submits them for remote execution. Every execution of the run() method will submit an independent job for remote execution that consists of the nodes reachable from the passed in node argument or entire graph if node is None. The run() method returns after the service created the job and will not wait for the job to finish if blocking is set to False. """ # Environment version information. It is passed to the service during a # a job submission and is used by the service to establish what features # are expected by the workers. BATCH_ENVIRONMENT_MAJOR_VERSION = '4' STREAMING_ENVIRONMENT_MAJOR_VERSION = '0' def __init__(self, cache=None, blocking=False): # Cache of CloudWorkflowStep protos generated while the runner # "executes" a pipeline. self._cache = cache if cache is not None else PValueCache() self.blocking = blocking self.result = None self._unique_step_id = 0 def _get_unique_step_name(self): self._unique_step_id += 1 return 's%s' % self._unique_step_id @staticmethod def poll_for_job_completion(runner, job_id): """Polls for the specified job to finish running (successfully or not).""" last_message_time = None last_message_id = None last_error_rank = float('-inf') last_error_msg = None last_job_state = None # How long to wait after pipeline failure for the error # message to show up giving the reason for the failure. # It typically takes about 30 seconds. final_countdown_timer_secs = 50.0 sleep_secs = 5.0 # Try to prioritize the user-level traceback, if any. def rank_error(msg): if 'work item was attempted' in msg: return -1 elif 'Traceback' in msg: return 1 else: return 0 while True: response = runner.dataflow_client.get_job(job_id) # If get() is called very soon after Create() the response may not contain # an initialized 'currentState' field. if response.currentState is not None: if response.currentState != last_job_state: logging.info('Job %s is in state %s', job_id, response.currentState) last_job_state = response.currentState if str(response.currentState) != 'JOB_STATE_RUNNING': # Stop checking for new messages on timeout, explanatory # message received, success, or a terminal job state caused # by the user that therefore doesn't require explanation. if (final_countdown_timer_secs <= 0.0 or last_error_msg is not None or str(response.currentState) == 'JOB_STATE_DONE' or str(response.currentState) == 'JOB_STATE_CANCELLED' or str(response.currentState) == 'JOB_STATE_UPDATED' or str(response.currentState) == 'JOB_STATE_DRAINED'): break # The job has failed; ensure we see any final error messages. sleep_secs = 1.0 # poll faster during the final countdown final_countdown_timer_secs -= sleep_secs time.sleep(sleep_secs) # Get all messages since beginning of the job run or since last message. page_token = None while True: messages, page_token = runner.dataflow_client.list_messages( job_id, page_token=page_token, start_time=last_message_time) for m in messages: if last_message_id is not None and m.id == last_message_id: # Skip the first message if it is the last message we got in the # previous round. This can happen because we use the # last_message_time as a parameter of the query for new messages. continue last_message_time = m.time last_message_id = m.id # Skip empty messages. if m.messageImportance is None: continue logging.info( '%s: %s: %s: %s', m.id, m.time, m.messageImportance, m.messageText) if str(m.messageImportance) == 'JOB_MESSAGE_ERROR': if rank_error(m.messageText) >= last_error_rank: last_error_rank = rank_error(m.messageText) last_error_msg = m.messageText if not page_token: break runner.result = DataflowPipelineResult(response) runner.last_error_msg = last_error_msg def run(self, pipeline): """Remotely executes entire pipeline or parts reachable from node.""" # Import here to avoid adding the dependency for local running scenarios. # pylint: disable=g-import-not-at-top from google.cloud.dataflow.internal import apiclient self.job = apiclient.Job(pipeline.options) # The superclass's run will trigger a traversal of all reachable nodes. super(DataflowPipelineRunner, self).run(pipeline) # Get a Dataflow API client and submit the job. standard_options = pipeline.options.view_as(StandardOptions) if standard_options.streaming: job_version = DataflowPipelineRunner.STREAMING_ENVIRONMENT_MAJOR_VERSION else: job_version = DataflowPipelineRunner.BATCH_ENVIRONMENT_MAJOR_VERSION self.dataflow_client = apiclient.DataflowApplicationClient( pipeline.options, job_version) self.result = DataflowPipelineResult( self.dataflow_client.create_job(self.job)) if self.blocking: thread = threading.Thread( target=DataflowPipelineRunner.poll_for_job_completion, args=(self, self.result.job_id())) # Mark the thread as a daemon thread so a keyboard interrupt on the main # thread will terminate everything. This is also the reason we will not # use thread.join() to wait for the polling thread. thread.daemon = True thread.start() while thread.isAlive(): time.sleep(5.0) if self.result.current_state() != PipelineState.DONE: raise DataflowRuntimeException( 'Dataflow pipeline failed:\n%s' % getattr(self, 'last_error_msg', None), self.result) return self.result def _get_typehint_based_encoding(self, typehint, window_coder): """Returns an encoding based on a typehint object.""" return self._get_cloud_encoding(self._get_coder(typehint, window_coder=window_coder)) def _get_coder(self, typehint, window_coder): """Returns a coder based on a typehint object.""" if window_coder: return coders.WindowedValueCoder( coders.registry.get_coder(typehint), coders.TimestampCoder(), window_coder) else: return coders.registry.get_coder(typehint) def _get_cloud_encoding(self, coder): """Returns an encoding based on a coder object.""" if not isinstance(coder, coders.Coder): raise TypeError('Coder object must inherit from coders.Coder: %s.' % str(coder)) return coder.as_cloud_object() def _get_side_input_encoding(self, input_encoding): """Returns an encoding for the output of a view transform. Args: input_encoding: encoding of current transform's input. Side inputs need this because the service will check that input and output types match. Returns: An encoding that matches the output and input encoding. This is essential for the View transforms introduced to produce side inputs to a ParDo. """ return { '@type': input_encoding['@type'], 'component_encodings': [input_encoding] } def _get_encoded_output_coder(self, transform_node, window_value=True): """Returns the cloud encoding of the coder for the output of a transform.""" if (len(transform_node.outputs) == 1 and transform_node.outputs[0].element_type is not None): # TODO(robertwb): Handle type hints for multi-output transforms. element_type = transform_node.outputs[0].element_type else: # TODO(silviuc): Remove this branch (and assert) when typehints are # propagated everywhere. Returning an 'Any' as type hint will trigger # usage of the fallback coder (i.e., cPickler). element_type = typehints.Any if window_value: window_coder = ( transform_node.outputs[0].windowing.windowfn.get_window_coder()) else: window_coder = None return self._get_typehint_based_encoding( element_type, window_coder=window_coder) def _add_step(self, step_kind, step_label, transform_node, side_tags=()): """Creates a Step object and adds it to the cache.""" # Import here to avoid adding the dependency for local running scenarios. # pylint: disable=g-import-not-at-top from google.cloud.dataflow.internal import apiclient step = apiclient.Step(step_kind, self._get_unique_step_name()) self.job.proto.steps.append(step.proto) step.add_property(PropertyNames.USER_NAME, step_label) # Cache the node/step association for the main output of the transform node. self._cache.cache_output(transform_node, None, step) # If side_tags is not () then this is a multi-output transform node and we # need to cache the (node, tag, step) for each of the tags used to access # the outputs. This is essential because the keys used to search in the # cache always contain the tag. for tag in side_tags: self._cache.cache_output(transform_node, tag, step) return step def run_Create(self, transform_node): transform = transform_node.transform step = self._add_step(TransformNames.CREATE_PCOLLECTION, transform_node.full_label, transform_node) # TODO(silviuc): Eventually use a coder based on typecoders. # Note that we base64-encode values here so that the service will accept # the values. element_coder = coders.PickleCoder() step.add_property( PropertyNames.ELEMENT, [base64.b64encode(element_coder.encode(v)) for v in transform.value]) # The service expects a WindowedValueCoder here, so we wrap the actual # encoding in a WindowedValueCoder. step.encoding = self._get_cloud_encoding( coders.WindowedValueCoder(element_coder)) step.add_property( PropertyNames.OUTPUT_INFO, [{PropertyNames.USER_NAME: ( '%s.%s' % (transform_node.full_label, PropertyNames.OUT)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: PropertyNames.OUT}]) def run_CreatePCollectionView(self, transform_node): step = self._add_step(TransformNames.COLLECTION_TO_SINGLETON, transform_node.full_label, transform_node) input_tag = transform_node.inputs[0].tag input_step = self._cache.get_pvalue(transform_node.inputs[0]) step.add_property( PropertyNames.PARALLEL_INPUT, {'@type': 'OutputReference', PropertyNames.STEP_NAME: input_step.proto.name, PropertyNames.OUTPUT_NAME: input_step.get_output(input_tag)}) step.encoding = self._get_side_input_encoding(input_step.encoding) step.add_property( PropertyNames.OUTPUT_INFO, [{PropertyNames.USER_NAME: ( '%s.%s' % (transform_node.full_label, PropertyNames.OUT)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: PropertyNames.OUT}]) def run_Flatten(self, transform_node): step = self._add_step(TransformNames.FLATTEN, transform_node.full_label, transform_node) inputs = [] for one_input in transform_node.inputs: input_step = self._cache.get_pvalue(one_input) inputs.append( {'@type': 'OutputReference', PropertyNames.STEP_NAME: input_step.proto.name, PropertyNames.OUTPUT_NAME: input_step.get_output(one_input.tag)}) step.add_property(PropertyNames.INPUTS, inputs) step.encoding = self._get_encoded_output_coder(transform_node) step.add_property( PropertyNames.OUTPUT_INFO, [{PropertyNames.USER_NAME: ( '%s.%s' % (transform_node.full_label, PropertyNames.OUT)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: PropertyNames.OUT}]) def apply_GroupByKey(self, transform, pcoll): coder = self._get_coder(pcoll.element_type or typehints.Any, None) if not coder.is_kv_coder(): raise ValueError(('Coder for the GroupByKey operation "%s" is not a ' 'key-value coder: %s.') % (transform.label, coder)) # TODO(robertwb): Update the coder itself if it changed. coders.registry.verify_deterministic( coder.key_coder(), 'GroupByKey operation "%s"' % transform.label) return pvalue.PCollection(pcoll.pipeline) def run_GroupByKey(self, transform_node): input_tag = transform_node.inputs[0].tag input_step = self._cache.get_pvalue(transform_node.inputs[0]) step = self._add_step( TransformNames.GROUP, transform_node.full_label, transform_node) step.add_property( PropertyNames.PARALLEL_INPUT, {'@type': 'OutputReference', PropertyNames.STEP_NAME: input_step.proto.name, PropertyNames.OUTPUT_NAME: input_step.get_output(input_tag)}) step.encoding = self._get_encoded_output_coder(transform_node) step.add_property( PropertyNames.OUTPUT_INFO, [{PropertyNames.USER_NAME: ( '%s.%s' % (transform_node.full_label, PropertyNames.OUT)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: PropertyNames.OUT}]) windowing = transform_node.transform.get_windowing( transform_node.inputs) step.add_property(PropertyNames.SERIALIZED_FN, pickler.dumps(windowing)) def run_ParDo(self, transform_node): transform = transform_node.transform input_tag = transform_node.inputs[0].tag input_step = self._cache.get_pvalue(transform_node.inputs[0]) # Attach side inputs. si_dict = {} si_tags_and_types = [] for side_pval in transform_node.side_inputs: assert isinstance(side_pval, PCollectionView) side_input_step = self._cache.get_pvalue(side_pval) si_label = side_input_step.step_name si_dict[si_label] = { '@type': 'OutputReference', PropertyNames.STEP_NAME: si_label, PropertyNames.OUTPUT_NAME: PropertyNames.OUT} # The label for the side input step will appear as a 'tag' property for # the side input source specification. Its type (singleton or iterator) # will also be used to read the entire source or just first element. si_tags_and_types.append((si_label, side_pval.__class__, side_pval._view_options())) # pylint: disable=protected-access # Now create the step for the ParDo transform being handled. step = self._add_step( TransformNames.DO, transform_node.full_label, transform_node, transform_node.transform.side_output_tags) fn_data = (transform.fn, transform.args, transform.kwargs, si_tags_and_types, transform_node.inputs[0].windowing) step.add_property(PropertyNames.SERIALIZED_FN, pickler.dumps(fn_data)) step.add_property( PropertyNames.PARALLEL_INPUT, {'@type': 'OutputReference', PropertyNames.STEP_NAME: input_step.proto.name, PropertyNames.OUTPUT_NAME: input_step.get_output(input_tag)}) # Add side inputs if any. step.add_property(PropertyNames.NON_PARALLEL_INPUTS, si_dict) # Generate description for main output and side outputs. The output names # will be 'out' for main output and 'out_' for a tagged output. # Using 'out' as a tag will not clash with the name for main since it will # be transformed into 'out_out' internally. outputs = [] step.encoding = self._get_encoded_output_coder(transform_node) # Add the main output to the description. outputs.append( {PropertyNames.USER_NAME: ( '%s.%s' % (transform_node.full_label, PropertyNames.OUT)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: PropertyNames.OUT}) for side_tag in transform.side_output_tags: # The assumption here is that side outputs will have the same typehint # and coder as the main output. This is certainly the case right now # but conceivably it could change in the future. outputs.append( {PropertyNames.USER_NAME: ( '%s.%s' % (transform_node.full_label, side_tag)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: ( '%s_%s' % (PropertyNames.OUT, side_tag))}) step.add_property(PropertyNames.OUTPUT_INFO, outputs) def apply_CombineValues(self, transform, pcoll): return pvalue.PCollection(pcoll.pipeline) def run_CombineValues(self, transform_node): transform = transform_node.transform input_tag = transform_node.inputs[0].tag input_step = self._cache.get_pvalue(transform_node.inputs[0]) step = self._add_step( TransformNames.COMBINE, transform_node.full_label, transform_node) # Combiner functions do not take deferred side-inputs (i.e. PValues) and # therefore the code to handle extra args/kwargs is simpler than for the # DoFn's of the ParDo transform. In the last, empty argument is where # side inputs information would go. fn_data = (transform.fn, transform.args, transform.kwargs, ()) step.add_property(PropertyNames.SERIALIZED_FN, pickler.dumps(fn_data)) step.add_property( PropertyNames.PARALLEL_INPUT, {'@type': 'OutputReference', PropertyNames.STEP_NAME: input_step.proto.name, PropertyNames.OUTPUT_NAME: input_step.get_output(input_tag)}) # Note that the accumulator must not have a WindowedValue encoding, while # the output of this step does in fact have a WindowedValue encoding. accumulator_encoding = self._get_encoded_output_coder(transform_node, window_value=False) output_encoding = self._get_encoded_output_coder(transform_node) step.encoding = output_encoding step.add_property(PropertyNames.ENCODING, accumulator_encoding) # Generate description for main output 'out.' outputs = [] # Add the main output to the description. outputs.append( {PropertyNames.USER_NAME: ( '%s.%s' % (transform_node.full_label, PropertyNames.OUT)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: PropertyNames.OUT}) step.add_property(PropertyNames.OUTPUT_INFO, outputs) def run_Read(self, transform_node): transform = transform_node.transform step = self._add_step( TransformNames.READ, transform_node.full_label, transform_node) # TODO(mairbek): refactor if-else tree to use registerable functions. # Initialize the source specific properties. if not hasattr(transform.source, 'format'): # If a format is not set, we assume the source to be a custom source. source_dict = dict() spec_dict = dict() spec_dict[names.SERIALIZED_SOURCE_KEY] = pickler.dumps(transform.source) spec_dict['@type'] = names.SOURCE_TYPE source_dict['spec'] = spec_dict step.add_property(PropertyNames.SOURCE_STEP_INPUT, source_dict) elif transform.source.format == 'text': step.add_property(PropertyNames.FILE_PATTERN, transform.source.path) elif transform.source.format == 'bigquery': # TODO(silviuc): Add table validation if transform.source.validate. if transform.source.table_reference is not None: step.add_property(PropertyNames.BIGQUERY_DATASET, transform.source.table_reference.datasetId) step.add_property(PropertyNames.BIGQUERY_TABLE, transform.source.table_reference.tableId) # If project owning the table was not specified then the project owning # the workflow (current project) will be used. if transform.source.table_reference.projectId is not None: step.add_property(PropertyNames.BIGQUERY_PROJECT, transform.source.table_reference.projectId) elif transform.source.query is not None: step.add_property(PropertyNames.BIGQUERY_QUERY, transform.source.query) else: raise ValueError('BigQuery source %r must specify either a table or' ' a query', transform.source) elif transform.source.format == 'pubsub': standard_options = ( transform_node.inputs[0].pipeline.options.view_as(StandardOptions)) if not standard_options.streaming: raise ValueError('PubSubSource is currently available for use only in ' 'streaming pipelines.') step.add_property(PropertyNames.PUBSUB_TOPIC, transform.source.topic) if transform.source.subscription: step.add_property(PropertyNames.PUBSUB_SUBSCRIPTION, transform.source.topic) if transform.source.id_label: step.add_property(PropertyNames.PUBSUB_ID_LABEL, transform.source.id_label) else: raise ValueError( 'Source %r has unexpected format %s.' % ( transform.source, transform.source.format)) if not hasattr(transform.source, 'format'): step.add_property(PropertyNames.FORMAT, names.SOURCE_FORMAT) else: step.add_property(PropertyNames.FORMAT, transform.source.format) if isinstance(transform.source, iobase.BoundedSource): coder = transform.source.default_output_coder() else: coder = transform.source.coder step.encoding = self._get_cloud_encoding(coder) step.add_property( PropertyNames.OUTPUT_INFO, [{PropertyNames.USER_NAME: ( '%s.%s' % (transform_node.full_label, PropertyNames.OUT)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: PropertyNames.OUT}]) def run__NativeWrite(self, transform_node): transform = transform_node.transform input_tag = transform_node.inputs[0].tag input_step = self._cache.get_pvalue(transform_node.inputs[0]) step = self._add_step( TransformNames.WRITE, transform_node.full_label, transform_node) # TODO(mairbek): refactor if-else tree to use registerable functions. # Initialize the sink specific properties. if transform.sink.format == 'text': # Note that it is important to use typed properties (@type/value dicts) # for non-string properties and also for empty strings. For example, # in the code below the num_shards must have type and also # file_name_suffix and shard_name_template (could be empty strings). step.add_property( PropertyNames.FILE_NAME_PREFIX, transform.sink.file_name_prefix, with_type=True) step.add_property( PropertyNames.FILE_NAME_SUFFIX, transform.sink.file_name_suffix, with_type=True) step.add_property( PropertyNames.SHARD_NAME_TEMPLATE, transform.sink.shard_name_template, with_type=True) if transform.sink.num_shards > 0: step.add_property( PropertyNames.NUM_SHARDS, transform.sink.num_shards, with_type=True) # TODO(silviuc): Implement sink validation. step.add_property(PropertyNames.VALIDATE_SINK, False, with_type=True) elif transform.sink.format == 'bigquery': # TODO(silviuc): Add table validation if transform.sink.validate. step.add_property(PropertyNames.BIGQUERY_DATASET, transform.sink.table_reference.datasetId) step.add_property(PropertyNames.BIGQUERY_TABLE, transform.sink.table_reference.tableId) # If project owning the table was not specified then the project owning # the workflow (current project) will be used. if transform.sink.table_reference.projectId is not None: step.add_property(PropertyNames.BIGQUERY_PROJECT, transform.sink.table_reference.projectId) step.add_property(PropertyNames.BIGQUERY_CREATE_DISPOSITION, transform.sink.create_disposition) step.add_property(PropertyNames.BIGQUERY_WRITE_DISPOSITION, transform.sink.write_disposition) if transform.sink.table_schema is not None: step.add_property( PropertyNames.BIGQUERY_SCHEMA, transform.sink.schema_as_json()) elif transform.sink.format == 'pubsub': standard_options = ( transform_node.inputs[0].pipeline.options.view_as(StandardOptions)) if not standard_options.streaming: raise ValueError('PubSubSink is currently available for use only in ' 'streaming pipelines.') step.add_property(PropertyNames.PUBSUB_TOPIC, transform.sink.topic) else: raise ValueError( 'Sink %r has unexpected format %s.' % ( transform.sink, transform.sink.format)) step.add_property(PropertyNames.FORMAT, transform.sink.format) step.encoding = self._get_cloud_encoding(transform.sink.coder) step.add_property(PropertyNames.ENCODING, step.encoding) step.add_property( PropertyNames.PARALLEL_INPUT, {'@type': 'OutputReference', PropertyNames.STEP_NAME: input_step.proto.name, PropertyNames.OUTPUT_NAME: input_step.get_output(input_tag)}) class DataflowPipelineResult(PipelineResult): """Represents the state of a pipeline run on the Dataflow service.""" def __init__(self, job): """Job is a Job message from the Dataflow API.""" self._job = job def job_id(self): return self._job.id def current_state(self): """Return the current state of the remote job. Returns: A PipelineState object. """ values_enum = dataflow_api.Job.CurrentStateValueValuesEnum api_jobstate_map = { values_enum.JOB_STATE_UNKNOWN: PipelineState.UNKNOWN, values_enum.JOB_STATE_STOPPED: PipelineState.STOPPED, values_enum.JOB_STATE_RUNNING: PipelineState.RUNNING, values_enum.JOB_STATE_DONE: PipelineState.DONE, values_enum.JOB_STATE_FAILED: PipelineState.FAILED, values_enum.JOB_STATE_CANCELLED: PipelineState.CANCELLED, values_enum.JOB_STATE_UPDATED: PipelineState.UPDATED, values_enum.JOB_STATE_DRAINING: PipelineState.DRAINING, values_enum.JOB_STATE_DRAINED: PipelineState.DRAINED, } return (api_jobstate_map[self._job.currentState] if self._job.currentState else PipelineState.UNKNOWN) def __str__(self): return '<%s %s %s>' % ( self.__class__.__name__, self.job_id(), self.current_state()) def __repr__(self): return '<%s %s at %s>' % (self.__class__.__name__, self._job, hex(id(self))) class DataflowRuntimeException(Exception): """Indicates an error has occurred in running this pipeline.""" def __init__(self, msg, result): super(DataflowRuntimeException, self).__init__(msg) self.result = result DataflowPythonSDK-0.2.5/google/cloud/dataflow/runners/direct_runner.py000066400000000000000000000313771272340251000261260ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """DirectPipelineRunner, executing on the local machine. The DirectPipelineRunner class implements what is called in Dataflow parlance the "direct runner". Such a runner executes the entire graph of transformations belonging to a pipeline on the local machine. """ from __future__ import absolute_import import collections import itertools import logging from google.cloud.dataflow import coders from google.cloud.dataflow import error from google.cloud.dataflow.io import fileio from google.cloud.dataflow.io import iobase from google.cloud.dataflow.pvalue import DictPCollectionView from google.cloud.dataflow.pvalue import EmptySideInput from google.cloud.dataflow.pvalue import IterablePCollectionView from google.cloud.dataflow.pvalue import ListPCollectionView from google.cloud.dataflow.pvalue import SingletonPCollectionView from google.cloud.dataflow.runners.common import DoFnRunner from google.cloud.dataflow.runners.common import DoFnState from google.cloud.dataflow.runners.runner import PipelineResult from google.cloud.dataflow.runners.runner import PipelineRunner from google.cloud.dataflow.runners.runner import PipelineState from google.cloud.dataflow.runners.runner import PValueCache from google.cloud.dataflow.transforms import DoFnProcessContext from google.cloud.dataflow.transforms.window import GlobalWindows from google.cloud.dataflow.transforms.window import WindowedValue from google.cloud.dataflow.typehints.typecheck import OutputCheckWrapperDoFn from google.cloud.dataflow.typehints.typecheck import TypeCheckError from google.cloud.dataflow.typehints.typecheck import TypeCheckWrapperDoFn from google.cloud.dataflow.utils import counters from google.cloud.dataflow.utils.options import TypeOptions class DirectPipelineRunner(PipelineRunner): """A local pipeline runner. The runner computes everything locally and does not make any attempt to optimize for time or space. """ def __init__(self, cache=None): # Cache of values computed while the runner executes a pipeline. self._cache = cache if cache is not None else PValueCache() self._counter_factory = counters.CounterFactory() # Element counts used for debugging footprint issues in the direct runner. # The values computed are used only for logging and do not take part in # any decision making logic. The key for the counter dictionary is either # the full label for the transform producing the elements or a tuple # (full label, output tag) for ParDo transforms since they can output values # on multiple outputs. self.debug_counters = {} self.debug_counters['element_counts'] = collections.Counter() @property def cache(self): return self._cache def get_pvalue(self, pvalue): """Gets the PValue's computed value from the runner's cache.""" try: return self._cache.get_pvalue(pvalue) except KeyError: raise error.PValueError('PValue is not computed.') def clear_pvalue(self, pvalue): """Removes a PValue from the runner's cache.""" self._cache.clear_pvalue(pvalue) def skip_if_cached(func): # pylint: disable=no-self-argument """Decorator to skip execution of a transform if value is cached.""" def func_wrapper(self, pvalue, *args, **kwargs): logging.debug('Current: Debug counters: %s', self.debug_counters) if self._cache.is_cached(pvalue): # pylint: disable=protected-access return else: func(self, pvalue, *args, **kwargs) return func_wrapper def run(self, pipeline): super(DirectPipelineRunner, self).run(pipeline) logging.info('Final: Debug counters: %s', self.debug_counters) return DirectPipelineResult(state=PipelineState.DONE, counter_factory=self._counter_factory) @skip_if_cached def run_CreatePCollectionView(self, transform_node): transform = transform_node.transform view = transform.view values = self._cache.get_pvalue(transform_node.inputs[0]) if isinstance(view, SingletonPCollectionView): has_default, default_value = view._view_options() # pylint: disable=protected-access if len(values) == 0: # pylint: disable=g-explicit-length-test if has_default: result = default_value else: result = EmptySideInput() elif len(values) == 1: # TODO(ccy): Figure out whether side inputs should ever be given as # windowed values result = values[0].value else: raise ValueError(('PCollection with more than one element accessed as ' 'a singleton view: %s.') % view) elif isinstance(view, IterablePCollectionView): result = [v.value for v in values] elif isinstance(view, ListPCollectionView): result = [v.value for v in values] elif isinstance(view, DictPCollectionView): result = dict(v.value for v in values) else: raise NotImplementedError self._cache.cache_output(transform_node, result) @skip_if_cached def run_ParDo(self, transform_node): transform = transform_node.transform # TODO(gildea): what is the appropriate object to attach the state to? context = DoFnProcessContext(label=transform.label, state=DoFnState(self._counter_factory)) side_inputs = [self._cache.get_pvalue(view) for view in transform_node.side_inputs] # TODO(robertwb): Do this type checking inside DoFnRunner to get it on # remote workers as well? options = transform_node.inputs[0].pipeline.options if options is not None and options.view_as(TypeOptions).runtime_type_check: transform.dofn = TypeCheckWrapperDoFn( transform.dofn, transform.get_type_hints()) # TODO(robertwb): Should this be conditionally done on the workers as well? transform.dofn = OutputCheckWrapperDoFn( transform.dofn, transform_node.full_label) class RecordingReceiverSet(object): def __init__(self, tag): self.tag = tag def output(self, element): results[self.tag].append(element) class TaggedReceivers(dict): def __missing__(self, key): return RecordingReceiverSet(key) results = collections.defaultdict(list) # Some tags may be empty. for tag in transform.side_output_tags: results[tag] = [] runner = DoFnRunner(transform.dofn, transform.args, transform.kwargs, side_inputs, transform_node.inputs[0].windowing, context, TaggedReceivers(), step_name=transform_node.full_label) runner.start() for v in self._cache.get_pvalue(transform_node.inputs[0]): runner.process(v) runner.finish() self._cache.cache_output(transform_node, []) for tag, value in results.items(): self.debug_counters['element_counts'][ (transform_node.full_label, tag)] += len(value) self._cache.cache_output(transform_node, tag, value) @skip_if_cached def run_GroupByKeyOnly(self, transform_node): result_dict = collections.defaultdict(list) # The input type of a GroupByKey will be KV[Any, Any] or more specific. kv_type_hint = transform_node.transform.get_type_hints().input_types[0] key_coder = coders.registry.get_coder(kv_type_hint[0].tuple_types[0]) for wv in self._cache.get_pvalue(transform_node.inputs[0]): if (isinstance(wv, WindowedValue) and isinstance(wv.value, collections.Iterable) and len(wv.value) == 2): k, v = wv.value # We use as key a string encoding of the key object to support keys # that are based on custom classes. This mimics also the remote # execution behavior where key objects are encoded before being written # to the shuffler system responsible for grouping. result_dict[key_coder.encode(k)].append(v) else: raise TypeCheckError('Input to GroupByKeyOnly must be a PCollection of ' 'windowed key-value pairs. Instead received: %r.' % wv) gbk_result = map( GlobalWindows.windowed_value, ((key_coder.decode(k), v) for k, v in result_dict.iteritems())) self.debug_counters['element_counts'][ transform_node.full_label] += len(gbk_result) self._cache.cache_output(transform_node, gbk_result) @skip_if_cached def run_Create(self, transform_node): transform = transform_node.transform create_result = [GlobalWindows.windowed_value(v) for v in transform.value] self.debug_counters['element_counts'][ transform_node.full_label] += len(create_result) self._cache.cache_output(transform_node, create_result) @skip_if_cached def run_Flatten(self, transform_node): flatten_result = list( itertools.chain.from_iterable( self._cache.get_pvalue(pc) for pc in transform_node.inputs)) self.debug_counters['element_counts'][ transform_node.full_label] += len(flatten_result) self._cache.cache_output(transform_node, flatten_result) @skip_if_cached def run_Read(self, transform_node): # TODO(chamikara) Implement a more generic way for passing PipelineOptions # to sources and sinks when using DirectRunner. source = transform_node.transform.source source.pipeline_options = transform_node.inputs[0].pipeline.options def read_values(reader): read_result = [GlobalWindows.windowed_value(e) for e in reader] self.debug_counters['element_counts'][ transform_node.full_label] += len(read_result) self._cache.cache_output(transform_node, read_result) if isinstance(source, iobase.BoundedSource): reader = source.read(None) read_values(reader) else: with source.reader() as reader: read_values(reader) @skip_if_cached def run__NativeWrite(self, transform_node): sink = transform_node.transform.sink if isinstance(sink, fileio.NativeTextFileSink): assert sink.num_shards in (0, 1) if sink.shard_name_template: sink.file_path += '-00000-of-00001' sink.file_path += sink.file_name_suffix sink.pipeline_options = transform_node.inputs[0].pipeline.options with sink.writer() as writer: for v in self._cache.get_pvalue(transform_node.inputs[0]): self.debug_counters['element_counts'][transform_node.full_label] += 1 writer.Write(v.value) class DirectPipelineResult(PipelineResult): """A DirectPipelineResult provides access to info about a pipeline.""" def __init__(self, state, counter_factory=None): super(DirectPipelineResult, self).__init__(state) self._counter_factory = counter_factory def aggregated_values(self, aggregator_or_name): return self._counter_factory.get_aggregator_values(aggregator_or_name) class EagerPipelineRunner(DirectPipelineRunner): is_eager = True def __init__(self): super(EagerPipelineRunner, self).__init__() self._seen_transforms = set() def run_transform(self, transform): if transform not in self._seen_transforms: self._seen_transforms.add(transform) super(EagerPipelineRunner, self).run_transform(transform) class DiskCachedPipelineRunner(DirectPipelineRunner): """A DirectPipelineRunner that uses a disk backed cache. DiskCachedPipelineRunner uses a temporary disk backed cache for running pipelines. This allows for running pipelines that will require more memory than it is available, however this comes with a performance cost due to disk IO. Memory requirement for DiskCachedPipelineRunner is approximately capped by the single transform in the pipeline that consumes and outputs largest total collection (i.e. inputs, side-inputs and outputs in aggregate). In the extreme case a where a transform will use all previous intermediate values as input, memory requirements for DiskCachedPipelineRunner will be the same as DirectPipelineRunner. """ def __init__(self): self._null_cache = () super(DiskCachedPipelineRunner, self).__init__(self._null_cache) def run(self, pipeline): try: self._cache = PValueCache(use_disk_backed_cache=True) return super(DirectPipelineRunner, self).run(pipeline) finally: del self._cache self._cache = self._null_cache @property def cache(self): raise NotImplementedError( 'DiskCachedPipelineRunner does not keep cache outside the scope of its ' 'run method.') DataflowPythonSDK-0.2.5/google/cloud/dataflow/runners/runner.py000066400000000000000000000256761272340251000246010ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """PipelineRunner, an abstract base runner object.""" from __future__ import absolute_import import logging import os import shelve import shutil import tempfile def create_runner(runner_name): """Creates a runner instance from a runner class name. Args: runner_name: Name of the pipeline runner. Possible values are: DirectPipelineRunner, DataflowPipelineRunner and BlockingDataflowPipelineRunner. Returns: A runner object. Raises: RuntimeError: if an invalid runner name is used. """ # pylint: disable=g-import-not-at-top if runner_name == 'DirectPipelineRunner': import google.cloud.dataflow.runners.direct_runner return google.cloud.dataflow.runners.direct_runner.DirectPipelineRunner() if runner_name == 'DiskCachedPipelineRunner': import google.cloud.dataflow.runners.direct_runner return google.cloud.dataflow.runners.direct_runner.DiskCachedPipelineRunner( ) if runner_name == 'EagerPipelineRunner': import google.cloud.dataflow.runners.direct_runner return google.cloud.dataflow.runners.direct_runner.EagerPipelineRunner() elif runner_name in ('DataflowPipelineRunner', 'BlockingDataflowPipelineRunner'): import google.cloud.dataflow.runners.dataflow_runner return google.cloud.dataflow.runners.dataflow_runner.DataflowPipelineRunner( blocking=runner_name == 'BlockingDataflowPipelineRunner') else: raise RuntimeError( 'Unexpected pipeline runner: %s. Valid values are ' 'DirectPipelineRunner, DataflowPipelineRunner, EagerPipelineRunner, or ' 'BlockingDataflowPipelineRunner.' % runner_name) class PipelineRunner(object): """A runner of a pipeline object. The base runner provides a run() method for visiting every node in the pipeline's DAG and executing the transforms computing the PValue in the node. It also provides a clear() method for visiting every node and clearing out the values contained in PValue objects produced during a run. A custom runner will typically provide implementations for some of the transform methods (ParDo, GroupByKey, Create, etc.). It may also provide a new implementation for clear_pvalue(), which is used to wipe out materialized values in order to reduce footprint. """ def run(self, pipeline): """Execute the entire pipeline or the sub-DAG reachable from a node.""" # Imported here to avoid circular dependencies. # pylint: disable=g-import-not-at-top from google.cloud.dataflow.pipeline import PipelineVisitor class RunVisitor(PipelineVisitor): def __init__(self, runner): self.runner = runner def visit_transform(self, transform_node): try: self.runner.run_transform(transform_node) except: logging.error('Error while visiting %s', transform_node.full_label) raise pipeline.visit(RunVisitor(self)) def clear(self, pipeline, node=None): """Clear all nodes or nodes reachable from node of materialized values. Args: pipeline: Pipeline object containing PValues to be cleared. node: Optional node in the Pipeline processing DAG. If specified only nodes reachable from this node will be cleared (ancestors of the node). This method is not intended (for now) to be called by users of Runner objects. It is a hook for future layers on top of the current programming model to control how much of the previously computed values are kept around. Presumably an interactivity layer will use it. The simplest way to change the behavior would be to define a runner that overwrites the clear_pvalue() method since this method (runner.clear) will visit all relevant nodes and call clear_pvalue on them. """ # Imported here to avoid circular dependencies. # pylint: disable=g-import-not-at-top from google.cloud.dataflow.pipeline import PipelineVisitor class ClearVisitor(PipelineVisitor): def __init__(self, runner): self.runner = runner def visit_value(self, value, _): self.runner.clear_pvalue(value) pipeline.visit(ClearVisitor(self), node=node) def apply(self, transform, input): """Runner callback for a pipeline.apply call. Args: transform: the transform to apply. input: transform's input (typically a PCollection). A concrete implementation of the Runner class may want to do custom pipeline construction for a given transform. To override the behavior for a transform class Xyz, implement an apply_Xyz method with this same signature. """ for cls in transform.__class__.mro(): m = getattr(self, 'apply_%s' % cls.__name__, None) if m: return m(transform, input) raise NotImplementedError( 'Execution of [%s] not implemented in runner %s.' % (transform, self)) def apply_PTransform(self, transform, input): # The base case of apply is to call the transform's apply. return transform.apply(input) def run_transform(self, transform_node): """Runner callback for a pipeline.run call. Args: transform_node: transform node for the transform to run. A concrete implementation of the Runner class must implement run_Abc for some class Abc in the method resolution order for every non-composite transform Xyz in the pipeline. """ for cls in transform_node.transform.__class__.mro(): m = getattr(self, 'run_%s' % cls.__name__, None) if m: return m(transform_node) raise NotImplementedError( 'Execution of [%s] not implemented in runner %s.' % ( transform_node.transform, self)) class PValueCache(object): """Local cache for arbitrary information computed for PValue objects.""" def __init__(self, use_disk_backed_cache=False): # Cache of values computed while a runner executes a pipeline. This is a # dictionary of PValues and their computed values. Note that in principle # the runner could contain PValues from several pipelines without clashes # since a PValue is associated with one and only one pipeline. The keys of # the dictionary are tuple of PValue instance addresses obtained using id() # and tag names converted to strings. self._use_disk_backed_cache = use_disk_backed_cache if use_disk_backed_cache: self._tempdir = tempfile.mkdtemp() self._cache = shelve.open(os.path.join(self._tempdir, 'shelve')) else: self._cache = {} def __del__(self): if self._use_disk_backed_cache: self._cache.close() shutil.rmtree(self._tempdir) def __len__(self): return len(self._cache) def to_cache_key(self, transform, tag): return str((id(transform), tag)) def _ensure_pvalue_has_real_producer(self, pvalue): """Ensure the passed-in PValue has the real_producer attribute. Args: pvalue: A PValue instance whose cached value is requested. During the runner's execution only the results of the primitive transforms are cached. Whenever we are looking for a PValue that is the output of a composite transform we need to find the output of its rightmost transform part. """ if not hasattr(pvalue, 'real_producer'): real_producer = pvalue.producer while real_producer.parts: real_producer = real_producer.parts[-1] pvalue.real_producer = real_producer def is_cached(self, pobj): from google.cloud.dataflow.pipeline import AppliedPTransform if isinstance(pobj, AppliedPTransform): transform = pobj tag = None else: self._ensure_pvalue_has_real_producer(pobj) transform = pobj.real_producer tag = pobj.tag return self.to_cache_key(transform, tag) in self._cache def cache_output(self, transform, tag_or_value, value=None): if value is None: value = tag_or_value tag = None else: tag = tag_or_value self._cache[ self.to_cache_key(transform, tag)] = [value, transform.refcounts[tag]] def get_pvalue(self, pvalue): """Gets the value associated with a PValue from the cache.""" self._ensure_pvalue_has_real_producer(pvalue) try: value_with_refcount = self._cache[self.key(pvalue)] value_with_refcount[1] -= 1 logging.debug('PValue computed by %s (tag %s): refcount: %d => %d', pvalue.real_producer.full_label, self.key(pvalue)[1], value_with_refcount[1] + 1, value_with_refcount[1]) if value_with_refcount[1] <= 0: self.clear_pvalue(pvalue) return value_with_refcount[0] except KeyError: if (pvalue.tag is not None and self.to_cache_key(pvalue.real_producer, None) in self._cache): # This is an undeclared, empty side output of a DoFn executed # in the local runner before this side output referenced. return [] else: raise def get_unwindowed_pvalue(self, pvalue): return [v.value for v in self.get_pvalue(pvalue)] def clear_pvalue(self, pvalue): """Removes a PValue from the cache.""" if self.is_cached(pvalue): del self._cache[self.key(pvalue)] def key(self, pobj): self._ensure_pvalue_has_real_producer(pobj) return self.to_cache_key(pobj.real_producer, pobj.tag) class PipelineState(object): """State of the Pipeline, as returned by PipelineResult.current_state(). This is meant to be the union of all the states any runner can put a pipeline in. Currently, it represents the values of the dataflow API JobState enum. """ UNKNOWN = 'UNKNOWN' # not specified STOPPED = 'STOPPED' # paused or not yet started RUNNING = 'RUNNING' # currently running DONE = 'DONE' # successfully completed (terminal state) FAILED = 'FAILED' # failed (terminal state) CANCELLED = 'CANCELLED' # explicitly cancelled (terminal state) UPDATED = 'UPDATED' # replaced by another job (terminal state) DRAINING = 'DRAINING' # still processing, no longer reading data DRAINED = 'DRAINED' # draining completed (terminal state) class PipelineResult(object): """A PipelineResult provides access to info about a pipeline.""" def __init__(self, state): self._state = state def current_state(self): """Return the current state of running the pipeline.""" return self._state # pylint: disable=unused-argument def aggregated_values(self, aggregator_or_name): """Return a dict of step names to values of the Aggregator.""" logging.warn('%s does not implement aggregated_values', self.__class__.__name__) return {} DataflowPythonSDK-0.2.5/google/cloud/dataflow/runners/runner_test.py000066400000000000000000000050261272340251000256230ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the PipelineRunner and DirectPipelineRunner classes. Note that PipelineRunner and DirectPipelineRunner functionality is tested in all the other unit tests. In this file we choose to test only aspects related to caching and clearing values that are not tested elsewhere. """ import unittest from google.cloud.dataflow.internal import apiclient from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.runners import create_runner from google.cloud.dataflow.runners import DataflowPipelineRunner from google.cloud.dataflow.runners import DirectPipelineRunner import google.cloud.dataflow.transforms as ptransform from google.cloud.dataflow.utils.options import PipelineOptions class RunnerTest(unittest.TestCase): def test_create_runner(self): self.assertTrue( isinstance(create_runner('DirectPipelineRunner'), DirectPipelineRunner)) self.assertTrue( isinstance(create_runner('DataflowPipelineRunner'), DataflowPipelineRunner)) self.assertTrue( isinstance(create_runner('BlockingDataflowPipelineRunner'), DataflowPipelineRunner)) self.assertRaises(RuntimeError, create_runner, 'xyz') def test_remote_runner_translation(self): remote_runner = DataflowPipelineRunner() p = Pipeline(remote_runner, options=PipelineOptions([ '--dataflow_endpoint=ignored', '--job_name=test-job', '--project=test-project', '--staging_location=ignored', '--temp_location=/dev/null', '--no_auth=True' ])) res = (p | ptransform.Create('create', [1, 2, 3]) | ptransform.FlatMap('do', lambda x: [(x, x)]) | ptransform.GroupByKey('gbk')) remote_runner.job = apiclient.Job(p.options) super(DataflowPipelineRunner, remote_runner).run(p) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/000077500000000000000000000000001272340251000234005ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/__init__.py000066400000000000000000000017461272340251000255210ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """PTransform and descendants.""" # pylint: disable=wildcard-import from google.cloud.dataflow.transforms import combiners from google.cloud.dataflow.transforms.aggregator import * from google.cloud.dataflow.transforms.core import * from google.cloud.dataflow.transforms.ptransform import * from google.cloud.dataflow.transforms.timeutil import * from google.cloud.dataflow.transforms.util import * DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/aggregator.py000066400000000000000000000065431272340251000261040ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Support for user-defined Aggregators. Aggregators allow a pipeline to have the workers do custom aggregation of statistics about the data processed. To update an aggregator's value, call aggregate_to() on the context passed to a DoFn. Example: import google.cloud.dataflow as df simple_counter = df.Aggregator('example-counter') class ExampleDoFn(df.DoFn): def process(self, context): context.aggregate_to(simple_counter, 1) ... The aggregators defined here show up in the UI as "Custom counters." You can also query the combined value(s) of an aggregator by calling aggregated_value() or aggregated_values() on the result of running a pipeline. """ from __future__ import absolute_import from google.cloud.dataflow.transforms import core class Aggregator(object): """A user-specified aggregator of statistics about pipeline data. Args: combine_fn: how to combine values input to the aggregation. It must be one of these arithmetic functions: - Python's built-in sum, min, max, any, and all. - df.combiners.MeanCombineFn() The default is sum of 64-bit ints. type: describes the type that will be accepted as input for aggregation; by default types appropriate to the combine_fn are accepted. Example uses:: import google.cloud.dataflow as df simple_counter = df.Aggregator('example-counter') complex_counter = df.Aggregator('other-counter', df.Mean(), float) """ def __init__(self, name, combine_fn=sum, input_type=int): combine_fn = core.CombineFn.maybe_from_callable(combine_fn).for_input_type( input_type) if not _is_supported_kind(combine_fn): raise ValueError( 'combine_fn %r (class %r) ' 'does not map to a supported aggregation kind' % (combine_fn, combine_fn.__class__)) self.name = name self.combine_fn = combine_fn self.input_type = input_type def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): """Internal helper function for both __str__ and __repr__.""" def get_name(thing): try: return thing.__name__ except AttributeError: return thing.__class__.__name__ combine_fn_str = get_name(self.combine_fn) input_arg = '(%s)' % get_name(self.input_type) if self.input_type else '' if combine_fn_str == 'sum' and not input_arg: combine_call = '' else: combine_call = ' %s%s' % (combine_fn_str, input_arg) return 'Aggregator %s%s' % (self.name, combine_call) def _is_supported_kind(combine_fn): # pylint: disable=g-import-not-at-top from google.cloud.dataflow.internal.apiclient import metric_translations return combine_fn.__class__ in metric_translations DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/aggregator_test.py000066400000000000000000000045671272340251000271470ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for Aggregator class.""" import unittest import google.cloud.dataflow as df from google.cloud.dataflow.transforms import combiners from google.cloud.dataflow.transforms.aggregator import Aggregator class AggregatorTest(unittest.TestCase): def test_str(self): basic = Aggregator('a-name') self.assertEqual('', str(basic)) for_max = Aggregator('max-name', max) self.assertEqual('', str(for_max)) for_float = Aggregator('f-name', sum, float) self.assertEqual('', str(for_float)) for_mean = Aggregator('m-name', combiners.MeanCombineFn(), float) self.assertEqual('', str(for_mean)) def test_aggregation(self): mean = combiners.MeanCombineFn() mean.__name__ = 'mean' counter_types = [ (sum, int, 6), (min, int, 0), (max, int, 3), (mean, int, 1), (sum, float, 6.0), (min, float, 0.0), (max, float, 3.0), (mean, float, 1.5), (any, int, True), (all, float, False), ] aggeregators = [Aggregator('%s_%s' % (f.__name__, t.__name__), f, t) for f, t, _ in counter_types] class UpdateAggregators(df.DoFn): def process(self, context): for a in aggeregators: context.aggregate_to(a, context.element) p = df.Pipeline('DirectPipelineRunner') p | df.Create([0, 1, 2, 3]) | df.ParDo(UpdateAggregators()) res = p.run() for (_, _, expected), a in zip(counter_types, aggeregators): actual = res.aggregated_values(a).values()[0] self.assertEqual(expected, actual) self.assertEqual(type(expected), type(actual)) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/combiners.py000066400000000000000000000410471272340251000257410ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A library of basic combiner PTransform subclasses.""" from __future__ import absolute_import import heapq import itertools import random from google.cloud.dataflow.transforms import core from google.cloud.dataflow.transforms import cy_combiners from google.cloud.dataflow.transforms import ptransform from google.cloud.dataflow.typehints import Any from google.cloud.dataflow.typehints import Dict from google.cloud.dataflow.typehints import KV from google.cloud.dataflow.typehints import List from google.cloud.dataflow.typehints import Tuple from google.cloud.dataflow.typehints import TypeVariable from google.cloud.dataflow.typehints import Union from google.cloud.dataflow.typehints import with_input_types from google.cloud.dataflow.typehints import with_output_types __all__ = [ 'Count', 'Mean', 'Sample', 'Top', 'ToDict', 'ToList', ] class Mean(object): """Combiners for computing arithmetic means of elements.""" class Globally(ptransform.PTransform): """combiners.Mean.Globally computes the arithmetic mean of the elements.""" def apply(self, pcoll): return pcoll | core.CombineGlobally(MeanCombineFn()) class PerKey(ptransform.PTransform): """combiners.Mean.PerKey finds the means of the values for each key.""" def apply(self, pcoll): return pcoll | core.CombinePerKey(MeanCombineFn()) # TODO(laolu): This type signature is overly restrictive. This should be # more general. @with_input_types(Union[float, int, long]) @with_output_types(float) class MeanCombineFn(core.CombineFn): """CombineFn for computing an arithmetic mean.""" def create_accumulator(self): return (0, 0) def add_input(self, (sum_, count), element): return sum_ + element, count + 1 def merge_accumulators(self, accumulators): sums, counts = zip(*accumulators) return sum(sums), sum(counts) def extract_output(self, (sum_, count)): if count == 0: return float('NaN') return sum_ / float(count) def for_input_type(self, input_type): if input_type is int: return cy_combiners.MeanInt64Fn() elif input_type is float: return cy_combiners.MeanFloatFn() else: return self class Count(object): """Combiners for counting elements.""" class Globally(ptransform.PTransform): """combiners.Count.Globally counts the total number of elements.""" def apply(self, pcoll): return pcoll | core.CombineGlobally(CountCombineFn()) class PerKey(ptransform.PTransform): """combiners.Count.PerKey counts how many elements each unique key has.""" def apply(self, pcoll): return pcoll | core.CombinePerKey(CountCombineFn()) class PerElement(ptransform.PTransform): """combiners.Count.PerElement counts how many times each element occurs.""" def apply(self, pcoll): paired_with_void_type = KV[pcoll.element_type, Any] return (pcoll | (core.Map('%s:PairWithVoid' % self.label, lambda x: (x, None)) .with_output_types(paired_with_void_type)) | core.CombinePerKey(CountCombineFn())) @with_input_types(Any) @with_output_types(int) class CountCombineFn(core.CombineFn): """CombineFn for computing PCollection size.""" def create_accumulator(self): return 0 def add_inputs(self, accumulator, elements): return accumulator + len(elements) def merge_accumulators(self, accumulators): return sum(accumulators) def extract_output(self, accumulator): return accumulator class Top(object): """Combiners for obtaining extremal elements.""" # pylint: disable=no-self-argument @ptransform.ptransform_fn def Of(label, pcoll, n, compare, *args, **kwargs): """Obtain a list of the compare-most N elements in a PCollection. This transform will retrieve the n greatest elements in the PCollection to which it is applied, where "greatest" is determined by the comparator function supplied as the compare argument. compare should be an implementation of "a < b" taking at least two arguments (a and b). Additional arguments and side inputs specified in the apply call become additional arguments to the comparator. Args: label: display label for transform processes. pcoll: PCollection to process. n: number of elements to extract from pcoll. compare: as described above. *args: as described above. **kwargs: as described above. """ return pcoll | core.CombineGlobally( label, TopCombineFn(n, compare), *args, **kwargs) @ptransform.ptransform_fn def PerKey(label, pcoll, n, compare, *args, **kwargs): """Identifies the compare-most N elements associated with each key. This transform will produce a PCollection mapping unique keys in the input PCollection to the n greatest elements with which they are associated, where "greatest" is determined by the comparator function supplied as the compare argument. compare should be an implementation of "a < b" taking at least two arguments (a and b). Additional arguments and side inputs specified in the apply call become additional arguments to the comparator. Args: label: display label for transform processes. pcoll: PCollection to process. n: number of elements to extract from pcoll. compare: as described above. *args: as described above. **kwargs: as described above. Raises: TypeCheckError: If the output type of the input PCollection is not compatible with KV[A, B]. """ return pcoll | core.CombinePerKey( label, TopCombineFn(n, compare), *args, **kwargs) @ptransform.ptransform_fn def Largest(label, pcoll, n): """Obtain a list of the greatest N elements in a PCollection.""" return pcoll | Top.Of(label, n, lambda a, b: a < b) @ptransform.ptransform_fn def Smallest(label, pcoll, n): """Obtain a list of the least N elements in a PCollection.""" return pcoll | Top.Of(label, n, lambda a, b: b < a) @ptransform.ptransform_fn def LargestPerKey(label, pcoll, n): """Identifies the N greatest elements associated with each key.""" return pcoll | Top.PerKey(label, n, lambda a, b: a < b) @ptransform.ptransform_fn def SmallestPerKey(label, pcoll, n): """Identifies the N least elements associated with each key.""" return pcoll | Top.PerKey(label, n, lambda a, b: b < a) T = TypeVariable('T') @with_input_types(T) @with_output_types(List[T]) class TopCombineFn(core.CombineFn): """CombineFn doing the combining for all of the Top transforms. The comparator function supplied as an argument to the apply call invoking TopCombineFn should be an implementation of "a < b" taking at least two arguments (a and b). Additional arguments and side inputs specified in the apply call become additional arguments to the comparator. """ # Actually pickling the comparison operators (including, often, their # entire globals) can be very expensive. Instead refer to them by index # in this dictionary, which is populated on construction (including # unpickling). compare_by_id = {} def __init__(self, n, compare, _compare_id=None): # pylint: disable=invalid-name self._n = n self._compare = compare self._compare_id = _compare_id or id(compare) TopCombineFn.compare_by_id[self._compare_id] = self._compare def __reduce_ex__(self, _): return TopCombineFn, (self._n, self._compare, self._compare_id) class _HeapItem(object): """A wrapper for values supporting arbitrary comparisons. The heap implementation supplied by Python is a min heap that always uses the __lt__ operator if one is available. This wrapper overloads __lt__, letting us specify arbitrary precedence for elements in the PCollection. """ def __init__(self, item, compare_id, *args, **kwargs): # item: wrapped item. # compare: an implementation of the pairwise < operator. # args, kwargs: extra arguments supplied to the compare function. self.item = item self.compare_id = compare_id self.args = args self.kwargs = kwargs def __lt__(self, other): return TopCombineFn.compare_by_id[self.compare_id]( self.item, other.item, *self.args, **self.kwargs) def create_accumulator(self, *args, **kwargs): return [] # Empty heap. def add_input(self, heap, element, *args, **kwargs): # Note that because heap is a min heap, heappushpop will discard incoming # elements that are lesser (according to compare) than those in the heap # (since that's what you would get if you pushed a small element on and # popped the smallest element off). So, filtering a collection with a # min-heap gives you the largest elements in the collection. item = self._HeapItem(element, self._compare_id, *args, **kwargs) if len(heap) < self._n: heapq.heappush(heap, item) else: heapq.heappushpop(heap, item) return heap def merge_accumulators(self, heaps, *args, **kwargs): heap = [] for e in itertools.chain(*heaps): if len(heap) < self._n: heapq.heappush(heap, e) else: heapq.heappushpop(heap, e) return heap def extract_output(self, heap, *args, **kwargs): # Items in the heap are heap-ordered. We put them in sorted order, but we # have to use the reverse order because the result is expected to go # from greatest to least (as defined by the supplied comparison function). return [e.item for e in sorted(heap, reverse=True)] # Python's pickling is broken for nested classes. _HeapItem = TopCombineFn._HeapItem # pylint: disable=protected-access class Largest(TopCombineFn): def __init__(self, n): super(Largest, self).__init__(n, lambda a, b: a < b) def default_label(self): return 'Largest(%s)' % self._n class Smallest(TopCombineFn): def __init__(self, n): super(Smallest, self).__init__(n, lambda a, b: b < a) def default_label(self): return 'Smallest(%s)' % self._n class Sample(object): """Combiners for sampling n elements without replacement.""" # pylint: disable=no-self-argument @ptransform.ptransform_fn def FixedSizeGlobally(label, pcoll, n): return pcoll | core.CombineGlobally(label, SampleCombineFn(n)) @ptransform.ptransform_fn def FixedSizePerKey(label, pcoll, n): return pcoll | core.CombinePerKey(label, SampleCombineFn(n)) T = TypeVariable('T') @with_input_types(T) @with_output_types(List[T]) class SampleCombineFn(core.CombineFn): """CombineFn for all Sample transforms.""" def __init__(self, n): super(SampleCombineFn, self).__init__() # Most of this combiner's work is done by a TopCombineFn. We could just # subclass TopCombineFn to make this class, but since sampling is not # really a kind of Top operation, we use a TopCombineFn instance as a # helper instead. self._top_combiner = TopCombineFn(n, lambda a, b: a < b) def create_accumulator(self): return self._top_combiner.create_accumulator() def add_input(self, heap, element): # Before passing elements to the Top combiner, we pair them with random # numbers. The elements with the n largest random number "keys" will be # selected for the output. return self._top_combiner.add_input(heap, (random.random(), element)) def merge_accumulators(self, heaps): return self._top_combiner.merge_accumulators(heaps) def extract_output(self, heap): # Here we strip off the random number keys we added in add_input. return [e for _, e in self._top_combiner.extract_output(heap)] class _TupleCombineFnBase(core.CombineFn): def __init__(self, *combiners): self._combiners = [core.CombineFn.maybe_from_callable(c) for c in combiners] def create_accumulator(self): return [c.create_accumulator() for c in self._combiners] def merge_accumulators(self, accumulators): return [c.merge_accumulators(a) for c, a in zip(self._combiners, zip(*accumulators))] def extract_output(self, accumulator): return tuple([c.extract_output(a) for c, a in zip(self._combiners, accumulator)]) class TupleCombineFn(_TupleCombineFnBase): def add_inputs(self, accumulator, elements): return [c.add_inputs(a, e) for c, a, e in zip(self._combiners, accumulator, zip(*elements))] def with_common_input(self): return SingleInputTupleCombineFn(*self._combiners) class SingleInputTupleCombineFn(_TupleCombineFnBase): def add_inputs(self, accumulator, elements): return [c.add_inputs(a, elements) for c, a in zip(self._combiners, accumulator)] class ToList(ptransform.PTransform): """A global CombineFn that condenses a PCollection into a single list.""" def __init__(self, label='ToList'): super(ToList, self).__init__(label) def apply(self, pcoll): return pcoll | core.CombineGlobally(self.label, ToListCombineFn()) T = TypeVariable('T') @with_input_types(T) @with_output_types(List[T]) class ToListCombineFn(core.CombineFn): """CombineFn for to_list.""" def create_accumulator(self): return [] def add_input(self, accumulator, element): accumulator.append(element) return accumulator def merge_accumulators(self, accumulators): return sum(accumulators, []) def extract_output(self, accumulator): return accumulator class ToDict(ptransform.PTransform): """A global CombineFn that condenses a PCollection into a single dict. PCollections should consist of 2-tuples, notionally (key, value) pairs. If multiple values are associated with the same key, only one of the values will be present in the resulting dict. """ def __init__(self, label='ToDict'): super(ToDict, self).__init__(label) def apply(self, pcoll): return pcoll | core.CombineGlobally(self.label, ToDictCombineFn()) K = TypeVariable('K') V = TypeVariable('V') @with_input_types(Tuple[K, V]) @with_output_types(Dict[K, V]) class ToDictCombineFn(core.CombineFn): """CombineFn for to_dict.""" def create_accumulator(self): return dict() def add_input(self, accumulator, element): key, value = element accumulator[key] = value return accumulator def merge_accumulators(self, accumulators): result = dict() for a in accumulators: result.update(a) return result def extract_output(self, accumulator): return accumulator def curry_combine_fn(fn, args, kwargs): if not args and not kwargs: return fn else: class CurriedFn(core.CombineFn): """CombineFn that applies extra arguments.""" def create_accumulator(self): return fn.create_accumulator(*args, **kwargs) def add_input(self, accumulator, element): return fn.add_input(accumulator, element, *args, **kwargs) def add_inputs(self, accumulator, elements): return fn.add_inputs(accumulator, elements, *args, **kwargs) def merge_accumulators(self, accumulators): return fn.merge_accumulators(accumulators, *args, **kwargs) def extract_output(self, accumulator): return fn.extract_output(accumulator, *args, **kwargs) def apply(self, elements): return fn.apply(elements, *args, **kwargs) return CurriedFn() class PhasedCombineFnExecutor(object): """Executor for phases of combine operations.""" def __init__(self, phase, fn, args, kwargs): self.combine_fn = curry_combine_fn(fn, args, kwargs) if phase == 'all': self.apply = self.full_combine elif phase == 'add': self.apply = self.add_only elif phase == 'merge': self.apply = self.merge_only elif phase == 'extract': self.apply = self.extract_only else: raise ValueError('Unexpected phase: %s' % phase) def full_combine(self, elements): # pylint: disable=invalid-name return self.combine_fn.apply(elements) def add_only(self, elements): # pylint: disable=invalid-name return self.combine_fn.add_inputs( self.combine_fn.create_accumulator(), elements) def merge_only(self, accumulators): # pylint: disable=invalid-name return self.combine_fn.merge_accumulators(accumulators) def extract_only(self, accumulator): # pylint: disable=invalid-name return self.combine_fn.extract_output(accumulator) DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/combiners_test.py000066400000000000000000000212621272340251000267750ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for our libraries of combine PTransforms.""" import unittest import google.cloud.dataflow as df from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.transforms import combiners import google.cloud.dataflow.transforms.combiners as combine from google.cloud.dataflow.transforms.core import CombineGlobally from google.cloud.dataflow.transforms.core import Create from google.cloud.dataflow.transforms.core import Map from google.cloud.dataflow.transforms.ptransform import PTransform from google.cloud.dataflow.transforms.util import assert_that, equal_to class CombineTest(unittest.TestCase): def test_builtin_combines(self): pipeline = Pipeline('DirectPipelineRunner') vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] mean = sum(vals) / float(len(vals)) size = len(vals) # First for global combines. pcoll = pipeline | Create('start', vals) result_mean = pcoll | combine.Mean.Globally('mean') result_count = pcoll | combine.Count.Globally('count') assert_that(result_mean, equal_to([mean]), label='assert:mean') assert_that(result_count, equal_to([size]), label='assert:size') # Again for per-key combines. pcoll = pipeline | Create('start-perkey', [('a', x) for x in vals]) result_key_mean = pcoll | combine.Mean.PerKey('mean-perkey') result_key_count = pcoll | combine.Count.PerKey('count-perkey') assert_that(result_key_mean, equal_to([('a', mean)]), label='key:mean') assert_that(result_key_count, equal_to([('a', size)]), label='key:size') pipeline.run() def test_top(self): pipeline = Pipeline('DirectPipelineRunner') # A parameter we'll be sharing with a custom comparator. names = {0: 'zo', 1: 'one', 2: 'twoo', 3: 'three', 5: 'fiiive', 6: 'sssssix', 9: 'nniiinne'} # First for global combines. pcoll = pipeline | Create('start', [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]) result_top = pcoll | combine.Top.Largest('top', 5) result_bot = pcoll | combine.Top.Smallest('bot', 4) result_cmp = pcoll | combine.Top.Of( 'cmp', 6, lambda a, b, names: len(names[a]) < len(names[b]), names) # Note parameter passed to comparator. assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top') assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot') assert_that(result_cmp, equal_to([[9, 6, 6, 5, 3, 2]]), label='assert:cmp') # Again for per-key combines. pcoll = pipeline | Create( 'start-perkey', [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]]) result_key_top = pcoll | combine.Top.LargestPerKey('top-perkey', 5) result_key_bot = pcoll | combine.Top.SmallestPerKey('bot-perkey', 4) result_key_cmp = pcoll | combine.Top.PerKey( 'cmp-perkey', 6, lambda a, b, names: len(names[a]) < len(names[b]), names) # Note parameter passed to comparator. assert_that(result_key_top, equal_to([('a', [9, 6, 6, 5, 3])]), label='key:top') assert_that(result_key_bot, equal_to([('a', [0, 1, 1, 1])]), label='key:bot') assert_that(result_key_cmp, equal_to([('a', [9, 6, 6, 5, 3, 2])]), label='key:cmp') pipeline.run() def test_top_shorthands(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | Create('start', [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]) result_top = pcoll | df.CombineGlobally('top', combiners.Largest(5)) result_bot = pcoll | df.CombineGlobally('bot', combiners.Smallest(4)) assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top') assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot') pcoll = pipeline | Create( 'start-perkey', [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]]) result_ktop = pcoll | df.CombinePerKey('top-perkey', combiners.Largest(5)) result_kbot = pcoll | df.CombinePerKey('bot-perkey', combiners.Smallest(4)) assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='k:top') assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='k:bot') pipeline.run() def test_sample(self): # First test global samples (lots of them). for ix in xrange(300): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | Create('start', [1, 1, 2, 2]) result = pcoll | combine.Sample.FixedSizeGlobally('sample-%d' % ix, 3) def matcher(): def match(actual): # There is always exactly one result. equal_to([1])([len(actual)]) # There are always exactly three samples in the result. equal_to([3])([len(actual[0])]) # Sampling is without replacement. num_ones = sum(1 for x in actual[0] if x == 1) num_twos = sum(1 for x in actual[0] if x == 2) equal_to([1, 2])([num_ones, num_twos]) return match assert_that(result, matcher()) pipeline.run() # Now test per-key samples. pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | Create( 'start-perkey', sum(([(i, 1), (i, 1), (i, 2), (i, 2)] for i in xrange(300)), [])) result = pcoll | combine.Sample.FixedSizePerKey('sample', 3) def matcher(): def match(actual): for _, samples in actual: equal_to([3])([len(samples)]) num_ones = sum(1 for x in samples if x == 1) num_twos = sum(1 for x in samples if x == 2) equal_to([1, 2])([num_ones, num_twos]) return match assert_that(result, matcher()) pipeline.run() def test_tuple_combine_fn(self): p = Pipeline('DirectPipelineRunner') result = ( p | Create([('a', 100, 0.0), ('b', 10, -1), ('c', 1, 100)]) | df.CombineGlobally(combine.TupleCombineFn(max, combine.MeanCombineFn(), sum)).without_defaults()) assert_that(result, equal_to([('c', 111.0 / 3, 99.0)])) p.run() def test_tuple_combine_fn_without_defaults(self): p = Pipeline('DirectPipelineRunner') result = ( p | Create([1, 1, 2, 3]) | df.CombineGlobally( combine.TupleCombineFn(min, combine.MeanCombineFn(), max) .with_common_input()).without_defaults()) assert_that(result, equal_to([(1, 7.0 / 4, 3)])) p.run() def test_to_list_and_to_dict(self): pipeline = Pipeline('DirectPipelineRunner') the_list = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] pcoll = pipeline | Create('start', the_list) result = pcoll | combine.ToList('to list') def matcher(expected): def match(actual): equal_to(expected[0])(actual[0]) return match assert_that(result, matcher([the_list])) pipeline.run() pipeline = Pipeline('DirectPipelineRunner') pairs = [(1, 2), (3, 4), (5, 6)] pcoll = pipeline | Create('start-pairs', pairs) result = pcoll | combine.ToDict('to dict') def matcher(): def match(actual): equal_to([1])([len(actual)]) equal_to(pairs)(actual[0].iteritems()) return match assert_that(result, matcher()) pipeline.run() def test_combine_globally_with_default(self): p = Pipeline('DirectPipelineRunner') assert_that(p | Create([]) | CombineGlobally(sum), equal_to([0])) p.run() def test_combine_globally_without_default(self): p = Pipeline('DirectPipelineRunner') result = p | Create([]) | CombineGlobally(sum).without_defaults() assert_that(result, equal_to([])) p.run() def test_combine_globally_with_default_side_input(self): class CombineWithSideInput(PTransform): def apply(self, pcoll): side = pcoll | CombineGlobally(sum).as_singleton_view() main = pcoll.pipeline | Create([None]) return main | Map(lambda _, s: s, side) p = Pipeline('DirectPipelineRunner') result1 = p | Create('label1', []) | CombineWithSideInput('L1') result2 = p | Create('label2', [1, 2, 3, 4]) | CombineWithSideInput('L2') assert_that(result1, equal_to([0]), label='r1') assert_that(result2, equal_to([10]), label='r2') p.run() if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/core.py000066400000000000000000001354021272340251000247070ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Core PTransform subclasses, such as FlatMap, GroupByKey, and Map.""" from __future__ import absolute_import import copy import uuid from google.cloud.dataflow import pvalue from google.cloud.dataflow import typehints from google.cloud.dataflow.coders import typecoders from google.cloud.dataflow.internal import util from google.cloud.dataflow.pvalue import AsIter from google.cloud.dataflow.pvalue import AsSingleton from google.cloud.dataflow.transforms import ptransform from google.cloud.dataflow.transforms import window from google.cloud.dataflow.transforms.ptransform import PTransform from google.cloud.dataflow.transforms.ptransform import ptransform_fn from google.cloud.dataflow.transforms.ptransform import PTransformWithSideInputs from google.cloud.dataflow.transforms.window import MIN_TIMESTAMP from google.cloud.dataflow.transforms.window import OutputTimeFn from google.cloud.dataflow.transforms.window import WindowedValue from google.cloud.dataflow.transforms.window import WindowFn from google.cloud.dataflow.typehints import Any from google.cloud.dataflow.typehints import get_type_hints from google.cloud.dataflow.typehints import is_consistent_with from google.cloud.dataflow.typehints import Iterable from google.cloud.dataflow.typehints import KV from google.cloud.dataflow.typehints import trivial_inference from google.cloud.dataflow.typehints import TypeCheckError from google.cloud.dataflow.typehints import Union from google.cloud.dataflow.typehints import WithTypeHints from google.cloud.dataflow.typehints.trivial_inference import element_type from google.cloud.dataflow.utils.options import TypeOptions class DoFnProcessContext(object): """A processing context passed to DoFn methods during execution. Most importantly, a DoFn.process method will access context.element to get the element it is supposed to process. Attributes: label: label of the ParDo whose element is being processed. element: element being processed (in process method only; always None in start_bundle and finish_bundle) timestamp: timestamp of the element (in process method only; always None in start_bundle and finish_bundle) windows: windows of the element (in process method only; always None in start_bundle and finish_bundle) state: a DoFnState object, which holds the runner's internal state for this element. For example, aggregator state is here. Not used by the pipeline code. """ def __init__(self, label, element=None, state=None): """Initialize a processing context object with an element and state. The element represents one value from a PCollection that will be accessed by a DoFn object during pipeline execution, and state is an arbitrary object where counters and other pipeline state information can be passed in. DoFnProcessContext objects are also used as inputs to PartitionFn instances. Args: label: label of the PCollection whose element is being processed. element: element of a PCollection being processed using this context. state: a DoFnState object with state to be passed in to the DoFn object. """ self.label = label self.state = state if element is not None: self.set_element(element) def set_element(self, windowed_value): if windowed_value is None: # Not currently processing an element. if hasattr(self, 'element'): del self.element del self.timestamp del self.windows else: self.element = windowed_value.value self.timestamp = windowed_value.timestamp self.windows = windowed_value.windows def aggregate_to(self, aggregator, input_value): """Provide a new input value for the aggregator. Args: aggregator: the aggregator to update input_value: the new value to input to the combine_fn of this aggregator. """ self.state.counter_for(aggregator).update(input_value) class DoFn(WithTypeHints): """A function object used by a transform with custom processing. The ParDo transform is such a transform. The ParDo.apply method will take an object of type DoFn and apply it to all elements of a PCollection object. In order to have concrete DoFn objects one has to subclass from DoFn and define the desired behavior (start_bundle/finish_bundle and process) or wrap a callable object using the CallableWrapperDoFn class. """ def default_label(self): return self.__class__.__name__ def infer_output_type(self, input_type): # TODO(robertwb): Side inputs types. # TODO(robertwb): Assert compatibility with input type hint? return self._strip_output_annotations( trivial_inference.infer_return_type(self.process, [input_type])) def start_bundle(self, context, *args, **kwargs): """Called before a bundle of elements is processed on a worker. Elements to be processed are split into bundles and distributed to workers. Before a worker calls process() on the first element of its bundle, it calls this method. Args: context: a DoFnProcessContext object *args: side inputs **kwargs: keyword side inputs """ pass def finish_bundle(self, context, *args, **kwargs): """Called after a bundle of elements is processed on a worker. Args: context: a DoFnProcessContext object *args: side inputs **kwargs: keyword side inputs """ pass def process(self, context, *args, **kwargs): """Called for each element of a pipeline. Args: context: a DoFnProcessContext object containing, among other attributes, the element to be processed. See the DoFnProcessContext documentation for details. *args: side inputs **kwargs: keyword side inputs """ raise NotImplementedError @staticmethod def from_callable(fn): return CallableWrapperDoFn(fn) def process_argspec_fn(self): """Returns the Python callable that will eventually be invoked. This should ideally be the user-level function that is called with the main and (if any) side inputs, and is used to relate the type hint parameters with the input parameters (e.g., by argument name). """ return self.process def _strip_output_annotations(self, type_hint): annotations = (window.TimestampedValue, window.WindowedValue, pvalue.SideOutputValue) # TODO(robertwb): These should be parameterized types that the # type inferencer understands. if (type_hint in annotations or trivial_inference.element_type(type_hint) in annotations): return Any else: return type_hint class CallableWrapperDoFn(DoFn): """A DoFn (function) object wrapping a callable object. The purpose of this class is to conveniently wrap simple functions and use them in transforms. """ def __init__(self, fn): """Initializes a CallableWrapperDoFn object wrapping a callable. Args: fn: A callable object. Raises: TypeError: if fn parameter is not a callable type. """ if not callable(fn): raise TypeError('Expected a callable object instead of: %r' % fn) self._fn = fn super(CallableWrapperDoFn, self).__init__() def __repr__(self): return 'CallableWrapperDoFn(%s)' % self._fn def default_type_hints(self): type_hints = get_type_hints(self._fn) # If the fn was a DoFn annotated with a type-hint that hinted a return # type compatible with Iterable[Any], then we strip off the outer # container type due to the 'flatten' portion of FlatMap. # TODO(robertwb): Should we require an iterable specification for FlatMap? if type_hints.output_types: args, kwargs = type_hints.output_types if len(args) == 1 and is_consistent_with(args[0], Iterable[Any]): type_hints = type_hints.copy() type_hints.set_output_types(element_type(args[0]), **kwargs) return type_hints def infer_output_type(self, input_type): return self._strip_output_annotations( trivial_inference.infer_return_type(self._fn, [input_type])) def process(self, context, *args, **kwargs): return self._fn(context.element, *args, **kwargs) def process_argspec_fn(self): return getattr(self._fn, '_argspec_fn', self._fn) class CombineFn(WithTypeHints): """A function object used by a Combine transform with custom processing. A CombineFn specifies how multiple values in all or part of a PCollection can be merged into a single value---essentially providing the same kind of information as the arguments to the Python "reduce" builtin (except for the input argument, which is an instance of CombineFnProcessContext). The combining process proceeds as follows: 1. Input values are partitioned into one or more batches. 2. For each batch, the create_accumulator method is invoked to create a fresh initial "accumulator" value representing the combination of zero values. 3. For each input value in the batch, the add_inputs method is invoked to combine more values with the accumulator for that batch. 4. The merge_accumulators method is invoked to combine accumulators from separate batches into a single combined output accumulator value, once all of the accumulators have had all the input value in their batches added to them. This operation is invoked repeatedly, until there is only one accumulator value left. 5. The extract_output operation is invoked on the final accumulator to get the output value. """ def default_label(self): return self.__class__.__name__ def create_accumulator(self, *args, **kwargs): """Return a fresh, empty accumulator for the combine operation. Args: *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ raise NotImplementedError(str(self)) def add_input(self, accumulator, element, *args, **kwargs): """Return result of folding element into accumulator. CombineFn implementors must override either add_input or add_inputs. Args: accumulator: the current accumulator element: the element to add *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ raise NotImplementedError(str(self)) def add_inputs(self, accumulator, elements, *args, **kwargs): """Returns the result of folding each element in elements into accumulator. This is provided in case the implementation affords more efficient bulk addition of elements. The default implementation simply loops over the inputs invoking add_input for each one. Args: accumulator: the current accumulator elements: the elements to add *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ for element in elements: accumulator = self.add_input(accumulator, element, *args, **kwargs) return accumulator def merge_accumulators(self, accumulators, *args, **kwargs): """Returns the result of merging several accumulators to a single accumulator value. Args: accumulators: the accumulators to merge *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ raise NotImplementedError(str(self)) def extract_output(self, accumulator, *args, **kwargs): """Return result of converting accumulator into the output value. Args: accumulator: the final accumulator value computed by this CombineFn for the entire input key or PCollection. *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ raise NotImplementedError(str(self)) def apply(self, elements, *args, **kwargs): """Returns result of applying this CombineFn to the input values. Args: elements: the set of values to combine. *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ return self.extract_output( self.add_inputs( self.create_accumulator(*args, **kwargs), elements, *args, **kwargs), *args, **kwargs) def for_input_type(self, input_type): """Returns a specialized implementation of self, if it exists. Otherwise, returns self. Args: input_type: the type of input elements. """ return self @staticmethod def from_callable(fn): return CallableWrapperCombineFn(fn) @staticmethod def maybe_from_callable(fn): return fn if isinstance(fn, CombineFn) else CallableWrapperCombineFn(fn) class CallableWrapperCombineFn(CombineFn): """A CombineFn (function) object wrapping a callable object. The purpose of this class is to conveniently wrap simple functions and use them in Combine transforms. """ _EMPTY = object() def __init__(self, fn): """Initializes a CallableFn object wrapping a callable. Args: fn: A callable object that reduces elements of an iterable to a single value (like the builtins sum and max). This callable must be capable of receiving the kind of values it generates as output in its input, and for best results, its operation must be commutative and associative. Raises: TypeError: if fn parameter is not a callable type. """ if not callable(fn): raise TypeError('Expected a callable object instead of: %r' % fn) super(CallableWrapperCombineFn, self).__init__() self._fn = fn def __repr__(self): return "CallableWrapperCombineFn(%s)" % self._fn def create_accumulator(self, *args, **kwargs): return self._EMPTY def add_input(self, accumulator, element, *args, **kwargs): if accumulator is self._EMPTY: return element else: return self._fn([accumulator, element], *args, **kwargs) def add_inputs(self, accumulator, elements, *args, **kwargs): if accumulator is self._EMPTY: return self._fn(elements, *args, **kwargs) elif isinstance(elements, (list, tuple)): return self._fn([accumulator] + elements, *args, **kwargs) else: def union(): yield accumulator for e in elements: yield e return self._fn(union(), *args, **kwargs) def merge_accumulators(self, accumulators, *args, **kwargs): # It's (weakly) assumed that self._fn is associative. return self._fn(accumulators, *args, **kwargs) def extract_output(self, accumulator, *args, **kwargs): return self._fn(()) if accumulator is self._EMPTY else accumulator def default_type_hints(self): fn_hints = get_type_hints(self._fn) if fn_hints.input_types is None: return fn_hints else: # fn(Iterable[V]) -> V becomes CombineFn(V) -> V input_args, input_kwargs = fn_hints.input_types if not input_args: if len(input_kwargs) == 1: input_args, input_kwargs = tuple(input_kwargs.values()), {} else: raise TypeError('Combiner input type must be specified positionally.') if not is_consistent_with(input_args[0], Iterable[Any]): raise TypeCheckError( 'All functions for a Combine PTransform must accept a ' 'single argument compatible with: Iterable[Any]. ' 'Instead a function with input type: %s was received.' % input_args[0]) input_args = (element_type(input_args[0]),) + input_args[1:] # TODO(robertwb): Assert output type is consistent with input type? hints = fn_hints.copy() hints.set_input_types(*input_args, **input_kwargs) return hints def for_input_type(self, input_type): # Avoid circular imports. from google.cloud.dataflow.transforms import cy_combiners if self._fn is any: return cy_combiners.AnyCombineFn() elif self._fn is all: return cy_combiners.AllCombineFn() else: known_types = { (sum, int): cy_combiners.SumInt64Fn(), (min, int): cy_combiners.MinInt64Fn(), (max, int): cy_combiners.MaxInt64Fn(), (sum, float): cy_combiners.SumFloatFn(), (min, float): cy_combiners.MinFloatFn(), (max, float): cy_combiners.MaxFloatFn(), } return known_types.get((self._fn, input_type), self) class PartitionFn(WithTypeHints): """A function object used by a Partition transform. A PartitionFn specifies how individual values in a PCollection will be placed into separate partitions, indexed by an integer. """ def default_label(self): return self.__class__.__name__ def partition_for(self, context, num_partitions, *args, **kwargs): """Specify which partition will receive this element. Args: context: A DoFnProcessContext containing an element of the input PCollection. num_partitions: Number of partitions, i.e., output PCollections. *args: optional parameters and side inputs. **kwargs: optional parameters and side inputs. Returns: An integer in [0, num_partitions). """ pass class CallableWrapperPartitionFn(PartitionFn): """A PartitionFn object wrapping a callable object. Instances of this class wrap simple functions for use in Partition operations. """ def __init__(self, fn): """Initializes a PartitionFn object wrapping a callable. Args: fn: A callable object, which should accept the following arguments: element - element to assign to a partition. num_partitions - number of output partitions. and may accept additional arguments and side inputs. Raises: TypeError: if fn is not a callable type. """ if not callable(fn): raise TypeError('Expected a callable object instead of: %r' % fn) self._fn = fn def partition_for(self, context, num_partitions, *args, **kwargs): return self._fn(context.element, num_partitions, *args, **kwargs) class ParDo(PTransformWithSideInputs): """A ParDo transform. Processes an input PCollection by applying a DoFn to each element and returning the accumulated results into an output PCollection. The type of the elements is not fixed as long as the DoFn can deal with it. In reality the type is restrained to some extent because the elements sometimes must be persisted to external storage. See the apply() method comments for a detailed description of all possible arguments. Note that the DoFn must return an iterable for each element of the input PCollection. An easy way to do this is to use the yield keyword in the process method. Args: label: name of this transform instance. Useful while monitoring and debugging a pipeline execution. pcoll: a PCollection to be processed. dofn: a DoFn object to be applied to each element of pcoll argument. *args: positional arguments passed to the dofn object. **kwargs: keyword arguments passed to the dofn object. Note that the positional and keyword arguments will be processed in order to detect PCollections that will be computed as side inputs to the transform. During pipeline execution whenever the DoFn object gets executed (its apply() method gets called) the PCollection arguments will be replaced by values from the PCollection in the exact positions where they appear in the argument lists. """ def __init__(self, fn_or_label, *args, **kwargs): super(ParDo, self).__init__(fn_or_label, *args, **kwargs) if not isinstance(self.fn, DoFn): raise TypeError('ParDo must be called with a DoFn instance.') def default_type_hints(self): return self.fn.get_type_hints() def infer_output_type(self, input_type): return trivial_inference.element_type( self.fn.infer_output_type(input_type)) def make_fn(self, fn): return fn if isinstance(fn, DoFn) else CallableWrapperDoFn(fn) def process_argspec_fn(self): return self.fn.process_argspec_fn() def apply(self, pcoll): self.side_output_tags = set() # TODO(robertwb): Change all uses of the dofn attribute to use fn instead. self.dofn = self.fn return pvalue.PCollection(pcoll.pipeline) def with_outputs(self, *tags, **main_kw): """Returns a tagged tuple allowing access to the outputs of a ParDo. The resulting object supports access to the PCollection associated with a tag (e.g., o.tag, o[tag]) and iterating over the available tags (e.g., for tag in o: ...). Args: *tags: if non-empty, list of valid tags. If a list of valid tags is given, it will be an error to use an undeclared tag later in the pipeline. **main_kw: dictionary empty or with one key 'main' defining the tag to be used for the main output (which will not have a tag associated with it). Returns: An object of type DoOutputsTuple that bundles together all the outputs of a ParDo transform and allows accessing the individual PCollections for each output using an object.tag syntax. Raises: TypeError: if the self object is not a PCollection that is the result of a ParDo transform. ValueError: if main_kw contains any key other than 'main'. """ main_tag = main_kw.pop('main', None) if main_kw: raise ValueError('Unexpected keyword arguments: %s' % main_kw.keys()) return _MultiParDo(self, tags, main_tag) class _MultiParDo(PTransform): def __init__(self, do_transform, tags, main_tag): super(_MultiParDo, self).__init__(do_transform.label) self._do_transform = do_transform self._tags = tags self._main_tag = main_tag def apply(self, pcoll): _ = pcoll | self._do_transform return pvalue.DoOutputsTuple( pcoll.pipeline, self._do_transform, self._tags, self._main_tag) def FlatMap(fn_or_label, *args, **kwargs): # pylint: disable=invalid-name """FlatMap is like ParDo except it takes a callable to specify the transformation. The callable must return an iterable for each element of the input PCollection. The elements of these iterables will be flattened into the output PCollection. Args: fn_or_label: name of this transform instance. Useful while monitoring and debugging a pipeline execution. *args: positional arguments passed to the transform callable. **kwargs: keyword arguments passed to the transform callable. Returns: A PCollection containing the Map outputs. Raises: TypeError: If the fn passed as argument is not a callable. Typical error is to pass a DoFn instance which is supported only for ParDo. """ if fn_or_label is None or isinstance(fn_or_label, str): label, fn, args = fn_or_label, args[0], args[1:] else: label, fn = None, fn_or_label if not callable(fn): raise TypeError( 'FlatMap can be used only with callable objects. ' 'Received %r instead for %s argument.' % (fn, 'first' if label is None else 'second')) if label is None: label = 'FlatMap(%s)' % ptransform.label_from_callable(fn) return ParDo(label, CallableWrapperDoFn(fn), *args, **kwargs) def Map(fn_or_label, *args, **kwargs): # pylint: disable=invalid-name """Map is like FlatMap except its callable returns only a single element. Args: fn_or_label: name of this transform instance. Useful while monitoring and debugging a pipeline execution. *args: positional arguments passed to the transform callable. **kwargs: keyword arguments passed to the transform callable. Returns: A PCollection containing the Map outputs. Raises: TypeError: If the fn passed as argument is not a callable. Typical error is to pass a DoFn instance which is supported only for ParDo. """ if isinstance(fn_or_label, str): label, fn, args = fn_or_label, args[0], args[1:] else: label, fn = None, fn_or_label if not callable(fn): raise TypeError( 'Map can be used only with callable objects. ' 'Received %r instead for %s argument.' % (fn, 'first' if label is None else 'second')) wrapper = lambda x, *args, **kwargs: [fn(x, *args, **kwargs)] # Proxy the type-hint information from the original function to this new # wrapped function. get_type_hints(wrapper).input_types = get_type_hints(fn).input_types output_hint = get_type_hints(fn).simple_output_type(label) if output_hint: get_type_hints(wrapper).set_output_types(typehints.Iterable[output_hint]) # pylint: disable=protected-access wrapper._argspec_fn = fn # pylint: enable=protected-access if label is None: label = 'Map(%s)' % ptransform.label_from_callable(fn) return FlatMap(label, wrapper, *args, **kwargs) def Filter(fn_or_label, *args, **kwargs): # pylint: disable=invalid-name """Filter is a FlatMap with its callable filtering out elements. Args: fn_or_label: name of this transform instance. Useful while monitoring and debugging a pipeline execution. *args: positional arguments passed to the transform callable. **kwargs: keyword arguments passed to the transform callable. Returns: A PCollection containing the Filter outputs. Raises: TypeError: If the fn passed as argument is not a callable. Typical error is to pass a DoFn instance which is supported only for FlatMap. """ if isinstance(fn_or_label, str): label, fn, args = fn_or_label, args[0], args[1:] else: label, fn = None, fn_or_label if not callable(fn): raise TypeError( 'Filter can be used only with callable objects. ' 'Received %r instead for %s argument.' % (fn, 'first' if label is None else 'second')) wrapper = lambda x, *args, **kwargs: [x] if fn(x, *args, **kwargs) else [] # Proxy the type-hint information from the function being wrapped, setting the # output type to be the same as the input type. get_type_hints(wrapper).input_types = get_type_hints(fn).input_types output_hint = get_type_hints(fn).simple_output_type(label) if (output_hint is None and get_type_hints(wrapper).input_types and get_type_hints(wrapper).input_types[0]): output_hint = get_type_hints(wrapper).input_types[0] if output_hint: get_type_hints(wrapper).set_output_types(typehints.Iterable[output_hint]) # pylint: disable=protected-access wrapper._argspec_fn = fn # pylint: enable=protected-access if label is None: label = 'Filter(%s)' % ptransform.label_from_callable(fn) return FlatMap(label, wrapper, *args, **kwargs) class CombineGlobally(PTransform): """A CombineGlobally transform. Reduces a PCollection to a single value by progressively applying a CombineFn to portions of the PCollection (and to intermediate values created thereby). See documentation in CombineFn for details on the specifics on how CombineFns are applied. Args: label: name of this transform instance. Useful while monitoring and debugging a pipeline execution. pcoll: a PCollection to be reduced into a single value. fn: a CombineFn object that will be called to progressively reduce the PCollection into single values, or a callable suitable for wrapping by CallableWrapperCombineFn. *args: positional arguments passed to the CombineFn object. **kwargs: keyword arguments passed to the CombineFn object. Raises: TypeError: If the output type of the input PCollection is not compatible with Iterable[A]. Returns: A single-element PCollection containing the main output of the Combine transform. Note that the positional and keyword arguments will be processed in order to detect PObjects that will be computed as side inputs to the transform. During pipeline execution whenever the CombineFn object gets executed (i.e., any of the CombineFn methods get called), the PObject arguments will be replaced by their actual value in the exact position where they appear in the argument lists. """ has_defaults = True as_view = False def __init__(self, label_or_fn, *args, **kwargs): if label_or_fn is None or isinstance(label_or_fn, str): label, fn, args = label_or_fn, args[0], args[1:] else: label, fn = None, label_or_fn super(CombineGlobally, self).__init__(label) self.fn = fn self.args = args self.kwargs = kwargs def default_label(self): return 'CombineGlobally(%s)' % ptransform.label_from_callable(self.fn) def clone(self, **extra_attributes): clone = copy.copy(self) clone.__dict__.update(extra_attributes) return clone def with_defaults(self, has_defaults=True): return self.clone(has_defaults=has_defaults) def without_defaults(self): return self.with_defaults(False) def as_singleton_view(self): return self.clone(as_view=True) def apply(self, pcoll): def add_input_types(transform): type_hints = self.get_type_hints() if type_hints.input_types: return transform.with_input_types(type_hints.input_types[0][0]) else: return transform combined = (pcoll | add_input_types(Map('KeyWithVoid', lambda v: (None, v)) .with_output_types(KV[None, pcoll.element_type])) | CombinePerKey('CombinePerKey', self.fn, *self.args, **self.kwargs) | Map('UnKey', lambda (k, v): v)) if not self.has_defaults and not self.as_view: return combined if self.has_defaults: combine_fn = ( self.fn if isinstance(self.fn, CombineFn) else CombineFn.from_callable(self.fn)) default_value = combine_fn.apply([], *self.args, **self.kwargs) else: default_value = pvalue._SINGLETON_NO_DEFAULT # pylint: disable=protected-access view = pvalue.AsSingleton(combined, default_value=default_value) if self.as_view: return view else: if pcoll.windowing.windowfn != window.GlobalWindows(): raise ValueError( "Default values are not yet supported in CombineGlobally() if the " "output PCollection is not windowed by GlobalWindows. " "Instead, use CombineGlobally().without_defaults() to output " "an empty PCollection if the input PCollection is empty, " "or CombineGlobally().as_singleton_view() to get the default " "output of the CombineFn if the input PCollection is empty.") def typed(transform): # TODO(robertwb): We should infer this. if combined.element_type: return transform.with_output_types(combined.element_type) else: return transform return (pcoll.pipeline | Create('DoOnce', [None]) | typed(Map('InjectDefault', lambda _, s: s, view))) @ptransform_fn def CombinePerKey(label, pcoll, fn, *args, **kwargs): # pylint: disable=invalid-name """A per-key Combine transform. Identifies sets of values associated with the same key in the input PCollection, then applies a CombineFn to condense those sets to single values. See documentation in CombineFn for details on the specifics on how CombineFns are applied. Args: label: name of this transform instance. Useful while monitoring and debugging a pipeline execution. pcoll: input pcollection. fn: instance of CombineFn to apply to all values under the same key in pcoll, or a callable whose signature is f(iterable, *args, **kwargs) (e.g., sum, max). *args: arguments and side inputs, passed directly to the CombineFn. **kwargs: arguments and side inputs, passed directly to the CombineFn. Returns: A PObject holding the result of the combine operation. """ return pcoll | GroupByKey() | CombineValues('Combine', fn, *args, **kwargs) # TODO(robertwb): Rename to CombineGroupedValues? class CombineValues(PTransformWithSideInputs): def make_fn(self, fn): return fn if isinstance(fn, CombineFn) else CombineFn.from_callable(fn) def apply(self, pcoll): args, kwargs = util.insert_values_in_args( self.args, self.kwargs, self.side_inputs) input_type = pcoll.element_type key_type = None if input_type is not None: key_type, _ = input_type.tuple_types runtime_type_check = ( pcoll.pipeline.options is not None and pcoll.pipeline.options.view_as(TypeOptions).runtime_type_check) return pcoll | ParDo( CombineValuesDoFn(key_type, self.fn, runtime_type_check), *args, **kwargs) class CombineValuesDoFn(DoFn): """DoFn for performing per-key Combine transforms.""" def __init__(self, input_pcoll_type, combinefn, runtime_type_check): super(CombineValuesDoFn, self).__init__() self.combinefn = combinefn self.runtime_type_check = runtime_type_check def process(self, p_context, *args, **kwargs): # Expected elements input to this DoFn are 2-tuples of the form # (key, iter), with iter an iterable of all the values associated with key # in the input PCollection. if self.runtime_type_check: # Apply the combiner in a single operation rather than artificially # breaking it up so that output type violations manifest as TypeCheck # errors rather than type errors. return [ (p_context.element[0], self.combinefn.apply(p_context.element[1], *args, **kwargs))] else: # Add the elements into three accumulators (for testing of merge). elements = p_context.element[1] accumulators = [] for k in range(3): if len(elements) <= k: break accumulators.append( self.combinefn.add_inputs( self.combinefn.create_accumulator(*args, **kwargs), elements[k::3], *args, **kwargs)) # Merge the accumulators. accumulator = self.combinefn.merge_accumulators( accumulators, *args, **kwargs) # Convert accumulator to the final result. return [(p_context.element[0], self.combinefn.extract_output(accumulator, *args, **kwargs))] def default_type_hints(self): hints = self.combinefn.get_type_hints().copy() if hints.input_types: K = typehints.TypeVariable('K') args, kwargs = hints.input_types args = (typehints.Tuple[K, typehints.Iterable[args[0]]],) + args[1:] hints.set_input_types(*args, **kwargs) else: K = typehints.Any if hints.output_types: main_output_type = hints.simple_output_type('') hints.set_output_types(typehints.Tuple[K, main_output_type]) return hints K = typehints.TypeVariable('K') V = typehints.TypeVariable('V') @typehints.with_input_types(typehints.KV[K, V]) @typehints.with_output_types(typehints.KV[K, typehints.Iterable[V]]) class GroupByKey(PTransform): """A group by key transform. Processes an input PCollection consisting of key/value pairs represented as a tuple pair. The result is a PCollection where values having a common key are grouped together. For example (a, 1), (b, 2), (a, 3) will result into (a, [1, 3]), (b, [2]). The implementation here is used only when run on the local direct runner. """ class ReifyWindows(DoFn): def process(self, context): try: k, v = context.element except TypeError: raise TypeCheckError('Input to GroupByKey must be a PCollection with ' 'elements compatible with KV[A, B]') return [(k, window.WindowedValue(v, context.timestamp, context.windows))] def infer_output_type(self, input_type): key_type, value_type = trivial_inference.key_value_types(input_type) return Iterable[KV[key_type, typehints.WindowedValue[value_type]]] class GroupAlsoByWindow(DoFn): # TODO(robertwb): Support combiner lifting. def __init__(self, windowing): super(GroupByKey.GroupAlsoByWindow, self).__init__() self.windowing = windowing def infer_output_type(self, input_type): key_type, windowed_value_iter_type = trivial_inference.key_value_types( input_type) value_type = windowed_value_iter_type.inner_type.inner_type return Iterable[KV[key_type, Iterable[value_type]]] def process(self, context): k, vs = context.element # pylint: disable=g-import-not-at-top from google.cloud.dataflow.transforms.trigger import InMemoryUnmergedState from google.cloud.dataflow.transforms.trigger import create_trigger_driver # pylint: enable=g-import-not-at-top driver = create_trigger_driver(self.windowing, True) state = InMemoryUnmergedState() # TODO(robertwb): Conditionally process in smaller chunks. for wvalue in driver.process_elements(state, vs, MIN_TIMESTAMP): yield wvalue.with_value((k, wvalue.value)) while state.timers: fired = state.get_and_clear_timers() for timer_window, (name, time_domain, fire_time) in fired: for wvalue in driver.process_timer( timer_window, name, time_domain, fire_time, state): yield wvalue.with_value((k, wvalue.value)) def apply(self, pcoll): # This code path is only used in the local direct runner. For Dataflow # runner execution, the GroupByKey transform is expanded on the service. input_type = pcoll.element_type if input_type is not None: # Initialize type-hints used below to enforce type-checking and to pass # downstream to further PTransforms. key_type, value_type = trivial_inference.key_value_types(input_type) typecoders.registry.verify_deterministic( typecoders.registry.get_coder(key_type), 'GroupByKey operation "%s"' % self.label) reify_output_type = KV[key_type, typehints.WindowedValue[value_type]] gbk_input_type = KV[key_type, Iterable[typehints.WindowedValue[value_type]]] gbk_output_type = KV[key_type, Iterable[value_type]] return (pcoll | (ParDo('reify_windows', self.ReifyWindows()) .with_output_types(reify_output_type)) | (GroupByKeyOnly('group_by_key') .with_input_types(reify_output_type) .with_output_types(gbk_input_type)) | (ParDo('group_by_window', self.GroupAlsoByWindow(pcoll.windowing)) .with_input_types(gbk_input_type) .with_output_types(gbk_output_type))) else: return (pcoll | ParDo('reify_windows', self.ReifyWindows()) | GroupByKeyOnly('group_by_key') | ParDo('group_by_window', self.GroupAlsoByWindow(pcoll.windowing))) K = typehints.TypeVariable('K') V = typehints.TypeVariable('V') @typehints.with_input_types(typehints.KV[K, V]) @typehints.with_output_types(typehints.KV[K, typehints.Iterable[V]]) class GroupByKeyOnly(PTransform): """A group by key transform, ignoring windows.""" def __init__(self, label=None): super(GroupByKeyOnly, self).__init__(label) def infer_output_type(self, input_type): key_type, value_type = trivial_inference.key_value_types(input_type) return KV[key_type, Iterable[value_type]] def apply(self, pcoll): self._check_pcollection(pcoll) return pvalue.PCollection(pcoll.pipeline) class Partition(PTransformWithSideInputs): """Split a PCollection into several partitions. Uses the specified PartitionFn to separate an input PCollection into the specified number of sub-PCollections. When apply()d, a Partition() PTransform requires the following: Args: partitionfn: a PartitionFn, or a callable with the signature described in CallableWrapperPartitionFn. n: number of output partitions. The result of this PTransform is a simple list of the output PCollections representing each of n partitions, in order. """ class ApplyPartitionFnFn(DoFn): """A DoFn that applies a PartitionFn.""" def process(self, context, partitionfn, n, *args, **kwargs): partition = partitionfn.partition_for(context, n, *args, **kwargs) if not 0 <= partition < n: raise ValueError( 'PartitionFn specified out-of-bounds partition index: ' '%d not in [0, %d)' % (partition, n)) # Each input is directed into the side output that corresponds to the # selected partition. yield pvalue.SideOutputValue(str(partition), context.element) def make_fn(self, fn): return fn if isinstance(fn, PartitionFn) else CallableWrapperPartitionFn(fn) def apply(self, pcoll): n = int(self.args[0]) return pcoll | ParDo( self.ApplyPartitionFnFn(), self.fn, *self.args, **self.kwargs).with_outputs(*[str(t) for t in range(n)]) class Windowing(object): def __init__(self, windowfn, triggerfn=None, accumulation_mode=None, output_time_fn=None): global AccumulationMode, DefaultTrigger # pylint: disable=g-import-not-at-top from google.cloud.dataflow.transforms.trigger import AccumulationMode, DefaultTrigger # pylint: enable=g-import-not-at-top if triggerfn is None: triggerfn = DefaultTrigger() if accumulation_mode is None: if triggerfn == DefaultTrigger(): accumulation_mode = AccumulationMode.DISCARDING else: raise ValueError( 'accumulation_mode must be provided for non-trivial triggers') self.windowfn = windowfn self.triggerfn = triggerfn self.accumulation_mode = accumulation_mode self.output_time_fn = output_time_fn or OutputTimeFn.OUTPUT_AT_EOW self._is_default = ( self.windowfn == window.GlobalWindows() and self.triggerfn == DefaultTrigger() and self.accumulation_mode == AccumulationMode.DISCARDING and self.output_time_fn == OutputTimeFn.OUTPUT_AT_EOW) def __repr__(self): return "Windowing(%s, %s, %s, %s)" % (self.windowfn, self.triggerfn, self.accumulation_mode, self.output_time_fn) def is_default(self): return self._is_default T = typehints.TypeVariable('T') @typehints.with_input_types(T) @typehints.with_output_types(T) class WindowInto(ParDo): # pylint: disable=g-wrong-blank-lines """A window transform assigning windows to each element of a PCollection. Transforms an input PCollection by applying a windowing function to each element. Each transformed element in the result will be a WindowedValue element with the same input value and timestamp, with its new set of windows determined by the windowing function. """ class WindowIntoFn(DoFn): """A DoFn that applies a WindowInto operation.""" def __init__(self, windowing): self.windowing = windowing def process(self, context): context = WindowFn.AssignContext(context.timestamp, element=context.element, existing_windows=context.windows) new_windows = self.windowing.windowfn.assign(context) yield WindowedValue(context.element, context.timestamp, new_windows) def __init__(self, *args, **kwargs): """Initializes a WindowInto transform. Args: *args: A tuple of position arguments. **kwargs: A dictionary of keyword arguments. The *args, **kwargs are expected to be (label, windowfn) or (windowfn). The optional trigger and accumulation_mode kwargs may also be provided. """ triggerfn = kwargs.pop('trigger', None) accumulation_mode = kwargs.pop('accumulation_mode', None) output_time_fn = kwargs.pop('output_time_fn', None) label, windowfn = self.parse_label_and_arg(args, kwargs, 'windowfn') self.windowing = Windowing(windowfn, triggerfn, accumulation_mode, output_time_fn) dofn = self.WindowIntoFn(self.windowing) super(WindowInto, self).__init__(label, dofn) def get_windowing(self, unused_inputs): return self.windowing def infer_output_type(self, input_type): return input_type def apply(self, pcoll): input_type = pcoll.element_type if input_type is not None: output_type = input_type self.with_input_types(input_type) self.with_output_types(output_type) return super(WindowInto, self).apply(pcoll) # Python's pickling is broken for nested classes. WindowIntoFn = WindowInto.WindowIntoFn class Flatten(PTransform): """Merges several PCollections into a single PCollection. Copies all elements in 0 or more PCollections into a single output PCollection. If there are no input PCollections, the resulting PCollection will be empty (but see also kwargs below). Args: label: name of this transform instance. Useful while monitoring and debugging a pipeline execution. **kwargs: Accepts a single named argument "pipeline", which specifies the pipeline that "owns" this PTransform. Ordinarily Flatten can obtain this information from one of the input PCollections, but if there are none (or if there's a chance there may be none), this argument is the only way to provide pipeline information and should be considered mandatory. """ def __init__(self, label=None, **kwargs): super(Flatten, self).__init__(label) self.pipeline = kwargs.pop('pipeline', None) if kwargs: raise ValueError('Unexpected keyword arguments: %s' % kwargs.keys()) def _extract_input_pvalues(self, pvalueish): try: pvalueish = tuple(pvalueish) except TypeError: raise ValueError('Input to Flatten must be an iterable.') return pvalueish, pvalueish def apply(self, pcolls): for pcoll in pcolls: self._check_pcollection(pcoll) return pvalue.PCollection(self.pipeline) def get_windowing(self, inputs): if not inputs: # TODO(robertwb): Return something compatible with every windowing? return Windowing(window.GlobalWindows()) else: return super(Flatten, self).get_windowing(inputs) class Create(PTransform): """A transform that creates a PCollection from an iterable.""" def __init__(self, *args, **kwargs): """Initializes a Create transform. Args: *args: A tuple of position arguments. **kwargs: A dictionary of keyword arguments. The *args, **kwargs are expected to be (label, value) or (value). """ label, value = self.parse_label_and_arg(args, kwargs, 'value') super(Create, self).__init__(label) if isinstance(value, basestring): raise TypeError('PTransform Create: Refusing to treat string as ' 'an iterable. (string=%r)' % value) elif isinstance(value, dict): value = value.items() self.value = tuple(value) def infer_output_type(self, unused_input_type): if not self.value: return Any else: return Union[[trivial_inference.instance_to_type(v) for v in self.value]] def apply(self, pbegin): assert isinstance(pbegin, pvalue.PBegin) self.pipeline = pbegin.pipeline return pvalue.PCollection(self.pipeline) def get_windowing(self, unused_inputs): return Windowing(window.GlobalWindows()) def Read(*args, **kwargs): from google.cloud.dataflow import io return io.Read(*args, **kwargs) def Write(*args, **kwargs): from google.cloud.dataflow import io return io.Write(*args, **kwargs) DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/cy_combiners.pxd000066400000000000000000000055061272340251000265770ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. cimport cython from libc.stdint cimport int64_t, INT64_MIN, INT64_MAX cdef double _NEG_INF, _POS_INF, _NAN cdef class CountAccumulator(object): cdef readonly int64_t value cpdef add_input(self, unused_element) @cython.locals(accumulator=CountAccumulator) cpdef merge(self, accumulators) cdef class SumInt64Accumulator(object): cdef readonly int64_t value cpdef add_input(self, int64_t element) @cython.locals(accumulator=SumInt64Accumulator) cpdef merge(self, accumulators) cdef class MinInt64Accumulator(object): cdef readonly int64_t value cpdef add_input(self, int64_t element) @cython.locals(accumulator=MinInt64Accumulator) cpdef merge(self, accumulators) cdef class MaxInt64Accumulator(object): cdef readonly int64_t value cpdef add_input(self, int64_t element) @cython.locals(accumulator=MaxInt64Accumulator) cpdef merge(self, accumulators) cdef class MeanInt64Accumulator(object): cdef readonly int64_t sum cdef readonly int64_t count cpdef add_input(self, int64_t element) @cython.locals(accumulator=MeanInt64Accumulator) cpdef merge(self, accumulators) cdef class SumDoubleAccumulator(object): cdef readonly double value cpdef add_input(self, double element) @cython.locals(accumulator=SumDoubleAccumulator) cpdef merge(self, accumulators) cdef class MinDoubleAccumulator(object): cdef readonly double value cpdef add_input(self, double element) @cython.locals(accumulator=MinDoubleAccumulator) cpdef merge(self, accumulators) cdef class MaxDoubleAccumulator(object): cdef readonly double value cpdef add_input(self, double element) @cython.locals(accumulator=MaxDoubleAccumulator) cpdef merge(self, accumulators) cdef class MeanDoubleAccumulator(object): cdef readonly double sum cdef readonly int64_t count cpdef add_input(self, double element) @cython.locals(accumulator=MeanDoubleAccumulator) cpdef merge(self, accumulators) cdef class AllAccumulator(object): cdef readonly bint value cpdef add_input(self, bint element) @cython.locals(accumulator=AllAccumulator) cpdef merge(self, accumulators) cdef class AnyAccumulator(object): cdef readonly bint value cpdef add_input(self, bint element) @cython.locals(accumulator=AnyAccumulator) cpdef merge(self, accumulators) DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/cy_combiners.py000066400000000000000000000157341272340251000264400ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A library of basic cythonized CombineFn subclasses.""" from __future__ import absolute_import from google.cloud.dataflow.transforms import core class AccumulatorCombineFn(core.CombineFn): # singleton? def create_accumulator(self): return self._accumulator_type() @staticmethod def add_input(accumulator, element): accumulator.add_input(element) return accumulator def merge_accumulators(self, accumulators): accumulator = self._accumulator_type() accumulator.merge(accumulators) return accumulator @staticmethod def extract_output(accumulator): return accumulator.extract_output() def __eq__(self, other): return (isinstance(other, AccumulatorCombineFn) and self._accumulator_type is other._accumulator_type) def __hash__(self): return hash(self._accumulator_type) _63 = 63 # Avoid large literals in C source code. globals()['INT64_MAX'] = 2**_63 - 1 globals()['INT64_MIN'] = -2**_63 class CountAccumulator(object): def __init__(self): self.value = 0 def add_input(self, unused_element): self.value += 1 def merge(self, accumulators): for accumulator in accumulators: self.value += accumulator.value def extract_output(self): return self.value class SumInt64Accumulator(object): def __init__(self): self.value = 0 def add_input(self, element): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) self.value += element def merge(self, accumulators): for accumulator in accumulators: self.value += accumulator.value def extract_output(self): if not INT64_MIN <= self.value <= INT64_MAX: self.value %= 2**64 if self.value >= INT64_MAX: self.value -= 2**64 return self.value class MinInt64Accumulator(object): def __init__(self): self.value = INT64_MAX def add_input(self, element): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) if element < self.value: self.value = element def merge(self, accumulators): for accumulator in accumulators: if accumulator.value < self.value: self.value = accumulator.value def extract_output(self): return self.value class MaxInt64Accumulator(object): def __init__(self): self.value = INT64_MIN def add_input(self, element): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) if element > self.value: self.value = element def merge(self, accumulators): for accumulator in accumulators: if accumulator.value > self.value: self.value = accumulator.value def extract_output(self): return self.value class MeanInt64Accumulator(object): def __init__(self): self.sum = 0 self.count = 0 def add_input(self, element): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) self.sum += element self.count += 1 def merge(self, accumulators): for accumulator in accumulators: self.sum += accumulator.sum self.count += accumulator.count def extract_output(self): if not INT64_MIN <= self.sum <= INT64_MAX: self.sum %= 2**64 if self.sum >= INT64_MAX: self.sum -= 2**64 return self.sum / self.count if self.count else _NAN class CountCombineFn(AccumulatorCombineFn): _accumulator_type = CountAccumulator class SumInt64Fn(AccumulatorCombineFn): _accumulator_type = SumInt64Accumulator class MinInt64Fn(AccumulatorCombineFn): _accumulator_type = MinInt64Accumulator class MaxInt64Fn(AccumulatorCombineFn): _accumulator_type = MaxInt64Accumulator class MeanInt64Fn(AccumulatorCombineFn): _accumulator_type = MeanInt64Accumulator _POS_INF = float('inf') _NEG_INF = float('-inf') _NAN = float('nan') class SumDoubleAccumulator(object): def __init__(self): self.value = 0 def add_input(self, element): element = float(element) self.value += element def merge(self, accumulators): for accumulator in accumulators: self.value += accumulator.value def extract_output(self): return self.value class MinDoubleAccumulator(object): def __init__(self): self.value = _POS_INF def add_input(self, element): element = float(element) if element < self.value: self.value = element def merge(self, accumulators): for accumulator in accumulators: if accumulator.value < self.value: self.value = accumulator.value def extract_output(self): return self.value class MaxDoubleAccumulator(object): def __init__(self): self.value = _NEG_INF def add_input(self, element): element = float(element) if element > self.value: self.value = element def merge(self, accumulators): for accumulator in accumulators: if accumulator.value > self.value: self.value = accumulator.value def extract_output(self): return self.value class MeanDoubleAccumulator(object): def __init__(self): self.sum = 0 self.count = 0 def add_input(self, element): element = float(element) self.sum += element self.count += 1 def merge(self, accumulators): for accumulator in accumulators: self.sum += accumulator.sum self.count += accumulator.count def extract_output(self): return self.sum / self.count if self.count else _NAN class SumFloatFn(AccumulatorCombineFn): _accumulator_type = SumDoubleAccumulator class MinFloatFn(AccumulatorCombineFn): _accumulator_type = MinDoubleAccumulator class MaxFloatFn(AccumulatorCombineFn): _accumulator_type = MaxDoubleAccumulator class MeanFloatFn(AccumulatorCombineFn): _accumulator_type = MeanDoubleAccumulator class AllAccumulator(object): def __init__(self): self.value = True def add_input(self, element): self.value &= not not element def merge(self, accumulators): for accumulator in accumulators: self.value &= accumulator.value def extract_output(self): return self.value class AnyAccumulator(object): def __init__(self): self.value = False def add_input(self, element): self.value |= not not element def merge(self, accumulators): for accumulator in accumulators: self.value |= accumulator.value def extract_output(self): return self.value class AnyCombineFn(AccumulatorCombineFn): _accumulator_type = AnyAccumulator class AllCombineFn(AccumulatorCombineFn): _accumulator_type = AllAccumulator DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/ptransform.py000066400000000000000000000617671272340251000261660ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """PTransform and descendants. A PTransform is an object describing (not executing) a computation. The actual execution semantics for a transform is captured by a runner object. A transform object always belongs to a pipeline object. A PTransform derived class needs to define the apply() method that describes how one or more PValues are created by the transform. The module defines a few standard transforms: FlatMap (parallel do), GroupByKey (group by key), etc. Note that the apply() methods for these classes contain code that will add nodes to the processing graph associated with a pipeline. As support for the FlatMap transform, the module also defines a DoFn class and wrapper class that allows lambda functions to be used as FlatMap processing functions. """ from __future__ import absolute_import import copy import inspect import operator import os import sys from google.cloud.dataflow import coders from google.cloud.dataflow import error from google.cloud.dataflow import pvalue from google.cloud.dataflow import typehints from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.internal import util from google.cloud.dataflow.typehints import getcallargs_forhints from google.cloud.dataflow.typehints import TypeCheckError from google.cloud.dataflow.typehints import validate_composite_type_param from google.cloud.dataflow.typehints import WithTypeHints from google.cloud.dataflow.typehints.trivial_inference import instance_to_type class _PValueishTransform(object): """Visitor for PValueish objects. A PValueish is a PValue, or list, tuple, dict of PValuesish objects. This visits a PValueish, contstructing a (possibly mutated) copy. """ def visit(self, node, *args): return getattr( self, 'visit_' + node.__class__.__name__, lambda x, *args: x)(node, *args) def visit_list(self, node, *args): return [self.visit(x, *args) for x in node] def visit_tuple(self, node, *args): return tuple(self.visit(x, *args) for x in node) def visit_dict(self, node, *args): return {key: self.visit(value, *args) for (key, value) in node.items()} class _SetInputPValues(_PValueishTransform): def visit(self, node, replacements): if id(node) in replacements: return replacements[id(node)] else: return super(_SetInputPValues, self).visit(node, replacements) class _MaterializedDoOutputsTuple(pvalue.DoOutputsTuple): def __init__(self, deferred, pvalue_cache): super(_MaterializedDoOutputsTuple, self).__init__( None, None, deferred._tags, deferred._main_tag) self._deferred = deferred self._pvalue_cache = pvalue_cache def __getitem__(self, tag): return self._pvalue_cache.get_unwindowed_pvalue(self._deferred[tag]) class _MaterializePValues(_PValueishTransform): def __init__(self, pvalue_cache): self._pvalue_cache = pvalue_cache def visit(self, node): if isinstance(node, pvalue.PValue): return self._pvalue_cache.get_unwindowed_pvalue(node) elif isinstance(node, pvalue.DoOutputsTuple): return _MaterializedDoOutputsTuple(node, self._pvalue_cache) else: return super(_MaterializePValues, self).visit(node) class GetPValues(_PValueishTransform): def visit(self, node, pvalues=None): if pvalues is None: pvalues = [] self.visit(node, pvalues) return pvalues elif isinstance(node, (pvalue.PValue, pvalue.DoOutputsTuple)): pvalues.append(node) else: super(GetPValues, self).visit(node, pvalues) class ZipPValues(_PValueishTransform): """Pairs each PValue in a pvalueish with a value in a parallel out sibling. Sibling should have the same nested structure as pvalueish. Leaves in sibling are expanded across nested pvalueish lists, tuples, and dicts. For example ZipPValues().visit({'a': pc1, 'b': (pc2, pc3)}, {'a': 'A', 'b', 'B'}) will return [('a', pc1, 'A'), ('b', pc2, 'B'), ('b', pc3, 'B')] """ def visit(self, pvalueish, sibling, pairs=None, context=None): if pairs is None: pairs = [] self.visit(pvalueish, sibling, pairs, context) return pairs elif isinstance(pvalueish, (pvalue.PValue, pvalue.DoOutputsTuple)): pairs.append((context, pvalueish, sibling)) else: super(ZipPValues, self).visit(pvalueish, sibling, pairs, context) def visit_list(self, pvalueish, sibling, pairs, context): if isinstance(sibling, (list, tuple)): for ix, (p, s) in enumerate(zip( pvalueish, list(sibling) + [None] * len(pvalueish))): self.visit(p, s, pairs, 'position %s' % ix) else: for p in pvalueish: self.visit(p, sibling, pairs, context) def visit_tuple(self, pvalueish, sibling, pairs, context): self.visit_list(pvalueish, sibling, pairs, context) def visit_dict(self, pvalueish, sibling, pairs, context): if isinstance(sibling, dict): for key, p in pvalueish.items(): self.visit(p, sibling.get(key), pairs, key) else: for p in pvalueish.values(): self.visit(p, sibling, pairs, context) class PTransform(WithTypeHints): """A transform object used to modify one or more PCollections. Subclasses must define an apply() method that will be used when the transform is applied to some arguments. Typical usage pattern will be: input | CustomTransform(...) The apply() method of the CustomTransform object passed in will be called with input as an argument. """ # By default, transforms don't have any side inputs. side_inputs = () # Used for nullary transforms. pipeline = None # Default is unset. _user_label = None def __init__(self, label=None): super(PTransform, self).__init__() self.label = label @property def label(self): return self._user_label or self.default_label() @label.setter def label(self, value): self._user_label = value def default_label(self): return self.__class__.__name__ @classmethod def parse_label_and_arg(cls, args, kwargs, arg_name): """Parses a tuple of positional arguments into label, arg_name. The function is used by functions that take a (label, arg_name) list of parameters and in which first label could be optional even if the arg_name is not passed as a keyword. More specifically the following calling patterns are allowed:: (value) ('label', value) (arg_name=value) ('label', arg_name=value) (value, label='label') (label='label', arg_name=value) Args: args: A tuple of position arguments. kwargs: A dictionary of keyword arguments. arg_name: The name of the second ergument. Returns: A (label, value) tuple. The label will be the one passed in or one derived from the class name. The value will the corresponding value for the arg_name argument. Raises: ValueError: If the label and value cannot be deduced from args and kwargs and also if the label is not a string. """ # TODO(robertwb): Fix to not silently drop extra arguments. kw_label = kwargs.get('label', None) kw_value = kwargs.get(arg_name, None) if kw_value is not None: value = kw_value else: value = args[1] if len(args) > 1 else args[0] if args else None if kw_label is not None: label = kw_label else: # We need to get a label from positional arguments. If we did not get a # keyword value for the arg_name either then expect that a one element # list will provide the value and the label will be derived from the class # name. num_args = len(args) if kw_value is None: label = args[0] if num_args >= 2 else cls.__name__ else: label = args[0] if num_args >= 1 else cls.__name__ if label is None or value is None or not isinstance(label, basestring): raise ValueError( '%s expects a (label, %s) or (%s) argument list ' 'instead of args=%s, kwargs=%s' % ( cls.__name__, arg_name, arg_name, args, kwargs)) return label, value def with_input_types(self, input_type_hint): """Annotates the input type of a PTransform with a type-hint. Args: input_type_hint: An instance of an allowed built-in type, a custom class, or an instance of a typehints.TypeConstraint. Raises: TypeError: If 'type_hint' is not a valid type-hint. See typehints.validate_composite_type_param for further details. Returns: A reference to the instance of this particular PTransform object. This allows chaining type-hinting related methods. """ validate_composite_type_param(input_type_hint, 'Type hints for a PTransform') return super(PTransform, self).with_input_types(input_type_hint) def with_output_types(self, type_hint): """Annotates the output type of a PTransform with a type-hint. Args: type_hint: An instance of an allowed built-in type, a custom class, or a typehints.TypeConstraint. Raises: TypeError: If 'type_hint' is not a valid type-hint. See typehints.validate_composite_type_param for further details. Returns: A reference to the instance of this particular PTransform object. This allows chaining type-hinting related methods. """ validate_composite_type_param(type_hint, 'Type hints for a PTransform') return super(PTransform, self).with_output_types(type_hint) def type_check_inputs(self, pvalueish): self.type_check_inputs_or_outputs(pvalueish, 'input') def infer_output_type(self, unused_input_type): return self.get_type_hints().simple_output_type(self.label) or typehints.Any def type_check_outputs(self, pvalueish): self.type_check_inputs_or_outputs(pvalueish, 'output') def type_check_inputs_or_outputs(self, pvalueish, input_or_output): hints = getattr(self.get_type_hints(), input_or_output + '_types') if not hints: return arg_hints, kwarg_hints = hints if arg_hints and kwarg_hints: raise TypeCheckError( 'PTransform cannot have both positional and keyword type hints ' 'without overriding %s._type_check_%s()' % ( self.__class__, input_or_output)) root_hint = ( arg_hints[0] if len(arg_hints) == 1 else arg_hints or kwarg_hints) for context, pvalue_, hint in ZipPValues().visit(pvalueish, root_hint): if pvalue_.element_type is None: # TODO(robertwb): It's a bug that we ever get here. (typecheck) continue if hint and not typehints.is_consistent_with(pvalue_.element_type, hint): at_context = ' %s %s' % (input_or_output, context) if context else '' raise TypeCheckError( '%s type hint violation at %s%s: expected %s, got %s' % ( input_or_output.title(), self.label, at_context, hint, pvalue_.element_type)) def clone(self, new_label): """Clones the current transform instance under a new label.""" transform = copy.copy(self) transform.label = new_label return transform def apply(self, input_or_inputs): raise NotImplementedError def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s(PTransform)%s%s%s' % ( self.__class__.__name__, ' label=[%s]' % self.label if (hasattr(self, 'label') and self.label) else '', ' inputs=%s' % str(self.inputs) if (hasattr(self, 'inputs') and self.inputs) else '', ' side_inputs=%s' % str(self.side_inputs) if self.side_inputs else '') def _check_pcollection(self, pcoll): if not isinstance(pcoll, pvalue.PCollection): raise error.TransformError('Expecting a PCollection argument.') if not pcoll.pipeline: raise error.TransformError('PCollection not part of a pipeline.') def get_windowing(self, inputs): """Returns the window function to be associated with transform's output. By default most transforms just return the windowing function associated with the input PCollection (or the first input if several). """ # TODO(robertwb): Assert all input WindowFns compatible. return inputs[0].windowing def __or__(self, right): """Used to compose PTransforms, e.g., ptransform1 | ptransform2.""" if isinstance(right, PTransform): return ChainedPTransform(self, right) else: return NotImplemented def __ror__(self, left): """Used to apply this PTransform to non-PValues, e.g., a tuple.""" pvalueish, pvalues = self._extract_input_pvalues(left) pipelines = [v.pipeline for v in pvalues if isinstance(v, pvalue.PValue)] if pvalues and not pipelines: deferred = False # pylint: disable=g-import-not-at-top from google.cloud.dataflow import pipeline from google.cloud.dataflow.utils.options import PipelineOptions # pylint: enable=g-import-not-at-top p = pipeline.Pipeline( 'DirectPipelineRunner', PipelineOptions(sys.argv)) else: if not pipelines: if self.pipeline is not None: p = self.pipeline else: raise ValueError('"%s" requires a pipeline to be specified ' 'as there are no deferred inputs.'% self.label) else: p = self.pipeline or pipelines[0] for pp in pipelines: if p != pp: raise ValueError( 'Mixing value from different pipelines not allowed.') deferred = not getattr(p.runner, 'is_eager', False) # pylint: disable=g-import-not-at-top from google.cloud.dataflow.transforms.core import Create # pylint: enable=g-import-not-at-top replacements = {id(v): p | Create('CreatePInput%s' % ix, v) for ix, v in enumerate(pvalues) if not isinstance(v, pvalue.PValue) and v is not None} pvalueish = _SetInputPValues().visit(pvalueish, replacements) self.pipeline = p result = p.apply(self, pvalueish) if deferred: return result else: # Get a reference to the runners internal cache, otherwise runner may # clean it after run. cache = p.runner.cache p.run() return _MaterializePValues(cache).visit(result) def _extract_input_pvalues(self, pvalueish): """Extract all the pvalues contained in the input pvalueish. Returns pvalueish as well as the flat inputs list as the input may have to be copied as inspection may be destructive. By default, recursively extracts tuple components and dict values. Generally only needs to be overriden for multi-input PTransforms. """ # pylint: disable=g-import-not-at-top from google.cloud.dataflow import pipeline # pylint: enable=g-import-not-at-top if isinstance(pvalueish, pipeline.Pipeline): pvalueish = pvalue.PBegin(pvalueish) def _dict_tuple_leaves(pvalueish): if isinstance(pvalueish, tuple): for a in pvalueish: for p in _dict_tuple_leaves(a): yield p elif isinstance(pvalueish, dict): for a in pvalueish.values(): for p in _dict_tuple_leaves(a): yield p else: yield pvalueish return pvalueish, tuple(_dict_tuple_leaves(pvalueish)) class ChainedPTransform(PTransform): def __init__(self, *parts): super(ChainedPTransform, self).__init__(label=self._chain_label(parts)) self._parts = parts def _chain_label(self, parts): return '|'.join(p.label for p in parts) def __or__(self, right): if isinstance(right, PTransform): # Create a flat list rather than a nested tree of composite # transforms for better monitoring, etc. return ChainedPTransform(*(self._parts + (right,))) else: return NotImplemented def apply(self, pval): return reduce(operator.or_, self._parts, pval) class PTransformWithSideInputs(PTransform): """A superclass for any PTransform (e.g. FlatMap or Combine) invoking user code. PTransforms like FlatMap invoke user-supplied code in some kind of package (e.g. a DoFn) and optionally provide arguments and side inputs to that code. This internal-use-only class contains common functionality for PTransforms that fit this model. """ def __init__(self, fn_or_label, *args, **kwargs): if fn_or_label is None or isinstance(fn_or_label, basestring): label = fn_or_label fn, args = args[0], args[1:] else: label = None fn = fn_or_label if isinstance(fn, type) and issubclass(fn, typehints.WithTypeHints): # Don't treat Fn class objects as callables. raise ValueError('Use %s() not %s.' % (fn.__name__, fn.__name__)) self.fn = self.make_fn(fn) # Now that we figure out the label, initialize the super-class. super(PTransformWithSideInputs, self).__init__(label=label) if (any([isinstance(v, pvalue.PCollection) for v in args]) or any([isinstance(v, pvalue.PCollection) for v in kwargs.itervalues()])): raise error.SideInputError( 'PCollection used directly as side input argument. Specify ' 'AsIter(pcollection) or AsSingleton(pcollection) to indicate how the ' 'PCollection is to be used.') self.args, self.kwargs, self.side_inputs = util.remove_objects_from_args( args, kwargs, pvalue.PCollectionView) self.raw_side_inputs = args, kwargs # Prevent name collisions with fns of the form ' at ...>' self._cached_fn = self.fn # Ensure fn and side inputs are picklable for remote execution. self.fn = pickler.loads(pickler.dumps(self.fn)) self.args = pickler.loads(pickler.dumps(self.args)) self.kwargs = pickler.loads(pickler.dumps(self.kwargs)) # For type hints, because loads(dumps(class)) != class. self.fn = self._cached_fn def with_input_types( self, input_type_hint, *side_inputs_arg_hints, **side_input_kwarg_hints): """Annotates the types of main inputs and side inputs for the PTransform. Args: input_type_hint: An instance of an allowed built-in type, a custom class, or an instance of a typehints.TypeConstraint. *side_inputs_arg_hints: A variable length argument composed of of an allowed built-in type, a custom class, or a typehints.TypeConstraint. **side_input_kwarg_hints: A dictionary argument composed of of an allowed built-in type, a custom class, or a typehints.TypeConstraint. Example of annotating the types of side-inputs: FlatMap().with_input_types(int, int, bool) Raises: TypeError: If 'type_hint' is not a valid type-hint. See typehints.validate_composite_type_param for further details. Returns: A reference to the instance of this particular PTransform object. This allows chaining type-hinting related methods. """ super(PTransformWithSideInputs, self).with_input_types(input_type_hint) for si in side_inputs_arg_hints: validate_composite_type_param(si, 'Type hints for a PTransform') for si in side_input_kwarg_hints.values(): validate_composite_type_param(si, 'Type hints for a PTransform') self.side_inputs_types = side_inputs_arg_hints return WithTypeHints.with_input_types( self, input_type_hint, *side_inputs_arg_hints, **side_input_kwarg_hints) def type_check_inputs(self, pvalueish): type_hints = self.get_type_hints().input_types if type_hints: args, kwargs = self.raw_side_inputs def element_type(side_input): if isinstance(side_input, pvalue.PCollectionView): return side_input.element_type else: return instance_to_type(side_input) arg_types = [pvalueish.element_type] + [element_type(v) for v in args] kwargs_types = {k: element_type(v) for (k, v) in kwargs.items()} argspec_fn = self.process_argspec_fn() bindings = getcallargs_forhints(argspec_fn, *arg_types, **kwargs_types) hints = getcallargs_forhints(argspec_fn, *type_hints[0], **type_hints[1]) for arg, hint in hints.items(): if arg.startswith('%unknown%'): continue if hint is None: continue if not typehints.is_consistent_with( bindings.get(arg, typehints.Any), hint): raise typehints.TypeCheckError( 'Type hint violation for \'%s\': requires %s but got %s for %s' % (self.label, hint, bindings[arg], arg)) def process_argspec_fn(self): """Returns an argspec of the function actually consuming the data. """ raise NotImplementedError def make_fn(self, fn): # TODO(silviuc): Add comment describing that this is meant to be overriden # by methods detecting callables and wrapping them in DoFns. return fn def default_label(self): return '%s(%s)' % (self.__class__.__name__, self.fn.default_label()) class CallablePTransform(PTransform): """A class wrapper for a function-based transform.""" def __init__(self, fn): # pylint: disable=super-init-not-called # This is a helper class for a function decorator. Only when the class # is called (and __call__ invoked) we will have all the information # needed to initialize the super class. self.fn = fn def __call__(self, *args, **kwargs): if args and args[0] is None: label, self._args = None, args[1:] elif args and isinstance(args[0], str): label, self._args = args[0], args[1:] else: label, self._args = None, args self._kwargs = kwargs # We know the label now, so initialize the super-class. super(CallablePTransform, self).__init__(label=label) return self def apply(self, pcoll): # Since the PTransform will be implemented entirely as a function # (once called), we need to pass through any type-hinting information that # may have been annotated via the .with_input_types() and # .with_output_types() methods. kwargs = dict(self._kwargs) args = tuple(self._args) try: if 'type_hints' in inspect.getargspec(self.fn).args: args = (self.get_type_hints(),) + args except TypeError: # Might not be a function. pass return self.fn(self.label, pcoll, *args, **kwargs) def default_label(self): if self._args: return '%s(%s)' % ( label_from_callable(self.fn), label_from_callable(self._args[0])) else: return label_from_callable(self.fn) def ptransform_fn(fn): """A decorator for a function-based PTransform. Args: fn: A function implementing a custom PTransform. Returns: A CallablePTransform instance wrapping the function-based PTransform. This wrapper provides an alternative, simpler way to define a PTransform. The standard method is to subclass from PTransform and override the apply() method. An equivalent effect can be obtained by defining a function that takes a label, an input PCollection and additional optional arguments and returns a resulting PCollection. For example:: @ptransform_fn def CustomMapper(label, pcoll, mapfn): return pcoll | ParDo(mapfn) The equivalent approach using PTransform subclassing:: class CustomMapper(PTransform): def apply(self, pcoll, mapfn): return pcoll | ParDo(mapfn) With either method the custom PTransform can be used in pipelines as if it were one of the "native" PTransforms:: result_pcoll = input_pcoll | CustomMapper('label', somefn) Note that for both solutions the underlying implementation of the pipe operator (i.e., `|`) will inject the pcoll argument in its proper place (first argument if no label was specified and second argument otherwise). """ return CallablePTransform(fn) def format_full_label(applied_transform, pending_transform): """Returns a fully formatted cumulative PTransform label. Args: applied_transform: An instance of an AppliedPTransform that has been fully applied prior to 'pending_transform'. pending_transform: An instance of PTransform that has yet to be applied to the Pipeline. Returns: A fully formatted PTransform label. Example: '/foo/bar/baz'. """ label = '/'.join([applied_transform.full_label, pending_transform.label]) # Remove leading backslash because the monitoring UI expects names that do not # start with such a character. return label if not label.startswith('/') else label[1:] def label_from_callable(fn): if hasattr(fn, 'default_label'): return fn.default_label() elif hasattr(fn, '__name__'): if fn.__name__ == '': return '' % ( os.path.basename(fn.func_code.co_filename), fn.func_code.co_firstlineno) else: return fn.__name__ else: return str(fn) DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/ptransform_test.py000066400000000000000000002145431272340251000272150ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the PTransform and descendants.""" from __future__ import absolute_import import operator import re import unittest import google.cloud.dataflow as df from google.cloud.dataflow.pipeline import Pipeline import google.cloud.dataflow.pvalue as pvalue import google.cloud.dataflow.transforms.combiners as combine from google.cloud.dataflow.transforms.ptransform import PTransform from google.cloud.dataflow.transforms.util import assert_that, equal_to import google.cloud.dataflow.typehints as typehints from google.cloud.dataflow.typehints import with_input_types from google.cloud.dataflow.typehints import with_output_types from google.cloud.dataflow.typehints.typehints_test import TypeHintTestCase from google.cloud.dataflow.utils.options import PipelineOptions from google.cloud.dataflow.utils.options import TypeOptions # Disable frequent lint warning due to pipe operator for chaining transforms. # pylint: disable=expression-not-assigned class PTransformTest(unittest.TestCase): def assertStartswith(self, msg, prefix): self.assertTrue(msg.startswith(prefix), '"%s" does not start with "%s"' % (msg, prefix)) def test_str(self): self.assertEqual('', str(PTransform())) pa = Pipeline('DirectPipelineRunner') res = pa | df.Create('a_label', [1, 2]) self.assertEqual('', str(res.producer.transform)) pc = Pipeline('DirectPipelineRunner') res = pc | df.Create('with_inputs', [1, 2]) inputs_tr = res.producer.transform inputs_tr.inputs = ('ci',) self.assertEqual( """""", str(inputs_tr)) pd = Pipeline('DirectPipelineRunner') res = pd | df.Create('with_sidei', [1, 2]) side_tr = res.producer.transform side_tr.side_inputs = (4,) self.assertEqual( '', str(side_tr)) inputs_tr.side_inputs = ('cs',) self.assertEqual( """""", str(inputs_tr)) def test_parse_label_and_arg(self): def fun(*args, **kwargs): return PTransform().parse_label_and_arg(args, kwargs, 'name') self.assertEqual(('PTransform', 'value'), fun('value')) self.assertEqual(('PTransform', 'value'), fun(name='value')) self.assertEqual(('label', 'value'), fun('label', 'value')) self.assertEqual(('label', 'value'), fun('label', name='value')) self.assertEqual(('label', 'value'), fun('value', label='label')) self.assertEqual(('label', 'value'), fun(name='value', label='label')) self.assertRaises(ValueError, fun) self.assertRaises(ValueError, fun, 0, 'value') self.assertRaises(ValueError, fun, label=0, name='value') self.assertRaises(ValueError, fun, other='value') with self.assertRaises(ValueError) as cm: fun(0, name='value') self.assertEqual( cm.exception.message, 'PTransform expects a (label, name) or (name) argument list ' 'instead of args=(0,), kwargs={\'name\': \'value\'}') def test_do_with_do_fn(self): class AddNDoFn(df.DoFn): def process(self, context, addon): return [context.element + addon] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [1, 2, 3]) result = pcoll | df.ParDo('do', AddNDoFn(), 10) assert_that(result, equal_to([11, 12, 13])) pipeline.run() def test_do_with_unconstructed_do_fn(self): class MyDoFn(df.DoFn): def process(self, context): pass pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [1, 2, 3]) with self.assertRaises(ValueError): pcoll | df.ParDo('do', MyDoFn) # Note the lack of ()'s def test_do_with_callable(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [1, 2, 3]) result = pcoll | df.FlatMap('do', lambda x, addon: [x + addon], 10) assert_that(result, equal_to([11, 12, 13])) pipeline.run() def test_do_with_side_input_as_arg(self): pipeline = Pipeline('DirectPipelineRunner') side = pipeline | df.Create('side', [10]) pcoll = pipeline | df.Create('start', [1, 2, 3]) result = pcoll | df.FlatMap( 'do', lambda x, addon: [x + addon], pvalue.AsSingleton(side)) assert_that(result, equal_to([11, 12, 13])) pipeline.run() def test_do_with_side_input_as_keyword_arg(self): pipeline = Pipeline('DirectPipelineRunner') side = pipeline | df.Create('side', [10]) pcoll = pipeline | df.Create('start', [1, 2, 3]) result = pcoll | df.FlatMap( 'do', lambda x, addon: [x + addon], addon=pvalue.AsSingleton(side)) assert_that(result, equal_to([11, 12, 13])) pipeline.run() def test_do_with_do_fn_returning_string_raises_warning(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', ['2', '9', '3']) pcoll | df.FlatMap('do', lambda x: x + '1') # Since the DoFn directly returns a string we should get an error warning # us. with self.assertRaises(typehints.TypeCheckError) as cm: pipeline.run() expected_error_prefix = ('Returning a str from a ParDo or FlatMap ' 'is discouraged.') self.assertStartswith(cm.exception.message, expected_error_prefix) def test_do_with_do_fn_returning_dict_raises_warning(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', ['2', '9', '3']) pcoll | df.FlatMap('do', lambda x: {x: '1'}) # Since the DoFn directly returns a dict we should get an error warning # us. with self.assertRaises(typehints.TypeCheckError) as cm: pipeline.run() expected_error_prefix = ('Returning a dict from a ParDo or FlatMap ' 'is discouraged.') self.assertStartswith(cm.exception.message, expected_error_prefix) def test_do_with_side_outputs_maintains_unique_name(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [1, 2, 3]) r1 = pcoll | df.FlatMap('a', lambda x: [x + 1]).with_outputs(main='m') r2 = pcoll | df.FlatMap('b', lambda x: [x + 2]).with_outputs(main='m') assert_that(r1.m, equal_to([2, 3, 4]), label='r1') assert_that(r2.m, equal_to([3, 4, 5]), label='r2') pipeline.run() def test_do_requires_do_fn_returning_iterable(self): # This function is incorrect because it returns an object that isn't an # iterable. def incorrect_par_do_fn(x): return x + 5 pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [2, 9, 3]) pcoll | df.FlatMap('do', incorrect_par_do_fn) # It's a requirement that all user-defined functions to a ParDo return # an iterable. with self.assertRaises(typehints.TypeCheckError) as cm: pipeline.run() expected_error_prefix = 'FlatMap and ParDo must return an iterable.' self.assertStartswith(cm.exception.message, expected_error_prefix) def test_do_fn_with_start_finish(self): class MyDoFn(df.DoFn): def start_bundle(self, c): yield 'start' def process(self, c): pass def finish_bundle(self, c): yield 'finish' pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [1, 2, 3]) result = pcoll | df.ParDo('do', MyDoFn()) # May have many bundles, but each has a start and finish. def matcher(): def match(actual): equal_to(['start', 'finish'])(list(set(actual))) equal_to([actual.count('start')])([actual.count('finish')]) return match assert_that(result, matcher()) pipeline.run() def test_filter(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [1, 2, 3, 4]) result = pcoll | df.Filter( 'filter', lambda x: x % 2 == 0) assert_that(result, equal_to([2, 4])) pipeline.run() class _MeanCombineFn(df.CombineFn): def create_accumulator(self): return (0, 0) def add_input(self, (sum_, count), element): return sum_ + element, count + 1 def merge_accumulators(self, accumulators): sums, counts = zip(*accumulators) return sum(sums), sum(counts) def extract_output(self, (sum_, count)): if not count: return float('nan') return sum_ / float(count) def test_combine_with_combine_fn(self): vals = [1, 2, 3, 4, 5, 6, 7] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', vals) result = pcoll | df.CombineGlobally('mean', self._MeanCombineFn()) assert_that(result, equal_to([sum(vals) / len(vals)])) pipeline.run() def test_combine_with_callable(self): vals = [1, 2, 3, 4, 5, 6, 7] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', vals) result = pcoll | df.CombineGlobally(sum) assert_that(result, equal_to([sum(vals)])) pipeline.run() def test_combine_with_side_input_as_arg(self): values = [1, 2, 3, 4, 5, 6, 7] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', values) divisor = pipeline | df.Create('divisor', [2]) result = pcoll | df.CombineGlobally( 'max', # Multiples of divisor only. lambda vals, d: max(v for v in vals if v % d == 0), pvalue.AsSingleton(divisor)).without_defaults() filt_vals = [v for v in values if v % 2 == 0] assert_that(result, equal_to([max(filt_vals)])) pipeline.run() def test_combine_per_key_with_combine_fn(self): vals_1 = [1, 2, 3, 4, 5, 6, 7] vals_2 = [2, 4, 6, 8, 10, 12, 14] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', ([('a', x) for x in vals_1] + [('b', x) for x in vals_2])) result = pcoll | df.CombinePerKey('mean', self._MeanCombineFn()) assert_that(result, equal_to([('a', sum(vals_1) / len(vals_1)), ('b', sum(vals_2) / len(vals_2))])) pipeline.run() def test_combine_per_key_with_callable(self): vals_1 = [1, 2, 3, 4, 5, 6, 7] vals_2 = [2, 4, 6, 8, 10, 12, 14] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', ([('a', x) for x in vals_1] + [('b', x) for x in vals_2])) result = pcoll | df.CombinePerKey(sum) assert_that(result, equal_to([('a', sum(vals_1)), ('b', sum(vals_2))])) pipeline.run() def test_combine_per_key_with_side_input_as_arg(self): vals_1 = [1, 2, 3, 4, 5, 6, 7] vals_2 = [2, 4, 6, 8, 10, 12, 14] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', ([('a', x) for x in vals_1] + [('b', x) for x in vals_2])) divisor = pipeline | df.Create('divisor', [2]) result = pcoll | df.CombinePerKey( lambda vals, d: max(v for v in vals if v % d == 0), pvalue.AsSingleton(divisor)) # Multiples of divisor only. m_1 = max(v for v in vals_1 if v % 2 == 0) m_2 = max(v for v in vals_2 if v % 2 == 0) assert_that(result, equal_to([('a', m_1), ('b', m_2)])) pipeline.run() def test_group_by_key(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create( 'start', [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]) result = pcoll | df.GroupByKey('group') assert_that(result, equal_to([(1, [1, 2, 3]), (2, [1, 2]), (3, [1])])) pipeline.run() def test_partition_with_partition_fn(self): class SomePartitionFn(df.PartitionFn): def partition_for(self, context, num_partitions, offset): return (context.element % 3) + offset pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [0, 1, 2, 3, 4, 5, 6, 7, 8]) # Attempt nominal partition operation. partitions = pcoll | df.Partition('part1', SomePartitionFn(), 4, 1) assert_that(partitions[0], equal_to([])) assert_that(partitions[1], equal_to([0, 3, 6]), label='p1') assert_that(partitions[2], equal_to([1, 4, 7]), label='p2') assert_that(partitions[3], equal_to([2, 5, 8]), label='p3') pipeline.run() # Check that a bad partition label will yield an error. For the # DirectPipelineRunner, this error manifests as an exception. pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [0, 1, 2, 3, 4, 5, 6, 7, 8]) partitions = pcoll | df.Partition('part2', SomePartitionFn(), 4, 10000) with self.assertRaises(ValueError): pipeline.run() def test_partition_with_callable(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', [0, 1, 2, 3, 4, 5, 6, 7, 8]) partitions = ( pcoll | df.Partition( 'part', lambda e, n, offset: (e % 3) + offset, 4, 1)) assert_that(partitions[0], equal_to([])) assert_that(partitions[1], equal_to([0, 3, 6]), label='p1') assert_that(partitions[2], equal_to([1, 4, 7]), label='p2') assert_that(partitions[3], equal_to([2, 5, 8]), label='p3') pipeline.run() def test_partition_followed_by_flatten_and_groupbykey(self): """Regression test for an issue with how partitions are handled.""" pipeline = Pipeline('DirectPipelineRunner') contents = [('aa', 1), ('bb', 2), ('aa', 2)] created = pipeline | df.Create('A', contents) partitioned = created | df.Partition('B', lambda x, n: len(x) % n, 3) flattened = partitioned | df.Flatten('C') grouped = flattened | df.GroupByKey('D') assert_that(grouped, equal_to([('aa', [1, 2]), ('bb', [2])])) pipeline.run() def test_flatten_pcollections(self): pipeline = Pipeline('DirectPipelineRunner') pcoll_1 = pipeline | df.Create('start_1', [0, 1, 2, 3]) pcoll_2 = pipeline | df.Create('start_2', [4, 5, 6, 7]) result = (pcoll_1, pcoll_2) | df.Flatten('flatten') assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7])) pipeline.run() def test_flatten_no_pcollections(self): pipeline = Pipeline('DirectPipelineRunner') with self.assertRaises(ValueError): () | df.Flatten('pipeline arg missing') result = () | df.Flatten('empty', pipeline=pipeline) assert_that(result, equal_to([])) pipeline.run() def test_flatten_pcollections_in_iterable(self): pipeline = Pipeline('DirectPipelineRunner') pcoll_1 = pipeline | df.Create('start_1', [0, 1, 2, 3]) pcoll_2 = pipeline | df.Create('start_2', [4, 5, 6, 7]) result = ([pcoll for pcoll in (pcoll_1, pcoll_2)] | df.Flatten('flatten')) assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7])) pipeline.run() def test_flatten_input_type_must_be_iterable(self): # Inputs to flatten *must* be an iterable. with self.assertRaises(ValueError): 4 | df.Flatten('flatten') def test_flatten_input_type_must_be_iterable_of_pcolls(self): # Inputs to flatten *must* be an iterable of PCollections. with self.assertRaises(TypeError): {'l': 'test'} | df.Flatten('flatten') with self.assertRaises(TypeError): set([1, 2, 3]) | df.Flatten('flatten') def test_co_group_by_key_on_list(self): pipeline = Pipeline('DirectPipelineRunner') pcoll_1 = pipeline | df.Create( 'start_1', [('a', 1), ('a', 2), ('b', 3), ('c', 4)]) pcoll_2 = pipeline | df.Create( 'start_2', [('a', 5), ('a', 6), ('c', 7), ('c', 8)]) result = (pcoll_1, pcoll_2) | df.CoGroupByKey('cgbk') assert_that(result, equal_to([('a', ([1, 2], [5, 6])), ('b', ([3], [])), ('c', ([4], [7, 8]))])) pipeline.run() def test_co_group_by_key_on_iterable(self): pipeline = Pipeline('DirectPipelineRunner') pcoll_1 = pipeline | df.Create( 'start_1', [('a', 1), ('a', 2), ('b', 3), ('c', 4)]) pcoll_2 = pipeline | df.Create( 'start_2', [('a', 5), ('a', 6), ('c', 7), ('c', 8)]) result = ([pc for pc in (pcoll_1, pcoll_2)] | df.CoGroupByKey('cgbk')) assert_that(result, equal_to([('a', ([1, 2], [5, 6])), ('b', ([3], [])), ('c', ([4], [7, 8]))])) pipeline.run() def test_co_group_by_key_on_dict(self): pipeline = Pipeline('DirectPipelineRunner') pcoll_1 = pipeline | df.Create( 'start_1', [('a', 1), ('a', 2), ('b', 3), ('c', 4)]) pcoll_2 = pipeline | df.Create( 'start_2', [('a', 5), ('a', 6), ('c', 7), ('c', 8)]) result = {'X': pcoll_1, 'Y': pcoll_2} | df.CoGroupByKey('cgbk') assert_that(result, equal_to([('a', {'X': [1, 2], 'Y': [5, 6]}), ('b', {'X': [3], 'Y': []}), ('c', {'X': [4], 'Y': [7, 8]})])) pipeline.run() def test_group_by_key_input_must_be_kv_pairs(self): pipeline = Pipeline('DirectPipelineRunner') pcolls = pipeline | df.Create('A', [1, 2, 3, 4, 5]) with self.assertRaises(typehints.TypeCheckError) as e: pcolls | df.GroupByKey('D') pipeline.run() self.assertStartswith( e.exception.message, 'Runtime type violation detected within ' 'ParDo(D/reify_windows): Input to GroupByKey must be ' 'a PCollection with elements compatible with KV[A, B]') def test_group_by_key_only_input_must_be_kv_pairs(self): pipeline = Pipeline('DirectPipelineRunner') pcolls = pipeline | df.Create('A', ['a', 'b', 'f']) with self.assertRaises(typehints.TypeCheckError) as cm: pcolls | df.GroupByKeyOnly('D') pipeline.run() expected_error_prefix = ('Input to GroupByKeyOnly must be a PCollection of ' 'windowed key-value pairs.') self.assertStartswith(cm.exception.message, expected_error_prefix) def test_keys_and_values(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create( 'start', [(3, 1), (2, 1), (1, 1), (3, 2), (2, 2), (3, 3)]) keys = pcoll.apply('keys', df.Keys()) vals = pcoll.apply('vals', df.Values()) assert_that(keys, equal_to([1, 2, 2, 3, 3, 3]), label='assert:keys') assert_that(vals, equal_to([1, 1, 1, 2, 2, 3]), label='assert:vals') pipeline.run() def test_kv_swap(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create( 'start', [(6, 3), (1, 2), (7, 1), (5, 2), (3, 2)]) result = pcoll.apply('swap', df.KvSwap()) assert_that(result, equal_to([(1, 7), (2, 1), (2, 3), (2, 5), (3, 6)])) pipeline.run() def test_remove_duplicates(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create( 'start', [6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel']) result = pcoll.apply('nodupes', df.RemoveDuplicates()) assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel'])) pipeline.run() def test_chained_ptransforms(self): pipeline = Pipeline('DirectPipelineRunner') t = (df.Map(lambda x: (x, 1)) | df.GroupByKey() | df.Map(lambda (x, ones): (x, sum(ones)))) result = pipeline | df.Create('start', ['a', 'a', 'b']) | t assert_that(result, equal_to([('a', 2), ('b', 1)])) pipeline.run() def test_apply_to_list(self): self.assertEqual([1, 2, 3], [0, 1, 2] | df.Map('add_one', lambda x: x + 1)) self.assertEqual([1], [0, 1, 2] | df.Filter('odd', lambda x: x % 2)) self.assertEqual([1, 2, 3, 100], ([1, 2, 3], [100]) | df.Flatten('flat')) join_input = ([('k', 'a')], [('k', 'b'), ('k', 'c')]) self.assertEqual([('k', (['a'], ['b', 'c']))], join_input | df.CoGroupByKey('join')) def test_multi_input_ptransform(self): class DisjointUnion(PTransform): def apply(self, pcollections): return (pcollections | df.Flatten() | df.Map(lambda x: (x, None)) | df.GroupByKey() | df.Map(lambda (x, _): x)) self.assertEqual([1, 2, 3], sorted(([1, 2], [2, 3]) | DisjointUnion())) def test_apply_to_crazy_pvaluish(self): class NestedFlatten(PTransform): """A PTransform taking and returning nested PValueish. Takes as input a list of dicts, and returns a dict with the corresponding values flattened. """ def _extract_input_pvalues(self, pvalueish): pvalueish = list(pvalueish) return pvalueish, sum([list(p.values()) for p in pvalueish], []) def apply(self, pcoll_dicts): keys = reduce(operator.or_, [set(p.keys()) for p in pcoll_dicts]) res = {} for k in keys: res[k] = [p[k] for p in pcoll_dicts if k in p] | df.Flatten(k) return res res = [{'a': [1, 2, 3]}, {'a': [4, 5, 6], 'b': ['x', 'y', 'z']}, {'a': [7, 8], 'b': ['x', 'y'], 'c': []}] | NestedFlatten() self.assertEqual(3, len(res)) self.assertEqual([1, 2, 3, 4, 5, 6, 7, 8], sorted(res['a'])) self.assertEqual(['x', 'x', 'y', 'y', 'z'], sorted(res['b'])) self.assertEqual([], sorted(res['c'])) @df.ptransform_fn def SamplePTransform(label, pcoll, context, *args, **kwargs): """Sample transform using the @ptransform_fn decorator.""" _ = label, args, kwargs map_transform = df.Map('ToPairs', lambda v: (v, None)) combine_transform = df.CombinePerKey('Group', lambda vs: None) keys_transform = df.Keys('RemoveDuplicates') context.extend([map_transform, combine_transform, keys_transform]) return pcoll | map_transform | combine_transform | keys_transform class PTransformLabelsTest(unittest.TestCase): class CustomTransform(df.PTransform): pardo = None def apply(self, pcoll): self.pardo = df.FlatMap('*do*', lambda x: [x + 1]) return pcoll | self.pardo def test_chained_ptransforms(self): """Tests that chaining gets proper nesting.""" pipeline = Pipeline('DirectPipelineRunner') map1 = df.Map('map1', lambda x: (x, 1)) gbk = df.GroupByKey('gbk') map2 = df.Map('map2', lambda (x, ones): (x, sum(ones))) t = (map1 | gbk | map2) result = pipeline | df.Create('start', ['a', 'a', 'b']) | t self.assertTrue('map1|gbk|map2/map1' in pipeline.applied_labels) self.assertTrue('map1|gbk|map2/gbk' in pipeline.applied_labels) self.assertTrue('map1|gbk|map2/map2' in pipeline.applied_labels) assert_that(result, equal_to([('a', 2), ('b', 1)])) pipeline.run() def test_apply_custom_transform_without_label(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('pcoll', [1, 2, 3]) custom = PTransformLabelsTest.CustomTransform() result = pipeline.apply(custom, pcoll) self.assertTrue('CustomTransform' in pipeline.applied_labels) self.assertTrue('CustomTransform/*do*' in pipeline.applied_labels) assert_that(result, equal_to([2, 3, 4])) pipeline.run() def test_apply_custom_transform_with_label(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('pcoll', [1, 2, 3]) custom = PTransformLabelsTest.CustomTransform('*custom*') result = pipeline.apply(custom, pcoll) self.assertTrue('*custom*' in pipeline.applied_labels) self.assertTrue('*custom*/*do*' in pipeline.applied_labels) assert_that(result, equal_to([2, 3, 4])) pipeline.run() def test_combine_without_label(self): vals = [1, 2, 3, 4, 5, 6, 7] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', vals) combine = df.CombineGlobally(sum) result = pcoll | combine self.assertTrue('CombineGlobally(sum)' in pipeline.applied_labels) assert_that(result, equal_to([sum(vals)])) pipeline.run() def test_apply_ptransform_using_decorator(self): pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('pcoll', [1, 2, 3]) context = [] sample = SamplePTransform('*sample*', context) _ = pcoll | sample self.assertTrue('*sample*' in pipeline.applied_labels) self.assertTrue('*sample*/ToPairs' in pipeline.applied_labels) self.assertTrue('*sample*/Group' in pipeline.applied_labels) self.assertTrue('*sample*/RemoveDuplicates' in pipeline.applied_labels) def test_combine_with_label(self): vals = [1, 2, 3, 4, 5, 6, 7] pipeline = Pipeline('DirectPipelineRunner') pcoll = pipeline | df.Create('start', vals) combine = df.CombineGlobally('*sum*', sum) result = pcoll | combine self.assertTrue('*sum*' in pipeline.applied_labels) assert_that(result, equal_to([sum(vals)])) pipeline.run() def check_label(self, ptransform, expected_label): pipeline = Pipeline('DirectPipelineRunner') pipeline | df.Create('start', [('a', 1)]) | ptransform actual_label = sorted(pipeline.applied_labels - {'start'})[0] self.assertEqual(expected_label, re.sub(r'\d{3,}', '#', actual_label)) def test_default_labels(self): self.check_label(df.Map(len), r'Map(len)') self.check_label(df.Map(lambda x: x), r'Map()') self.check_label(df.FlatMap(list), r'FlatMap(list)') self.check_label(df.Filter(sum), r'Filter(sum)') self.check_label(df.CombineGlobally(sum), r'CombineGlobally(sum)') self.check_label(df.CombinePerKey(sum), r'CombinePerKey(sum)') class MyDoFn(df.DoFn): def process(self, context): pass self.check_label(df.ParDo(MyDoFn()), r'ParDo(MyDoFn)') class PTransformTypeCheckTestCase(TypeHintTestCase): def assertStartswith(self, msg, prefix): self.assertTrue(msg.startswith(prefix), '"%s" does not start with "%s"' % (msg, prefix)) def setUp(self): self.p = Pipeline(options=PipelineOptions([])) def test_do_fn_pipeline_pipeline_type_check_satisfied(self): @with_input_types(int, int) @with_output_types(typehints.List[int]) class AddWithFive(df.DoFn): def process(self, context, five): return [context.element + five] d = (self.p | df.Create('t', [1, 2, 3]).with_output_types(int) | df.ParDo('add', AddWithFive(), 5)) assert_that(d, equal_to([6, 7, 8])) self.p.run() def test_do_fn_pipeline_pipeline_type_check_violated(self): @with_input_types(str, str) @with_output_types(typehints.List[str]) class ToUpperCaseWithPrefix(df.DoFn): def process(self, context, prefix): return [prefix + context.element.upper()] with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('t', [1, 2, 3]).with_output_types(int) | df.ParDo('upper', ToUpperCaseWithPrefix(), 'hello')) self.assertEqual("Type hint violation for 'upper': " "requires but got for context", e.exception.message) def test_do_fn_pipeline_runtime_type_check_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True @with_input_types(int, int) @with_output_types(int) class AddWithNum(df.DoFn): def process(self, context, num): return [context.element + num] d = (self.p | df.Create('t', [1, 2, 3]).with_output_types(int) | df.ParDo('add', AddWithNum(), 5)) assert_that(d, equal_to([6, 7, 8])) self.p.run() def test_do_fn_pipeline_runtime_type_check_violated(self): self.p.options.view_as(TypeOptions).runtime_type_check = True @with_input_types(int, int) @with_output_types(typehints.List[int]) class AddWithNum(df.DoFn): def process(self, context, num): return [context.element + num] with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('t', ['1', '2', '3']).with_output_types(str) | df.ParDo('add', AddWithNum(), 5)) self.p.run() self.assertEqual("Type hint violation for 'add': " "requires but got for context", e.exception.message) def test_pardo_does_not_type_check_using_type_hint_decorators(self): @with_input_types(a=int) @with_output_types(typehints.List[str]) def int_to_str(a): return [str(a)] # The function above is expecting an int for its only parameter. However, it # will receive a str instead, which should result in a raised exception. with self.assertRaises(typehints.TypeCheckError) as e: c = (self.p | df.Create('s', ['b', 'a', 'r']).with_output_types(str) | df.FlatMap('to str', int_to_str)) self.assertEqual("Type hint violation for 'to str': " "requires but got for a", e.exception.message) def test_pardo_properly_type_checks_using_type_hint_decorators(self): @with_input_types(a=str) @with_output_types(typehints.List[str]) def to_all_upper_case(a): return [a.upper()] # If this type-checks than no error should be raised. d = (self.p | df.Create('t', ['t', 'e', 's', 't']).with_output_types(str) | df.FlatMap('case', to_all_upper_case)) assert_that(d, equal_to(['T', 'E', 'S', 'T'])) self.p.run() # Output type should have been recognized as 'str' rather than List[str] to # do the flatten part of FlatMap. self.assertEqual(str, d.element_type) def test_pardo_does_not_type_check_using_type_hint_methods(self): # The first ParDo outputs pcoll's of type int, however the second ParDo is # expecting pcoll's of type str instead. with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('s', ['t', 'e', 's', 't']).with_output_types(str) | (df.FlatMap('score', lambda x: [1] if x == 't' else [2]) .with_input_types(str).with_output_types(int)) | (df.FlatMap('upper', lambda x: [x.upper()]) .with_input_types(str).with_output_types(str))) self.assertEqual("Type hint violation for 'upper': " "requires but got for x", e.exception.message) def test_pardo_properly_type_checks_using_type_hint_methods(self): # Pipeline should be created successfully without an error d = (self.p | df.Create('s', ['t', 'e', 's', 't']).with_output_types(str) | df.FlatMap('dup', lambda x: [x + x]) .with_input_types(str).with_output_types(str) | df.FlatMap('upper', lambda x: [x.upper()]) .with_input_types(str).with_output_types(str)) assert_that(d, equal_to(['TT', 'EE', 'SS', 'TT'])) self.p.run() def test_map_does_not_type_check_using_type_hints_methods(self): # The transform before 'Map' has indicated that it outputs PCollections with # int's, while Map is expecting one of str. with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('s', [1, 2, 3, 4]).with_output_types(int) | df.Map('upper', lambda x: x.upper()).with_input_types(str).with_output_types(str)) self.assertEqual("Type hint violation for 'upper': " "requires but got for x", e.exception.message) def test_map_properly_type_checks_using_type_hints_methods(self): # No error should be raised if this type-checks properly. d = (self.p | df.Create('s', [1, 2, 3, 4]).with_output_types(int) | df.Map('to_str', lambda x: str(x)).with_input_types(int).with_output_types(str)) assert_that(d, equal_to(['1', '2', '3', '4'])) self.p.run() def test_map_does_not_type_check_using_type_hints_decorator(self): @with_input_types(s=str) @with_output_types(str) def upper(s): return s.upper() # Hinted function above expects a str at pipeline construction. # However, 'Map' should detect that Create has hinted an int instead. with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('s', [1, 2, 3, 4]).with_output_types(int) | df.Map('upper', upper)) self.assertEqual("Type hint violation for 'upper': " "requires but got for s", e.exception.message) def test_map_properly_type_checks_using_type_hints_decorator(self): @with_input_types(a=bool) @with_output_types(int) def bool_to_int(a): return int(a) # If this type-checks than no error should be raised. d = (self.p | df.Create('bools', [True, False, True]).with_output_types(bool) | df.Map('to_ints', bool_to_int)) assert_that(d, equal_to([1, 0, 1])) self.p.run() def test_filter_does_not_type_check_using_type_hints_method(self): # Filter is expecting an int but instead looks to the 'left' and sees a str # incoming. with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('strs', ['1', '2', '3', '4', '5']).with_output_types(str) | df.Map('lower', lambda x: x.lower()).with_input_types(str).with_output_types(str) | df.Filter('below 3', lambda x: x < 3).with_input_types(int)) self.assertEqual("Type hint violation for 'below 3': " "requires but got for x", e.exception.message) def test_filter_type_checks_using_type_hints_method(self): # No error should be raised if this type-checks properly. d = (self.p | df.Create('strs', ['1', '2', '3', '4', '5']).with_output_types(str) | df.Map('to int', lambda x: int(x)).with_input_types(str).with_output_types(int) | df.Filter('below 3', lambda x: x < 3).with_input_types(int)) assert_that(d, equal_to([1, 2])) self.p.run() def test_filter_does_not_type_check_using_type_hints_decorator(self): @with_input_types(a=float) def more_than_half(a): return a > 0.50 # Func above was hinted to only take a float, yet an int will be passed. with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('ints', [1, 2, 3, 4]).with_output_types(int) | df.Filter('half', more_than_half)) self.assertEqual("Type hint violation for 'half': " "requires but got for a", e.exception.message) def test_filter_type_checks_using_type_hints_decorator(self): @with_input_types(b=int) def half(b): import random return bool(random.choice([0, 1])) # Filter should deduce that it returns the same type that it takes. (self.p | df.Create('str', range(5)).with_output_types(int) | df.Filter('half', half) | df.Map('to bool', lambda x: bool(x)).with_input_types(int).with_output_types(bool)) def test_group_by_key_only_output_type_deduction(self): d = (self.p | df.Create('str', ['t', 'e', 's', 't']).with_output_types(str) | (df.Map('pair', lambda x: (x, ord(x))) .with_output_types(typehints.KV[str, str])) | df.GroupByKeyOnly('O')) # Output type should correctly be deduced. # GBK-only should deduce that KV[A, B] is turned into KV[A, Iterable[B]]. self.assertCompatible(typehints.KV[str, typehints.Iterable[str]], d.element_type) def test_group_by_key_output_type_deduction(self): d = (self.p | df.Create('str', range(20)).with_output_types(int) | (df.Map('pair negative', lambda x: (x % 5, -x)) .with_output_types(typehints.KV[int, int])) | df.GroupByKey('T')) # Output type should correctly be deduced. # GBK should deduce that KV[A, B] is turned into KV[A, Iterable[B]]. self.assertCompatible(typehints.KV[int, typehints.Iterable[int]], d.element_type) def test_group_by_key_only_does_not_type_check(self): # GBK will be passed raw int's here instead of some form of KV[A, B]. with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('s', [1, 2, 3]).with_output_types(int) | df.GroupByKeyOnly('F')) self.assertEqual("Input type hint violation at F: " "expected Tuple[TypeVariable[K], TypeVariable[V]], " "got ", e.exception.message) def test_group_by_does_not_type_check(self): # Create is returning a List[int, str], rather than a KV[int, str] that is # aliased to Tuple[int, str]. with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | (df.Create('s', range(5)) .with_output_types(typehints.Iterable[int])) | df.GroupByKey('T')) self.assertEqual("Input type hint violation at T: " "expected Tuple[TypeVariable[K], TypeVariable[V]], " "got Iterable[int]", e.exception.message) def test_pipeline_checking_pardo_insufficient_type_information(self): self.p.options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED' # Type checking is enabled, but 'Create' doesn't pass on any relevant type # information to the ParDo. with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('nums', range(5)) | df.FlatMap('mod dup', lambda x: (x % 2, x))) self.assertEqual('Pipeline type checking is enabled, however no output ' 'type-hint was found for the PTransform Create(nums)', e.exception.message) def test_pipeline_checking_gbk_insufficient_type_information(self): self.p.options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED' # Type checking is enabled, but 'Map' doesn't pass on any relevant type # information to GBK-only. with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('nums', range(5)).with_output_types(int) | df.Map('mod dup', lambda x: (x % 2, x)) | df.GroupByKeyOnly('G')) self.assertEqual('Pipeline type checking is enabled, however no output ' 'type-hint was found for the PTransform ' 'ParDo(mod dup)', e.exception.message) def test_disable_pipeline_type_check(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False # The pipeline below should raise a TypeError, however pipeline type # checking was disabled above. (self.p | df.Create('t', [1, 2, 3]).with_output_types(int) | df.Map('lower', lambda x: x.lower()).with_input_types(str).with_output_types(str)) def test_run_time_type_checking_enabled_type_violation(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True @with_output_types(str) @with_input_types(x=int) def int_to_string(x): return str(x) # Function above has been type-hinted to only accept an int. But in the # pipeline execution it'll be passed a string due to the output of Create. (self.p | df.Create('t', ['some_string']) | df.Map('to str', int_to_string)) with self.assertRaises(typehints.TypeCheckError) as e: self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within ParDo(to str): " "Type-hint for argument: 'x' violated. " "Expected an instance of , " "instead found some_string, an instance of .") def test_run_time_type_checking_enabled_types_satisfied(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True @with_output_types(typehints.KV[int, str]) @with_input_types(x=str) def group_with_upper_ord(x): return (ord(x.upper()) % 5, x) # Pipeline checking is off, but the above function should satisfy types at # run-time. result = (self.p | df.Create('t', ['t', 'e', 's', 't', 'i', 'n', 'g']).with_output_types(str) | df.Map('gen keys', group_with_upper_ord) | df.GroupByKey('O')) assert_that(result, equal_to([(1, ['g']), (3, ['s', 'i', 'n']), (4, ['t', 'e', 't'])])) self.p.run() def test_pipeline_checking_satisfied_but_run_time_types_violate(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True @with_output_types(typehints.KV[bool, int]) @with_input_types(a=int) def is_even_as_key(a): # Simulate a programming error, should be: return (a % 2 == 0, a) # However this returns KV[int, int] return (a % 2, a) (self.p | df.Create('nums', range(5)).with_output_types(int) | df.Map('is even', is_even_as_key) | df.GroupByKey('parity')) # Although all the types appear to be correct when checked at pipeline # construction. Runtime type-checking should detect the 'is_even_as_key' is # returning Tuple[int, int], instead of Tuple[bool, int]. with self.assertRaises(typehints.TypeCheckError) as e: self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within ParDo(is even): " "Tuple[bool, int] hint type-constraint violated. " "The type of element #0 in the passed tuple is incorrect. " "Expected an instance of type bool, " "instead received an instance of type int.") def test_pipeline_checking_satisfied_run_time_checking_satisfied(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False @with_output_types(typehints.KV[bool, int]) @with_input_types(a=int) def is_even_as_key(a): # The programming error in the above test-case has now been fixed. # Everything should properly type-check. return (a % 2 == 0, a) result = (self.p | df.Create('nums', range(5)).with_output_types(int) | df.Map('is even', is_even_as_key) | df.GroupByKey('parity')) assert_that(result, equal_to([(False, [1, 3]), (True, [0, 2, 4])])) self.p.run() def test_pipeline_runtime_checking_violation_simple_type_input(self): self.p.options.view_as(TypeOptions).runtime_type_check = True self.p.options.view_as(TypeOptions).pipeline_type_check = False # The type-hinted applied via the 'with_input_types()' method indicates the # ParDo should receive an instance of type 'str', however an 'int' will be # passed instead. with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('n', [1, 2, 3]) | (df.FlatMap('to int', lambda x: [int(x)]) .with_input_types(str).with_output_types(int)) ) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within ParDo(to int): " "Type-hint for argument: 'x' violated. " "Expected an instance of , " "instead found 1, an instance of .") def test_pipeline_runtime_checking_violation_composite_type_input(self): self.p.options.view_as(TypeOptions).runtime_type_check = True self.p.options.view_as(TypeOptions).pipeline_type_check = False with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('n', [(1, 3.0), (2, 4.9), (3, 9.5)]) | (df.FlatMap('add', lambda (x, y): [x + y]) .with_input_types(typehints.Tuple[int, int]).with_output_types(int)) ) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within ParDo(add): " "Type-hint for argument: 'y' violated. " "Expected an instance of , " "instead found 3.0, an instance of .") def test_pipeline_runtime_checking_violation_simple_type_output(self): self.p.options.view_as(TypeOptions).runtime_type_check = True self.p.options.view_as(TypeOptions).pipeline_type_check = False # The type-hinted applied via the 'returns()' method indicates the ParDo # should output an instance of type 'int', however a 'float' will be # generated instead. print "HINTS", df.FlatMap( 'to int', lambda x: [float(x)]).with_input_types(int).with_output_types( int).get_type_hints() with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('n', [1, 2, 3]) | (df.FlatMap('to int', lambda x: [float(x)]) .with_input_types(int).with_output_types(int)) ) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within " "ParDo(to int): " "According to type-hint expected output should be " "of type . Instead, received '1.0', " "an instance of type .") def test_pipeline_runtime_checking_violation_composite_type_output(self): self.p.options.view_as(TypeOptions).runtime_type_check = True self.p.options.view_as(TypeOptions).pipeline_type_check = False # The type-hinted applied via the 'returns()' method indicates the ParDo # should return an instance of type: Tuple[float, int]. However, an instance # of 'int' will be generated instead. with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('n', [(1, 3.0), (2, 4.9), (3, 9.5)]) | (df.FlatMap('swap', lambda (x, y): [x + y]) .with_input_types(typehints.Tuple[int, float]) .with_output_types(typehints.Tuple[float, int])) ) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within " "ParDo(swap): Tuple type constraint violated. " "Valid object instance must be of type 'tuple'. Instead, " "an instance of 'float' was received.") def test_pipline_runtime_checking_violation_with_side_inputs_decorator(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True @with_output_types(int) @with_input_types(a=int, b=int) def add(a, b): return a + b with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('t', [1, 2, 3, 4]) | df.Map('add 1', add, 1.0)) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within ParDo(add 1): " "Type-hint for argument: 'b' violated. " "Expected an instance of , " "instead found 1.0, an instance of .") def test_pipline_runtime_checking_violation_with_side_inputs_via_method(self): self.p.options.view_as(TypeOptions).runtime_type_check = True self.p.options.view_as(TypeOptions).pipeline_type_check = False with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('t', [1, 2, 3, 4]) | (df.Map('add 1', lambda x, one: x + one, 1.0) .with_input_types(int, int) .with_output_types(float))) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within ParDo(add 1): " "Type-hint for argument: 'one' violated. " "Expected an instance of , " "instead found 1.0, an instance of .") def test_combine_properly_pipeline_type_checks_using_decorator(self): @with_output_types(int) @with_input_types(ints=typehints.Iterable[int]) def sum_ints(ints): return sum(ints) d = (self.p | df.Create('t', [1, 2, 3]).with_output_types(int) | df.CombineGlobally('sum', sum_ints)) self.assertEqual(int, d.element_type) assert_that(d, equal_to([6])) self.p.run() def test_combine_func_type_hint_does_not_take_iterable_using_decorator(self): @with_output_types(int) @with_input_types(a=int) def bad_combine(a): 5 + a with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('m', [1, 2, 3]).with_output_types(int) | df.CombineGlobally('add', bad_combine)) self.assertEqual( "All functions for a Combine PTransform must accept a " "single argument compatible with: Iterable[Any]. " "Instead a function with input type: was received.", e.exception.message) def test_combine_pipeline_type_propagation_using_decorators(self): @with_output_types(int) @with_input_types(ints=typehints.Iterable[int]) def sum_ints(ints): return sum(ints) @with_output_types(typehints.List[int]) @with_input_types(n=int) def range_from_zero(n): return list(range(n+1)) d = (self.p | df.Create('t', [1, 2, 3]).with_output_types(int) | df.CombineGlobally('sum', sum_ints) | df.ParDo('range', range_from_zero)) self.assertEqual(int, d.element_type) assert_that(d, equal_to([0, 1, 2, 3, 4, 5, 6])) self.p.run() def test_combine_runtime_type_check_satisfied_using_decorators(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False @with_output_types(int) @with_input_types(ints=typehints.Iterable[int]) def iter_mul(ints): return reduce(operator.mul, ints, 1) d = (self.p | df.Create('k', [5, 5, 5, 5]).with_output_types(int) | df.CombineGlobally('mul', iter_mul)) assert_that(d, equal_to([625])) self.p.run() def test_combine_runtime_type_check_violation_using_decorators(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True # Combine fn is returning the incorrect type @with_output_types(int) @with_input_types(ints=typehints.Iterable[int]) def iter_mul(ints): return str(reduce(operator.mul, ints, 1)) with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('k', [5, 5, 5, 5]).with_output_types(int) | df.CombineGlobally('mul', iter_mul)) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within " "ParDo(mul/CombinePerKey/Combine/ParDo(CombineValuesDoFn)): " "Tuple[TypeVariable[K], int] hint type-constraint violated. " "The type of element #1 in the passed tuple is incorrect. " "Expected an instance of type int, " "instead received an instance of type str.") def test_combine_pipeline_type_check_using_methods(self): d = (self.p | df.Create('s', ['t', 'e', 's', 't']).with_output_types(str) | (df.CombineGlobally('concat', lambda s: ''.join(s)) .with_input_types(str).with_output_types(str))) def matcher(expected): def match(actual): equal_to(expected)(list(actual[0])) return match assert_that(d, matcher('estt')) self.p.run() def test_combine_runtime_type_check_using_methods(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('s', range(5)).with_output_types(int) | (df.CombineGlobally('sum', lambda s: sum(s)) .with_input_types(int).with_output_types(int))) assert_that(d, equal_to([10])) self.p.run() def test_combine_pipeline_type_check_violation_using_methods(self): with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('e', range(3)).with_output_types(int) | (df.CombineGlobally('sort join', lambda s: ''.join(sorted(s))) .with_input_types(str).with_output_types(str))) self.assertEqual("Input type hint violation at sort join: " "expected , got ", e.exception.message) def test_combine_runtime_type_check_violation_using_methods(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('e', range(3)).with_output_types(int) | (df.CombineGlobally('sort join', lambda s: ''.join(sorted(s))) .with_input_types(str).with_output_types(str))) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within " "ParDo(sort join/KeyWithVoid): " "Type-hint for argument: 'v' violated. " "Expected an instance of , " "instead found 0, an instance of .") def test_combine_insufficient_type_hint_information(self): self.p.options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED' with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('e', range(3)).with_output_types(int) | df.CombineGlobally('sort join', lambda s: ''.join(sorted(s))) | df.Map('f', lambda x: x + 1)) self.assertEqual( 'Pipeline type checking is enabled, ' 'however no output type-hint was found for the PTransform ' 'ParDo(sort join/CombinePerKey/Combine/ParDo(CombineValuesDoFn))', e.exception.message) def test_mean_globally_pipeline_checking_satisfied(self): d = (self.p | df.Create('c', range(5)).with_output_types(int) | combine.Mean.Globally('mean')) self.assertTrue(d.element_type is float) assert_that(d, equal_to([2.0])) self.p.run() def test_mean_globally_pipeline_checking_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('c', ['test']).with_output_types(str) | combine.Mean.Globally('mean')) self.assertEqual("Type hint violation for 'ParDo(CombineValuesDoFn)': " "requires Tuple[TypeVariable[K], " "Iterable[Union[float, int, long]]] " "but got Tuple[None, Iterable[str]] for p_context", e.exception.message) def test_mean_globally_runtime_checking_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('c', range(5)).with_output_types(int) | combine.Mean.Globally('mean')) self.assertTrue(d.element_type is float) assert_that(d, equal_to([2.0])) self.p.run() def test_mean_globally_runtime_checking_violated(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('c', ['t', 'e', 's', 't']).with_output_types(str) | combine.Mean.Globally('mean')) self.p.run() self.assertEqual("Runtime type violation detected for transform input " "when executing ParDoFlatMap(Combine): Tuple[Any, " "Iterable[Union[int, float]]] hint type-constraint " "violated. The type of element #1 in the passed tuple " "is incorrect. Iterable[Union[int, float]] hint " "type-constraint violated. The type of element #0 in " "the passed Iterable is incorrect: Union[int, float] " "type-constraint violated. Expected an instance of one " "of: ('int', 'float'), received str instead.", e.exception.message) def test_mean_per_key_pipeline_checking_satisfied(self): d = (self.p | df.Create('c', range(5)).with_output_types(int) | (df.Map('even group', lambda x: (not x % 2, x)) .with_output_types(typehints.KV[bool, int])) | combine.Mean.PerKey('even mean')) self.assertCompatible(typehints.KV[bool, float], d.element_type) assert_that(d, equal_to([(False, 2.0), (True, 2.0)])) self.p.run() def test_mean_per_key_pipeline_checking_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('e', map(str, range(5))).with_output_types(str) | (df.Map('upper pair', lambda x: (x.upper(), x)) .with_output_types(typehints.KV[str, str])) | combine.Mean.PerKey('even mean')) self.p.run() self.assertEqual("Type hint violation for 'ParDo(CombineValuesDoFn)': " "requires Tuple[TypeVariable[K], " "Iterable[Union[float, int, long]]] " "but got Tuple[str, Iterable[str]] for p_context", e.exception.message) def test_mean_per_key_runtime_checking_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('c', range(5)).with_output_types(int) | (df.Map('odd group', lambda x: (bool(x % 2), x)) .with_output_types(typehints.KV[bool, int])) | combine.Mean.PerKey('odd mean')) self.assertCompatible(typehints.KV[bool, float], d.element_type) assert_that(d, equal_to([(False, 2.0), (True, 2.0)])) self.p.run() def test_mean_per_key_runtime_checking_violated(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False self.p.options.view_as(TypeOptions).runtime_type_check = True with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('c', range(5)).with_output_types(int) | (df.Map('odd group', lambda x: (x, str(bool(x % 2)))) .with_output_types(typehints.KV[int, str])) | combine.Mean.PerKey('odd mean')) self.p.run() self.assertStartswith( e.exception.message, "Runtime type violation detected within " "ParDo(odd mean/CombinePerKey(MeanCombineFn)/" "Combine/ParDo(CombineValuesDoFn)): " "Type-hint for argument: 'p_context' violated: " "Tuple[TypeVariable[K], Iterable[Union[float, int, long]]]" " hint type-constraint violated. " "The type of element #1 in the passed tuple is incorrect. " "Iterable[Union[float, int, long]] " "hint type-constraint violated. The type of element #0 " "in the passed Iterable is incorrect: " "Union[float, int, long] type-constraint violated. " "Expected an instance of one of: " "('float', 'int', 'long'), received str instead.") def test_count_globally_pipeline_type_checking_satisfied(self): d = (self.p | df.Create('p', range(5)).with_output_types(int) | combine.Count.Globally('count int')) self.assertTrue(d.element_type is int) assert_that(d, equal_to([5])) self.p.run() def test_count_globally_runtime_type_checking_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('p', range(5)).with_output_types(int) | combine.Count.Globally('count int')) self.assertTrue(d.element_type is int) assert_that(d, equal_to([5])) self.p.run() def test_count_perkey_pipeline_type_checking_satisfied(self): d = (self.p | df.Create('p', range(5)).with_output_types(int) | (df.Map('even group', lambda x: (not x % 2, x)) .with_output_types(typehints.KV[bool, int])) | combine.Count.PerKey('count int')) self.assertCompatible(typehints.KV[bool, int], d.element_type) assert_that(d, equal_to([(False, 2), (True, 3)])) self.p.run() def test_count_perkey_pipeline_type_checking_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('p', range(5)).with_output_types(int) | combine.Count.PerKey('count int')) self.assertEqual("Input type hint violation at GroupByKey: " "expected Tuple[TypeVariable[K], TypeVariable[V]], " "got ", e.exception.message) def test_count_perkey_runtime_type_checking_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('c', ['t', 'e', 's', 't']).with_output_types(str) | df.Map('dup key', lambda x: (x, x)).with_output_types(typehints.KV[str, str]) | combine.Count.PerKey('count dups')) self.assertCompatible(typehints.KV[str, int], d.element_type) assert_that(d, equal_to([('e', 1), ('s', 1), ('t', 2)])) self.p.run() def test_count_perelement_pipeline_type_checking_satisfied(self): d = (self.p | df.Create('w', [1, 1, 2, 3]).with_output_types(int) | combine.Count.PerElement('count elems')) self.assertCompatible(typehints.KV[int, int], d.element_type) assert_that(d, equal_to([(1, 2), (2, 1), (3, 1)])) self.p.run() def test_count_perelement_pipeline_type_checking_violated(self): self.p.options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED' with self.assertRaises(typehints.TypeCheckError) as e: (self.p | df.Create('f', [1, 1, 2, 3]) | combine.Count.PerElement('count elems')) self.assertEqual('Pipeline type checking is enabled, however no output ' 'type-hint was found for the PTransform ' 'Create(f)', e.exception.message) def test_count_perelement_runtime_type_checking_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('w', [True, True, False, True, True]).with_output_types(bool) | combine.Count.PerElement('count elems')) self.assertCompatible(typehints.KV[bool, int], d.element_type) assert_that(d, equal_to([(False, 1), (True, 4)])) self.p.run() def test_top_of_pipeline_checking_satisfied(self): d = (self.p | df.Create('n', range(5, 11)).with_output_types(int) | combine.Top.Of('top 3', 3, lambda x, y: x < y)) self.assertCompatible(typehints.Iterable[int], d.element_type) assert_that(d, equal_to([[10, 9, 8]])) self.p.run() def test_top_of_runtime_checking_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('n', list('testing')).with_output_types(str) | combine.Top.Of('acii top', 3, lambda x, y: x < y)) self.assertCompatible(typehints.Iterable[str], d.element_type) assert_that(d, equal_to([['t', 't', 's']])) self.p.run() def test_per_key_pipeline_checking_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('n', range(100)).with_output_types(int) | df.Map('num + 1', lambda x: x + 1).with_output_types(int) | combine.Top.PerKey('top mod', 1, lambda a, b: a < b)) self.assertEqual("Input type hint violation at GroupByKey: " "expected Tuple[TypeVariable[K], TypeVariable[V]], " "got ", e.exception.message) def test_per_key_pipeline_checking_satisfied(self): d = (self.p | df.Create('n', range(100)).with_output_types(int) | (df.Map('group mod 3', lambda x: (x % 3, x)) .with_output_types(typehints.KV[int, int])) | combine.Top.PerKey('top mod', 1, lambda a, b: a < b)) self.assertCompatible(typehints.Tuple[int, typehints.Iterable[int]], d.element_type) assert_that(d, equal_to([(0, [99]), (1, [97]), (2, [98])])) self.p.run() def test_per_key_runtime_checking_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('n', range(21)) | (df.Map('group mod 3', lambda x: (x % 3, x)) .with_output_types(typehints.KV[int, int])) | combine.Top.PerKey('top mod', 1, lambda a, b: a < b)) self.assertCompatible(typehints.KV[int, typehints.Iterable[int]], d.element_type) assert_that(d, equal_to([(0, [18]), (1, [19]), (2, [20])])) self.p.run() def test_sample_globally_pipeline_satisfied(self): d = (self.p | df.Create('m', [2, 2, 3, 3]).with_output_types(int) | combine.Sample.FixedSizeGlobally('sample', 3)) self.assertCompatible(typehints.Iterable[int], d.element_type) def matcher(expected_len): def match(actual): equal_to([expected_len])([len(actual[0])]) return match assert_that(d, matcher(3)) self.p.run() def test_sample_globally_runtime_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('m', [2, 2, 3, 3]).with_output_types(int) | combine.Sample.FixedSizeGlobally('sample', 2)) self.assertCompatible(typehints.Iterable[int], d.element_type) def matcher(expected_len): def match(actual): equal_to([expected_len])([len(actual[0])]) return match assert_that(d, matcher(2)) self.p.run() def test_sample_per_key_pipeline_satisfied(self): d = (self.p | (df.Create('m', [(1, 2), (1, 2), (2, 3), (2, 3)]) .with_output_types(typehints.KV[int, int])) | combine.Sample.FixedSizePerKey('sample', 2)) self.assertCompatible(typehints.KV[int, typehints.Iterable[int]], d.element_type) def matcher(expected_len): def match(actual): for _, sample in actual: equal_to([expected_len])([len(sample)]) return match assert_that(d, matcher(2)) self.p.run() def test_sample_per_key_runtime_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | (df.Create('m', [(1, 2), (1, 2), (2, 3), (2, 3)]) .with_output_types(typehints.KV[int, int])) | combine.Sample.FixedSizePerKey('sample', 1)) self.assertCompatible(typehints.KV[int, typehints.Iterable[int]], d.element_type) def matcher(expected_len): def match(actual): for _, sample in actual: equal_to([expected_len])([len(sample)]) return match assert_that(d, matcher(1)) self.p.run() def test_to_list_pipeline_check_satisfied(self): d = (self.p | df.Create('c', (1, 2, 3, 4)).with_output_types(int) | combine.ToList('to list')) self.assertCompatible(typehints.List[int], d.element_type) def matcher(expected): def match(actual): equal_to(expected)(actual[0]) return match assert_that(d, matcher([1, 2, 3, 4])) self.p.run() def test_to_list_runtime_check_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | df.Create('c', list('test')).with_output_types(str) | combine.ToList('to list')) self.assertCompatible(typehints.List[str], d.element_type) def matcher(expected): def match(actual): equal_to(expected)(actual[0]) return match assert_that(d, matcher(['e', 's', 't', 't'])) self.p.run() def test_to_dict_pipeline_check_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: d = (self.p | df.Create('d', [1, 2, 3, 4]).with_output_types(int) | combine.ToDict('to dict')) self.assertEqual("Type hint violation for 'ParDo(CombineValuesDoFn)': " "requires Tuple[TypeVariable[K], " "Iterable[Tuple[TypeVariable[K], TypeVariable[V]]]] " "but got Tuple[None, Iterable[int]] for p_context", e.exception.message) def test_to_dict_pipeline_check_satisfied(self): d = (self.p | df.Create( 'd', [(1, 2), (3, 4)]).with_output_types(typehints.Tuple[int, int]) | combine.ToDict('to dict')) self.assertCompatible(typehints.Dict[int, int], d.element_type) assert_that(d, equal_to([{1: 2, 3: 4}])) self.p.run() def test_to_dict_runtime_check_satisfied(self): self.p.options.view_as(TypeOptions).runtime_type_check = True d = (self.p | (df.Create('d', [('1', 2), ('3', 4)]) .with_output_types(typehints.Tuple[str, int])) | combine.ToDict('to dict')) self.assertCompatible(typehints.Dict[str, int], d.element_type) assert_that(d, equal_to([{'1': 2, '3': 4}])) self.p.run() def test_runtime_type_check_python_type_error(self): self.p.options.view_as(TypeOptions).runtime_type_check = True with self.assertRaises(TypeError) as e: (self.p | df.Create('t', [1, 2, 3]).with_output_types(int) | df.Map('len', lambda x: len(x)).with_output_types(int)) self.p.run() # Our special type-checking related TypeError shouldn't have been raised. # Instead the above pipeline should have triggered a regular Python runtime # TypeError. self.assertEqual("object of type 'int' has no len() [while running 'len']", e.exception.message) self.assertFalse(isinstance(e, typehints.TypeCheckError)) def test_pardo_type_inference(self): self.assertEqual(int, df.Filter(lambda x: False).infer_output_type(int)) self.assertEqual(typehints.Tuple[str, int], df.Map(lambda x: (x, 1)).infer_output_type(str)) def test_gbk_type_inference(self): self.assertEqual( typehints.Tuple[str, typehints.Iterable[int]], df.core.GroupByKeyOnly().infer_output_type(typehints.KV[str, int])) def test_pipeline_inference(self): created = self.p | df.Create('c', ['a', 'b', 'c']) mapped = created | df.Map('pair with 1', lambda x: (x, 1)) grouped = mapped | df.GroupByKey() self.assertEqual(str, created.element_type) self.assertEqual(typehints.KV[str, int], mapped.element_type) self.assertEqual(typehints.KV[str, typehints.Iterable[int]], grouped.element_type) def test_inferred_bad_kv_type(self): with self.assertRaises(typehints.TypeCheckError) as e: _ = (self.p | df.Create('t', ['a', 'b', 'c']) | df.Map('ungroupable', lambda x: (x, 0, 1.0)) | df.GroupByKey()) self.assertEqual('Input type hint violation at GroupByKey: ' 'expected Tuple[TypeVariable[K], TypeVariable[V]], ' 'got Tuple[str, int, float]', e.exception.message) def test_type_inference_command_line_flag_toggle(self): self.p.options.view_as(TypeOptions).pipeline_type_check = False x = self.p | df.Create('t', [1, 2, 3, 4]) self.assertIsNone(x.element_type) self.p.options.view_as(TypeOptions).pipeline_type_check = True x = self.p | df.Create('m', [1, 2, 3, 4]) self.assertEqual(int, x.element_type) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/sideinputs.py000066400000000000000000000116401272340251000261430ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Internal side input transforms and implementations. Important: this module is an implementation detail and should not be used directly by pipeline writers. Instead, users should use the helper methods AsSingleton, AsIter, AsList and AsDict in google.cloud.dataflow.pvalue. """ from __future__ import absolute_import from google.cloud.dataflow import pvalue from google.cloud.dataflow import typehints from google.cloud.dataflow.transforms.ptransform import PTransform class CreatePCollectionView(PTransform): """Transform to materialize a given PCollectionView in the pipeline. Important: this transform is an implementation detail and should not be used directly by pipeline writers. """ def __init__(self, view): self.view = view super(CreatePCollectionView, self).__init__() def infer_output_type(self, input_type): # TODO(ccy): Figure out if we want to create a new type of type hint, i.e., # typehints.View[...]. return input_type def apply(self, pcoll): return self.view class ViewAsSingleton(PTransform): """Transform to view PCollection as a singleton PCollectionView. Important: this transform is an implementation detail and should not be used directly by pipeline writers. Use pvalue.AsSingleton(...) instead. """ def __init__(self, has_default, default_value, label=None): if label: label = 'ViewAsSingleton(%s)' % label super(ViewAsSingleton, self).__init__(label=label) self.has_default = has_default self.default_value = default_value def apply(self, pcoll): self._check_pcollection(pcoll) input_type = pcoll.element_type output_type = input_type return (pcoll | CreatePCollectionView( pvalue.SingletonPCollectionView( pcoll.pipeline, self.has_default, self.default_value)) .with_input_types(input_type) .with_output_types(output_type)) class ViewAsIterable(PTransform): """Transform to view PCollection as an iterable PCollectionView. Important: this transform is an implementation detail and should not be used directly by pipeline writers. Use pvalue.AsIter(...) instead. """ def __init__(self, label=None): if label: label = 'ViewAsIterable(%s)' % label super(ViewAsIterable, self).__init__(label=label) def apply(self, pcoll): self._check_pcollection(pcoll) input_type = pcoll.element_type output_type = typehints.Iterable[input_type] return (pcoll | CreatePCollectionView( pvalue.IterablePCollectionView(pcoll.pipeline)) .with_input_types(input_type) .with_output_types(output_type)) class ViewAsList(PTransform): """Transform to view PCollection as a list PCollectionView. Important: this transform is an implementation detail and should not be used directly by pipeline writers. Use pvalue.AsList(...) instead. """ def __init__(self, label=None): if label: label = 'ViewAsList(%s)' % label super(ViewAsList, self).__init__(label=label) def apply(self, pcoll): self._check_pcollection(pcoll) input_type = pcoll.element_type output_type = typehints.List[input_type] return (pcoll | CreatePCollectionView(pvalue.ListPCollectionView(pcoll.pipeline)) .with_input_types(input_type) .with_output_types(output_type)) K = typehints.TypeVariable('K') V = typehints.TypeVariable('V') @typehints.with_input_types(typehints.Tuple[K, V]) @typehints.with_output_types(typehints.Dict[K, V]) class ViewAsDict(PTransform): # pylint: disable=g-wrong-blank-lines """Transform to view PCollection as a dict PCollectionView. Important: this transform is an implementation detail and should not be used directly by pipeline writers. Use pvalue.AsDict(...) instead. """ def __init__(self, label=None): if label: label = 'ViewAsDict(%s)' % label super(ViewAsDict, self).__init__(label=label) def apply(self, pcoll): self._check_pcollection(pcoll) input_type = pcoll.element_type key_type, value_type = ( typehints.trivial_inference.key_value_types(input_type)) output_type = typehints.Dict[key_type, value_type] return (pcoll | CreatePCollectionView( pvalue.DictPCollectionView(pcoll.pipeline)) .with_input_types(input_type) .with_output_types(output_type)) DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/timeutil.py000066400000000000000000000220321272340251000256050ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Time and timer utilities.""" from __future__ import absolute_import from abc import ABCMeta from abc import abstractmethod import datetime import sys class Timestamp(object): """Represents a Unix second timestamp with microsecond granularity. Can be treated in common timestamp arithmetic operations as a numeric type. Internally stores a time interval as an int of microseconds. This strategy is necessary since floating point values lose precision when storing values, especially after arithmetic operations (for example, 10000000 % 0.1 evaluates to 0.0999999994448885). """ def __init__(self, seconds=0, micros=0): self.micros = int(seconds * 1000000) + int(micros) @staticmethod def of(seconds): """Return the Timestamp for the given number of seconds. If the input is already a Timestamp, the input itself will be returned. Args: seconds: Number of seconds as int, float or Timestamp. Returns: Corresponding Timestamp object. """ if isinstance(seconds, Duration): raise TypeError('Can\'t interpret %s as Timestamp.' % seconds) if isinstance(seconds, Timestamp): return seconds return Timestamp(seconds) def __repr__(self): micros = self.micros sign = '' if micros < 0: sign = '-' micros = -micros int_part = micros / 1000000 frac_part = micros % 1000000 if frac_part: return 'Timestamp(%s%d.%06d)' % (sign, int_part, frac_part) else: return 'Timestamp(%s%d)' % (sign, int_part) def to_utc_datetime(self): epoch = datetime.datetime.utcfromtimestamp(0) # We can't easily construct a datetime object from microseconds, so we # create one at the epoch and add an appropriate timedelta interval. return epoch + datetime.timedelta(microseconds=self.micros) def isoformat(self): # Append 'Z' for UTC timezone. return self.to_utc_datetime().isoformat() + 'Z' def __float__(self): # Note that the returned value may have lost precision. return float(self.micros) / 1000000 def __int__(self): # Note that the returned value may have lost precision. return self.micros / 1000000 def __cmp__(self, other): # Allow comparisons between Duration and Timestamp values. if not isinstance(other, Duration): other = Timestamp.of(other) return cmp(self.micros, other.micros) def __hash__(self): return hash(self.micros) def __add__(self, other): other = Duration.of(other) return Timestamp(micros=self.micros + other.micros) def __radd__(self, other): return self + other def __sub__(self, other): other = Duration.of(other) return Timestamp(micros=self.micros - other.micros) def __mod__(self, other): other = Duration.of(other) return Duration(micros=self.micros % other.micros) MIN_TIMESTAMP = Timestamp(micros=-sys.maxint - 1) MAX_TIMESTAMP = Timestamp(micros=sys.maxint) class Duration(object): """Represents a second duration with microsecond granularity. Can be treated in common arithmetic operations as a numeric type. Internally stores a time interval as an int of microseconds. This strategy is necessary since floating point values lose precision when storing values, especially after arithmetic operations (for example, 10000000 % 0.1 evaluates to 0.0999999994448885). """ def __init__(self, seconds=0, micros=0): self.micros = int(seconds * 1000000) + int(micros) @staticmethod def of(seconds): """Return the Duration for the given number of seconds since Unix epoch. If the input is already a Duration, the input itself will be returned. Args: seconds: Number of seconds as int, float or Duration. Returns: Corresponding Duration object. """ if isinstance(seconds, Timestamp): raise TypeError('Can\'t interpret %s as Duration.' % seconds) if isinstance(seconds, Duration): return seconds return Duration(seconds) def __repr__(self): micros = self.micros sign = '' if micros < 0: sign = '-' micros = -micros int_part = micros / 1000000 frac_part = micros % 1000000 if frac_part: return 'Duration(%s%d.%06d)' % (sign, int_part, frac_part) else: return 'Duration(%s%d)' % (sign, int_part) def __float__(self): # Note that the returned value may have lost precision. return float(self.micros) / 1000000 def __int__(self): # Note that the returned value may have lost precision. return self.micros / 1000000 def __cmp__(self, other): # Allow comparisons between Duration and Timestamp values. if not isinstance(other, Timestamp): other = Duration.of(other) return cmp(self.micros, other.micros) def __hash__(self): return hash(self.micros) def __neg__(self): return Duration(micros=-self.micros) def __add__(self, other): if isinstance(other, Timestamp): return other + self other = Duration.of(other) return Duration(micros=self.micros + other.micros) def __radd__(self, other): return self + other def __sub__(self, other): other = Duration.of(other) return Duration(micros=self.micros - other.micros) def __rsub__(self, other): return -(self - other) def __mul__(self, other): other = Duration.of(other) return Duration(micros=self.micros * other.micros / 1000000) def __rmul__(self, other): return self * other def __mod__(self, other): other = Duration.of(other) return Duration(micros=self.micros % other.micros) class TimeDomain(object): """Time domain for streaming timers.""" WATERMARK = 'WATERMARK' REAL_TIME = 'REAL_TIME' DEPENDENT_REAL_TIME = 'DEPENDENT_REAL_TIME' @staticmethod def from_string(domain): if domain in (TimeDomain.WATERMARK, TimeDomain.REAL_TIME, TimeDomain.DEPENDENT_REAL_TIME): return domain raise ValueError('Unknown time domain: %s' % domain) class OutputTimeFnImpl(object): """Implementation of OutputTimeFn.""" __metaclass__ = ABCMeta @abstractmethod def assign_output_time(self, window, input_timestamp): pass @abstractmethod def combine(self, output_timestamp, other_output_timestamp): pass def combine_all(self, merging_timestamps): """Apply combine to list of timestamps.""" combined_output_time = None for output_time in merging_timestamps: if combined_output_time is None: combined_output_time = output_time else: combined_output_time = self.combine( combined_output_time, output_time) return combined_output_time def merge(self, unused_result_window, merging_timestamps): """Default to returning the result of combine_all.""" return self.combine_all(merging_timestamps) class DependsOnlyOnWindow(OutputTimeFnImpl): """OutputTimeFnImpl that only depends on the window.""" __metaclass__ = ABCMeta def combine(self, output_timestamp, other_output_timestamp): return output_timestamp def merge(self, result_window, unused_merging_timestamps): # Since we know that the result only depends on the window, we can ignore # the given timestamps. return self.assign_output_time(result_window, None) class OutputAtEarliestInputTimestampImpl(OutputTimeFnImpl): """OutputTimeFnImpl outputting at earliest input timestamp.""" def assign_output_time(self, window, input_timestamp): return input_timestamp def combine(self, output_timestamp, other_output_timestamp): """Default to returning the earlier of two timestamps.""" return min(output_timestamp, other_output_timestamp) class OutputAtEarliestTransformedInputTimestampImpl(OutputTimeFnImpl): """OutputTimeFnImpl outputting at earliest input timestamp.""" def __init__(self, window_fn): self.window_fn = window_fn def assign_output_time(self, window, input_timestamp): return self.window_fn.get_transformed_output_time(window, input_timestamp) def combine(self, output_timestamp, other_output_timestamp): return min(output_timestamp, other_output_timestamp) class OutputAtLatestInputTimestampImpl(OutputTimeFnImpl): """OutputTimeFnImpl outputting at latest input timestamp.""" def assign_output_time(self, window, input_timestamp): return input_timestamp def combine(self, output_timestamp, other_output_timestamp): return max(output_timestamp, other_output_timestamp) class OutputAtEndOfWindowImpl(DependsOnlyOnWindow): """OutputTimeFnImpl outputting at end of window.""" def assign_output_time(self, window, unused_input_timestamp): return window.end DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/timeutil_test.py000066400000000000000000000147671272340251000266640ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for time utilities.""" from __future__ import absolute_import import unittest from google.cloud.dataflow.transforms.timeutil import Duration from google.cloud.dataflow.transforms.timeutil import Timestamp class TimestampTest(unittest.TestCase): def test_of(self): interval = Timestamp(123) self.assertEqual(id(interval), id(Timestamp.of(interval))) self.assertEqual(interval, Timestamp.of(123.0)) with self.assertRaises(TypeError): Timestamp.of(Duration(10)) def test_precision(self): self.assertEqual(Timestamp(10000000) % 0.1, 0) self.assertEqual(Timestamp(10000000) % 0.05, 0) self.assertEqual(Timestamp(10000000) % 0.000005, 0) self.assertEqual(Timestamp(10000000) % Duration(0.1), 0) self.assertEqual(Timestamp(10000000) % Duration(0.05), 0) self.assertEqual(Timestamp(10000000) % Duration(0.000005), 0) def test_utc_timestamp(self): self.assertEqual(Timestamp(10000000).isoformat(), '1970-04-26T17:46:40Z') self.assertEqual(Timestamp(10000000.000001).isoformat(), '1970-04-26T17:46:40.000001Z') self.assertEqual(Timestamp(1458343379.123456).isoformat(), '2016-03-18T23:22:59.123456Z') def test_arithmetic(self): # Supported operations. self.assertEqual(Timestamp(123) + 456, 579) self.assertEqual(Timestamp(123) + Duration(456), 579) self.assertEqual(456 + Timestamp(123), 579) self.assertEqual(Duration(456) + Timestamp(123), 579) self.assertEqual(Timestamp(123) - 456, -333) self.assertEqual(Timestamp(123) - Duration(456), -333) self.assertEqual(Timestamp(1230) % 456, 318) self.assertEqual(Timestamp(1230) % Duration(456), 318) # Check that direct comparison of Timestamp and Duration is allowed. self.assertTrue(Duration(123) == Timestamp(123)) self.assertTrue(Timestamp(123) == Duration(123)) self.assertFalse(Duration(123) == Timestamp(1230)) self.assertFalse(Timestamp(123) == Duration(1230)) # Check return types. self.assertEqual((Timestamp(123) + 456).__class__, Timestamp) self.assertEqual((Timestamp(123) + Duration(456)).__class__, Timestamp) self.assertEqual((456 + Timestamp(123)).__class__, Timestamp) self.assertEqual((Duration(456) + Timestamp(123)).__class__, Timestamp) self.assertEqual((Timestamp(123) - 456).__class__, Timestamp) self.assertEqual((Timestamp(123) - Duration(456)).__class__, Timestamp) self.assertEqual((Timestamp(1230) % 456).__class__, Duration) self.assertEqual((Timestamp(1230) % Duration(456)).__class__, Duration) # Unsupported operations. with self.assertRaises(TypeError): self.assertEqual(Timestamp(123) * 456, 56088) with self.assertRaises(TypeError): self.assertEqual(Timestamp(123) * Duration(456), 56088) with self.assertRaises(TypeError): self.assertEqual(456 * Timestamp(123), 56088) with self.assertRaises(TypeError): self.assertEqual(Duration(456) * Timestamp(123), 56088) with self.assertRaises(TypeError): self.assertEqual(456 - Timestamp(123), 333) with self.assertRaises(TypeError): self.assertEqual(Duration(456) - Timestamp(123), 333) with self.assertRaises(TypeError): self.assertEqual(-Timestamp(123), -123) with self.assertRaises(TypeError): self.assertEqual(-Timestamp(123), -Duration(123)) with self.assertRaises(TypeError): self.assertEqual(1230 % Timestamp(456), 318) with self.assertRaises(TypeError): self.assertEqual(Duration(1230) % Timestamp(456), 318) def test_sort_order(self): self.assertEqual( [-63, Timestamp(-3), 2, 9, Timestamp(292.3), 500], sorted([9, 2, Timestamp(-3), Timestamp(292.3), -63, 500])) self.assertEqual( [4, 5, Timestamp(6), Timestamp(7), 8, 9], sorted([9, 8, Timestamp(7), Timestamp(6), 5, 4])) def test_str(self): self.assertEqual('Timestamp(1.234567)', str(Timestamp(1.234567))) self.assertEqual('Timestamp(-1.234567)', str(Timestamp(-1.234567))) self.assertEqual('Timestamp(-999999999.900000)', str(Timestamp(-999999999.9))) self.assertEqual('Timestamp(999999999)', str(Timestamp(999999999))) self.assertEqual('Timestamp(-999999999)', str(Timestamp(-999999999))) class DurationTest(unittest.TestCase): def test_of(self): interval = Duration(123) self.assertEqual(id(interval), id(Duration.of(interval))) self.assertEqual(interval, Duration.of(123.0)) with self.assertRaises(TypeError): Duration.of(Timestamp(10)) def test_precision(self): self.assertEqual(Duration(10000000) % 0.1, 0) self.assertEqual(Duration(10000000) % 0.05, 0) self.assertEqual(Duration(10000000) % 0.000005, 0) def test_arithmetic(self): self.assertEqual(Duration(123) + 456, 579) self.assertEqual(456 + Duration(123), 579) self.assertEqual(Duration(123) * 456, 56088) self.assertEqual(456 * Duration(123), 56088) self.assertEqual(Duration(123) - 456, -333) self.assertEqual(456 - Duration(123), 333) self.assertEqual(-Duration(123), -123) def test_sort_order(self): self.assertEqual( [-63, Duration(-3), 2, 9, Duration(292.3), 500], sorted([9, 2, Duration(-3), Duration(292.3), -63, 500])) self.assertEqual( [4, 5, Duration(6), Duration(7), 8, 9], sorted([9, 8, Duration(7), Duration(6), 5, 4])) def test_str(self): self.assertEqual('Duration(1.234567)', str(Duration(1.234567))) self.assertEqual('Duration(-1.234567)', str(Duration(-1.234567))) self.assertEqual('Duration(-999999999.900000)', str(Duration(-999999999.9))) self.assertEqual('Duration(999999999)', str(Duration(999999999))) self.assertEqual('Duration(-999999999)', str(Duration(-999999999))) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/trigger.py000066400000000000000000000755551272340251000254360ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Support for Dataflow triggers. Triggers control when in processing time windows get emitted. """ from abc import ABCMeta from abc import abstractmethod import collections import copy from google.cloud.dataflow.coders import observable from google.cloud.dataflow.transforms import combiners from google.cloud.dataflow.transforms import core from google.cloud.dataflow.transforms.timeutil import MAX_TIMESTAMP from google.cloud.dataflow.transforms.timeutil import MIN_TIMESTAMP from google.cloud.dataflow.transforms.timeutil import TimeDomain from google.cloud.dataflow.transforms.window import GlobalWindow from google.cloud.dataflow.transforms.window import OutputTimeFn from google.cloud.dataflow.transforms.window import WindowedValue from google.cloud.dataflow.transforms.window import WindowFn class AccumulationMode(object): """Controls what to do with data when a trigger fires multiple times. """ DISCARDING = 1 ACCUMULATING = 2 # TODO(robertwb): Provide retractions of previous outputs. # RETRACTING = 3 class StateTag(object): """An identifier used to store and retrieve typed, combinable state. The given tag must be unique for this stage. If CombineFn is None then all elements will be returned as a list, otherwise the given CombineFn will be applied (possibly incrementally and eagerly) when adding elements. """ __metaclass__ = ABCMeta def __init__(self, tag): self.tag = tag class ValueStateTag(StateTag): """StateTag pointing to an element.""" def __repr__(self): return 'ValueStateTag(%s, %s)' % (self.tag, self.combine_fn) def with_prefix(self, prefix): return ValueStateTag(prefix + self.tag) class CombiningValueStateTag(StateTag): """StateTag pointing to an element, accumulated with a combiner.""" # TODO(robertwb): Also store the coder (perhaps extracted from the combine_fn) def __init__(self, tag, combine_fn): super(CombiningValueStateTag, self).__init__(tag) if not combine_fn: raise ValueError('combine_fn must be specified.') if not isinstance(combine_fn, core.CombineFn): combine_fn = core.CombineFn.from_callable(combine_fn) self.combine_fn = combine_fn def __repr__(self): return 'CombiningValueStateTag(%s, %s)' % (self.tag, self.combine_fn) def with_prefix(self, prefix): return CombiningValueStateTag(prefix + self.tag, self.combine_fn) class ListStateTag(StateTag): """StateTag pointing to a list of elements.""" def __init__(self, tag): super(ListStateTag, self).__init__(tag) def __repr__(self): return 'ListStateTag(%s)' % self.tag def with_prefix(self, prefix): return ListStateTag(prefix + self.tag) class WatermarkHoldStateTag(StateTag): def __init__(self, tag, output_time_fn_impl): super(WatermarkHoldStateTag, self).__init__(tag) self.output_time_fn_impl = output_time_fn_impl def __repr__(self): return 'WatermarkHoldStateTag(%s, %s)' % (self.tag, self.output_time_fn_impl) def with_prefix(self, prefix): return WatermarkHoldStateTag(prefix + self.tag, self.output_time_fn_impl) # pylint: disable=unused-argument # TODO(robertwb): Provisional API, Java likely to change as well. class TriggerFn(object): """A TriggerFn determines when window (panes) are emitted. See https://cloud.google.com/dataflow/model/triggers. """ __metaclass__ = ABCMeta @abstractmethod def on_element(self, element, window, context): """Called when a new element arrives in a window. Args: element: the element being added window: the window to which the element is being added context: a context (e.g. a TriggerContext instance) for managing state and setting timers """ pass @abstractmethod def on_merge(self, to_be_merged, merge_result, context): """Called when multiple windows are merged. Args: to_be_merged: the set of windows to be merged merge_result: the window into which the windows are being merged context: a context (e.g. a TriggerContext instance) for managing state and setting timers """ pass @abstractmethod def should_fire(self, watermark, window, context): """Whether this trigger should cause the window to fire. Args: watermark: (a lower bound on) the watermark of the system window: the window whose trigger is being considered context: a context (e.g. a TriggerContext instance) for managing state and setting timers Returns: whether this trigger should cause a firing """ pass @abstractmethod def on_fire(self, watermark, window, context): """Called when a trigger actually fires. Args: watermark: (a lower bound on) the watermark of the system window: the window whose trigger is being fired context: a context (e.g. a TriggerContext instance) for managing state and setting timers Returns: whether this trigger is finished """ pass @abstractmethod def reset(self, window, context): """Clear any state and timers used by this TriggerFn.""" pass # pylint: enable=unused-argument class DefaultTrigger(TriggerFn): """Semantically Repeatedly(AfterWatermark()), but more optimized.""" def __init__(self): pass def __repr__(self): return 'DefaultTrigger()' def on_element(self, element, window, context): context.set_timer('', TimeDomain.WATERMARK, window.end) def on_merge(self, to_be_merged, merge_result, context): # Note: Timer clearing solely an optimization. for window in to_be_merged: if window.end != merge_result.end: context.clear_timer('', TimeDomain.WATERMARK) def should_fire(self, watermark, window, context): return watermark >= window.end def on_fire(self, watermark, window, context): return False def reset(self, window, context): context.clear_timer('', TimeDomain.WATERMARK) def __eq__(self, other): return type(self) == type(other) class AfterWatermark(TriggerFn): """Fire exactly once when the watermark passes the end of the window. Args: early: if not None, a speculative trigger to repeatedly evaluate before the watermark passes the end of the window late: if not None, a speculative trigger to repeatedly evaluate after the watermark passes the end of the window """ LATE_TAG = CombiningValueStateTag('is_late', any) def __init__(self, early=None, late=None): self.early = Repeatedly(early) if early else None self.late = Repeatedly(late) if late else None def __repr__(self): qualifiers = [] if self.early: qualifiers.append('early=%s' % self.early) if self.late: qualifiers.append('late=%s', self.late) return 'AfterWatermark(%s)' % ', '.join(qualifiers) def is_late(self, context): return self.late and context.get_state(self.LATE_TAG) def on_element(self, element, window, context): if self.is_late(context): self.late.on_element(element, window, NestedContext(context, 'late')) else: context.set_timer('', TimeDomain.WATERMARK, window.end) if self.early: self.early.on_element(element, window, NestedContext(context, 'early')) def on_merge(self, to_be_merged, merge_result, context): # TODO(robertwb): Figure out whether the 'rewind' semantics could be used # here. if self.is_late(context): self.late.on_merge( to_be_merged, merge_result, NestedContext(context, 'late')) else: # Note: Timer clearing solely an optimization. for window in to_be_merged: if window.end != merge_result.end: context.clear_timer('', TimeDomain.WATERMARK) if self.early: self.early.on_merge( to_be_merged, merge_result, NestedContext(context, 'early')) def should_fire(self, watermark, window, context): if self.is_late(context): return self.late.should_fire( watermark, window, NestedContext(context, 'late')) elif watermark >= window.end: return True elif self.early: return self.early.should_fire( watermark, window, NestedContext(context, 'early')) else: return False def on_fire(self, watermark, window, context): if self.is_late(context): return self.late.on_fire( watermark, window, NestedContext(context, 'late')) elif watermark >= window.end: context.add_state(self.LATE_TAG, True) return not self.late elif self.early: self.early.on_fire(watermark, window, NestedContext(context, 'early')) return False def reset(self, window, context): if self.late: context.clear_state(self.LATE_TAG) if self.early: self.early.reset(window, NestedContext(context, 'early')) if self.late: self.late.reset(window, NestedContext(context, 'late')) def __eq__(self, other): return (type(self) == type(other) and self.early == other.early and self.late == other.late) def __hash__(self): return hash((type(self), self.early, self.late)) class AfterCount(TriggerFn): """Fire when there are at least count elements in this window pane.""" COUNT_TAG = CombiningValueStateTag('count', combiners.CountCombineFn()) def __init__(self, count): self.count = count def __repr__(self): return 'AfterCount(%s)' % self.count def on_element(self, element, window, context): context.add_state(self.COUNT_TAG, 1) def on_merge(self, to_be_merged, merge_result, context): # states automatically merged pass def should_fire(self, watermark, window, context): return context.get_state(self.COUNT_TAG) >= self.count def on_fire(self, watermark, window, context): return True def reset(self, window, context): context.clear_state(self.COUNT_TAG) class Repeatedly(TriggerFn): """Repeatedly invoke the given trigger, never finishing.""" def __init__(self, underlying): self.underlying = underlying def __repr__(self): return 'Repeatedly(%s)' % self.underlying def on_element(self, element, window, context): # get window from context? self.underlying.on_element(element, window, context) def on_merge(self, to_be_merged, merge_result, context): self.underlying.on_merge(to_be_merged, merge_result, context) def should_fire(self, watermark, window, context): return self.underlying.should_fire(watermark, window, context) def on_fire(self, watermark, window, context): if self.underlying.on_fire(watermark, window, context): self.underlying.reset(window, context) return False def reset(self, window, context): self.underlying.reset(window, context) class ParallelTriggerFn(TriggerFn): __metaclass__ = ABCMeta def __init__(self, *triggers): self.triggers = triggers def __repr__(self): return '%s(%s)' % (self.__class__.__name__, ', '.join(str(t) for t in self.triggers)) @abstractmethod def combine_op(self, trigger_results): pass def on_element(self, element, window, context): for ix, trigger in enumerate(self.triggers): trigger.on_element(element, window, self._sub_context(context, ix)) def on_merge(self, to_be_merged, merge_result, context): for ix, trigger in enumerate(self.triggers): trigger.on_merge( to_be_merged, merge_result, self._sub_context(context, ix)) def should_fire(self, watermark, window, context): return self.combine_op( trigger.should_fire(watermark, window, self._sub_context(context, ix)) for ix, trigger in enumerate(self.triggers)) def on_fire(self, watermark, window, context): finished = [] for ix, trigger in enumerate(self.triggers): nested_context = self._sub_context(context, ix) if trigger.should_fire(watermark, window, nested_context): finished.append(trigger.on_fire(watermark, window, nested_context)) return self.combine_op(finished) def reset(self, window, context): for ix, trigger in enumerate(self.triggers): trigger.reset(window, self._sub_context(context, ix)) @staticmethod def _sub_context(context, index): return NestedContext(context, '%d/' % index) class AfterFirst(ParallelTriggerFn): """Fires when any subtrigger fires. Also finishes when any subtrigger finishes. """ combine_op = any class AfterAll(ParallelTriggerFn): """Fires when all subtriggers have fired. Also finishes when all subtriggers have finished. """ combine_op = all class AfterEach(TriggerFn): INDEX_TAG = CombiningValueStateTag('index', ( lambda indices: 0 if not indices else max(indices))) def __init__(self, *triggers): self.triggers = triggers def __repr__(self): return '%s(%s)' % (self.__class__.__name__, ', '.join(str(t) for t in self.triggers)) def on_element(self, element, window, context): ix = context.get_state(self.INDEX_TAG) if ix < len(self.triggers): self.triggers[ix].on_element( element, window, self._sub_context(context, ix)) def on_merge(self, to_be_merged, merge_result, context): # This takes the furthest window on merging. # TODO(robertwb): Revisit this when merging windows logic is settled for # all possible merging situations. ix = context.get_state(self.INDEX_TAG) if ix < len(self.triggers): self.triggers[ix].on_merge( to_be_merged, merge_result, self._sub_context(context, ix)) def should_fire(self, watermark, window, context): ix = context.get_state(self.INDEX_TAG) if ix < len(self.triggers): return self.triggers[ix].should_fire( watermark, window, self._sub_context(context, ix)) def on_fire(self, watermark, window, context): ix = context.get_state(self.INDEX_TAG) if ix < len(self.triggers): if self.triggers[ix].on_fire( watermark, window, self._sub_context(context, ix)): ix += 1 context.add_state(self.INDEX_TAG, ix) return ix == len(self.triggers) def reset(self, window, context): context.clear_state(self.INDEX_TAG) for ix, trigger in enumerate(self.triggers): trigger.reset(window, self._sub_context(context, ix)) @staticmethod def _sub_context(context, index): return NestedContext(context, '%d/' % index) class OrFinally(AfterFirst): def __init__(self, body_trigger, exit_trigger): super(OrFinally, self).__init__(body_trigger, exit_trigger) class TriggerContext(object): def __init__(self, outer, window): self._outer = outer self._window = window def set_timer(self, name, time_domain, timestamp): self._outer.set_timer(self._window, name, time_domain, timestamp) def clear_timer(self, name, time_domain): self._outer.clear_timer(self._window, name, time_domain) def add_state(self, tag, value): self._outer.add_state(self._window, tag, value) def get_state(self, tag): return self._outer.get_state(self._window, tag) def clear_state(self, tag): return self._outer.clear_state(self._window, tag) class NestedContext(object): """Namespaced context useful for defining composite triggers.""" def __init__(self, outer, prefix): self._outer = outer self._prefix = prefix def set_timer(self, name, time_domain, timestamp): self._outer.set_timer(self._prefix + name, time_domain, timestamp) def clear_timer(self, name, time_domain): self._outer.clear_timer(self._prefix + name, time_domain) def add_state(self, tag, value): self._outer.add_state(tag.with_prefix(self._prefix), value) def get_state(self, tag): return self._outer.get_state(tag.with_prefix(self._prefix)) def clear_state(self, tag): self._outer.clear_state(tag.with_prefix(self._prefix)) # pylint: disable=unused-argument class SimpleState(object): """Basic state storage interface used for triggering. Only timers must hold the watermark (by their timestamp). """ __metaclass__ = ABCMeta @abstractmethod def set_timer(self, window, name, time_domain, timestamp): pass @abstractmethod def get_window(self, window_id): pass @abstractmethod def clear_timer(self, window, name, time_domain): pass @abstractmethod def add_state(self, window, tag, value): pass @abstractmethod def get_state(self, window, tag): pass @abstractmethod def clear_state(self, window, tag): pass def at(self, window): return TriggerContext(self, window) class UnmergedState(SimpleState): """State suitable for use in TriggerDriver. This class must be implemented by each backend. """ @abstractmethod def set_global_state(self, tag, value): pass @abstractmethod def get_global_state(self, tag, default=None): pass # pylint: enable=unused-argument class MergeableStateAdapter(SimpleState): """Wraps an UnmergedState, tracking merged windows.""" # TODO(robertwb): A similar indirection could be used for sliding windows # or other window_fns when a single element typically belongs to many windows. WINDOW_IDS = ValueStateTag('window_ids') def __init__(self, raw_state): self.raw_state = raw_state self.window_ids = self.raw_state.get_global_state(self.WINDOW_IDS, {}) self.counter = None def set_timer(self, window, name, time_domain, timestamp): self.raw_state.set_timer(self._get_id(window), name, time_domain, timestamp) def clear_timer(self, window, name, time_domain): for window_id in self._get_ids(window): self.raw_state.clear_timer(window_id, name, time_domain) def add_state(self, window, tag, value): if isinstance(tag, ValueStateTag): raise ValueError( 'Merging requested for non-mergeable state tag: %r.' % tag) self.raw_state.add_state(self._get_id(window), tag, value) def get_state(self, window, tag): values = [self.raw_state.get_state(window_id, tag) for window_id in self._get_ids(window)] if isinstance(tag, ValueStateTag): raise ValueError( 'Merging requested for non-mergeable state tag: %r.' % tag) elif isinstance(tag, CombiningValueStateTag): # TODO(robertwb): Strip combine_fn.extract_output from raw_state tag. if not values: accumulator = tag.combine_fn.create_accumulator() elif len(values) == 1: accumulator = values[0] else: accumulator = tag.combine_fn.merge_accumulators(values) # TODO(robertwb): Store the merged value in the first tag. return tag.combine_fn.extract_output(accumulator) elif isinstance(tag, ListStateTag): return [v for vs in values for v in vs] elif isinstance(tag, WatermarkHoldStateTag): return tag.output_time_fn_impl.combine_all(values) else: raise ValueError('Invalid tag.', tag) def clear_state(self, window, tag): for window_id in self._get_ids(window): self.raw_state.clear_state(window_id, tag) if tag is None: del self.window_ids[window] self._persist_window_ids() def merge(self, to_be_merged, merge_result): for window in to_be_merged: if window != merge_result: if window in self.window_ids: if merge_result in self.window_ids: merge_window_ids = self.window_ids[merge_result] else: merge_window_ids = self.window_ids[merge_result] = [] merge_window_ids.extend(self.window_ids.pop(window)) self._persist_window_ids() def known_windows(self): return self.window_ids.keys() def get_window(self, window_id): for window, ids in self.window_ids.items(): if window_id in ids: return window raise ValueError('No window for %s' % window_id) def _get_id(self, window): if window in self.window_ids: return self.window_ids[window][0] else: window_id = self._get_next_counter() self.window_ids[window] = [window_id] self._persist_window_ids() return window_id def _get_ids(self, window): return self.window_ids.get(window, []) def _get_next_counter(self): if not self.window_ids: self.counter = 0 elif self.counter is None: self.counter = max(k for ids in self.window_ids.values() for k in ids) self.counter += 1 return self.counter def _persist_window_ids(self): self.raw_state.set_global_state(self.WINDOW_IDS, self.window_ids) def __repr__(self): return '\n\t'.join([repr(self.window_ids)] + repr(self.raw_state).split('\n')) def create_trigger_driver(windowing, is_batch=False, phased_combine_fn=None): """Create the TriggerDriver for the given windowing and options.""" # TODO(robertwb): We can do more if we know elements are in timestamp # sorted order. if windowing.is_default() and is_batch: driver = DefaultGlobalBatchTriggerDriver() else: driver = GeneralTriggerDriver(windowing) if phased_combine_fn: # TODO(ccy): Refactor GeneralTriggerDriver to combine values eagerly using # the known phased_combine_fn here. driver = CombiningTriggerDriver(phased_combine_fn, driver) return driver class TriggerDriver(object): """Breaks a series of bundle and timer firings into window (pane)s.""" __metaclass__ = ABCMeta @abstractmethod def process_elements(self, state, windowed_values, output_watermark): pass @abstractmethod def process_timer(self, window_id, name, time_domain, timestamp, state): pass class DefaultGlobalBatchTriggerDriver(TriggerDriver): """Breaks a bundles into window (pane)s according to the default triggering. """ GLOBAL_WINDOW_TUPLE = (GlobalWindow(),) def __init__(self): pass def process_elements(self, state, windowed_values, unused_output_watermark): if isinstance(windowed_values, list): unwindowed = [wv.value for wv in windowed_values] else: class UnwindowedValues(observable.ObservableMixin): def __iter__(self): for wv in windowed_values: unwindowed_value = wv.value self.notify_observers(unwindowed_value) yield unwindowed_value def __repr__(self): return '' % windowed_values unwindowed = UnwindowedValues() yield WindowedValue(unwindowed, MIN_TIMESTAMP, self.GLOBAL_WINDOW_TUPLE) def process_timer(self, window_id, name, time_domain, timestamp, state): raise TypeError('Triggers never set or called for batch default windowing.') class CombiningTriggerDriver(TriggerDriver): """Uses a phased_combine_fn to process output of wrapped TriggerDriver.""" def __init__(self, phased_combine_fn, underlying): self.phased_combine_fn = phased_combine_fn self.underlying = underlying def process_elements(self, state, windowed_values, output_watermark): uncombined = self.underlying.process_elements(state, windowed_values, output_watermark) for output in uncombined: yield output.with_value(self.phased_combine_fn.apply(output.value)) def process_timer(self, window_id, name, time_domain, timestamp, state): uncombined = self.underlying.process_timer(window_id, name, time_domain, timestamp, state) for output in uncombined: yield output.with_value(self.phased_combine_fn.apply(output.value)) class GeneralTriggerDriver(TriggerDriver): """Breaks a series of bundle and timer firings into window (pane)s. Suitable for all variants of Windowing. """ ELEMENTS = ListStateTag('elements') TOMBSTONE = CombiningValueStateTag('tombstone', combiners.CountCombineFn()) def __init__(self, windowing): self.window_fn = windowing.windowfn self.output_time_fn_impl = OutputTimeFn.get_impl(windowing.output_time_fn, self.window_fn) # pylint: disable=invalid-name self.WATERMARK_HOLD = WatermarkHoldStateTag('watermark', self.output_time_fn_impl) # pylint: enable=invalid-name self.trigger_fn = windowing.triggerfn self.accumulation_mode = windowing.accumulation_mode self.is_merging = True def process_elements(self, state, windowed_values, output_watermark): if self.is_merging: state = MergeableStateAdapter(state) windows_to_elements = collections.defaultdict(list) for wv in windowed_values: for window in wv.windows: windows_to_elements[window].append((wv.value, wv.timestamp)) # First handle merging. if self.is_merging: old_windows = set(state.known_windows()) all_windows = old_windows.union(windows_to_elements.keys()) if all_windows != old_windows: merged_away = {} class TriggerMergeContext(WindowFn.MergeContext): def merge(_, to_be_merged, merge_result): for window in to_be_merged: if window != merge_result: merged_away[window] = merge_result state.merge(to_be_merged, merge_result) self.trigger_fn.on_merge( to_be_merged, merge_result, state.at(merge_result)) self.window_fn.merge(TriggerMergeContext(all_windows)) merged_windows_to_elements = collections.defaultdict(list) for window, values in windows_to_elements.items(): while window in merged_away: window = merged_away[window] merged_windows_to_elements[window].extend(values) windows_to_elements = merged_windows_to_elements for window in merged_away: state.clear_state(window, self.WATERMARK_HOLD) # Next handle element adding. for window, elements in windows_to_elements.items(): if state.get_state(window, self.TOMBSTONE): continue # Add watermark hold. # TODO(ccy): Add late data and garbage-collection hold support. output_time = self.output_time_fn_impl.merge( window, (element_output_time for element_output_time in (self.output_time_fn_impl.assign_output_time(window, timestamp) for unused_value, timestamp in elements) if element_output_time >= output_watermark)) if output_time is not None: state.add_state(window, self.WATERMARK_HOLD, output_time) context = state.at(window) for value, unused_timestamp in elements: state.add_state(window, self.ELEMENTS, value) self.trigger_fn.on_element(value, window, context) # Maybe fire this window. watermark = MIN_TIMESTAMP if self.trigger_fn.should_fire(watermark, window, context): finished = self.trigger_fn.on_fire(watermark, window, context) yield self._output(window, finished, state) def process_timer(self, window_id, unused_name, time_domain, timestamp, state): if self.is_merging: state = MergeableStateAdapter(state) window = state.get_window(window_id) if state.get_state(window, self.TOMBSTONE): return if time_domain == TimeDomain.WATERMARK: if not self.is_merging or window in state.known_windows(): context = state.at(window) if self.trigger_fn.should_fire(timestamp, window, context): finished = self.trigger_fn.on_fire(timestamp, window, context) yield self._output(window, finished, state) else: raise Exception('Unexpected time domain: %s' % time_domain) def _output(self, window, finished, state): """Output window and clean up if appropriate.""" values = state.get_state(window, self.ELEMENTS) if finished: # TODO(robertwb): allowed lateness state.clear_state(window, self.ELEMENTS) state.add_state(window, self.TOMBSTONE, 1) elif self.accumulation_mode == AccumulationMode.DISCARDING: state.clear_state(window, self.ELEMENTS) timestamp = state.get_state(window, self.WATERMARK_HOLD) if timestamp is None: # If no watermark hold was set, output at end of window. timestamp = window.end else: state.clear_state(window, self.WATERMARK_HOLD) return WindowedValue(values, timestamp, (window,)) class InMemoryUnmergedState(UnmergedState): """In-memory implementation of UnmergedState. Used for batch and testing. """ def __init__(self, defensive_copy=True): # TODO(robertwb): Skip defensive_copy in production if it's too expensive. self.timers = collections.defaultdict(dict) self.state = collections.defaultdict(lambda: collections.defaultdict(list)) self.global_state = {} self.defensive_copy = defensive_copy def set_global_state(self, tag, value): assert isinstance(tag, ValueStateTag) if self.defensive_copy: value = copy.deepcopy(value) self.global_state[tag.tag] = value def get_global_state(self, tag, default=None): return self.global_state.get(tag.tag, default) def set_timer(self, window, name, time_domain, timestamp): self.timers[window][(name, time_domain)] = timestamp def clear_timer(self, window, name, time_domain): self.timers[window].pop((name, time_domain), None) def get_window(self, window_id): return window_id def add_state(self, window, tag, value): if self.defensive_copy: value = copy.deepcopy(value) if isinstance(tag, ValueStateTag): self.state[window][tag.tag] = value elif isinstance(tag, CombiningValueStateTag): self.state[window][tag.tag].append(value) elif isinstance(tag, ListStateTag): self.state[window][tag.tag].append(value) elif isinstance(tag, WatermarkHoldStateTag): self.state[window][tag.tag].append(value) else: raise ValueError('Invalid tag.', tag) def get_state(self, window, tag): values = self.state[window][tag.tag] if isinstance(tag, ValueStateTag): return values elif isinstance(tag, CombiningValueStateTag): return tag.combine_fn.apply(values) elif isinstance(tag, ListStateTag): return values elif isinstance(tag, WatermarkHoldStateTag): return tag.output_time_fn_impl.combine_all(values) else: raise ValueError('Invalid tag.', tag) def clear_state(self, window, tag): self.state[window].pop(tag.tag, None) if not self.state[window]: self.state.pop(window, None) def get_and_clear_timers(self, watermark=MAX_TIMESTAMP): expired = [] for window, timers in list(self.timers.items()): for (name, time_domain), timestamp in list(timers.items()): if timestamp <= watermark: expired.append((window, (name, time_domain, timestamp))) del timers[(name, time_domain)] if not timers: del self.timers[window] return expired def __repr__(self): state_str = '\n'.join('%s: %s' % (key, dict(state)) for key, state in self.state.items()) return 'timers: %s\nstate: %s' % (dict(self.timers), state_str) DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/trigger_test.py000066400000000000000000000464271272340251000264710ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the triggering classes.""" import collections import os.path import unittest import yaml import google.cloud.dataflow as df from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.transforms.core import Windowing from google.cloud.dataflow.transforms.trigger import AccumulationMode from google.cloud.dataflow.transforms.trigger import AfterAll from google.cloud.dataflow.transforms.trigger import AfterCount from google.cloud.dataflow.transforms.trigger import AfterEach from google.cloud.dataflow.transforms.trigger import AfterFirst from google.cloud.dataflow.transforms.trigger import AfterWatermark from google.cloud.dataflow.transforms.trigger import DefaultTrigger from google.cloud.dataflow.transforms.trigger import GeneralTriggerDriver from google.cloud.dataflow.transforms.trigger import InMemoryUnmergedState from google.cloud.dataflow.transforms.trigger import Repeatedly from google.cloud.dataflow.transforms.util import assert_that, equal_to from google.cloud.dataflow.transforms.window import FixedWindows from google.cloud.dataflow.transforms.window import IntervalWindow from google.cloud.dataflow.transforms.window import MIN_TIMESTAMP from google.cloud.dataflow.transforms.window import OutputTimeFn from google.cloud.dataflow.transforms.window import Sessions from google.cloud.dataflow.transforms.window import TimestampedValue from google.cloud.dataflow.transforms.window import WindowedValue from google.cloud.dataflow.transforms.window import WindowFn class CustomTimestampingFixedWindowsWindowFn(FixedWindows): """WindowFn for testing custom timestamping.""" def get_transformed_output_time(self, unused_window, input_timestamp): return input_timestamp + 100 class TriggerTest(unittest.TestCase): def run_trigger_simple(self, window_fn, trigger_fn, accumulation_mode, timestamped_data, expected_panes, *groupings, **kwargs): late_data = kwargs.pop('late_data', []) assert not kwargs def bundle_data(data, size): bundle = [] for timestamp, elem in data: windows = window_fn.assign(WindowFn.AssignContext(timestamp, elem)) bundle.append(WindowedValue(elem, timestamp, windows)) if len(bundle) == size: yield bundle bundle = [] if bundle: yield bundle if not groupings: groupings = [1] for group_by in groupings: bundles = [] bundle = [] for timestamp, elem in timestamped_data: windows = window_fn.assign(WindowFn.AssignContext(timestamp, elem)) bundle.append(WindowedValue(elem, timestamp, windows)) if len(bundle) == group_by: bundles.append(bundle) bundle = [] bundles.append(bundle) self.run_trigger(window_fn, trigger_fn, accumulation_mode, bundle_data(timestamped_data, group_by), bundle_data(late_data, group_by), expected_panes) def run_trigger(self, window_fn, trigger_fn, accumulation_mode, bundles, late_bundles, expected_panes): actual_panes = collections.defaultdict(list) driver = GeneralTriggerDriver( Windowing(window_fn, trigger_fn, accumulation_mode)) state = InMemoryUnmergedState() for bundle in bundles: for wvalue in driver.process_elements(state, bundle, MIN_TIMESTAMP): window, = wvalue.windows actual_panes[window].append(set(wvalue.value)) while state.timers: for timer_window, (name, time_domain, timestamp) in ( state.get_and_clear_timers()): for wvalue in driver.process_timer( timer_window, name, time_domain, timestamp, state): window, = wvalue.windows actual_panes[window].append(set(wvalue.value)) for bundle in late_bundles: for wvalue in driver.process_elements(state, bundle, MIN_TIMESTAMP): window, = wvalue.windows actual_panes[window].append(set(wvalue.value)) while state.timers: for timer_window, (name, time_domain, timestamp) in ( state.get_and_clear_timers()): for wvalue in driver.process_timer( timer_window, name, time_domain, timestamp, state): window, = wvalue.windows actual_panes[window].append(set(wvalue.value)) self.assertEqual(expected_panes, actual_panes) def test_fixed_watermark(self): self.run_trigger_simple( FixedWindows(10), # pyformat break AfterWatermark(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (13, 'c')], {IntervalWindow(0, 10): [set('ab')], IntervalWindow(10, 20): [set('c')]}, 1, 2, 3) def test_fixed_watermark_with_early(self): self.run_trigger_simple( FixedWindows(10), # pyformat break AfterWatermark(early=AfterCount(2)), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(0, 10): [set('ab'), set('abc')]}, 2) self.run_trigger_simple( FixedWindows(10), # pyformat break AfterWatermark(early=AfterCount(2)), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(0, 10): [set('abc'), set('abc')]}, 3) def test_fixed_watermark_with_early_late(self): self.run_trigger_simple( FixedWindows(100), # pyformat break AfterWatermark(early=AfterCount(3), late=AfterCount(2)), AccumulationMode.DISCARDING, zip(range(9), 'abcdefghi'), {IntervalWindow(0, 100): [ set('abcd'), set('efgh'), # early set('i'), # on time set('vw'), set('xy') # late ]}, 2, late_data=zip(range(5), 'vwxyz')) def test_sessions_watermark_with_early_late(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterWatermark(early=AfterCount(2), late=AfterCount(1)), AccumulationMode.ACCUMULATING, [(1, 'a'), (15, 'b'), (7, 'c'), (30, 'd')], {IntervalWindow(1, 25): [ set('abc'), # early set('abc'), # on time set('abcxy') # late ], IntervalWindow(30, 40): [ set('d'), # on time ], IntervalWindow(1, 40): [ set('abcdxyz') # late ], }, 2, late_data=[(1, 'x'), (2, 'y'), (21, 'z')]) def test_fixed_after_count(self): self.run_trigger_simple( FixedWindows(10), # pyformat break AfterCount(2), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c'), (11, 'z')], {IntervalWindow(0, 10): [set('ab')]}, 1, 2) self.run_trigger_simple( FixedWindows(10), # pyformat break AfterCount(2), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c'), (11, 'z')], {IntervalWindow(0, 10): [set('abc')]}, 3, 4) def test_fixed_after_first(self): self.run_trigger_simple( FixedWindows(10), # pyformat break AfterFirst(AfterCount(2), AfterWatermark()), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(0, 10): [set('ab')]}, 1, 2) self.run_trigger_simple( FixedWindows(10), # pyformat break AfterFirst(AfterCount(5), AfterWatermark()), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(0, 10): [set('abc')]}, 1, 2, late_data=[(1, 'x'), (2, 'y'), (3, 'z')]) def test_repeatedly_after_first(self): self.run_trigger_simple( FixedWindows(100), # pyformat break Repeatedly(AfterFirst(AfterCount(3), AfterWatermark())), AccumulationMode.ACCUMULATING, zip(range(7), 'abcdefg'), {IntervalWindow(0, 100): [ set('abc'), set('abcdef'), set('abcdefg'), set('abcdefgx'), set('abcdefgxy'), set('abcdefgxyz')]}, 1, late_data=zip(range(3), 'xyz')) def test_sessions_after_all(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterAll(AfterCount(2), AfterWatermark()), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(1, 13): [set('abc')]}, 1, 2) self.run_trigger_simple( Sessions(10), # pyformat break AfterAll(AfterCount(5), AfterWatermark()), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(1, 13): [set('abcxy')]}, 1, 2, late_data=[(1, 'x'), (2, 'y'), (3, 'z')]) def test_sessions_default(self): self.run_trigger_simple( Sessions(10), # pyformat break DefaultTrigger(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b')], {IntervalWindow(1, 12): [set('ab')]}, 1, 2) self.run_trigger_simple( Sessions(10), # pyformat break AfterWatermark(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (15, 'c'), (16, 'd'), (30, 'z'), (9, 'e'), (10, 'f'), (30, 'y')], {IntervalWindow(1, 26): [set('abcdef')], IntervalWindow(30, 40): [set('yz')]}, 1, 2, 3, 4, 5, 6) def test_sessions_watermark(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterWatermark(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b')], {IntervalWindow(1, 12): [set('ab')]}, 1, 2) self.run_trigger_simple( Sessions(10), # pyformat break AfterWatermark(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (15, 'c'), (16, 'd'), (30, 'z'), (9, 'e'), (10, 'f'), (30, 'y')], {IntervalWindow(1, 26): [set('abcdef')], IntervalWindow(30, 40): [set('yz')]}, 1, 2, 3, 4, 5, 6) def test_sessions_after_count(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterCount(2), AccumulationMode.ACCUMULATING, [(1, 'a'), (15, 'b'), (6, 'c'), (30, 's'), (31, 't'), (50, 'z'), (50, 'y')], {IntervalWindow(1, 25): [set('abc')], IntervalWindow(30, 41): [set('st')], IntervalWindow(50, 60): [set('yz')]}, 1, 2, 3) def test_sessions_repeatedly_after_count(self): self.run_trigger_simple( Sessions(10), # pyformat break Repeatedly(AfterCount(2)), AccumulationMode.ACCUMULATING, [(1, 'a'), (15, 'b'), (6, 'c'), (2, 'd'), (7, 'e')], {IntervalWindow(1, 25): [set('abc'), set('abcde')]}, 1, 3) self.run_trigger_simple( Sessions(10), # pyformat break Repeatedly(AfterCount(2)), AccumulationMode.DISCARDING, [(1, 'a'), (15, 'b'), (6, 'c'), (2, 'd'), (7, 'e')], {IntervalWindow(1, 25): [set('abc'), set('de')]}, 1, 3) def test_sessions_after_each(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterEach(AfterCount(2), AfterCount(3)), AccumulationMode.ACCUMULATING, zip(range(10), 'abcdefghij'), {IntervalWindow(0, 11): [set('ab')], IntervalWindow(0, 15): [set('abcdef')]}, 2) self.run_trigger_simple( Sessions(10), # pyformat break Repeatedly(AfterEach(AfterCount(2), AfterCount(3))), AccumulationMode.ACCUMULATING, zip(range(10), 'abcdefghij'), {IntervalWindow(0, 11): [set('ab')], IntervalWindow(0, 15): [set('abcdef')], IntervalWindow(0, 17): [set('abcdefgh')]}, 2) class TriggerPipelineTest(unittest.TestCase): def test_after_count(self): p = Pipeline('DirectPipelineRunner') result = (p | df.Create([1, 2, 3, 4, 5, 10, 11]) | df.FlatMap(lambda t: [('A', t), ('B', t + 5)]) | df.Map(lambda (k, t): TimestampedValue((k, t), t)) | df.WindowInto(FixedWindows(10), trigger=AfterCount(3), accumulation_mode=AccumulationMode.DISCARDING) | df.GroupByKey() | df.Map(lambda (k, v): ('%s-%s' % (k, len(v)), set(v)))) assert_that(result, equal_to( { 'A-5': {1, 2, 3, 4, 5}, # A-10, A-11 never emitted due to AfterCount(3) never firing. 'B-4': {6, 7, 8, 9}, 'B-3': {10, 15, 16}, }.iteritems())) class TranscriptTest(unittest.TestCase): # We must prepend an underscore to this name so that the open-source unittest # runner does not execute this method directly as a test. @classmethod def _create_test(cls, spec): counter = 0 name = spec.get('name', 'unnamed') unique_name = 'test_' + name while hasattr(cls, unique_name): counter += 1 unique_name = 'test_%s_%d' % (name, counter) setattr(cls, unique_name, lambda self: self._run_log_test(spec)) # We must prepend an underscore to this name so that the open-source unittest # runner does not execute this method directly as a test. @classmethod def _create_tests(cls, transcript_filename): for spec in yaml.load_all(open(transcript_filename)): cls._create_test(spec) def _run_log_test(self, spec): if 'error' in spec: self.assertRaisesRegexp( AssertionError, spec['error'], self._run_log, spec) else: self._run_log(spec) def _run_log(self, spec): def parse_int_list(s): """Parses strings like '[1, 2, 3]'.""" s = s.strip() assert s[0] == '[' and s[-1] == ']', s if not s[1:-1].strip(): return [] else: return [int(x) for x in s[1:-1].split(',')] def split_args(s): """Splits 'a, b, [c, d]' into ['a', 'b', '[c, d]'].""" args = [] start = 0 depth = 0 for ix in xrange(len(s)): c = s[ix] if c in '({[': depth += 1 elif c in ')}]': depth -= 1 elif c == ',' and depth == 0: args.append(s[start:ix].strip()) start = ix + 1 assert depth == 0, s args.append(s[start:].strip()) return args def parse(s, names): """Parse (recursive) 'Foo(arg, kw=arg)' for Foo in the names dict.""" s = s.strip() if s in names: return names[s] elif s[0] == '[': return parse_int_list(s) elif '(' in s: assert s[-1] == ')', s callee = parse(s[:s.index('(')], names) posargs = [] kwargs = {} for arg in split_args(s[s.index('(') + 1:-1]): if '=' in arg: kw, value = arg.split('=', 1) kwargs[kw] = parse(value, names) else: posargs.append(parse(arg, names)) return callee(*posargs, **kwargs) else: try: return int(s) except ValueError: raise ValueError('Unknown function: %s' % s) def parse_fn(s, names): """Like parse(), but implicitly calls no-arg constructors.""" fn = parse(s, names) if isinstance(fn, type): return fn() else: return fn # pylint: disable=g-import-not-at-top from google.cloud.dataflow.transforms import window as window_module from google.cloud.dataflow.transforms import trigger as trigger_module # pylint: enable=g-import-not-at-top window_fn_names = dict(window_module.__dict__) window_fn_names.update({'CustomTimestampingFixedWindowsWindowFn': CustomTimestampingFixedWindowsWindowFn}) trigger_names = {'Default': DefaultTrigger} trigger_names.update(trigger_module.__dict__) window_fn = parse_fn(spec.get('window_fn', 'GlobalWindows'), window_fn_names) trigger_fn = parse_fn(spec.get('trigger_fn', 'Default'), trigger_names) accumulation_mode = getattr( AccumulationMode, spec.get('accumulation_mode', 'ACCUMULATING').upper()) output_time_fn = getattr( OutputTimeFn, spec.get('output_time_fn', 'OUTPUT_AT_EOW').upper()) allowed_lateness = float(spec.get('allowed_lateness', '-inf')) driver = GeneralTriggerDriver( Windowing(window_fn, trigger_fn, accumulation_mode, output_time_fn)) state = InMemoryUnmergedState() output = [] watermark = MIN_TIMESTAMP def fire_timers(): to_fire = state.get_and_clear_timers(watermark) while to_fire: for timer_window, (name, time_domain, t_timestamp) in to_fire: for wvalue in driver.process_timer( timer_window, name, time_domain, t_timestamp, state): window, = wvalue.windows output.append({'window': [window.start, window.end - 1], 'values': sorted(wvalue.value), 'timestamp': wvalue.timestamp}) to_fire = state.get_and_clear_timers(watermark) for line in spec['transcript']: action, params = line.items()[0] if action != 'expect': # Fail if we have output that was not expected in the transcript. self.assertEquals( [], output, msg='Unexpected output: %s before %s' % (output, line)) if action == 'input': bundle = [ WindowedValue(t, t, window_fn.assign(WindowFn.AssignContext(t, t))) for t in params] output = [{'window': [wvalue.windows[0].start, wvalue.windows[0].end - 1], 'values': sorted(wvalue.value), 'timestamp': wvalue.timestamp} for wvalue in driver.process_elements(state, bundle, watermark)] fire_timers() elif action == 'watermark': watermark = params fire_timers() elif action == 'expect': for expected_output in params: for candidate in output: if all(candidate[k] == expected_output[k] for k in candidate if k in expected_output): output.remove(candidate) break else: self.fail('Unmatched output %s in %s' % (expected_output, output)) elif action == 'state': # TODO(robertwb): Implement once we support allowed lateness. pass else: self.fail('Unknown action: ' + action) # Fail if we have output that was not expected in the transcript. self.assertEquals([], output, msg='Unexpected output: %s' % output) TRANSCRIPT_TEST_FILE = os.path.join(os.path.dirname(__file__), 'trigger_transcripts.yaml') if os.path.exists(TRANSCRIPT_TEST_FILE): TranscriptTest._create_tests(TRANSCRIPT_TEST_FILE) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/trigger_transcripts.yaml000066400000000000000000000134241272340251000303670ustar00rootroot00000000000000name: fixed_default window_fn: FixedWindows(10) # Python names/syntax, unless otherwise noted. trigger_fn: Default # Same. Empty () may be omitted. transcript: # Ordered list of events. - input: [1, 2, 3, 10, 11] # The elements are the timestamps. - watermark: 25 - expect: # Every expected output from the last action. - {window: [0, 9], values: [1, 2, 3]} - {window: [10, 19], values: [10, 11]} # Partial match on attributes OK. --- name: fixed_default_late_data window_fn: FixedWindows(10) trigger_fn: Default output_time_fn: OUTPUT_AT_EOW transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 10, final: false} - {window: [10, 19], values: [10, 11], timestamp: 20} - {window: [20, 29], values: [25], timestamp: 30, late: false} - input: [7] - expect: - {window: [0, 9], values: [1, 2, 3, 7], timestamp: 10, late: true} --- name: output_time_fn_earliest window_fn: FixedWindows(10) trigger_fn: Default output_time_fn: OUTPUT_AT_EARLIEST transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 1, final: false} - {window: [10, 19], values: [10, 11], timestamp: 10} - {window: [20, 29], values: [25], timestamp: 25, late: false} --- name: output_time_fn_latest window_fn: FixedWindows(10) trigger_fn: Default output_time_fn: OUTPUT_AT_LATEST transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 3, final: false} - {window: [10, 19], values: [10, 11], timestamp: 11} - {window: [20, 29], values: [25], timestamp: 25, late: false} --- # Test that custom timestamping is not invoked. name: output_time_fn_custom_timestamping_eow window_fn: CustomTimestampingFixedWindowsWindowFn(10) trigger_fn: Default output_time_fn: OUTPUT_AT_EOW transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 10, final: false} - {window: [10, 19], values: [10, 11], timestamp: 20} - {window: [20, 29], values: [25], timestamp: 30, late: false} --- # Test that custom timestamping is not invoked. name: output_time_fn_custom_timestamping_earliest window_fn: CustomTimestampingFixedWindowsWindowFn(10) trigger_fn: Default output_time_fn: OUTPUT_AT_EARLIEST transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 1, final: false} - {window: [10, 19], values: [10, 11], timestamp: 10} - {window: [20, 29], values: [25], timestamp: 25, late: false} --- # Test that custom timestamping is in fact invoked. name: output_time_fn_custom_timestamping_earliest window_fn: CustomTimestampingFixedWindowsWindowFn(10) trigger_fn: Default output_time_fn: OUTPUT_AT_EARLIEST_TRANSFORMED transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 101, final: false} - {window: [10, 19], values: [10, 11], timestamp: 110} - {window: [20, 29], values: [25], timestamp: 125, late: false} --- name: early_late_sessions window_fn: Sessions(10) trigger_fn: AfterWatermark(early=AfterCount(2), late=AfterCount(3)) output_time_fn: OUTPUT_AT_EOW transcript: - input: [1, 2, 3] - expect: - {window: [1, 12], values: [1, 2, 3], timestamp: 13, early: true} - input: [4] # no output - input: [5] - expect: - {window: [1, 14], values: [1, 2, 3, 4, 5], timestamp: 14, early: true} - input: [6] - watermark: 100 - expect: - {window: [1, 15], values:[1, 2, 3, 4, 5, 6], timestamp: 16, final: true} - input: [1] - input: [3, 4] - expect: - {window: [1, 15], values: [1, 1, 2, 3, 3, 4, 4, 5, 6], timestamp: 16} --- name: garbage_collection window_fn: FixedWindows(10) trigger_fn: AfterCount(2) output_time_fn: OUTPUT_AT_EOW allowed_lateness: 10 accumulation_mode: discarding transcript: - input: [1, 2, 3, 10, 11, 25] - expect: - {window: [0, 9], timestamp: 10} - {window: [10, 19], timestamp: 20} - state: present: [[20, 29]] absent: [[0, 9]] tombstone: [[10, 19]] --- name: known_late_data_watermark window_fn: FixedWindows(10) trigger_fn: Default output_time_fn: OUTPUT_AT_EARLIEST transcript: - watermark: 5 - input: [2, 3, 7, 8] - watermark: 11 - expect: - {window: [0, 9], values: [2, 3, 7, 8], timestamp: 7} --- name: known_late_data_no_watermark_hold_possible window_fn: FixedWindows(10) trigger_fn: Default output_time_fn: OUTPUT_AT_EARLIEST transcript: - watermark: 8 - input: [2, 3, 7] - watermark: 11 - expect: - {window: [0, 9], values: [2, 3, 7], timestamp: 10} # These next examples test that bad/incomplete transcripts are rejected. --- name: bad_output error: Unmatched output windowfn: FixedWindows(10) transcript: - input: [1, 2, 3] - expect: - {window: [0, 9], values: [1, 2, 3]} # bad - watermark: 100 --- name: bad_expected_values error: Unmatched output window_fn: FixedWindows(10) transcript: - input: [1, 2, 3] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2]} # bad values --- name: bad_expected_window error: Unmatched output window_fn: FixedWindows(10) transcript: - input: [1, 2, 3] - watermark: 100 - expect: - {window: [0, 19], values: [1, 2, 3]} # bad window --- name: missing_output error: Unexpected output window_fn: FixedWindows(10) transcript: - input: [1, 2, 3] - watermark: 100 # missing output - watermark: 200 --- name: missing_output_at_end error: Unexpected output window_fn: FixedWindows(10) transcript: - input: [1, 2, 3] - watermark: 100 # missing output DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/util.py000066400000000000000000000200411272340251000247240ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Simple utility PTransforms. """ from __future__ import absolute_import import collections import operator from google.cloud.dataflow.pvalue import AsIter as AllOf from google.cloud.dataflow.transforms.core import CombinePerKey, Create, Flatten, GroupByKey, Map from google.cloud.dataflow.transforms.ptransform import PTransform from google.cloud.dataflow.transforms.ptransform import ptransform_fn __all__ = [ 'CoGroupByKey', 'Keys', 'KvSwap', 'RemoveDuplicates', 'Values', 'assert_that', 'equal_to', 'is_empty', ] class CoGroupByKey(PTransform): """Groups results across several PCollections by key. Given an input dict mapping serializable keys (called "tags") to 0 or more PCollections of (key, value) tuples, e.g.:: {'pc1': pcoll1, 'pc2': pcoll2, 33333: pcoll3} creates a single output PCollection of (key, value) tuples whose keys are the unique input keys from all inputs, and whose values are dicts mapping each tag to an iterable of whatever values were under the key in the corresponding PCollection:: ('some key', {'pc1': ['value 1 under "some key" in pcoll1', 'value 2 under "some key" in pcoll1'], 'pc2': [], 33333: ['only value under "some key" in pcoll3']}) Note that pcoll2 had no values associated with "some key". CoGroupByKey also works for tuples, lists, or other flat iterables of PCollections, in which case the values of the resulting PCollections will be tuples whose nth value is the list of values from the nth PCollection---conceptually, the "tags" are the indices into the input. Thus, for this input:: (pcoll1, pcoll2, pcoll3) the output PCollection's value for "some key" is:: ('some key', (['value 1 under "some key" in pcoll1', 'value 2 under "some key" in pcoll1'], [], ['only value under "some key" in pcoll3'])) Args: label: name of this transform instance. Useful while monitoring and debugging a pipeline execution. **kwargs: Accepts a single named argument "pipeline", which specifies the pipeline that "owns" this PTransform. Ordinarily CoGroupByKey can obtain this information from one of the input PCollections, but if there are none (or if there's a chance there may be none), this argument is the only way to provide pipeline information, and should be considered mandatory. """ def __init__(self, label=None, **kwargs): super(CoGroupByKey, self).__init__(label) self.pipeline = kwargs.pop('pipeline', None) if kwargs: raise ValueError('Unexpected keyword arguments: %s' % kwargs.keys()) def _extract_input_pvalues(self, pvalueish): try: # If this works, it's a dict. return pvalueish, tuple(pvalueish.viewvalues()) except AttributeError: pcolls = tuple(pvalueish) return pcolls, pcolls def apply(self, pcolls): """Performs CoGroupByKey on argument pcolls; see class docstring.""" # For associating values in K-V pairs with the PCollections they came from. def _pair_tag_with_value((key, value), tag): return (key, (tag, value)) # Creates the key, value pairs for the output PCollection. Values are either # lists or dicts (per the class docstring), initialized by the result of # result_ctor(result_ctor_arg). def _merge_tagged_vals_under_key((key, grouped), result_ctor, result_ctor_arg): result_value = result_ctor(result_ctor_arg) for tag, value in grouped: result_value[tag].append(value) return (key, result_value) try: # If pcolls is a dict, we turn it into (tag, pcoll) pairs for use in the # general-purpose code below. The result value constructor creates dicts # whose keys are the tags. result_ctor_arg = pcolls.keys() result_ctor = lambda tags: dict((tag, []) for tag in tags) pcolls = pcolls.items() except AttributeError: # Otherwise, pcolls is a list/tuple, so we turn it into (index, pcoll) # pairs. The result value constructor makes tuples with len(pcolls) slots. pcolls = list(enumerate(pcolls)) result_ctor_arg = len(pcolls) result_ctor = lambda size: tuple([] for _ in xrange(size)) # Check input PCollections for PCollection-ness, and that they all belong # to the same pipeline. for _, pcoll in pcolls: self._check_pcollection(pcoll) if self.pipeline: assert pcoll.pipeline == self.pipeline return ([pcoll | Map('pair_with_%s' % tag, _pair_tag_with_value, tag) for tag, pcoll in pcolls] | Flatten(pipeline=self.pipeline) | GroupByKey() | Map(_merge_tagged_vals_under_key, result_ctor, result_ctor_arg)) def Keys(label='Keys'): # pylint: disable=invalid-name """Produces a PCollection of first elements of 2-tuples in a PCollection.""" return Map(label, lambda (k, v): k) def Values(label='Values'): # pylint: disable=invalid-name """Produces a PCollection of second elements of 2-tuples in a PCollection.""" return Map(label, lambda (k, v): v) def KvSwap(label='KvSwap'): # pylint: disable=invalid-name """Produces a PCollection reversing 2-tuples in a PCollection.""" return Map(label, lambda (k, v): (v, k)) @ptransform_fn def RemoveDuplicates(label, pcoll): # pylint: disable=invalid-name """Produces a PCollection containing the unique elements of a PCollection.""" return (pcoll | Map('%s:ToPairs' % label, lambda v: (v, None)) | CombinePerKey('%s:Group' % label, lambda vs: None) | Keys('%s:RemoveDuplicates' % label)) class DataflowAssertException(Exception): """Exception raised by matcher classes used by assert_that transform.""" pass # Note that equal_to always sorts the expected and actual since what we # compare are PCollections for which there is no guaranteed order. # However the sorting does not go beyond top level therefore [1,2] and [2,1] # are considered equal and [[1,2]] and [[2,1]] are not. # TODO(silviuc): Add contains_in_any_order-style matchers. def equal_to(expected): expected = list(expected) def _equal(actual): sorted_expected = sorted(expected) sorted_actual = sorted(actual) if sorted_expected != sorted_actual: raise DataflowAssertException( 'Failed assert: %r == %r' % (sorted_expected, sorted_actual)) return _equal def is_empty(): def _empty(actual): if actual: raise DataflowAssertException( 'Failed assert: [] == %r' % actual) return _empty def assert_that(actual, matcher, label='assert_that'): """A PTransform that checks a PCollection has an expected value. Note that assert_that should be used only for testing pipelines since the check relies on materializing the entire PCollection being checked. Args: actual: A PCollection. matcher: A matcher function taking as argument the actual value of a materialized PCollection. The matcher validates this actual value against expectations and raises DataflowAssertException if they are not met. label: Optional string label. This is needed in case several assert_that transforms are introduced in the same pipeline. Returns: Ignored. """ def match(_, actual): matcher(actual) class AssertThat(PTransform): def apply(self, pipeline): return pipeline | Create('singleton', [None]) | Map(match, AllOf(actual)) def default_label(self): return label actual.pipeline | AssertThat() DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/window.py000066400000000000000000000306551272340251000252720ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Windowing concepts. A WindowInto transform logically divides up or groups the elements of a PCollection into finite windows according to a windowing function (derived from WindowFn). The output of WindowInto contains the same elements as input, but they have been logically assigned to windows. The next GroupByKey(s) transforms, including one within a composite transform, will group by the combination of keys and windows. Windowing a PCollection allows chunks of it to be processed individually, before the entire PCollection is available. This is especially important for PCollection(s) with unbounded size, since the full PCollection is never available at once, since more data is continually arriving. For PCollection(s) with a bounded size (aka. conventional batch mode), by default, all data is implicitly in a single window (see GlobalWindows), unless WindowInto is applied. For example, a simple form of windowing divides up the data into fixed-width time intervals, using FixedWindows. Seconds are used as the time unit for the built-in windowing primitives here. Integer or floating point seconds can be passed to these primitives. Internally, seconds, with microsecond granularity, are stored as timeutil.Timestamp and timeutil.Duration objects. This is done to avoid precision errors that would occur with floating point representations. Custom windowing function classes can be created, by subclassing from WindowFn. """ from __future__ import absolute_import from google.cloud.dataflow import coders from google.cloud.dataflow.transforms import timeutil from google.cloud.dataflow.transforms.timeutil import Duration from google.cloud.dataflow.transforms.timeutil import MAX_TIMESTAMP from google.cloud.dataflow.transforms.timeutil import MIN_TIMESTAMP from google.cloud.dataflow.transforms.timeutil import Timestamp # TODO(ccy): revisit naming and semantics once Java Apache Beam finalizes their # behavior. class OutputTimeFn(object): """Determines how output timestamps of grouping operations are assigned.""" OUTPUT_AT_EOW = 'OUTPUT_AT_EOW' OUTPUT_AT_EARLIEST = 'OUTPUT_AT_EARLIEST' OUTPUT_AT_LATEST = 'OUTPUT_AT_LATEST' OUTPUT_AT_EARLIEST_TRANSFORMED = 'OUTPUT_AT_EARLIEST_TRANSFORMED' @staticmethod def get_impl(output_time_fn, window_fn): if output_time_fn == OutputTimeFn.OUTPUT_AT_EOW: return timeutil.OutputAtEndOfWindowImpl() elif output_time_fn == OutputTimeFn.OUTPUT_AT_EARLIEST: return timeutil.OutputAtEarliestInputTimestampImpl() elif output_time_fn == OutputTimeFn.OUTPUT_AT_LATEST: return timeutil.OutputAtLatestInputTimestampImpl() elif output_time_fn == OutputTimeFn.OUTPUT_AT_EARLIEST_TRANSFORMED: return timeutil.OutputAtEarliestTransformedInputTimestampImpl(window_fn) else: raise ValueError('Invalid OutputTimeFn: %s.' % output_time_fn) class WindowFn(object): """An abstract windowing function defining a basic assign and merge.""" class AssignContext(object): """Context passed to WindowFn.assign().""" def __init__(self, timestamp, element=None, existing_windows=None): self.timestamp = Timestamp.of(timestamp) self.element = element self.existing_windows = existing_windows def assign(self, assign_context): """Associates a timestamp and set of windows to an element.""" raise NotImplementedError class MergeContext(object): """Context passed to WindowFn.merge() to perform merging, if any.""" def __init__(self, windows): self.windows = list(windows) def merge(self, to_be_merged, merge_result): raise NotImplementedError def merge(self, merge_context): """Returns a window that is the result of merging a set of windows.""" raise NotImplementedError def get_window_coder(self): return coders.PickleCoder() def get_transformed_output_time(self, window, input_timestamp): # pylint: disable=unused-argument """Given input time and output window, returns output time for window. If OutputTimeFn.OUTPUT_AT_EARLIEST_TRANSFORMED is used in the Windowing, the output timestamp for the given window will be the earliest of the timestamps returned by get_transformed_output_time() for elements of the window. Arguments: window: Output window of element. input_timestamp: Input timestamp of element as a timeutil.Timestamp object. Returns: Transformed timestamp. """ # By default, just return the input timestamp. return input_timestamp class BoundedWindow(object): """A window for timestamps in range (-infinity, end). Attributes: end: End of window. """ def __init__(self, end): self.end = Timestamp.of(end) def __cmp__(self, other): # Order first by endpoint, then arbitrarily. return cmp(self.end, other.end) or cmp(hash(self), hash(other)) def __eq__(self, other): raise NotImplementedError def __hash__(self): return hash(self.end) def __repr__(self): return '[?, %s)' % float(self.end) class IntervalWindow(BoundedWindow): """A window for timestamps in range [start, end). Attributes: start: Start of window as seconds since Unix epoch. end: End of window as seconds since Unix epoch. """ def __init__(self, start, end): super(IntervalWindow, self).__init__(end) self.start = Timestamp.of(start) def __hash__(self): return hash((self.start, self.end)) def __eq__(self, other): return self.start == other.start and self.end == other.end def __repr__(self): return '[%s, %s)' % (float(self.start), float(self.end)) def intersects(self, other): return other.start < self.end or self.start < other.end def union(self, other): return IntervalWindow( min(self.start, other.start), max(self.end, other.end)) class WindowedValue(object): """A windowed value having a value, a timestamp and set of windows. Attributes: value: The underlying value of a windowed value. timestamp: Timestamp associated with the value as seconds since Unix epoch. windows: A set (iterable) of window objects for the value. The window object are descendants of the BoundedWindow class. """ def __init__(self, value, timestamp, windows): self.value = value self.timestamp = Timestamp.of(timestamp) self.windows = windows def __repr__(self): return '(%s, %s, %s)' % ( repr(self.value), 'MIN_TIMESTAMP' if self.timestamp == MIN_TIMESTAMP else 'MAX_TIMESTAMP' if self.timestamp == MAX_TIMESTAMP else float(self.timestamp), self.windows) def __hash__(self): return hash((self.value, self.timestamp, self.windows)) def __eq__(self, other): return (type(self) == type(other) and self.value == other.value and self.timestamp == other.timestamp and self.windows == other.windows) def with_value(self, new_value): return WindowedValue(new_value, self.timestamp, self.windows) class TimestampedValue(object): """A timestamped value having a value and a timestamp. Attributes: value: The underlying value. timestamp: Timestamp associated with the value as seconds since Unix epoch. """ def __init__(self, value, timestamp): self.value = value self.timestamp = Timestamp.of(timestamp) class GlobalWindow(BoundedWindow): """The default window into which all data is placed (via GlobalWindows).""" _instance = None def __new__(cls): if cls._instance is None: cls._instance = super(GlobalWindow, cls).__new__(cls) return cls._instance def __init__(self): super(GlobalWindow, self).__init__(MAX_TIMESTAMP) self.start = MIN_TIMESTAMP def __repr__(self): return 'GlobalWindow' def __hash__(self): return hash(type(self)) def __eq__(self, other): # Global windows are always and only equal to each other. return self is other or type(self) is type(other) class GlobalWindows(WindowFn): """A windowing function that assigns everything to one global window.""" @classmethod def windowed_value(cls, value, timestamp=MIN_TIMESTAMP): return WindowedValue(value, timestamp, [GlobalWindow()]) def assign(self, assign_context): return [GlobalWindow()] def merge(self, merge_context): pass # No merging. def get_window_coder(self): return coders.SingletonCoder(GlobalWindow()) def __hash__(self): return hash(type(self)) def __eq__(self, other): # Global windowfn is always and only equal to each other. return self is other or type(self) is type(other) def __ne__(self, other): return not self == other class FixedWindows(WindowFn): """A windowing function that assigns each element to one time interval. The attributes size and offset determine in what time interval a timestamp will be slotted. The time intervals have the following formula: [N * size + offset, (N + 1) * size + offset) Attributes: size: Size of the window as seconds. offset: Offset of this window as seconds since Unix epoch. Windows start at t=N * size + offset where t=0 is the epoch. The offset must be a value in range [0, size). If it is not it will be normalized to this range. """ def __init__(self, size, offset=0): if size <= 0: raise ValueError('The size parameter must be strictly positive.') self.size = Duration.of(size) self.offset = Timestamp.of(offset) % self.size def assign(self, context): timestamp = context.timestamp start = timestamp - (timestamp - self.offset) % self.size return [IntervalWindow(start, start + self.size)] def merge(self, merge_context): pass # No merging. class SlidingWindows(WindowFn): """A windowing function that assigns each element to a set of sliding windows. The attributes size and offset determine in what time interval a timestamp will be slotted. The time intervals have the following formula: [N * period + offset, N * period + offset + size) Attributes: size: Size of the window as seconds. period: Period of the windows as seconds. offset: Offset of this window as seconds since Unix epoch. Windows start at t=N * period + offset where t=0 is the epoch. The offset must be a value in range [0, period). If it is not it will be normalized to this range. """ def __init__(self, size, period, offset=0): if size <= 0: raise ValueError('The size parameter must be strictly positive.') self.size = Duration.of(size) self.period = Duration.of(period) self.offset = Timestamp.of(offset) % size def assign(self, context): timestamp = context.timestamp start = timestamp - (timestamp - self.offset) % self.period return [IntervalWindow(Timestamp.of(s), Timestamp.of(s) + self.size) for s in range(start, start - self.size, -self.period)] def merge(self, merge_context): pass # No merging. class Sessions(WindowFn): """A windowing function that groups elements into sessions. A session is defined as a series of consecutive events separated by a specified gap size. Attributes: gap_size: Size of the gap between windows as floating-point seconds. """ def __init__(self, gap_size): if gap_size <= 0: raise ValueError('The size parameter must be strictly positive.') self.gap_size = Duration.of(gap_size) def assign(self, context): timestamp = context.timestamp return [IntervalWindow(timestamp, timestamp + self.gap_size)] def merge(self, merge_context): to_merge = [] for w in sorted(merge_context.windows, key=lambda w: w.start): if to_merge: if end > w.start: to_merge.append(w) if w.end > end: end = w.end else: if len(to_merge) > 1: merge_context.merge(to_merge, IntervalWindow(to_merge[0].start, end)) to_merge = [w] end = w.end else: to_merge = [w] end = w.end if len(to_merge) > 1: merge_context.merge(to_merge, IntervalWindow(to_merge[0].start, end)) DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/window_test.py000066400000000000000000000201261272340251000263210ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the windowing classes.""" import unittest from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.transforms import CombinePerKey from google.cloud.dataflow.transforms import combiners from google.cloud.dataflow.transforms import core from google.cloud.dataflow.transforms import Create from google.cloud.dataflow.transforms import GroupByKey from google.cloud.dataflow.transforms import Map from google.cloud.dataflow.transforms import window from google.cloud.dataflow.transforms import WindowInto from google.cloud.dataflow.transforms.util import assert_that, equal_to from google.cloud.dataflow.transforms.window import FixedWindows from google.cloud.dataflow.transforms.window import IntervalWindow from google.cloud.dataflow.transforms.window import Sessions from google.cloud.dataflow.transforms.window import SlidingWindows from google.cloud.dataflow.transforms.window import TimestampedValue from google.cloud.dataflow.transforms.window import WindowedValue from google.cloud.dataflow.transforms.window import WindowFn def context(element, timestamp, windows): return WindowFn.AssignContext(timestamp, element, windows) sort_values = Map(lambda (k, vs): (k, sorted(vs))) class ReifyWindowsFn(core.DoFn): def process(self, context): key, values = context.element for window in context.windows: yield "%s @ %s" % (key, window), values reify_windows = core.ParDo(ReifyWindowsFn()) class WindowTest(unittest.TestCase): def test_fixed_windows(self): # Test windows with offset: 2, 7, 12, 17, ... windowfn = window.FixedWindows(size=5, offset=2) self.assertEqual([window.IntervalWindow(7, 12)], windowfn.assign(context('v', 7, []))) self.assertEqual([window.IntervalWindow(7, 12)], windowfn.assign(context('v', 11, []))) self.assertEqual([window.IntervalWindow(12, 17)], windowfn.assign(context('v', 12, []))) # Test windows without offset: 0, 5, 10, 15, ... windowfn = window.FixedWindows(size=5) self.assertEqual([window.IntervalWindow(5, 10)], windowfn.assign(context('v', 5, []))) self.assertEqual([window.IntervalWindow(5, 10)], windowfn.assign(context('v', 9, []))) self.assertEqual([window.IntervalWindow(10, 15)], windowfn.assign(context('v', 10, []))) # Test windows with offset out of range. windowfn = window.FixedWindows(size=5, offset=12) self.assertEqual([window.IntervalWindow(7, 12)], windowfn.assign(context('v', 11, []))) def test_sliding_windows_assignment(self): windowfn = SlidingWindows(size=15, period=5, offset=2) expected = [IntervalWindow(7, 22), IntervalWindow(2, 17), IntervalWindow(-3, 12)] self.assertEqual(expected, windowfn.assign(context('v', 7, []))) self.assertEqual(expected, windowfn.assign(context('v', 8, []))) self.assertEqual(expected, windowfn.assign(context('v', 11, []))) def test_sessions_merging(self): windowfn = Sessions(10) def merge(*timestamps): windows = [windowfn.assign(context(None, t, [])) for t in timestamps] running = set() class TestMergeContext(WindowFn.MergeContext): def __init__(self): super(TestMergeContext, self).__init__(running) def merge(self, to_be_merged, merge_result): for w in to_be_merged: if w in running: running.remove(w) running.add(merge_result) for ws in windows: running.update(ws) windowfn.merge(TestMergeContext()) windowfn.merge(TestMergeContext()) return sorted(running) self.assertEqual([IntervalWindow(2, 12)], merge(2)) self.assertEqual([IntervalWindow(2, 12), IntervalWindow(19, 29)], merge(2, 19)) self.assertEqual([IntervalWindow(2, 19)], merge(2, 9)) self.assertEqual([IntervalWindow(2, 19)], merge(9, 2)) self.assertEqual([IntervalWindow(2, 19), IntervalWindow(19, 29)], merge(2, 9, 19)) self.assertEqual([IntervalWindow(2, 19), IntervalWindow(19, 29)], merge(19, 9, 2)) self.assertEqual([IntervalWindow(2, 25)], merge(2, 15, 10)) def timestamped_key_values(self, pipeline, key, *timestamps): return (pipeline | Create('start', timestamps) | Map(lambda x: WindowedValue((key, x), x, []))) def test_sliding_windows(self): p = Pipeline('DirectPipelineRunner') pcoll = self.timestamped_key_values(p, 'key', 1, 2, 3) result = (pcoll | WindowInto('w', SlidingWindows(period=2, size=4)) | GroupByKey() | reify_windows) expected = [('key @ [-2.0, 2.0)', [1]), ('key @ [0.0, 4.0)', [1, 2, 3]), ('key @ [2.0, 6.0)', [2, 3])] assert_that(result, equal_to(expected)) p.run() def test_sessions(self): p = Pipeline('DirectPipelineRunner') pcoll = self.timestamped_key_values(p, 'key', 1, 2, 3, 20, 35, 27) result = (pcoll | WindowInto('w', Sessions(10)) | GroupByKey() | sort_values | reify_windows) expected = [('key @ [1.0, 13.0)', [1, 2, 3]), ('key @ [20.0, 45.0)', [20, 27, 35])] assert_that(result, equal_to(expected)) p.run() def test_timestamped_value(self): p = Pipeline('DirectPipelineRunner') result = (p | Create('start', [(k, k) for k in range(10)]) | Map(lambda (x, t): TimestampedValue(x, t)) | WindowInto('w', FixedWindows(5)) | Map(lambda v: ('key', v)) | GroupByKey()) assert_that(result, equal_to([('key', [0, 1, 2, 3, 4]), ('key', [5, 6, 7, 8, 9])])) p.run() def test_timestamped_with_combiners(self): p = Pipeline('DirectPipelineRunner') result = (p # Create some initial test values. | Create('start', [(k, k) for k in range(10)]) # The purpose of the WindowInto transform is to establish a # FixedWindows windowing function for the PCollection. # It does not bucket elements into windows since the timestamps # from Create are not spaced 5 ms apart and very likely they all # fall into the same window. | WindowInto('w', FixedWindows(5)) # Generate timestamped values using the values as timestamps. # Now there are values 5 ms apart and since Map propagates the # windowing function from input to output the output PCollection # will have elements falling into different 5ms windows. | Map(lambda (x, t): TimestampedValue(x, t)) # We add a 'key' to each value representing the index of the # window. This is important since there is no guarantee of # order for the elements of a PCollection. | Map(lambda v: (v / 5, v))) # Sum all elements associated with a key and window. Although it # is called CombinePerKey it is really CombinePerKeyAndWindow the # same way GroupByKey is really GroupByKeyAndWindow. sum_per_window = result | CombinePerKey(sum) # Compute mean per key and window. mean_per_window = result | combiners.Mean.PerKey() assert_that(sum_per_window, equal_to([(0, 10), (1, 35)]), label='assert:sum') assert_that(mean_per_window, equal_to([(0, 2.0), (1, 7.0)]), label='assert:mean') p.run() if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/transforms/write_ptransform_test.py000066400000000000000000000123211272340251000304150ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the write transform.""" import logging import unittest import google.cloud.dataflow as df from google.cloud.dataflow.io import iobase from google.cloud.dataflow.pipeline import Pipeline from google.cloud.dataflow.transforms.ptransform import PTransform from google.cloud.dataflow.transforms.util import assert_that, is_empty from google.cloud.dataflow.utils.options import PipelineOptions class _TestSink(iobase.Sink): STATE_UNSTARTED, STATE_INITIALIZED, STATE_OPENED, STATE_FINALIZED = 0, 1, 2, 3 TEST_INIT_RESULT = 'test_init_result' def __init__(self, return_init_result=True, return_write_results=True): self.state = _TestSink.STATE_UNSTARTED self.last_writer = None self.return_init_result = return_init_result self.return_write_results = return_write_results def initialize_write(self): assert self.state == _TestSink.STATE_UNSTARTED self.state = _TestSink.STATE_INITIALIZED if self.return_init_result: return _TestSink.TEST_INIT_RESULT def finalize_write(self, init_result, writer_results): assert (self.state == _TestSink.STATE_OPENED or self.state == _TestSink.STATE_INITIALIZED) self.state = _TestSink.STATE_FINALIZED self.init_result_at_finalize = init_result self.write_results_at_finalize = writer_results def open_writer(self, init_result, uid): assert self.state == _TestSink.STATE_INITIALIZED self.state = _TestSink.STATE_OPENED writer = _TestWriter(init_result, uid, self.return_write_results) self.last_writer = writer return writer class _TestWriter(iobase.Writer): STATE_UNSTARTED, STATE_WRITTEN, STATE_CLOSED = 0, 1, 2 TEST_WRITE_RESULT = 'test_write_result' def __init__(self, init_result, uid, return_write_results=True): self.state = _TestWriter.STATE_UNSTARTED self.init_result = init_result self.uid = uid self.write_output = [] self.return_write_results = return_write_results def close(self): assert self.state == _TestWriter.STATE_WRITTEN self.state = _TestWriter.STATE_CLOSED if self.return_write_results: return _TestWriter.TEST_WRITE_RESULT def write(self, value): if self.write_output: assert self.state == _TestWriter.STATE_WRITTEN else: assert self.state == _TestWriter.STATE_UNSTARTED self.state = _TestWriter.STATE_WRITTEN self.write_output.append(value) class WriteToTestSink(PTransform): def __init__(self, return_init_result=True, return_write_results=True): self.return_init_result = return_init_result self.return_write_results = return_write_results self.last_sink = None self.label = 'write_to_test_sink' def apply(self, pcoll): self.last_sink = _TestSink(return_init_result=self.return_init_result, return_write_results=self.return_write_results) return pcoll | df.io.Write(self.last_sink) class WriteTest(unittest.TestCase): DATA = ['some data', 'more data', 'another data', 'yet another data'] def _run_write_test(self, data, return_init_result=True, return_write_results=True): write_to_test_sink = WriteToTestSink(return_init_result, return_write_results) p = Pipeline(options=PipelineOptions([])) result = p | df.Create('start', data) | write_to_test_sink assert_that(result, is_empty()) p.run() sink = write_to_test_sink.last_sink self.assertIsNotNone(sink) self.assertEqual(sink.state, _TestSink.STATE_FINALIZED) if data: self.assertIsNotNone(sink.last_writer) self.assertEqual(sink.last_writer.state, _TestWriter.STATE_CLOSED) self.assertEqual(sink.last_writer.write_output, data) if return_init_result: self.assertEqual(sink.last_writer.init_result, _TestSink.TEST_INIT_RESULT) self.assertEqual(sink.init_result_at_finalize, _TestSink.TEST_INIT_RESULT) self.assertIsNotNone(sink.last_writer.uid) if return_write_results: self.assertEqual(sink.write_results_at_finalize, [_TestWriter.TEST_WRITE_RESULT]) else: self.assertIsNone(sink.last_writer) def test_write(self): self._run_write_test(WriteTest.DATA) def test_write_with_empty_pcollection(self): data = [] self._run_write_test(data) def test_write_no_init_result(self): self._run_write_test(WriteTest.DATA, return_init_result=False) def test_write_no_write_results(self): self._run_write_test(WriteTest.DATA, return_write_results=False) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/000077500000000000000000000000001272340251000232315ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/__init__.py000066400000000000000000000014661272340251000253510ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A package defining the syntax and decorator semantics for type-hints.""" # pylint: disable=wildcard-import from google.cloud.dataflow.typehints.typehints import * from google.cloud.dataflow.typehints.decorators import * DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/decorators.py000066400000000000000000000432441272340251000257570ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Type hinting decorators allowing static or runtime type-checking for the SDK. This module defines decorators which utilize the type-hints defined in 'type_hints.py' to allow annotation of the types of function arguments and return values. Type-hints for functions are annotated using two separate decorators. One is for type-hinting the types of function arguments, the other for type-hinting the function return value. Type-hints can either be specified in the form of positional arguments:: @with_input_types(int, int) def add(a, b): return a + b Keyword arguments:: @with_input_types(a=int, b=int) def add(a, b): return a + b Or even a mix of both:: @with_input_types(int, b=int) def add(a, b): return a + b Example usage for type-hinting arguments only:: @with_input_types(s=str) def to_lower(a): return a.lower() Example usage for type-hinting return values only:: @with_output_types(Tuple[int, bool]) def compress_point(ec_point): return ec_point.x, ec_point.y < 0 Example usage for type-hinting both arguments and return values:: @with_input_types(a=int) @with_output_types(str) def int_to_str(a): return str(a) Type-hinting a function with arguments that unpack tuples are also supported. As an example, such a function would be defined as:: def foo((a, b)): ... The valid type-hint for such as function looks like the following:: @with_input_types(a=int, b=int) def foo((a, b)): ... Notice that we hint the type of each unpacked argument independently, rather than hinting the type of the tuple as a whole (Tuple[int, int]). Optionally, type-hints can be type-checked at runtime. To toggle this behavior this module defines two functions: 'enable_run_time_type_checking' and 'disable_run_time_type_checking'. NOTE: for this toggle behavior to work properly it must appear at the top of the module where all functions are defined, or before importing a module containing type-hinted functions. """ import inspect import types from google.cloud.dataflow.typehints import check_constraint from google.cloud.dataflow.typehints import CompositeTypeHintError from google.cloud.dataflow.typehints import SimpleTypeHintError from google.cloud.dataflow.typehints import typehints from google.cloud.dataflow.typehints import validate_composite_type_param # This is missing in the builtin types module. str.upper is arbitrary, any # method on a C-implemented type will do. # pylint: disable=invalid-name _MethodDescriptorType = type(str.upper) # pylint: enable=invalid-name # Monkeypatch inspect.getargspec to allow passing non-function objects. # This is needed to use higher-level functions such as getcallargs. _original_getargspec = inspect.getargspec def getargspec(func): try: return _original_getargspec(func) except TypeError: if isinstance(func, (type, types.ClassType)): argspec = getargspec(func.__init__) del argspec.args[0] return argspec elif callable(func): try: return _original_getargspec(func.__call__) except TypeError: # Return an ArgSpec with at least one positional argument, # and any number of other (positional or keyword) arguments # whose name won't match any real agument. # Arguments with the %unknown% prefix will be ignored in the type # checking code. return inspect.ArgSpec( ['_'], '%unknown%varargs', '%unknown%keywords', ()) else: raise inspect.getargspec = getargspec class IOTypeHints(object): """Encapsulates all type hint information about a Dataflow construct. This should primarily be used via the WithTypeHints mixin class, though may also be attached to other objects (such as Python functions). """ __slots__ = ('input_types', 'output_types') def __init__(self, input_types=None, output_types=None): self.input_types = input_types self.output_types = output_types def set_input_types(self, *args, **kwargs): self.input_types = args, kwargs def set_output_types(self, *args, **kwargs): self.output_types = args, kwargs def simple_output_type(self, context): if self.output_types: args, kwargs = self.output_types if len(args) != 1 or kwargs: raise TypeError('Expected simple output type hint for %s' % context) return args[0] def copy(self): return IOTypeHints(self.input_types, self.output_types) def with_defaults(self, hints): if not hints: return self elif not self: return hints else: return IOTypeHints(self.input_types or hints.input_types, self.output_types or hints.output_types) def __nonzero__(self): return bool(self.input_types or self.output_types) def __repr__(self): return 'IOTypeHints[inputs=%s, outputs=%s]' % ( self.input_types, self.output_types) class WithTypeHints(object): """A mixin class that provides the ability to set and retrieve type hints. """ def __init__(self, *unused_args, **unused_kwargs): self._type_hints = IOTypeHints() def _get_or_create_type_hints(self): # __init__ may have not been called try: return self._type_hints except AttributeError: self._type_hints = IOTypeHints() return self._type_hints def get_type_hints(self): return (self._get_or_create_type_hints() .with_defaults(self.default_type_hints()) .with_defaults(get_type_hints(self.__class__))) def default_type_hints(self): return None def with_input_types(self, *arg_hints, **kwarg_hints): self._get_or_create_type_hints().set_input_types(*arg_hints, **kwarg_hints) return self def with_output_types(self, *arg_hints, **kwarg_hints): self._get_or_create_type_hints().set_output_types(*arg_hints, **kwarg_hints) return self class TypeCheckError(Exception): pass def _positional_arg_hints(arg, hints): """Returns the type of a (possibly tuple-packed) positional argument. E.g. for lambda ((a, b), c): None the single positional argument is (as returned by inspect) [[a, b], c] which should have type Tuple[Tuple[Int, Any], float] when applied to the type hints {a: int, b: Any, c: float}. """ if isinstance(arg, list): return typehints.Tuple[[_positional_arg_hints(a, hints) for a in arg]] else: return hints.get(arg, typehints.Any) def _unpack_positional_arg_hints(arg, hint): """Unpacks the given hint according to the nested structure of arg. For example, if arg is [[a, b], c] and hint is Tuple[Any, int], than this function would return ((Any, Any), int) so it can be used in conjunction with inspect.getcallargs. """ if isinstance(arg, list): tuple_constraint = typehints.Tuple[[typehints.Any] * len(arg)] if not typehints.is_consistent_with(hint, tuple_constraint): raise typehints.TypeCheckError( 'Bad tuple arguments for %s: expected %s, got %s' % ( arg, tuple_constraint, hint)) if isinstance(hint, typehints.TupleConstraint): return tuple(_unpack_positional_arg_hints(a, t) for a, t in zip(arg, hint.tuple_types)) else: return (typehints.Any,) * len(arg) else: return hint def getcallargs_forhints(func, *typeargs, **typekwargs): """Like inspect.getcallargs, but understands that Tuple[] and an Any unpack. """ argspec = inspect.getargspec(func) # Turn Tuple[x, y] into (x, y) so getcallargs can do the proper unpacking. packed_typeargs = [_unpack_positional_arg_hints(arg, hint) for (arg, hint) in zip(argspec.args, typeargs)] packed_typeargs += list(typeargs[len(packed_typeargs):]) try: callargs = inspect.getcallargs(func, *packed_typeargs, **typekwargs) except TypeError, e: raise TypeCheckError(e) if argspec.defaults: # Declare any default arguments to be Any. for k, var in enumerate(reversed(argspec.args)): if k >= len(argspec.defaults): break if callargs.get(var, None) is argspec.defaults[-k]: callargs[var] = typehints.Any # Patch up varargs and keywords if argspec.varargs: callargs[argspec.varargs] = typekwargs.get( argspec.varargs, typehints.Tuple[typehints.Any, ...]) if argspec.keywords: # TODO(robertwb): Consider taking the union of key and value types. callargs[argspec.keywords] = typekwargs.get( argspec.keywords, typehints.Dict[typehints.Any, typehints.Any]) return callargs def get_type_hints(fn): """Gets the type hint associated with an arbitrary object fn. Always returns a valid IOTypeHints object, creating one if necissary. """ # pylint: disable=protected-access if not hasattr(fn, '_type_hints'): try: fn._type_hints = IOTypeHints() except (AttributeError, TypeError): # Can't add arbitrary attributes to this object, # but might have some restrictions anyways... hints = IOTypeHints() if isinstance(fn, _MethodDescriptorType): hints.set_input_types(fn.__objclass__) return hints return fn._type_hints # pylint: enable=protected-access def with_input_types(*positional_hints, **keyword_hints): """A decorator that type-checks defined type-hints with passed func arguments. All type-hinted arguments can be specified using positional arguments, keyword arguments, or a mix of both. Additionaly, all function arguments must be type-hinted in totality if even one parameter is type-hinted. Once fully decorated, if the arguments passed to the resulting function violate the type-hint constraints defined, a TypeCheckError detailing the error will be raised. To be used as:: * @with_input_types(s=str) # just @with_input_types(str) will work too. def upper(s): return s.upper() Or:: * @with_input_types(ls=List[Tuple[int, int]) def increment(ls): [(i + 1, j + 1) for (i,j) in ls] Args: *positional_hints: Positional type-hints having identical order as the function's formal arguments. Values for this argument must either be a built-in Python type or an instance of a TypeContraint created by 'indexing' a CompositeTypeHint instance with a type parameter. **keyword_hints: Keyword arguments mirroring the names of the parameters to the decorated functions. The value of each keyword argument must either be one of the allowed built-in Python types, a custom class, or an instance of a TypeContraint created by 'indexing' a CompositeTypeHint instance with a type parameter. Raises: ValueError: If not all function arguments have corresponding type-hints specified. Or if the inner wrapper function isn't passed a function object. TypeCheckError: If the any of the passed type-hint constraints are not a type or TypeContraint instance. Returns: The original function decorated such that it enforces type-hint constraints for all received function arguments. """ def annotate(f): if isinstance(f, types.FunctionType): for t in list(positional_hints) + list(keyword_hints.values()): validate_composite_type_param( t, error_msg_prefix='All type hint arguments') get_type_hints(f).set_input_types(*positional_hints, **keyword_hints) return f return annotate def with_output_types(*return_type_hint, **kwargs): """A decorator that type-checks defined type-hints for return values(s). This decorator will type-check the return value(s) of the decorated function. Only a single type-hint is accepted to specify the return type of the return value. If the function to be decorated has multiple return values, then one should use: 'Tuple[type_1, type_2]' to annotate the types of the return values. If the ultimate return value for the function violates the specified type-hint a TypeCheckError will be raised detailing the type-constraint violation. This decorator is intended to be used like:: * @with_output_types(Set[Coordinate]) def parse_ints(ints): .... return [Coordinate.from_int(i) for i in ints] Or with a simple type-hint:: * @with_output_types(bool) def negate(p): return not p if p else p Args: *return_type_hint: A type-hint specifying the proper return type of the function. This argument should either be a built-in Python type or an instance of a 'TypeConstraint' created by 'indexing' a 'CompositeTypeHint'. **kwargs: Not used. Raises: ValueError: If any kwarg parameters are passed in, or the length of 'return_type_hint' is greater than 1. Or if the inner wrapper function isn't passed a function object. TypeCheckError: If the 'return_type_hint' object is in invalid type-hint. Returns: The original function decorated such that it enforces type-hint constraints for all return values. """ if kwargs: raise ValueError("All arguments for the 'returns' decorator must be " "positional arguments.") if len(return_type_hint) != 1: raise ValueError("'returns' accepts only a single positional argument. In " "order to specify multiple return types, use the 'Tuple' " "type-hint.") return_type_hint = return_type_hint[0] validate_composite_type_param( return_type_hint, error_msg_prefix='All type hint arguments' ) def annotate(f): get_type_hints(f).set_output_types(return_type_hint) return f return annotate def _check_instance_type( type_constraint, instance, var_name=None, verbose=False): """A helper function to report type-hint constraint violations. Args: type_constraint: An instance of a 'TypeConstraint' or a built-in Python type. instance: The candidate object which will be checked by to satisfy 'type_constraint'. var_name: If 'instance' is an argument, then the actual name for the parameter in the original function definition. Raises: TypeCheckError: If 'instance' fails to meet the type-constraint of 'type_constraint'. """ hint_type = ( "argument: '%s'" % var_name if var_name is not None else 'return type') try: check_constraint(type_constraint, instance) except SimpleTypeHintError: if verbose: verbose_instance = '%s, ' % instance else: verbose_instance = '' raise TypeCheckError('Type-hint for %s violated. Expected an ' 'instance of %s, instead found %san instance of %s.' % (hint_type, type_constraint, verbose_instance, type(instance))) except CompositeTypeHintError as e: raise TypeCheckError('Type-hint for %s violated: %s' % (hint_type, e)) def _interleave_type_check(type_constraint, var_name=None): """Lazily type-check the type-hint for a lazily generated sequence type. This function can be applied as a decorator or called manually in a curried manner: * @_interleave_type_check(List[int]) def gen(): yield 5 or * gen = _interleave_type_check(Tuple[int, int], 'coord_gen')(gen) As a result, all type-checking for the passed generator will occur at 'yield' time. This way, we avoid having to depleat the generator in order to type-check it. Args: type_constraint: An instance of a TypeConstraint. The output yielded of 'gen' will be type-checked according to this type constraint. var_name: The variable name binded to 'gen' if type-checking a function argument. Used solely for templating in error message generation. Returns: A function which takes a generator as an argument and returns a wrapped version of the generator that interleaves type-checking at 'yield' iteration. If the generator received is already wrapped, then it is simply returned to avoid nested wrapping. """ def wrapper(gen): if isinstance(gen, GeneratorWrapper): return gen else: return GeneratorWrapper( gen, lambda x: _check_instance_type(type_constraint, x, var_name) ) return wrapper class GeneratorWrapper(object): """A wrapper around a generator, allows execution of a callback per yield. Additionally, wrapping a generator with this class allows one to assign arbitary attributes to a generator object just as with a function object. Attributes: internal_gen: A instance of a generator object. As part of 'step' of the generator, the yielded object will be passed to 'interleave_func'. interleave_func: A callback accepting a single argument. This function will be called with the result of each yielded 'step' in the internal generator. """ def __init__(self, gen, interleave_func): self.internal_gen = gen self.interleave_func = interleave_func def __getattr__(self, attr): # TODO(laolu): May also want to intercept 'send' in the future if we move to # a GeneratorHint with 3 type-params: # * Generator[send_type, return_type, yield_type] if attr == '__next__': return self.__next__() elif attr == '__iter__': return self.__iter__() else: return getattr(self.internal_gen, attr) def next(self): next_val = next(self.internal_gen) self.interleave_func(next_val) return next_val def __iter__(self): while True: x = next(self.internal_gen) self.interleave_func(x) yield x DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/opcodes.py000066400000000000000000000170151272340251000252430ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Defines the actions various bytecodes have on the frame. Each function here corresponds to a bytecode documented in https://docs.python.org/2/library/dis.html. The first argument is a (mutable) FrameState object, the second the integer opcode argument. Bytecodes with more complicated behavior (e.g. modifying the program counter) are handled inline rather than here. """ import types from trivial_inference import union, element_type, Const, BoundMethod import typehints from typehints import Any, Dict, Iterable, List, Tuple, Union def pop_one(state, unused_arg): del state.stack[-1:] def pop_two(state, unused_arg): del state.stack[-2:] def pop_three(state, unused_arg): del state.stack[-3:] def push_value(v): def pusher(state, unused_arg): state.stack.append(v) return pusher def nop(unused_state, unused_arg): pass def pop_top(state, unused_arg): state.stack.pop() def rot_n(state, n): state.stack[-n:] = [state.stack[-1]] + state.stack[-n:-1] def rot_two(state, unused_arg): rot_n(state, 2) def rot_three(state, unused_arg): rot_n(state, 3) def rot_four(state, unused_arg): rot_n(state, 4) def dup_top(state, unused_arg): state.stack.append(state.stack[-1]) def unary(state, unused_arg): state.stack[-1] = Const.unwrap(state.stack[-1]) unary_positive = unary_negative = unary_invert = unary def unary_not(state, unused_arg): state.stack[-1] = bool def unary_convert(state, unused_arg): state.stack[-1] = str def get_iter(state, unused_arg): state.stack.append(Iterable[element_type(state.stack.pop())]) def symmetric_binary_op(state, unused_arg): # TODO(robertwb): This may not be entirely correct... b, a = state.stack.pop(), state.stack.pop() if a == b: state.stack.append(a) elif type(a) == type(b) and isinstance(a, typehints.SequenceTypeConstraint): state.stack.append(type(a)(union(element_type(a), element_type(b)))) else: state.stack.append(Any) # Except for int ** -int binary_power = inplace_power = symmetric_binary_op binary_multiply = inplace_multiply = symmetric_binary_op binary_divide = inplace_divide = symmetric_binary_op binary_floor_divide = inplace_floor_divide = symmetric_binary_op def binary_true_divide(state, unused_arg): u = union(state.stack.pop(), state.stack.pop) if u == int: state.stack.append(float) else: state.stack.append(u) inplace_true_divide = binary_true_divide binary_modulo = inplace_modulo = symmetric_binary_op # TODO(robertwb): Tuple add. binary_add = inplace_add = symmetric_binary_op binary_subtract = inplace_subtract = symmetric_binary_op def binary_subscr(state, unused_arg): tos = state.stack.pop() if tos in (str, unicode): out = tos else: out = element_type(tos) state.stack.append(out) # As far as types are concerned. binary_lshift = inplace_lshift = binary_rshift = inplace_rshift = pop_top binary_and = inplace_and = symmetric_binary_op binary_xor = inplace_xor = symmetric_binary_op binary_or = inpalce_or = symmetric_binary_op # As far as types are concerned. slice_0 = nop slice_1 = slice_2 = pop_top slice_3 = pop_two store_slice_0 = store_slice_1 = store_slice_2 = store_slice_3 = nop delete_slice_0 = delete_slice_1 = delete_slice_2 = delete_slice_3 = nop def store_subscr(unused_state, unused_args): # TODO(robertwb): Update element/value type of iterable/dict. pass binary_divide = binary_floor_divide = binary_modulo = symmetric_binary_op binary_divide = binary_floor_divide = binary_modulo = symmetric_binary_op binary_divide = binary_floor_divide = binary_modulo = symmetric_binary_op # print_expr print_item = pop_top # print_item_to print_newline = nop # print_newline_to # break_loop # continue_loop def list_append(state, arg): state.stack[-arg] = List[Union[element_type(state.stack[-arg]), Const.unwrap(state.stack.pop())]] load_locals = push_value(Dict[str, Any]) # return_value # yield_value # import_star exec_stmt = pop_three # pop_block # end_finally build_class = pop_three # setup_with # with_cleanup # store_name # delete_name def unpack_sequence(state, arg): t = state.stack.pop() if isinstance(t, Const): try: unpacked = [Const(ti) for ti in t.value] if len(unpacked) != arg: unpacked = [Any] * arg except TypeError: unpacked = [Any] * arg elif (isinstance(t, typehints.TupleHint.TupleConstraint) and len(t.tuple_types) == arg): unpacked = list(t.tuple_types) else: unpacked = [element_type(t)] * arg state.stack += reversed(unpacked) def dup_topx(state, arg): state.stack += state[-arg:] store_attr = pop_top delete_attr = nop store_global = pop_top delete_global = nop def load_const(state, arg): state.stack.append(state.const_type(arg)) load_name = push_value(Any) def build_tuple(state, arg): if arg == 0: state.stack.append(Tuple[()]) else: state.stack[-arg:] = [Tuple[[Const.unwrap(t) for t in state.stack[-arg:]]]] def build_list(state, arg): if arg == 0: state.stack.append(List[Union[()]]) else: state.stack[-arg:] = [List[reduce(union, state.stack[-arg:], Union[()])]] build_map = push_value(Dict[Any, Any]) def load_attr(state, arg): o = state.stack.pop() name = state.get_name(arg) if isinstance(o, Const) and hasattr(o.value, name): state.stack.append(Const(getattr(o.value, name))) elif (isinstance(o, (type, types.ClassType)) and isinstance(getattr(o, name, None), types.MethodType)): state.stack.append(Const(BoundMethod(getattr(o, name)))) else: state.stack.append(Any) def compare_op(state, unused_arg): # Could really be anything... state.stack[-2:] = [bool] def import_name(state, unused_arg): state.stack[-2:] = [Any] import_from = push_value(Any) # jump # for_iter def load_global(state, arg): state.stack.append(state.get_global(arg)) # setup_loop # setup_except # setup_finally store_map = pop_two def load_fast(state, arg): state.stack.append(state.vars[arg]) def store_fast(state, arg): state.vars[arg] = state.stack.pop() def delete_fast(state, arg): state.vars[arg] = Any # really an error def load_closure(state, unused_arg): state.stack.append(Any) # really a Cell def load_deref(state, arg): state.stack.append(state.closure_type(arg)) # raise_varargs def call_function(state, arg, has_var=False, has_kw=False): # TODO(robertwb): Recognize builtins and dataflow objects # (especially special return values). pop_count = (arg & 0xF) + (arg & 0xF0) / 8 + 1 + has_var + has_kw state.stack[-pop_count:] = [Any] def make_function(state, arg): state.stack[-arg - 1:] = [Any] # a callable def make_closure(state, arg): state.stack[-arg - 2:] = [Any] # a callable def build_slice(state, arg): state.stack[-arg:] = [Any] # a slice object def call_function_var(state, arg): call_function(state, arg, has_var=True) def call_function_kw(state, arg): call_function(state, arg, has_kw=True) def call_function_var_wk(state, arg): call_function(state, arg, has_var=True, has_kw=True) DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/trivial_inference.py000066400000000000000000000267511272340251000273060ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Trivial type inference for simple functions. """ import __builtin__ import collections import dis import pprint import sys import types from google.cloud.dataflow.typehints import Any from google.cloud.dataflow.typehints import typehints class TypeInferenceError(ValueError): """Error to raise when type inference failed.""" pass def instance_to_type(o): """Given a Python object o, return the corresponding type hint. """ t = type(o) if o is None: # TODO(robertwb): Eliminate inconsistent use of None vs. NoneType. return None elif t not in typehints.DISALLOWED_PRIMITIVE_TYPES: if t == types.InstanceType: return o.__class__ elif t == BoundMethod: return types.MethodType else: return t elif t == tuple: return typehints.Tuple[[instance_to_type(item) for item in o]] elif t == list: return typehints.List[ typehints.Union[[instance_to_type(item) for item in o]] ] elif t == set: return typehints.Set[ typehints.Union[[instance_to_type(item) for item in o]] ] elif t == dict: return typehints.Dict[ typehints.Union[[instance_to_type(k) for k, v in o.items()]], typehints.Union[[instance_to_type(v) for k, v in o.items()]], ] else: raise TypeInferenceError('Unknown forbidden type: %s' % t) def union_list(xs, ys): assert len(xs) == len(ys) return [union(x, y) for x, y in zip(xs, ys)] class Const(object): def __init__(self, value): self.value = value self.type = instance_to_type(value) def __eq__(self, other): return isinstance(other, Const) and self.value == other.value def __hash__(self): return hash(self.value) def __repr__(self): return 'Const[%s]' % str(self.value)[:100] @staticmethod def unwrap(x): if isinstance(x, Const): return x.type else: return x @staticmethod def unwrap_all(xs): return [Const.unwrap(x) for x in xs] class FrameState(object): """Stores the state of the frame at a particular point of execution. """ def __init__(self, f, local_vars=None, stack=()): self.f = f self.co = f.func_code self.vars = list(local_vars) self.stack = list(stack) def __eq__(self, other): return self.__dict__ == other.__dict__ def copy(self): return FrameState(self.f, self.vars, self.stack) def const_type(self, i): return Const(self.co.co_consts[i]) def closure_type(self, i): ncellvars = len(self.co.co_cellvars) if i < ncellvars: return Any else: return Const(self.f.func_closure[i - ncellvars].cell_contents) def get_global(self, i): name = self.get_name(i) if name in self.f.func_globals: return Const(self.f.func_globals[name]) if name in __builtin__.__dict__: return Const(__builtin__.__dict__[name]) else: return Any def get_name(self, i): return self.co.co_names[i] def __repr__(self): return 'Stack: %s Vars: %s' % (self.stack, self.vars) def __or__(self, other): if self is None: return other.copy() elif other is None: return self.copy() else: return FrameState(self.f, union_list(self.vars, other.vars), union_list( self.stack, other.stack)) def __ror__(self, left): return self | left def union(a, b): """Returns the union of two types or Const values. """ if a == b: return a elif not a: return b elif not b: return a a = Const.unwrap(a) b = Const.unwrap(b) # TODO(robertwb): Work this into the Union code in a more generic way. if type(a) == type(b) and element_type(a) == typehints.Union[()]: return b elif type(a) == type(b) and element_type(b) == typehints.Union[()]: return a else: return typehints.Union[a, b] def element_type(hint): """Returns the element type of a composite type. """ hint = Const.unwrap(hint) if isinstance(hint, typehints.SequenceTypeConstraint): return hint.inner_type elif isinstance(hint, typehints.TupleHint.TupleConstraint): return typehints.Union[hint.tuple_types] else: return Any def key_value_types(kv_type): """Returns the key and value type of a KV type. """ # TODO(robertwb): Unions of tuples, etc. # TODO(robertwb): Assert? if (isinstance(kv_type, typehints.TupleHint.TupleConstraint) and len(kv_type.tuple_types) == 2): return kv_type.tuple_types return Any, Any known_return_types = {len: int, hash: int,} class BoundMethod(object): """Used to create a bound method when we only know the type of the instance. """ def __init__(self, unbound): self.unbound = unbound def hashable(c): try: hash(c) return True except TypeError: return False def infer_return_type(c, input_types, debug=False, depth=5): """Analyses a callable to deduce its return type. Args: f: A Python function object to infer the return type of. input_types: A sequence of inputs corresponding to the input types. debug: Whether to print verbose debugging information. Returns: A TypeConstraint that that the return value of this function will (likely) satisfy given the specified inputs. """ try: if hashable(c) and c in known_return_types: return known_return_types[c] elif isinstance(c, types.FunctionType): return infer_return_type_func(c, input_types, debug, depth) elif isinstance(c, types.MethodType): if c.im_self is not None: input_types = [Const(c.im_self)] + input_types return infer_return_type_func(c.im_func, input_types, debug, depth) elif isinstance(c, BoundMethod): input_types = [c.unbound.im_class] + input_types return infer_return_type_func(c.unbound.im_func, input_types, debug, depth) elif isinstance(c, (type, types.ClassType)): if c in typehints.DISALLOWED_PRIMITIVE_TYPES: return { list: typehints.List[Any], set: typehints.Set[Any], tuple: typehints.Tuple[Any, ...], dict: typehints.Dict[Any, Any] }[c] else: return c else: return Any except TypeInferenceError: return Any except Exception: if debug: sys.stdout.flush() raise else: return Any def infer_return_type_func(f, input_types, debug=False, depth=0): """Analyses a function to deduce its return type. Args: f: A Python function object to infer the return type of. input_types: A sequence of inputs corresponding to the input types. debug: Whether to print verbose debugging information. Returns: A TypeConstraint that that the return value of this function will (likely) satisfy given the specified inputs. Raises: TypeInferenceError: if no type can be inferred. """ if debug: print print f, id(f), input_types import opcodes simple_ops = dict((k.upper(), v) for k, v in opcodes.__dict__.items()) co = f.func_code code = co.co_code end = len(code) pc = 0 extended_arg = 0 free = None yields = set() returns = set() # TODO(robertwb): Default args via inspect module. local_vars = list(input_types) + [typehints.Union[()]] * (len(co.co_varnames) - len(input_types)) state = FrameState(f, local_vars) states = collections.defaultdict(lambda: None) jumps = collections.defaultdict(int) last_pc = -1 while pc < end: start = pc op = ord(code[pc]) if debug: print '-->' if pc == last_pc else ' ', print repr(pc).rjust(4), print dis.opname[op].ljust(20), pc += 1 if op >= dis.HAVE_ARGUMENT: arg = ord(code[pc]) + ord(code[pc + 1]) * 256 + extended_arg extended_arg = 0 pc += 2 if op == dis.EXTENDED_ARG: extended_arg = arg * 65536L if debug: print str(arg).rjust(5), if op in dis.hasconst: print '(' + repr(co.co_consts[arg]) + ')', elif op in dis.hasname: print '(' + co.co_names[arg] + ')', elif op in dis.hasjrel: print '(to ' + repr(pc + arg) + ')', elif op in dis.haslocal: print '(' + co.co_varnames[arg] + ')', elif op in dis.hascompare: print '(' + dis.cmp_op[arg] + ')', elif op in dis.hasfree: if free is None: free = co.co_cellvars + co.co_freevars print '(' + free[arg] + ')', # Acutally emulate the op. if state is None and states[start] is None: # No control reaches here (yet). if debug: print continue state |= states[start] opname = dis.opname[op] jmp = jmp_state = None if opname.startswith('CALL_FUNCTION'): standard_args = (arg & 0xF) + (arg & 0xF0) / 8 var_args = 'VAR' in opname kw_args = 'KW' in opname pop_count = standard_args + var_args + kw_args + 1 if depth <= 0: return_type = Any elif arg & 0xF0: # TODO(robertwb): Handle this case. return_type = Any elif isinstance(state.stack[-pop_count], Const): # TODO(robertwb): Handle this better. if var_args or kw_args: state.stack[-1] = Any state.stack[-var_args - kw_args] = Any inputs = [] if pop_count == 1 else state.stack[1 - pop_count:] return_type = infer_return_type(state.stack[-pop_count].value, state.stack[1 - pop_count:], debug=debug, depth=depth - 1) else: return_type = Any state.stack[-pop_count:] = [return_type] elif opname in simple_ops: simple_ops[opname](state, arg) elif opname == 'RETURN_VALUE': returns.add(state.stack[-1]) state = None elif opname == 'YIELD_VALUE': yields.add(state.stack[-1]) elif opname == 'JUMP_FORWARD': jmp = pc + arg jmp_state = state state = None elif opname == 'JUMP_ABSOLUTE': jmp = arg jmp_state = state state = None elif opname in ('POP_JUMP_IF_TRUE', 'POP_JUMP_IF_FALSE'): state.stack.pop() jmp = arg jmp_state = state.copy() elif opname in ('JUMP_IF_TRUE_OR_POP', 'JUMP_IF_FALSE_OR_POP'): jmp = arg jmp_state = state.copy() state.stack.pop() elif opname == 'FOR_ITER': jmp = pc + arg jmp_state = state.copy() jmp_state.stack.pop() state.stack.append(element_type(state.stack[-1])) else: raise TypeInferenceError('unable to handle %s' % opname) if jmp is not None: # TODO(robertwb): Is this guerenteed to converge? new_state = states[jmp] | jmp_state if jmp < pc and new_state != states[jmp] and jumps[pc] < 5: jumps[pc] += 1 pc = jmp states[jmp] = new_state if debug: print print state pprint.pprint(dict(item for item in states.items() if item[1])) if yields: result = typehints.Iterable[reduce(union, Const.unwrap_all(yields))] else: result = reduce(union, Const.unwrap_all(returns)) if debug: print f, id(f), input_types, '->', result return result DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/trivial_inference_test.py000066400000000000000000000107661272340251000303440ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for google.cloud.dataflow.typehints.trivial_inference.""" import unittest from google.cloud.dataflow.typehints import trivial_inference from google.cloud.dataflow.typehints import typehints global_int = 1 class TrivialInferenceTest(unittest.TestCase): def assertReturnType(self, expected, f, inputs=()): self.assertEquals(expected, trivial_inference.infer_return_type(f, inputs)) def testIdentity(self): self.assertReturnType(int, lambda x: x, [int]) def testTuples(self): self.assertReturnType( typehints.Tuple[typehints.Tuple[()], int], lambda x: ((), x), [int]) self.assertReturnType( typehints.Tuple[str, int, float], lambda x: (x, 0, 1.0), [str]) def testUnpack(self): def reverse((a, b)): return b, a any_tuple = typehints.Tuple[typehints.Any, typehints.Any] self.assertReturnType( typehints.Tuple[int, float], reverse, [typehints.Tuple[float, int]]) self.assertReturnType( typehints.Tuple[int, int], reverse, [typehints.Tuple[int, ...]]) self.assertReturnType( typehints.Tuple[int, int], reverse, [typehints.List[int]]) self.assertReturnType( typehints.Tuple[typehints.Union[int, float, str], typehints.Union[int, float, str]], reverse, [typehints.Tuple[int, float, str]]) self.assertReturnType(any_tuple, reverse, [typehints.Any]) self.assertReturnType(typehints.Tuple[int, float], reverse, [trivial_inference.Const((1.0, 1))]) self.assertReturnType(any_tuple, reverse, [trivial_inference.Const((1, 2, 3))]) def testListComprehension(self): self.assertReturnType( typehints.List[int], lambda xs: [x for x in xs], [typehints.Tuple[int, ...]]) def testTupleListComprehension(self): self.assertReturnType( typehints.List[int], lambda xs: [x for x in xs], [typehints.Tuple[int, int, int]]) self.assertReturnType( typehints.List[typehints.Union[int, float]], lambda xs: [x for x in xs], [typehints.Tuple[int, float]]) def testGenerator(self): def foo(x, y): yield x yield y self.assertReturnType(typehints.Iterable[int], foo, [int, int]) self.assertReturnType( typehints.Iterable[typehints.Union[int, float]], foo, [int, float]) def testBinOp(self): self.assertReturnType(int, lambda a, b: a + b, [int, int]) self.assertReturnType( typehints.Any, lambda a, b: a + b, [int, typehints.Any]) self.assertReturnType( typehints.List[typehints.Union[int, str]], lambda a, b: a + b, [typehints.List[int], typehints.List[str]]) def testCall(self): f = lambda x, *args: x self.assertReturnType( typehints.Tuple[int, float], lambda: (f(1), f(2.0, 3))) def testClosure(self): x = 1 y = 1.0 self.assertReturnType(typehints.Tuple[int, float], lambda: (x, y)) def testGlobals(self): self.assertReturnType(int, lambda: global_int) def testBuiltins(self): self.assertReturnType(int, lambda x: len(x), [typehints.Any]) def testGetAttr(self): self.assertReturnType( typehints.Tuple[str, typehints.Any], lambda: (typehints.__doc__, typehints.fake)) def testMethod(self): class A(object): def m(self, x): return x self.assertReturnType(int, lambda: A().m(3)) self.assertReturnType(float, lambda: A.m(A(), 3.0)) def testAlwaysReturnsEarly(self): def some_fn(v): if v: return 1 else: return 2 self.assertReturnType(int, some_fn) def testDict(self): self.assertReturnType( typehints.Dict[typehints.Any, typehints.Any], lambda: {}) def testDictComprehension(self): # Just ensure it doesn't crash. fields = [] self.assertReturnType( typehints.Any, lambda row: {f: row[f] for f in fields}, [typehints.Any]) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/typecheck.py000066400000000000000000000144721272340251000255720ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Runtime type checking support.""" import collections import inspect import sys import types from google.cloud.dataflow.pvalue import SideOutputValue from google.cloud.dataflow.transforms.core import DoFn from google.cloud.dataflow.transforms.window import WindowedValue from google.cloud.dataflow.typehints import check_constraint from google.cloud.dataflow.typehints import CompositeTypeHintError from google.cloud.dataflow.typehints import GeneratorWrapper from google.cloud.dataflow.typehints import SimpleTypeHintError from google.cloud.dataflow.typehints import TypeCheckError from google.cloud.dataflow.typehints.decorators import _check_instance_type from google.cloud.dataflow.typehints.decorators import getcallargs_forhints class TypeCheckWrapperDoFn(DoFn): """A wrapper around a DoFn which performs type-checking of input and output. """ def __init__(self, dofn, type_hints, label=None): super(TypeCheckWrapperDoFn, self).__init__() self._dofn = dofn self._label = label self._process_fn = self._dofn.process_argspec_fn() if type_hints.input_types: input_args, input_kwargs = type_hints.input_types self._input_hints = getcallargs_forhints( self._process_fn, *input_args, **input_kwargs) else: self._input_hints = None # TODO(robertwb): Actually extract this. self.context_var = 'context' # TODO(robertwb): Multi-output. self._output_type_hint = type_hints.simple_output_type(label) def start_bundle(self, context, *args, **kwargs): return self._type_check_result( self._dofn.start_bundle(context, *args, **kwargs)) def finish_bundle(self, context, *args, **kwargs): return self._type_check_result( self._dofn.finish_bundle(context, *args, **kwargs)) def process(self, context, *args, **kwargs): if self._input_hints: actual_inputs = inspect.getcallargs( self._process_fn, context.element, *args, **kwargs) for var, hint in self._input_hints.items(): if hint is actual_inputs[var]: # self parameter continue var_name = var + '.element' if var == self.context_var else var _check_instance_type(hint, actual_inputs[var], var_name, True) return self._type_check_result(self._dofn.process(context, *args, **kwargs)) def _type_check_result(self, transform_results): if self._output_type_hint is None or transform_results is None: return transform_results def type_check_output(o): # TODO(robertwb): Multi-output. x = o.value if isinstance(o, (SideOutputValue, WindowedValue)) else o self._type_check(self._output_type_hint, x, is_input=False) # If the return type is a generator, then we will need to interleave our # type-checking with its normal iteration so we don't deplete the # generator initially just by type-checking its yielded contents. if isinstance(transform_results, types.GeneratorType): return GeneratorWrapper(transform_results, type_check_output) else: for o in transform_results: type_check_output(o) return transform_results def _type_check(self, type_constraint, datum, is_input): """Typecheck a PTransform related datum according to a type constraint. This function is used to optionally type-check either an input or an output to a PTransform. Args: type_constraint: An instance of a typehints.TypeContraint, one of the white-listed builtin Python types, or a custom user class. datum: An instance of a Python object. is_input: True if 'datum' is an input to a PTransform's DoFn. False otherwise. Raises: TypeError: If 'datum' fails to type-check according to 'type_constraint'. """ datum_type = 'input' if is_input else 'output' try: check_constraint(type_constraint, datum) except CompositeTypeHintError as e: raise TypeCheckError, e.message, sys.exc_info()[2] except SimpleTypeHintError: error_msg = ("According to type-hint expected %s should be of type %s. " "Instead, received '%s', an instance of type %s." % (datum_type, type_constraint, datum, type(datum))) raise TypeCheckError, error_msg, sys.exc_info()[2] class OutputCheckWrapperDoFn(DoFn): """A DoFn that verifies against common errors in the output type.""" def __init__(self, dofn, full_label): self.dofn = dofn self.full_label = full_label def run(self, method, context, args, kwargs): try: result = method(context, *args, **kwargs) except TypeCheckError as e: error_msg = ('Runtime type violation detected within ParDo(%s): ' '%s' % (self.full_label, e)) raise TypeCheckError, error_msg, sys.exc_info()[2] else: return self._check_type(result) def start_bundle(self, context, *args, **kwargs): return self.run(self.dofn.start_bundle, context, args, kwargs) def finish_bundle(self, context, *args, **kwargs): return self.run(self.dofn.finish_bundle, context, args, kwargs) def process(self, context, *args, **kwargs): return self.run(self.dofn.process, context, args, kwargs) def _check_type(self, output): if output is None: return output elif isinstance(output, (dict, basestring)): object_type = type(output).__name__ raise TypeCheckError('Returning a %s from a ParDo or FlatMap is ' 'discouraged. Please use list("%s") if you really ' 'want this behavior.' % (object_type, output)) elif not isinstance(output, collections.Iterable): raise TypeCheckError('FlatMap and ParDo must return an ' 'iterable. %s was returned instead.' % type(output)) return output DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/typed_pipeline_test.py000066400000000000000000000214331272340251000276570ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the type-hint objects and decorators.""" import inspect import unittest import google.cloud.dataflow as df from google.cloud.dataflow import pvalue from google.cloud.dataflow import typehints from google.cloud.dataflow.transforms.util import assert_that, equal_to from google.cloud.dataflow.typehints import WithTypeHints from google.cloud.dataflow.utils.options import OptionsContext from google.cloud.dataflow.utils.options import PipelineOptions # These test often construct a pipeline as value | PTransform to test side # effects (e.g. errors). # pylint: disable=expression-not-assigned class MainInputTest(unittest.TestCase): def test_bad_main_input(self): @typehints.with_input_types(str, int) def repeat(s, times): return s * times with self.assertRaises(typehints.TypeCheckError): [1, 2, 3] | df.Map(repeat, 3) def test_non_function(self): result = ['a', 'bb', 'c'] | df.Map(str.upper) self.assertEqual(['A', 'BB', 'C'], sorted(result)) result = ['xa', 'bbx', 'xcx'] | df.Map(str.strip, 'x') self.assertEqual(['a', 'bb', 'c'], sorted(result)) result = ['1', '10', '100'] | df.Map(int) self.assertEqual([1, 10, 100], sorted(result)) result = ['1', '10', '100'] | df.Map(int, 16) self.assertEqual([1, 16, 256], sorted(result)) with self.assertRaises(typehints.TypeCheckError): [1, 2, 3] | df.Map(str.upper) def test_loose_bounds(self): @typehints.with_input_types(typehints.Union[int, float, long]) @typehints.with_output_types(basestring) def format_number(x): return '%g' % x result = [1, 2, 3] | df.Map(format_number) self.assertEqual(['1', '2', '3'], sorted(result)) def test_typed_dofn_class(self): @typehints.with_input_types(int) @typehints.with_output_types(str) class MyDoFn(df.DoFn): def process(self, context): return [str(context.element)] result = [1, 2, 3] | df.ParDo(MyDoFn()) self.assertEqual(['1', '2', '3'], sorted(result)) with self.assertRaises(typehints.TypeCheckError): ['a', 'b', 'c'] | df.ParDo(MyDoFn()) with self.assertRaises(typehints.TypeCheckError): [1, 2, 3] | (df.ParDo(MyDoFn()) | df.ParDo('again', MyDoFn())) def test_typed_dofn_instance(self): class MyDoFn(df.DoFn): def process(self, context): return [str(context.element)] my_do_fn = MyDoFn().with_input_types(int).with_output_types(str) result = [1, 2, 3] | df.ParDo(my_do_fn) self.assertEqual(['1', '2', '3'], sorted(result)) with self.assertRaises(typehints.TypeCheckError): ['a', 'b', 'c'] | df.ParDo(my_do_fn) with self.assertRaises(typehints.TypeCheckError): [1, 2, 3] | (df.ParDo(my_do_fn) | df.ParDo('again', my_do_fn)) class SideInputTest(unittest.TestCase): def _run_repeat_test(self, repeat): self._run_repeat_test_good(repeat) self._run_repeat_test_bad(repeat) @OptionsContext(pipeline_type_check=True) def _run_repeat_test_good(self, repeat): # As a positional argument. result = ['a', 'bb', 'c'] | df.Map(repeat, 3) self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result)) # As a keyword argument. result = ['a', 'bb', 'c'] | df.Map(repeat, times=3) self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result)) def _run_repeat_test_bad(self, repeat): # Various mismatches. with self.assertRaises(typehints.TypeCheckError): ['a', 'bb', 'c'] | df.Map(repeat, 'z') with self.assertRaises(typehints.TypeCheckError): ['a', 'bb', 'c'] | df.Map(repeat, times='z') with self.assertRaises(typehints.TypeCheckError): ['a', 'bb', 'c'] | df.Map(repeat, 3, 4) if not inspect.getargspec(repeat).defaults: with self.assertRaises(typehints.TypeCheckError): ['a', 'bb', 'c'] | df.Map(repeat) def test_basic_side_input_hint(self): @typehints.with_input_types(str, int) def repeat(s, times): return s * times self._run_repeat_test(repeat) def test_keyword_side_input_hint(self): @typehints.with_input_types(str, times=int) def repeat(s, times): return s * times self._run_repeat_test(repeat) def test_default_typed_hint(self): @typehints.with_input_types(str, int) def repeat(s, times=3): return s * times self._run_repeat_test(repeat) def test_default_untyped_hint(self): @typehints.with_input_types(str) def repeat(s, times=3): return s * times # No type checking on dfault arg. self._run_repeat_test_good(repeat) @OptionsContext(pipeline_type_check=True) def test_varargs_side_input_hint(self): @typehints.with_input_types(str, int) def repeat(s, *times): return s * times[0] result = ['a', 'bb', 'c'] | df.Map(repeat, 3) self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result)) # TODO(robertwb): Support partially defined varargs. # with self.assertRaises(typehints.TypeCheckError): # ['a', 'bb', 'c'] | df.Map(repeat, 'z') def test_deferred_side_inputs(self): @typehints.with_input_types(str, int) def repeat(s, times): return s * times p = df.Pipeline(options=PipelineOptions([])) main_input = p | df.Create(['a', 'bb', 'c']) side_input = p | df.Create('side', [3]) result = main_input | df.Map(repeat, pvalue.AsSingleton(side_input)) assert_that(result, equal_to(['aaa', 'bbbbbb', 'ccc'])) p.run() bad_side_input = p | df.Create('bad_side', ['z']) with self.assertRaises(typehints.TypeCheckError): main_input | df.Map('again', repeat, pvalue.AsSingleton(bad_side_input)) def test_deferred_side_input_iterable(self): @typehints.with_input_types(str, typehints.Iterable[str]) def concat(glue, items): return glue.join(sorted(items)) p = df.Pipeline(options=PipelineOptions([])) main_input = p | df.Create(['a', 'bb', 'c']) side_input = p | df.Create('side', ['x', 'y', 'z']) result = main_input | df.Map(concat, pvalue.AsIter(side_input)) assert_that(result, equal_to(['xayaz', 'xbbybbz', 'xcycz'])) p.run() bad_side_input = p | df.Create('bad_side', [1, 2, 3]) with self.assertRaises(typehints.TypeCheckError): main_input | df.Map('fail', concat, pvalue.AsIter(bad_side_input)) class CustomTransformTest(unittest.TestCase): class CustomTransform(df.PTransform): def _extract_input_pvalues(self, pvalueish): return pvalueish, (pvalueish['in0'], pvalueish['in1']) def apply(self, pvalueish): return {'out0': pvalueish['in0'], 'out1': pvalueish['in1']} # TODO(robertwb): (typecheck) Make these the default? def with_input_types(self, *args, **kwargs): return WithTypeHints.with_input_types(self, *args, **kwargs) def with_output_types(self, *args, **kwargs): return WithTypeHints.with_output_types(self, *args, **kwargs) test_input = {'in0': ['a', 'b', 'c'], 'in1': [1, 2, 3]} def check_output(self, result): self.assertEqual(['a', 'b', 'c'], sorted(result['out0'])) self.assertEqual([1, 2, 3], sorted(result['out1'])) def test_custom_transform(self): self.check_output(self.test_input | self.CustomTransform()) def test_keyword_type_hints(self): self.check_output( self.test_input | self.CustomTransform().with_input_types( in0=str, in1=int)) self.check_output( self.test_input | self.CustomTransform().with_input_types(in0=str)) self.check_output( self.test_input | self.CustomTransform().with_output_types( out0=str, out1=int)) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_input_types(in0=int) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_output_types(out0=int) def test_flat_type_hint(self): # Type hint is applied to both. ({'in0': ['a', 'b', 'c'], 'in1': ['x', 'y', 'z']} | self.CustomTransform().with_input_types(str)) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_input_types(str) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_input_types(int) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_output_types(int) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/typehints.py000066400000000000000000001054541272340251000256430ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Syntax and semantics for type-hinting custom-functions/PTransforms in the SDK. This module defines type-hinting objects and the corresponding syntax for type-hinting function arguments, function return types, or PTransform object themselves. TypeHint's defined in the module can be used to implement either static or run-time type-checking in regular Python code. Type-hints are defined by 'indexing' a type-parameter into a defined CompositeTypeHint instance: * 'List[int]'. Valid type-hints are partitioned into two categories: simple, and composite. Simple type hints are type hints based on a subset of Python primitive types: int, bool, float, str, object, None, and bytes. No other primitive types are allowed. Composite type-hints are reserved for hinting the types of container-like Python objects such as 'list'. Composite type-hints can be parameterized by an inner simple or composite type-hint, using the 'indexing' syntax. In order to avoid conflicting with the namespace of the built-in container types, when specifying this category of type-hints, the first letter should capitalized. The following composite type-hints are permitted. NOTE: 'T' can be any of the type-hints listed or a simple Python type: * Any * Union[T, T, T] * Optional[T] * Tuple[T, T] * Tuple[T, ...] * List[T] * KV[T, T] * Dict[T, T] * Set[T] * Iterable[T] * Iterator[T] * Generator[T] Type-hints can be nested, allowing one to define type-hints for complex types: * 'List[Tuple[int, int, str]] In addition, type-hints can be used to implement run-time type-checking via the 'type_check' method on each TypeConstraint. """ import collections import copy import types # A set of the built-in Python types we don't support, guiding the users # to templated (upper-case) versions instead. DISALLOWED_PRIMITIVE_TYPES = (list, set, tuple, dict) class SimpleTypeHintError(TypeError): pass class CompositeTypeHintError(TypeError): pass class GetitemConstructor(type): """A metaclass that makes Cls[arg] an alias for Cls(arg).""" def __getitem__(cls, arg): return cls(arg) class TypeConstraint(object): """The base-class for all created type-constraints defined below. A TypeConstraint is the result of parameterizing a CompositeTypeHint with with one of the allowed Python types or another CompositeTypeHint. It binds and enforces a specific version of a generalized TypeHint. """ def _consistent_with_check_(self, sub): """Returns whether sub is consistent with self. Has the same relationship to is_consistent_with() as __subclasscheck__ does for issubclass(). Not meant to be called directly; call is_consistent_with(sub, self) instead. Implementation may assume that maybe_sub_type is not Any and has been normalized. """ raise NotImplementedError def type_check(self, instance): """Determines if the type of 'instance' satisfies this type constraint. Args: instance: An instance of a Python object. Raises: TypeError: The passed 'instance' doesn't satisfy this TypeConstraint. Subclasses of TypeConstraint are free to raise any of the subclasses of TypeError defined above, depending on the manner of the type hint error. All TypeConstraint sub-classes must define this method in other for the class object to be created. """ raise NotImplementedError def match_type_variables(self, unused_concrete_type): return {} def bind_type_variables(self, unused_bindings): return self def _inner_types(self): """Iterates over the inner types of the composite type.""" return [] def visit(self, visitor, visitor_arg): """Visitor method to visit all inner types of a composite type. Args: visitor: A callable invoked for all nodes in the type tree comprising a composite type. The visitor will be called with the node visited and the visitor argument specified here. visitor_arg: Visitor callback second argument. """ visitor(self, visitor_arg) for t in self._inner_types(): if isinstance(t, TypeConstraint): t.visit(visitor, visitor_arg) else: visitor(t, visitor_arg) def match_type_variables(type_constraint, concrete_type): if isinstance(type_constraint, TypeConstraint): return type_constraint.match_type_variables(concrete_type) else: return {} def bind_type_variables(type_constraint, bindings): if isinstance(type_constraint, TypeConstraint): return type_constraint.bind_type_variables(bindings) else: return type_constraint class SequenceTypeConstraint(TypeConstraint): """A common base-class for all sequence related type-constraint classes. A sequence is defined as an arbitrary length homogeneous container type. Type hints which fall under this category include: List[T], Set[T], Iterable[T], and Tuple[T, ...]. Sub-classes may need to override '_consistent_with_check_' if a particular sequence requires special handling with respect to type compatibility. Attributes: inner_type: The type which every element in the sequence should be an instance of. """ def __init__(self, inner_type, sequence_type): self.inner_type = inner_type self._sequence_type = sequence_type def __eq__(self, other): return (isinstance(other, SequenceTypeConstraint) and type(self) == type(other) and self.inner_type == other.inner_type) def __hash__(self): return hash(self.inner_type) ^ 13 * hash(type(self)) def _inner_types(self): yield self.inner_type def _consistent_with_check_(self, sub): return (isinstance(sub, self.__class__) and is_consistent_with(sub.inner_type, self.inner_type)) def type_check(self, sequence_instance): if not isinstance(sequence_instance, self._sequence_type): raise CompositeTypeHintError( "%s type-constraint violated. Valid object instance " "must be of type '%s'. Instead, an instance of '%s' " "was received." % (self._sequence_type.__name__.title(), self._sequence_type.__name__.lower(), sequence_instance.__class__.__name__)) for index, elem in enumerate(sequence_instance): try: check_constraint(self.inner_type, elem) except SimpleTypeHintError as e: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element #%s in ' 'the passed %s is incorrect. Expected an instance of type %s, ' 'instead received an instance of type %s.' % (repr(self), index, _unified_repr(self._sequence_type), _unified_repr(self.inner_type), elem.__class__.__name__)) except CompositeTypeHintError as e: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element #%s in ' 'the passed %s is incorrect: %s' % (repr(self), index, self._sequence_type.__name__, e)) def match_type_variables(self, concrete_type): if isinstance(concrete_type, SequenceTypeConstraint): return match_type_variables(self.inner_type, concrete_type.inner_type) else: return {} def bind_type_variables(self, bindings): bound_inner_type = bind_type_variables(self.inner_type, bindings) if bound_inner_type == self.inner_type: return self else: bound_self = copy.copy(self) bound_self.inner_type = bound_inner_type return bound_self class CompositeTypeHint(object): """The base-class for all created type-hint classes defined below. CompositeTypeHint's serve primarily as TypeConstraint factories. They are only required to define a single method: '__getitem__' which should return a parameterized TypeConstraint, that can be used to enforce static or run-time type-checking. '__getitem__' is used as a factory function in order to provide a familiar API for defining type-hints. The ultimate result is that one will be able to use: CompositeTypeHint[type_parameter] to create a type-hint object that behaves like any other Python object. This allows one to create 'type-aliases' by assigning the returned type-hints to a variable. * Example: 'Coordinates = List[Tuple[int, int]]' """ def __getitem___(self, py_type): """Given a type creates a TypeConstraint instance parameterized by the type. This function serves as a factory function which creates TypeConstraint instances. Additionally, implementations by sub-classes should perform any sanity checking of the passed types in this method in order to rule-out disallowed behavior. Such as, attempting to create a TypeConstraint whose parameterized type is actually an object instance. Args: py_type: An instance of a Python type or TypeConstraint. Returns: An instance of a custom TypeConstraint for this CompositeTypeHint. Raises: TypeError: If the passed type violates any contraints for this particular TypeHint. """ raise NotImplementedError def validate_composite_type_param(type_param, error_msg_prefix): """Determines if an object is a valid type parameter to a CompositeTypeHint. Implements sanity checking to disallow things like: * List[1, 2, 3] or Dict[5]. Args: type_param: An object instance. error_msg_prefix: A string prefix used to format an error message in the case of an exception. Raises: TypeError: If the passed 'type_param' is not a valid type parameter for a CompositeTypeHint. """ # Must either be a TypeConstraint instance or a basic Python type. is_not_type_constraint = ( not isinstance(type_param, (type, types.ClassType, TypeConstraint)) and type_param is not None) is_forbidden_type = (isinstance(type_param, type) and type_param in DISALLOWED_PRIMITIVE_TYPES) if is_not_type_constraint or is_forbidden_type: raise TypeError('%s must be a non-sequence, a type, or a TypeConstraint. %s' ' is an instance of %s.' % (error_msg_prefix, type_param, type_param.__class__.__name__)) def _unified_repr(o): """Given an object return a qualified name for the object. This function closely mirrors '__qualname__' which was introduced in Python 3.3. It is used primarily to format types or object instances for error messages. Args: o: An instance of a TypeConstraint or a type. Returns: A qualified name for the passed Python object fit for string formatting. """ return repr(o) if isinstance(o, (TypeConstraint, types.NoneType)) else o.__name__ def check_constraint(type_constraint, object_instance): """Determine if the passed type instance satisfies the TypeConstraint. When examining a candidate type for constraint satisfaction in 'type_check', all CompositeTypeHint's eventually call this function. This function may end up being called recursively if the hinted type of a CompositeTypeHint is another CompositeTypeHint. Args: type_constraint: An instance of a TypeConstraint or a built-in Python type. object_instance: An object instance. Raises: SimpleTypeHintError: If 'type_constraint' is a one of the allowed primitive Python types and 'object_instance' isn't an instance of this type. CompositeTypeHintError: If 'type_constraint' is a TypeConstraint object and 'object_instance' does not satisfy its constraint. """ if type_constraint is None and object_instance is None: return elif isinstance(type_constraint, TypeConstraint): type_constraint.type_check(object_instance) elif type_constraint is None: # TODO(robertwb): Fix uses of None for Any. pass elif not isinstance(type_constraint, type): raise RuntimeError("bad type: %s" % (type_constraint,)) elif not isinstance(object_instance, type_constraint): raise SimpleTypeHintError class AnyTypeConstraint(TypeConstraint): """An Any type-hint. Any is intended to be used as a "don't care" when hinting the types of function arguments or return types. All other TypeConstraint's are equivalent to 'Any', and its 'type_check' method is a no-op. """ def __repr__(self): return 'Any' def type_check(self, instance): pass class TypeVariable(AnyTypeConstraint): def __init__(self, name): self.name = name def __repr__(self): return 'TypeVariable[%s]' % self.name def match_type_variables(self, concrete_type): return {self: concrete_type} def bind_type_variables(self, bindings): return bindings.get(self, self) class UnionHint(CompositeTypeHint): """A Union type-hint. Union[X, Y] accepts instances of type X OR type Y. Duplicate type parameters are ignored. Additonally, Nested Union hints will be flattened out. For example: * Union[Union[str, int], bool] -> Union[str, int, bool] A candidate type instance satisfies a UnionConstraint if it is an instance of any of the parameterized 'union_types' for a Union. Union[X] is disallowed, and all type parameters will be sanity checked to ensure compatibility with nested type-hints. When comparing two Union hints, ordering is enforced before comparison. * Union[int, str] == Union[str, int] """ class UnionConstraint(TypeConstraint): def __init__(self, union_types): self.union_types = set(union_types) def __eq__(self, other): return (isinstance(other, UnionHint.UnionConstraint) and self.union_types == other.union_types) def __hash__(self): return 1 + sum(hash(t) for t in self.union_types) def __repr__(self): # Sorting the type name strings simplifies unit tests. return 'Union[%s]' % (', '.join(sorted(_unified_repr(t) for t in self.union_types))) def _inner_types(self): for t in self.union_types: yield t def _consistent_with_check_(self, sub): if isinstance(sub, UnionConstraint): # A union type is compatible if every possible type is compatible. # E.g. Union[A, B, C] > Union[A, B]. return all(is_consistent_with(elem, self) for elem in sub.union_types) else: # Other must be compatible with at least one of this union's subtypes. # E.g. Union[A, B, C] > T if T > A or T > B or T > C. return any(is_consistent_with(sub, elem) for elem in self.union_types) def type_check(self, instance): error_msg = '' for t in self.union_types: try: check_constraint(t, instance) return except TypeError as e: error_msg = str(e) continue raise CompositeTypeHintError( '%s type-constraint violated. Expected an instance of one of: %s, ' 'received %s instead.%s' % (repr(self), tuple(sorted(_unified_repr(t) for t in self.union_types)), instance.__class__.__name__, error_msg)) def __getitem__(self, type_params): if not isinstance(type_params, (collections.Sequence, set)): raise TypeError('Cannot create Union without a sequence of types.') # Flatten nested Union's and duplicated repeated type hints. params = set() for t in type_params: validate_composite_type_param( t, error_msg_prefix='All parameters to a Union hint' ) if isinstance(t, self.UnionConstraint): params |= t.union_types else: params.add(t) if Any in params: return Any elif len(params) == 1: return iter(params).next() else: return self.UnionConstraint(params) UnionConstraint = UnionHint.UnionConstraint class OptionalHint(UnionHint): """An Option type-hint. Optional[X] accepts instances of X or None. The Optional[X] factory function proxies to Union[X, None] """ def __getitem__(self, py_type): # A single type must have been passed. if isinstance(py_type, collections.Sequence): raise TypeError('An Option type-hint only accepts a single type ' 'parameter.') return Union[py_type, None] class TupleHint(CompositeTypeHint): """A Tuple type-hint. Tuple can accept 1 or more type-hint parameters. Tuple[X, Y] represents a tuple of *exactly* two elements, with the first being of type 'X' and the second an instance of type 'Y'. * (1, 2) satisfies Tuple[int, int] Additionally, one is able to type-hint an arbitary length, homogeneous tuple by passing the Ellipsis (...) object as the second parameter. As an example, Tuple[str, ...] indicates a tuple of any length with each element being an instance of 'str'. """ class TupleSequenceConstraint(SequenceTypeConstraint): def __init__(self, type_param): super(TupleHint.TupleSequenceConstraint, self).__init__(type_param, tuple) def __repr__(self): return 'Tuple[%s, ...]' % _unified_repr(self.inner_type) def _consistent_with_check_(self, sub): if isinstance(sub, TupleConstraint): # E.g. Tuple[A, B] < Tuple[C, ...] iff A < C and B < C. return all(is_consistent_with(elem, self.inner_type) for elem in sub.tuple_types) else: return super(TupleSequenceConstraint, self)._consistent_with_check_(sub) class TupleConstraint(TypeConstraint): def __init__(self, type_params): self.tuple_types = tuple(type_params) def __eq__(self, other): return (isinstance(other, TupleHint.TupleConstraint) and self.tuple_types == other.tuple_types) def __hash__(self): return hash(self.tuple_types) def __repr__(self): return 'Tuple[%s]' % (', '.join(_unified_repr(t) for t in self.tuple_types)) def _inner_types(self): for t in self.tuple_types: yield t def _consistent_with_check_(self, sub): return (isinstance(sub, self.__class__) and len(sub.tuple_types) == len(self.tuple_types) and all(is_consistent_with(sub_elem, elem) for sub_elem, elem in zip(sub.tuple_types, self.tuple_types))) def type_check(self, tuple_instance): if not isinstance(tuple_instance, tuple): raise CompositeTypeHintError( "Tuple type constraint violated. Valid object instance must be of " "type 'tuple'. Instead, an instance of '%s' was received." % tuple_instance.__class__.__name__) if len(tuple_instance) != len(self.tuple_types): raise CompositeTypeHintError( 'Passed object instance is of the proper type, but differs in ' 'length from the hinted type. Expected a tuple of length %s, ' 'received a tuple of length %s.' % (len(self.tuple_types), len(tuple_instance))) for type_pos, (expected, actual) in enumerate(zip(self.tuple_types, tuple_instance)): try: check_constraint(expected, actual) continue except SimpleTypeHintError: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element #%s in ' 'the passed tuple is incorrect. Expected an instance of ' 'type %s, instead received an instance of type %s.' % (repr(self), type_pos, _unified_repr(expected), actual.__class__.__name__)) except CompositeTypeHintError as e: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element #%s in ' 'the passed tuple is incorrect. %s' % (repr(self), type_pos, e)) def match_type_variables(self, concrete_type): bindings = {} if isinstance(concrete_type, TupleConstraint): for a, b in zip(self.tuple_types, concrete_type.tuple_types): bindings.update(match_type_variables(a, b)) return bindings def bind_type_variables(self, bindings): bound_tuple_types = tuple( bind_type_variables(t, bindings) for t in self.tuple_types) if bound_tuple_types == self.tuple_types: return self else: return Tuple[bound_tuple_types] def __getitem__(self, type_params): ellipsis = False if not isinstance(type_params, collections.Iterable): # Special case for hinting tuples with arity-1. type_params = (type_params,) if type_params and type_params[-1] == Ellipsis: if len(type_params) != 2: raise TypeError('Ellipsis can only be used to type-hint an arbitrary ' 'length tuple of containing a single type: ' 'Tuple[A, ...].') # Tuple[A, ...] indicates an arbitary length homogeneous tuple. type_params = type_params[:1] ellipsis = True for t in type_params: validate_composite_type_param( t, error_msg_prefix='All parameters to a Tuple hint' ) if ellipsis: return self.TupleSequenceConstraint(type_params[0]) else: return self.TupleConstraint(type_params) TupleConstraint = TupleHint.TupleConstraint TupleSequenceConstraint = TupleHint.TupleSequenceConstraint class ListHint(CompositeTypeHint): """A List type-hint. List[X] represents an instance of a list populated by a single homogeneous type. The parameterized type 'X' can either be a built-in Python type or an instance of another TypeConstraint. * ['1', '2', '3'] satisfies List[str] """ class ListConstraint(SequenceTypeConstraint): def __init__(self, list_type): super(ListHint.ListConstraint, self).__init__(list_type, list) def __repr__(self): return 'List[%s]' % _unified_repr(self.inner_type) def __getitem__(self, t): validate_composite_type_param(t, error_msg_prefix='Parameter to List hint') return self.ListConstraint(t) ListConstraint = ListHint.ListConstraint class KVHint(CompositeTypeHint): """A KV type-hint, represents a Key-Value pair of a particular type. Internally, KV[X, Y] proxies to Tuple[X, Y]. A KV type-hint accepts only accepts exactly two type-parameters. The first represents the required key-type and the second the required value-type. """ def __getitem__(self, type_params): if not isinstance(type_params, tuple): raise TypeError('Parameter to KV type-hint must be a tuple of types: ' 'KV[.., ..].') if len(type_params) != 2: raise TypeError( 'Length of parameters to a KV type-hint must be exactly 2. Passed ' 'parameters: %s, have a length of %s.' % (type_params, len(type_params)) ) return Tuple[type_params] def key_value_types(kv): """Returns the key and value type of a KV type-hint. Args: kv: An instance of a TypeConstraint sub-class. Returns: A tuple: (key_type, value_type) if the passed type-hint is an instance of a KV type-hint, and (Any, Any) otherwise. """ if isinstance(kv, TupleHint.TupleConstraint): return kv.tuple_types return Any, Any class DictHint(CompositeTypeHint): """A Dict type-hint. Dict[K, V] Represents a dictionary where all keys are of a particular type and all values are of another (possible the same) type. """ class DictConstraint(TypeConstraint): def __init__(self, key_type, value_type): self.key_type = key_type self.value_type = value_type def __repr__(self): return 'Dict[%s, %s]' % (_unified_repr(self.key_type), _unified_repr(self.value_type)) def __eq__(self, other): return (type(self) == type(other) and self.key_type == other.key_type and self.value_type == other.value_type) def __hash__(self): return hash((type(self), self.key_type, self.value_type)) def _inner_types(self): yield self.key_type yield self.value_type def _consistent_with_check_(self, sub): return (isinstance(sub, self.__class__) and is_consistent_with(sub.key_type, self.key_type) and is_consistent_with(sub.key_type, self.key_type)) def _raise_hint_exception_or_inner_exception(self, is_key, incorrect_instance, inner_error_message=''): incorrect_type = 'values' if not is_key else 'keys' hinted_type = self.value_type if not is_key else self.key_type if inner_error_message: raise CompositeTypeHintError( '%s hint %s-type constraint violated. All %s should be of type ' '%s. Instead: %s' % (repr(self), incorrect_type[:-1], incorrect_type, _unified_repr(hinted_type), inner_error_message) ) else: raise CompositeTypeHintError( '%s hint %s-type constraint violated. All %s should be of ' 'type %s. Instead, %s is of type %s.' % (repr(self), incorrect_type[:-1], incorrect_type, _unified_repr(hinted_type), incorrect_instance, incorrect_instance.__class__.__name__) ) def type_check(self, dict_instance): if not isinstance(dict_instance, dict): raise CompositeTypeHintError( 'Dict type-constraint violated. All passed instances must be of ' 'type dict. %s is of type %s.' % (dict_instance, dict_instance.__class__.__name__)) for key, value in dict_instance.iteritems(): try: check_constraint(self.key_type, key) except CompositeTypeHintError as e: self._raise_hint_exception_or_inner_exception(True, key, str(e)) except SimpleTypeHintError: self._raise_hint_exception_or_inner_exception(True, key) try: check_constraint(self.value_type, value) except CompositeTypeHintError as e: self._raise_hint_exception_or_inner_exception(False, value, str(e)) except SimpleTypeHintError: self._raise_hint_exception_or_inner_exception(False, value) def match_type_variables(self, concrete_type): if isinstance(concrete_type, DictConstraint): bindings = {} bindings.update( match_type_variables(self.key_type, concrete_type.key_type)) bindings.update( match_type_variables(self.value_type, concrete_type.value_type)) return bindings else: return {} def bind_type_variables(self, bindings): bound_key_type = bind_type_variables(self.key_type, bindings) bound_value_type = bind_type_variables(self.value_type, bindings) if (bound_key_type, self.key_type) == (bound_value_type, self.value_type): return self else: return Dict[bound_key_type, bound_value_type] def __getitem__(self, type_params): # Type param must be a (k, v) pair. if not isinstance(type_params, tuple): raise TypeError('Parameter to Dict type-hint must be a tuple of types: ' 'Dict[.., ..].') if len(type_params) != 2: raise TypeError( 'Length of parameters to a Dict type-hint must be exactly 2. Passed ' 'parameters: %s, have a length of %s.' % (type_params, len(type_params)) ) key_type, value_type = type_params validate_composite_type_param( key_type, error_msg_prefix='Key-type parameter to a Dict hint' ) validate_composite_type_param( value_type, error_msg_prefix='Value-type parameter to a Dict hint' ) return self.DictConstraint(key_type, value_type) DictConstraint = DictHint.DictConstraint class SetHint(CompositeTypeHint): """A Set type-hint. Set[X] defines a type-hint for a set of homogeneous types. 'X' may be either a built-in Python type or a another nested TypeConstraint. """ class SetTypeConstraint(SequenceTypeConstraint): def __init__(self, type_param): super(SetHint.SetTypeConstraint, self).__init__(type_param, set) def __repr__(self): return 'Set[%s]' % _unified_repr(self.inner_type) def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to a Set hint' ) return self.SetTypeConstraint(type_param) SetTypeConstraint = SetHint.SetTypeConstraint class IterableHint(CompositeTypeHint): """An Iterable type-hint. Iterable[X] defines a type-hint for an object implementing an '__iter__' method which yields objects which are all of the same type. """ class IterableTypeConstraint(SequenceTypeConstraint): def __init__(self, iter_type): super(IterableHint.IterableTypeConstraint, self).__init__( iter_type, collections.Iterable) def __repr__(self): return 'Iterable[%s]' % _unified_repr(self.inner_type) def _consistent_with_check_(self, sub): if isinstance(sub, SequenceTypeConstraint): return is_consistent_with(sub.inner_type, self.inner_type) elif isinstance(sub, TupleConstraint): if not sub.tuple_types: # The empty tuple is consistent with Iterator[T] for any T. return True else: # Each element in the hetrogenious tuple must be consistent with # the iterator type. # E.g. Tuple[A, B] < Iterable[C] if A < C and B < C. return all(is_consistent_with(elem, self.inner_type) for elem in sub.tuple_types) else: return False def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to an Iterable hint' ) return self.IterableTypeConstraint(type_param) IterableTypeConstraint = IterableHint.IterableTypeConstraint class IteratorHint(CompositeTypeHint): """An Iterator type-hint. Iterator[X] defines a type-hint for an object implementing both '__iter__' and a 'next' method which yields objects which are all of the same type. Type checking a type-hint of this type is deferred in order to avoid depleting the underlying lazily generated sequence. See decorators.interleave_type_check for further information. """ class IteratorTypeConstraint(TypeConstraint): def __init__(self, t): self.yielded_type = t def __repr__(self): return 'Iterator[%s]' % _unified_repr(self.yielded_type) def _inner_types(self): yield self.yielded_type def _consistent_with_check_(self, sub): return (isinstance(sub, self.__class__) and is_consistent_with(sub.yielded_type, self.yielded_type)) def type_check(self, instance): # Special case for lazy types, we only need to enforce the underlying # type. This avoid having to compute the entirety of the generator/iter. try: check_constraint(self.yielded_type, instance) return except CompositeTypeHintError as e: raise CompositeTypeHintError( '%s hint type-constraint violated: %s' % (repr(self), str(e))) except SimpleTypeHintError: raise CompositeTypeHintError( '%s hint type-constraint violated. Expected a iterator of type %s. ' 'Instead received a iterator of type %s.' % (repr(self), _unified_repr(self.yielded_type), instance.__class__.__name__)) def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to an Iterator hint' ) return self.IteratorTypeConstraint(type_param) IteratorTypeConstraint = IteratorHint.IteratorTypeConstraint class WindowedTypeConstraint(TypeConstraint): """A type constraint for WindowedValue objects. Mostly for internal use. Attributes: inner_type: The type which the element should be an instance of. """ __metaclass__ = GetitemConstructor def __init__(self, inner_type): self.inner_type = inner_type def __eq__(self, other): return (isinstance(other, WindowedTypeConstraint) and self.inner_type == other.inner_type) def __hash__(self): return hash(self.inner_type) ^ 13 * hash(type(self)) def _inner_types(self): yield self.inner_type def _consistent_with_check_(self, sub): return (isinstance(sub, self.__class__) and is_consistent_with(sub.inner_type, self.inner_type)) def type_check(self, instance): from google.cloud.dataflow.transforms import window if not isinstance(instance, window.WindowedValue): raise CompositeTypeHintError( "Window type-constraint violated. Valid object instance " "must be of type 'WindowedValue'. Instead, an instance of '%s' " "was received." % (instance.__class__.__name__)) try: check_constraint(self.inner_type, instance.value) except (CompositeTypeHintError, SimpleTypeHintError) as e: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element in ' 'is incorrect. Expected an instance of type %s, ' 'instead received an instance of type %s.' % (repr(self), _unified_repr(self.inner_type), elem.__class__.__name__)) class GeneratorHint(IteratorHint): pass # Create the actual instances for all defined type-hints above. Any = AnyTypeConstraint() Union = UnionHint() Optional = OptionalHint() Tuple = TupleHint() List = ListHint() KV = KVHint() Dict = DictHint() Set = SetHint() Iterable = IterableHint() Iterator = IteratorHint() Generator = GeneratorHint() WindowedValue = WindowedTypeConstraint _KNOWN_PRIMITIVE_TYPES = { dict: Dict[Any, Any], list: List[Any], tuple: Tuple[Any, ...], set: Set[Any], # Using None for the NoneType is a common convention. None: type(None), } def normalize(x): if x in _KNOWN_PRIMITIVE_TYPES: return _KNOWN_PRIMITIVE_TYPES[x] else: return x def is_consistent_with(sub, base): """Returns whether the type a is consistent with b. This is accordig to the terminology of PEP 483/484. This relationship is neither symmetric nor transitive, but a good mnemonic to keep in mind is that is_consistent_with(a, b) is roughly equivalent to the issubclass(a, b) relation, but also handles the special Any type as well as type parameterization. """ if sub == base: # Common special case. return True if isinstance(sub, AnyTypeConstraint) or isinstance(base, AnyTypeConstraint): return True sub = normalize(sub) base = normalize(base) if isinstance(base, TypeConstraint): return base._consistent_with_check_(sub) elif isinstance(sub, TypeConstraint): # Nothing but object lives above any type constraints. return base == object else: return issubclass(sub, base) DataflowPythonSDK-0.2.5/google/cloud/dataflow/typehints/typehints_test.py000066400000000000000000001140341272340251000266740ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the type-hint objects and decorators.""" import functools import inspect import unittest import google.cloud.dataflow.typehints as typehints from google.cloud.dataflow.typehints import Any from google.cloud.dataflow.typehints import is_consistent_with from google.cloud.dataflow.typehints import Tuple from google.cloud.dataflow.typehints import TypeCheckError from google.cloud.dataflow.typehints import Union from google.cloud.dataflow.typehints import with_input_types from google.cloud.dataflow.typehints import with_output_types from google.cloud.dataflow.typehints.decorators import _check_instance_type from google.cloud.dataflow.typehints.decorators import _interleave_type_check from google.cloud.dataflow.typehints.decorators import _positional_arg_hints from google.cloud.dataflow.typehints.decorators import get_type_hints from google.cloud.dataflow.typehints.decorators import getcallargs_forhints def check_or_interleave(hint, value, var): if hint is None: return value elif isinstance(hint, typehints.IteratorHint.IteratorTypeConstraint): return _interleave_type_check(hint, var)(value) else: _check_instance_type(hint, value, var) return value def check_type_hints(f): @functools.wraps(f) def wrapper(*args, **kwargs): hints = get_type_hints(f) if hints.input_types: input_hints = getcallargs_forhints( f, *hints.input_types[0], **hints.input_types[1]) inputs = inspect.getcallargs(f, *args, **kwargs) for var, hint in input_hints.items(): value = inputs[var] new_value = check_or_interleave(hint, value, var) if new_value is not value: if var in kwargs: kwargs[var] = new_value else: args = list(args) for ix, pvar in enumerate(inspect.getargspec(f).args): if pvar == var: args[ix] = new_value break else: raise NotImplementedError('Iterable in nested argument %s' % var) res = f(*args, **kwargs) return check_or_interleave(hints.simple_output_type('typecheck'), res, None) return wrapper class DummyTestClass1(object): pass class DummyTestClass2(object): pass class SuperClass(object): pass class SubClass(SuperClass): pass class TypeHintTestCase(unittest.TestCase): def assertCompatible(self, base, sub): # pylint: disable=invalid-name self.assertTrue( is_consistent_with(sub, base), '%s is not consistent with %s' % (sub, base)) def assertNotCompatible(self, base, sub): # pylint: disable=invalid-name self.assertFalse( is_consistent_with(sub, base), '%s is consistent with %s' % (sub, base)) class AnyTypeConstraintTestCase(TypeHintTestCase): def test_any_compatibility(self): self.assertCompatible(typehints.Any, typehints.List[int]) self.assertCompatible(typehints.Any, DummyTestClass1) self.assertCompatible(typehints.Union[int, bool], typehints.Any) self.assertCompatible(typehints.Optional[int], typehints.Any) self.assertCompatible(typehints.Tuple[int], typehints.Any) self.assertCompatible(typehints.KV[int, str], typehints.Any) self.assertCompatible(typehints.Dict[int, bool], typehints.Any) self.assertCompatible(typehints.Set[int], typehints.Any) self.assertCompatible(typehints.Iterable[int], typehints.Any) self.assertCompatible(typehints.Iterator[int], typehints.Any) self.assertCompatible(typehints.Generator[int], typehints.Any) self.assertCompatible(object, typehints.Any) self.assertCompatible(typehints.Any, object) def test_repr(self): self.assertEqual('Any', repr(typehints.Any)) def test_type_check(self): # This test passes if the type_check call does not raise any exception. typehints.Any.type_check(1) typehints.Any.type_check([1, 2, 3]) typehints.Any.type_check(DummyTestClass1()) class UnionHintTestCase(TypeHintTestCase): def test_getitem_must_be_valid_type_param_cant_be_object_instance(self): with self.assertRaises(TypeError) as e: typehints.Union[5] self.assertEqual('Cannot create Union without a sequence of types.', e.exception.message) def test_getitem_must_be_valid_type_param(self): t = [2, 3] with self.assertRaises(TypeError) as e: typehints.Union[t] self.assertEqual('All parameters to a Union hint must be a non-sequence, ' 'a type, or a TypeConstraint. 2 is an instance of int.', e.exception.message) def test_getitem_duplicates_ignored(self): # Types should be de-duplicated. hint = typehints.Union[int, int, str] self.assertEqual(len(hint.union_types), 2) def test_getitem_nested_unions_flattened(self): # The two Union's should be merged into 1. hint = typehints.Union[typehints.Union[int, str], typehints.Union[float, bool]] self.assertTrue(len(hint.union_types) == 4) self.assertTrue(all(t in hint.union_types for t in (int, str, float, bool))) def test_union_hint_compatibility(self): self.assertCompatible(typehints.Union[int, float], int) self.assertCompatible(typehints.Union[int, str], typehints.Union[str, int]) self.assertCompatible(typehints.Union[int, float, str], typehints.Union[str, int]) self.assertCompatible( typehints.Union[DummyTestClass1, str], typehints.Union[DummyTestClass1, str]) self.assertCompatible(typehints.Union[int, str], typehints.Union[str, typehints.Union[int, str]]) self.assertNotCompatible(typehints.Union[float, bool], typehints.Union[int, bool]) self.assertNotCompatible(typehints.Union[bool, str], typehints.Union[float, bool, int]) def test_nested_compatibility(self): self.assertCompatible(Union[int, Tuple[Any, int]], Tuple[int, int]) self.assertCompatible(Union[int, Tuple[Any, Any]], Union[Tuple[int, Any], Tuple[Any, int]]) self.assertCompatible(Union[int, SuperClass], SubClass) self.assertCompatible(Union[int, float, SuperClass], Union[int, SubClass]) self.assertNotCompatible(Union[int, SubClass], SuperClass) self.assertNotCompatible(Union[int, float, SubClass], Union[int, SuperClass]) self.assertNotCompatible(Union[int, SuperClass], Union[int, float, SubClass]) def test_union_hint_repr(self): hint = typehints.Union[DummyTestClass1, str] self.assertIn( str(hint), # Uses frozen set internally, so order not guaranteed. ['Union[str, DummyTestClass1]', 'Union[DummyTestClass1, str]'] ) def test_union_hint_enforcement_composite_type_in_union(self): o = DummyTestClass1() hint = typehints.Union[int, DummyTestClass1] self.assertIsNone(hint.type_check(4)) self.assertIsNone(hint.type_check(o)) def test_union_hint_enforcement_part_of_union(self): hint = typehints.Union[int, str] self.assertIsNone(hint.type_check(5)) self.assertIsNone(hint.type_check('test')) def test_union_hint_enforcement_not_part_of_union(self): hint = typehints.Union[int, float] with self.assertRaises(TypeError) as e: hint.type_check('test') self.assertEqual("Union[float, int] type-constraint violated. Expected an " "instance of one of: ('float', 'int'), received str " "instead.", e.exception.message) class OptionalHintTestCase(TypeHintTestCase): def test_getitem_sequence_not_allowed(self): with self.assertRaises(TypeError) as e: typehints.Optional[int, str] self.assertTrue(e.exception.message.startswith( 'An Option type-hint only accepts a single type parameter.')) def test_getitem_proxy_to_union(self): hint = typehints.Optional[int] self.assertTrue(isinstance(hint, typehints.UnionHint.UnionConstraint)) class TupleHintTestCase(TypeHintTestCase): def test_getitem_invalid_ellipsis_type_param(self): error_msg = ('Ellipsis can only be used to type-hint an arbitrary length ' 'tuple of containing a single type: Tuple[A, ...].') with self.assertRaises(TypeError) as e: typehints.Tuple[int, int, ...] self.assertEqual(error_msg, e.exception.message) with self.assertRaises(TypeError) as e: typehints.Tuple[...] self.assertEqual(error_msg, e.exception.message) def test_getitem_params_must_be_type_or_constraint(self): expected_error_prefix = 'All parameters to a Tuple hint must be' with self.assertRaises(TypeError) as e: typehints.Tuple[5, [1, 3]] self.assertTrue(e.exception.message.startswith(expected_error_prefix)) with self.assertRaises(TypeError) as e: typehints.Tuple[list, dict] self.assertTrue(e.exception.message.startswith(expected_error_prefix)) def test_compatibility_arbitrary_length(self): self.assertNotCompatible( typehints.Tuple[int, int], typehints.Tuple[int, ...]) self.assertCompatible( typehints.Tuple[int, ...], typehints.Tuple[int, int]) self.assertCompatible( typehints.Tuple[Any, ...], typehints.Tuple[int, float]) self.assertCompatible( typehints.Tuple[SuperClass, ...], typehints.Tuple[SubClass, SuperClass]) self.assertCompatible(typehints.Iterable[int], typehints.Tuple[int, ...]) self.assertCompatible(typehints.Iterable[SuperClass], typehints.Tuple[SubClass, ...]) def test_compatibility(self): self.assertCompatible(typehints.Tuple[int, str], typehints.Tuple[int, str]) self.assertCompatible(typehints.Tuple[int, Any], typehints.Tuple[int, str]) self.assertCompatible(typehints.Tuple[int, str], typehints.Tuple[int, Any]) self.assertCompatible(typehints.Tuple[typehints.Union[int, str], bool], typehints.Tuple[typehints.Union[int, str], bool]) self.assertCompatible(typehints.Tuple[typehints.Union[str, int], int], typehints.Tuple[typehints.Union[int, str], int]) self.assertCompatible(typehints.Tuple[SuperClass, int], typehints.Tuple[SubClass, int]) self.assertNotCompatible(typehints.Tuple[int, int], typehints.Tuple[int, int, int]) def test_raw_tuple(self): self.assertCompatible(tuple, typehints.Tuple[int]) self.assertCompatible(tuple, typehints.Tuple[int, float]) self.assertCompatible(tuple, typehints.Tuple[int, ...]) def test_repr(self): hint = typehints.Tuple[int, str, float] self.assertEqual('Tuple[int, str, float]', str(hint)) hint = typehints.Tuple[DummyTestClass1, DummyTestClass2] self.assertEqual('Tuple[DummyTestClass1, DummyTestClass2]', str(hint)) hint = typehints.Tuple[float, ...] self.assertEqual('Tuple[float, ...]', str(hint)) def test_type_check_must_be_tuple(self): hint = typehints.Tuple[int, str] expected_error_prefix = 'Tuple type constraint violated. Valid object' invalid_instances = ([1, 2, 3], {4: 'f'}, 9, 'test', None) for t in invalid_instances: with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertTrue(e.exception.message.startswith(expected_error_prefix)) def test_type_check_must_have_same_arity(self): # A 2-tuple of ints. hint = typehints.Tuple[int, int] t = (1, 2, 3) with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertEqual('Passed object instance is of the proper type, but ' 'differs in length from the hinted type. Expected a ' 'tuple of length 2, received a tuple of length 3.', e.exception.message) def test_type_check_invalid_simple_types(self): hint = typehints.Tuple[str, bool] with self.assertRaises(TypeError) as e: hint.type_check((4, False)) self.assertEqual('Tuple[str, bool] hint type-constraint violated. The ' 'type of element #0 in the passed tuple is incorrect.' ' Expected an instance of type str, instead received ' 'an instance of type int.', e.exception.message) def test_type_check_invalid_composite_type(self): hint = typehints.Tuple[DummyTestClass1, DummyTestClass2] t = (DummyTestClass2(), DummyTestClass1()) with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertEqual('Tuple[DummyTestClass1, DummyTestClass2] hint ' 'type-constraint violated. The type of element #0 in the ' 'passed tuple is incorrect. Expected an instance of type ' 'DummyTestClass1, instead received an instance of type ' 'DummyTestClass2.', e.exception.message) def test_type_check_valid_simple_types(self): hint = typehints.Tuple[float, bool] self.assertIsNone(hint.type_check((4.3, True))) hint = typehints.Tuple[int] self.assertIsNone(hint.type_check((1,))) def test_type_check_valid_composite_types(self): hint = typehints.Tuple[typehints.Tuple[int, str], typehints.Tuple[int, bool]] self.assertIsNone(hint.type_check( ((4, 'test'), (4, True)) )) def test_type_check_valid_simple_type_arbitrary_length(self): hint = typehints.Tuple[int, ...] t = (1, 2, 3, 4) self.assertIsNone(hint.type_check(t)) def test_type_check_valid_composite_type_arbitrary_length(self): hint = typehints.Tuple[typehints.List[str], ...] t = (['h', 'e'], ['l', 'l'], ['o']) self.assertIsNone(hint.type_check(t)) def test_type_check_invalid_simple_type_arbitrary_length(self): hint = typehints.Tuple[str, ...] t = ('t', 'e', 5, 't') with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertEqual('Tuple[str, ...] hint type-constraint violated. The type ' 'of element #2 in the passed tuple is incorrect. Expected ' 'an instance of type str, instead received an instance of ' 'type int.', e.exception.message) def test_type_check_invalid_composite_type_arbitrary_length(self): hint = typehints.Tuple[typehints.List[int], ...] t = ([1, 2], 'e', 's', 't') with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertEqual("Tuple[List[int], ...] hint type-constraint violated. The " "type of element #1 in the passed tuple is incorrect: " "List type-constraint violated. Valid object instance " "must be of type 'list'. Instead, an instance of 'str' " "was received.", e.exception.message) class ListHintTestCase(TypeHintTestCase): def test_getitem_invalid_composite_type_param(self): with self.assertRaises(TypeError) as e: typehints.List[4] def test_list_constraint_compatibility(self): hint1 = typehints.List[typehints.Tuple[int, str]] hint2 = typehints.List[typehints.Tuple[float, bool]] self.assertCompatible(hint1, hint1) self.assertNotCompatible(hint1, hint2) self.assertCompatible(typehints.List[SuperClass], typehints.List[SubClass]) def test_list_repr(self): hint = ( typehints.List[typehints.Tuple[DummyTestClass1, DummyTestClass2]] ) self.assertEqual('List[Tuple[DummyTestClass1, DummyTestClass2]]', repr(hint)) def test_enforce_list_type_constraint_valid_simple_type(self): hint = typehints.List[int] self.assertIsNone(hint.type_check([1, 2, 3])) def test_enforce_list_type_constraint_valid_composite_type(self): hint = typehints.List[DummyTestClass1] l = [DummyTestClass1(), DummyTestClass1()] self.assertIsNone(hint.type_check(l)) def test_enforce_list_type_constraint_invalid_simple_type(self): hint = typehints.List[int] l = ['f', 'd', 'm'] with self.assertRaises(TypeError) as e: hint.type_check(l) self.assertEqual('List[int] hint type-constraint violated. The type of ' 'element #0 in the passed list is incorrect. Expected an ' 'instance of type int, instead received an instance of ' 'type str.', e.exception.message) def test_enforce_list_type_constraint_invalid_composite_type(self): hint = typehints.List[typehints.Tuple[int, int]] l = [('f', 1), ('m', 5)] with self.assertRaises(TypeError) as e: hint.type_check(l) self.assertEqual('List[Tuple[int, int]] hint type-constraint violated.' ' The type of element #0 in the passed list is ' 'incorrect: Tuple[int, int] hint type-constraint ' 'violated. The type of element #0 in the passed tuple' ' is incorrect. Expected an instance of type int, ' 'instead received an instance of type str.', e.exception.message) class KVHintTestCase(TypeHintTestCase): def test_getitem_param_must_be_tuple(self): with self.assertRaises(TypeError) as e: typehints.KV[4] self.assertEqual('Parameter to KV type-hint must be a tuple of types: ' 'KV[.., ..].', e.exception.message) def test_getitem_param_must_have_length_2(self): with self.assertRaises(TypeError) as e: typehints.KV[int, str, bool] self.assertEqual("Length of parameters to a KV type-hint must be " "exactly 2. Passed parameters: (, , ), have a length of 3.", e.exception.message) def test_getitem_proxy_to_tuple(self): hint = typehints.KV[int, str] self.assertTrue(isinstance(hint, typehints.Tuple.TupleConstraint)) def test_enforce_kv_type_constraint(self): hint = typehints.KV[str, typehints.Tuple[int, int, int]] t = ('test', (1, 2, 3)) self.assertIsNone(hint.type_check(t)) class DictHintTestCase(TypeHintTestCase): def test_getitem_param_must_be_tuple(self): with self.assertRaises(TypeError) as e: typehints.Dict[4] self.assertEqual('Parameter to Dict type-hint must be a tuple of ' 'types: Dict[.., ..].', e.exception.message) def test_getitem_param_must_have_length_2(self): with self.assertRaises(TypeError) as e: typehints.Dict[float, int, bool] self.assertEqual("Length of parameters to a Dict type-hint must be " "exactly 2. Passed parameters: (, , ), have a length of 3.", e.exception.message) def test_key_type_must_be_valid_composite_param(self): with self.assertRaises(TypeError) as e: typehints.Dict[list, int] def test_value_type_must_be_valid_composite_param(self): with self.assertRaises(TypeError) as e: typehints.Dict[str, 5] def test_compatibility(self): hint1 = typehints.Dict[int, str] hint2 = typehints.Dict[bool, int] hint3 = typehints.Dict[int, typehints.List[typehints.Tuple[str, str, str]]] self.assertCompatible(hint1, hint1) self.assertCompatible(hint3, hint3) self.assertNotCompatible(hint3, 4) self.assertNotCompatible(hint2, hint1) def test_repr(self): hint3 = typehints.Dict[int, typehints.List[typehints.Tuple[str, str, str]]] self.assertEqual('Dict[int, List[Tuple[str, str, str]]]', repr(hint3)) def test_type_checks_not_dict(self): hint = typehints.Dict[int, str] l = [1, 2] with self.assertRaises(TypeError) as e: hint.type_check(l) self.assertEqual('Dict type-constraint violated. All passed instances ' 'must be of type dict. [1, 2] is of type list.', e.exception.message) def test_type_check_invalid_key_type(self): hint = typehints.Dict[typehints.Tuple[int, int, int], typehints.List[str] ] d = {(1, 2): ['m', '1', '2', '3']} with self.assertRaises((TypeError, TypeError)) as e: hint.type_check(d) self.assertEqual('Dict[Tuple[int, int, int], List[str]] hint key-type ' 'constraint violated. All keys should be of type ' 'Tuple[int, int, int]. Instead: Passed object ' 'instance is of the proper type, but differs in ' 'length from the hinted type. Expected a tuple of ' 'length 3, received a tuple of length 2.', e.exception.message) def test_type_check_invalid_value_type(self): hint = typehints.Dict[str, typehints.Dict[int, str]] d = {'f': [1, 2, 3]} with self.assertRaises(TypeError) as e: hint.type_check(d) self.assertEqual('Dict[str, Dict[int, str]] hint value-type constraint' ' violated. All values should be of type ' 'Dict[int, str]. Instead: Dict type-constraint ' 'violated. All passed instances must be of type dict.' ' [1, 2, 3] is of type list.', e.exception.message) def test_type_check_valid_simple_type(self): hint = typehints.Dict[int, str] d = {4: 'f', 9: 'k'} self.assertIsNone(hint.type_check(d)) def test_type_check_valid_composite_type(self): hint = typehints.Dict[typehints.Tuple[str, str], typehints.List[int]] d = {('f', 'k'): [1, 2, 3], ('m', 'r'): [4, 6, 9]} self.assertIsNone(hint.type_check(d)) def test_match_type_variables(self): S = typehints.TypeVariable('S') # pylint: disable=invalid-name T = typehints.TypeVariable('T') # pylint: disable=invalid-name hint = typehints.Dict[S, T] self.assertEqual({S: int, T: str}, hint.match_type_variables(typehints.Dict[int, str])) class SetHintTestCase(TypeHintTestCase): def test_getitem_invalid_composite_type_param(self): with self.assertRaises(TypeError) as e: typehints.Set[list] self.assertEqual("Parameter to a Set hint must be a non-sequence, a " "type, or a TypeConstraint. is an " "instance of type.", e.exception.message) def test_compatibility(self): hint1 = typehints.Set[typehints.List[str]] hint2 = typehints.Set[typehints.Tuple[int, int]] self.assertCompatible(hint1, hint1) self.assertNotCompatible(hint2, hint1) def test_repr(self): hint = typehints.Set[typehints.List[bool]] self.assertEqual('Set[List[bool]]', repr(hint)) def test_type_check_must_be_set(self): hint = typehints.Set[str] with self.assertRaises(TypeError) as e: hint.type_check(4) self.assertEqual("Set type-constraint violated. Valid object instance " "must be of type 'set'. Instead, an instance of 'int'" " was received.", e.exception.message) def test_type_check_invalid_elem_type(self): hint = typehints.Set[float] with self.assertRaises(TypeError): hint.type_check('f') def test_type_check_valid_elem_simple_type(self): hint = typehints.Set[str] s = set(['f', 'm', 'k']) self.assertIsNone(hint.type_check(s)) def test_type_check_valid_elem_composite_type(self): hint = typehints.Set[typehints.Union[int, str]] s = set([9, 'm', 'k']) self.assertIsNone(hint.type_check(s)) class IterableHintTestCase(TypeHintTestCase): def test_getitem_invalid_composite_type_param(self): with self.assertRaises(TypeError) as e: typehints.Iterable[5] self.assertEqual('Parameter to an Iterable hint must be a ' 'non-sequence, a type, or a TypeConstraint. 5 is ' 'an instance of int.', e.exception.message) def test_compatibility(self): self.assertCompatible(typehints.Iterable[int], typehints.List[int]) self.assertCompatible(typehints.Iterable[int], typehints.Set[int]) self.assertCompatible(typehints.Iterable[typehints.Any], typehints.List[typehints.Tuple[int, bool]]) self.assertCompatible(typehints.Iterable[int], typehints.Iterable[int]) self.assertCompatible(typehints.Iterable[typehints.Union[int, str]], typehints.Iterable[typehints.Union[int, str]]) self.assertNotCompatible(typehints.Iterable[str], typehints.Iterable[bool]) self.assertCompatible(typehints.Iterable[int], typehints.List[int]) self.assertCompatible(typehints.Iterable[int], typehints.Set[int]) self.assertCompatible(typehints.Iterable[typehints.Any], typehints.List[typehints.Tuple[int, bool]]) def test_tuple_compatibility(self): self.assertCompatible(typehints.Iterable[int], typehints.Tuple[int, ...]) self.assertCompatible(typehints.Iterable[SuperClass], typehints.Tuple[SubClass, ...]) self.assertCompatible(typehints.Iterable[int], typehints.Tuple[int, int]) self.assertCompatible(typehints.Iterable[Any], typehints.Tuple[int, float]) self.assertCompatible(typehints.Iterable[typehints.Union[int, float]], typehints.Tuple[int, ...]) self.assertCompatible(typehints.Iterable[typehints.Union[int, float]], typehints.Tuple[int, float]) self.assertCompatible(typehints.Iterable[typehints.Union[int, float]], typehints.Tuple[int, float, int]) def test_repr(self): hint = typehints.Iterable[typehints.Set[str]] self.assertEqual('Iterable[Set[str]]', repr(hint)) def test_type_check_must_be_iterable(self): with self.assertRaises(TypeError) as e: hint = typehints.Iterable[int] hint.type_check(5) self.assertEqual("Iterable type-constraint violated. Valid object " "instance must be of type 'iterable'. Instead, an " "instance of 'int' was received.", e.exception.message) def test_type_check_violation_invalid_simple_type(self): hint = typehints.Iterable[float] l = set([1, 2, 3, 4]) with self.assertRaises(TypeError): hint.type_check(l) def test_type_check_violation_valid_simple_type(self): hint = typehints.Iterable[str] l = ('t', 'e', 's', 't') self.assertIsNone(hint.type_check(l)) def test_type_check_violation_invalid_composite_type(self): hint = typehints.Iterable[typehints.List[int]] l = ([['t', 'e'], ['s', 't']]) with self.assertRaises(TypeError) as e: hint.type_check(l) def test_type_check_violation_valid_composite_type(self): hint = typehints.Iterable[typehints.List[int]] l = ([[1, 2], [3, 4, 5]]) self.assertIsNone(hint.type_check(l)) class TestGeneratorWrapper(TypeHintTestCase): def test_functions_as_regular_generator(self): def count(n): for i in range(n): yield i l = [] interleave_func = lambda x: l.append(x) wrapped_gen = typehints.GeneratorWrapper(count(4), interleave_func) # Should function as a normal generator. self.assertEqual(0, next(wrapped_gen)) self.assertEqual((1, 2, 3), tuple(wrapped_gen)) # Interleave function should have been called each time. self.assertEqual([0, 1, 2, 3], l) class GeneratorHintTestCase(TypeHintTestCase): def test_repr(self): hint = typehints.Iterator[typehints.Set[str]] self.assertEqual('Iterator[Set[str]]', repr(hint)) def test_compatibility(self): self.assertCompatible(typehints.Iterator[int], typehints.Iterator[int]) self.assertNotCompatible(typehints.Iterator[str], typehints.Iterator[float]) def test_generator_return_hint_invalid_yield_type(self): @check_type_hints @with_output_types(typehints.Iterator[int]) def all_upper(s): for e in s: yield e.upper() with self.assertRaises(TypeCheckError) as e: next(all_upper('hello')) self.assertEqual('Type-hint for return type violated: Iterator[int] ' 'hint type-constraint violated. Expected a iterator ' 'of type int. Instead received a iterator of type ' 'str.', e.exception.message) def test_generator_argument_hint_invalid_yield_type(self): def wrong_yield_gen(): for e in ['a', 'b']: yield e @check_type_hints @with_input_types(a=typehints.Iterator[int]) def increment(a): return [e + 1 for e in a] with self.assertRaises(TypeCheckError) as e: increment(wrong_yield_gen()) self.assertEqual("Type-hint for argument: 'a' violated: Iterator[int] " "hint type-constraint violated. Expected a iterator " "of type int. Instead received a iterator of type " "str.", e.exception.message) class TakesDecoratorTestCase(TypeHintTestCase): def test_must_be_primitive_type_or_constraint(self): with self.assertRaises(TypeError) as e: t = [1, 2] @with_input_types(a=t) def foo(a): pass self.assertEqual('All type hint arguments must be a non-sequence, a ' 'type, or a TypeConstraint. [1, 2] is an instance of ' 'list.', e.exception.message) with self.assertRaises(TypeError) as e: t = 5 @check_type_hints @with_input_types(a=t) def foo(a): pass self.assertEqual('All type hint arguments must be a non-sequence, a type, ' 'or a TypeConstraint. 5 is an instance of int.', e.exception.message) def test_basic_type_assertion(self): @check_type_hints @with_input_types(a=int) def foo(a): return a + 1 with self.assertRaises(TypeCheckError) as e: m = 'a' foo(m) self.assertEqual("Type-hint for argument: 'a' violated. Expected an " "instance of , instead found an " "instance of .", e.exception.message) def test_composite_type_assertion(self): @check_type_hints @with_input_types(a=typehints.List[int]) def foo(a): a.append(1) return a with self.assertRaises(TypeCheckError) as e: m = ['f', 'f'] foo(m) self.assertEqual("Type-hint for argument: 'a' violated: List[int] hint " "type-constraint violated. The type of element #0 in " "the passed list is incorrect. Expected an instance of " "type int, instead received an instance of type str.", e.exception.message) def test_valid_simple_type_arguments(self): @with_input_types(a=str) def upper(a): return a.upper() # Type constraints should pass, and function will be evaluated as normal. self.assertEqual('M', upper('m')) def test_any_argument_type_hint(self): @check_type_hints @with_input_types(a=typehints.Any) def foo(a): return 4 self.assertEqual(4, foo('m')) def test_valid_mix_positional_and_keyword_arguments(self): @check_type_hints @with_input_types(typehints.List[int], elem=typehints.List[int]) def combine(container, elem): return container + elem self.assertEqual([1, 2, 3], combine([1, 2], [3])) def test_invalid_only_positional_arguments(self): @check_type_hints @with_input_types(int, int) def sub(a, b): return a - b with self.assertRaises(TypeCheckError) as e: sub(1, 'two') self.assertEqual("Type-hint for argument: 'b' violated. Expected an " "instance of , instead found an instance " "of .", e.exception.message) def test_valid_only_positional_arguments(self): @with_input_types(int, int) def add(a, b): return a + b self.assertEqual(3, add(1, 2)) class ReturnsDecoratorTestCase(TypeHintTestCase): def test_no_kwargs_accepted(self): with self.assertRaises(ValueError): @with_output_types(m=int) def foo(): return 5 def test_must_be_primitive_type_or_type_constraint(self): with self.assertRaises(TypeError): @with_output_types(5) def foo(): pass with self.assertRaises(TypeError): @with_output_types([1, 2]) def foo(): pass def test_must_be_single_return_type(self): with self.assertRaises(ValueError): @with_output_types(int, str) def foo(): return 4, 'f' def test_type_check_violation(self): @check_type_hints @with_output_types(int) def foo(a): return 'test' with self.assertRaises(TypeCheckError) as e: foo(4) self.assertEqual("Type-hint for return type violated. Expected an " "instance of , instead found an instance " "of .", e.exception.message) def test_type_check_simple_type(self): @with_output_types(str) def upper(a): return a.upper() self.assertEqual('TEST', upper('test')) def test_type_check_composite_type(self): @with_output_types(typehints.List[typehints.Tuple[int, int]]) def bar(): return [(i, i+1) for i in range(5)] self.assertEqual([(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)], bar()) def test_any_return_type_hint(self): @with_output_types(typehints.Any) def bar(): return 'foo' self.assertEqual('foo', bar()) class CombinedReturnsAndTakesTestCase(TypeHintTestCase): def test_enable_and_disable_type_checking_takes(self): @with_input_types(a=int) def int_to_str(a): return str(a) # The function call below violates the argument type-hint above, but won't # result in an exception since run-time type-checking was disabled above. self.assertEqual('a', int_to_str('a')) # Must re-define since the conditional is in the (maybe)wrapper. @check_type_hints @with_input_types(a=int) def int_to_str(a): return str(a) # With run-time type checking enabled once again the same call-atttempt # should result in a TypeCheckError. with self.assertRaises(TypeCheckError): int_to_str('a') def test_enable_and_disable_type_checking_returns(self): @with_output_types(str) def int_to_str(a): return a # The return value of the function above violates the return-type # type-hint above, but won't result in an exception since run-time # type-checking was disabled above. self.assertEqual(9, int_to_str(9)) # Must re-define since the conditional is in the (maybe)wrapper. @check_type_hints @with_output_types(str) def int_to_str(a): return a # With type-checking enabled once again we should get a TypeCheckError here. with self.assertRaises(TypeCheckError): int_to_str(9) def test_valid_mix_pos_and_keyword_with_both_orders(self): @with_input_types(str, start=int) @with_output_types(str) def to_upper_with_slice(string, start): return string.upper()[start:] self.assertEqual('ELLO', to_upper_with_slice('hello', 1)) def test_simple_takes_and_returns_hints(self): @check_type_hints @with_output_types(str) @with_input_types(a=str) def to_lower(a): return a.lower() # Return type and argument type satisfied, should work as normal. self.assertEqual('m', to_lower('M')) # Invalid argument type should raise a TypeCheckError with self.assertRaises(TypeCheckError): to_lower(5) @check_type_hints @with_output_types(str) @with_input_types(a=str) def to_lower(a): return 9 # Modified function now has an invalid return type. with self.assertRaises(TypeCheckError): to_lower('a') def test_composite_takes_and_returns_hints(self): @check_type_hints @with_input_types(it=typehints.List[int]) @with_output_types(typehints.List[typehints.Tuple[int, int]]) def expand_ints(it): return [(i, i + 1) for i in it] # Return type and argument type satisfied, should work as normal. self.assertEqual([(0, 1), (1, 2), (2, 3)], expand_ints(list(range(3)))) # Invalid argument, list of str instead of int. with self.assertRaises(TypeCheckError): expand_ints('t e s t'.split()) @check_type_hints @with_output_types(typehints.List[typehints.Tuple[int, int]]) @with_input_types(it=typehints.List[int]) def expand_ints(it): return [str(i) for i in it] # Modified function now has invalid return type. with self.assertRaises(TypeCheckError): expand_ints(list(range(2))) class DecoratorHelpers(TypeHintTestCase): def test_hint_helper(self): self.assertTrue(typehints.is_consistent_with(Any, int)) self.assertTrue(typehints.is_consistent_with(int, Any)) self.assertTrue(typehints.is_consistent_with(str, object)) self.assertFalse(typehints.is_consistent_with(object, str)) self.assertTrue(typehints.is_consistent_with(str, Union[str, int])) self.assertFalse(typehints.is_consistent_with(Union[str, int], str)) def test_positional_arg_hints(self): self.assertEquals(typehints.Any, _positional_arg_hints('x', {})) self.assertEquals(int, _positional_arg_hints('x', {'x': int})) self.assertEquals(typehints.Tuple[int, typehints.Any], _positional_arg_hints(['x', 'y'], {'x': int})) def test_getcallargs_forhints(self): func = lambda a, (b, c), *d: None self.assertEquals( {'a': Any, 'b': Any, 'c': Any, 'd': Tuple[Any, ...]}, getcallargs_forhints(func, *[Any, Any])) self.assertEquals( {'a': Any, 'b': Any, 'c': Any, 'd': Tuple[Any, ...]}, getcallargs_forhints(func, *[Any, Any, Any, int])) self.assertEquals( {'a': int, 'b': str, 'c': Any, 'd': Tuple[Any, ...]}, getcallargs_forhints(func, *[int, Tuple[str, Any]])) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/000077500000000000000000000000001272340251000223425ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/__init__.py000066400000000000000000000014361272340251000244570ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A package containing utilities.""" # We must import path here to support the pattern of referencing utils.path # without needing to explicitly import google.cloud.dataflow.utils.path. import path DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/counters.pxd000066400000000000000000000016571272340251000247320ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. # cython: profile=True # cython: overflowcheck=True cdef class Counter(object): cdef readonly object name cdef readonly object combine_fn cdef readonly object accumulator cdef readonly object _add_input cpdef bint update(self, value) except -1 cdef class AccumulatorCombineFnCounter(Counter): cdef readonly object _fast_add_input DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/counters.py000066400000000000000000000136461272340251000245700ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. # cython: profile=False # cython: overflowcheck=True """Counters collect the progress of the Worker for reporting to the service.""" import threading from google.cloud.dataflow.transforms import cy_combiners class Counter(object): """A counter aggregates a series of values. The aggregation kind of the Counter is specified when the Counter is created. The values aggregated must be of an appropriate for the aggregation used. Aggregations supported are listed in the code. (The aggregated value will be reported to the Dataflow service.) Do not create directly; call CounterFactory.get_counter instead. Attributes: name: the name of the counter, a string aggregation_kind: one of the aggregation kinds defined by this class. total: the total size of all the items passed to update() elements: the number of times update() was called """ # Handy references to common counters. SUM = cy_combiners.SumInt64Fn() MEAN = cy_combiners.MeanInt64Fn() def __init__(self, name, combine_fn): """Creates a Counter object. Args: name: the name of this counter. Typically has three parts: "step-output-counter". combine_fn: the CombineFn to use for aggregation """ self.name = name self.combine_fn = combine_fn self.accumulator = combine_fn.create_accumulator() self._add_input = self.combine_fn.add_input def update(self, value): self.accumulator = self._add_input(self.accumulator, value) def value(self): return self.combine_fn.extract_output(self.accumulator) def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s %s %s' % (self.name, self.combine_fn.__class__.__name__, self.value()) class AccumulatorCombineFnCounter(Counter): """Counter optimized for a mutating accumulator that holds all the logic.""" def __init__(self, name, combine_fn): assert isinstance(combine_fn, cy_combiners.AccumulatorCombineFn) super(AccumulatorCombineFnCounter, self).__init__(name, combine_fn) self._fast_add_input = self.accumulator.add_input def update(self, value): self._fast_add_input(value) # Counters that represent Accumulators have names starting with this USER_COUNTER_PREFIX = 'user-' class CounterFactory(object): """Keeps track of unique counters.""" def __init__(self): self.counters = {} # Lock to be acquired when accessing the counters map. self._lock = threading.Lock() def get_counter(self, name, combine_fn): """Returns a counter with the requested name. Passing in the same name will return the same counter; the combine_fn must agree. Args: name: the name of this counter. Typically has three parts: "step-output-counter". combine_fn: the CombineFn to use for aggregation Returns: A new or existing counter with the requested name. """ with self._lock: counter = self.counters.get(name, None) if counter: assert counter.combine_fn == combine_fn else: if isinstance(combine_fn, cy_combiners.AccumulatorCombineFn): counter = AccumulatorCombineFnCounter(name, combine_fn) else: counter = Counter(name, combine_fn) self.counters[name] = counter return counter def get_aggregator_counter(self, step_name, aggregator): """Returns an AggregationCounter for this step's aggregator. Passing in the same values will return the same counter. Args: step_name: the name of this step. aggregator: an Aggregator object. Returns: A new or existing counter. """ return self.get_counter( '%s%s-%s' % (USER_COUNTER_PREFIX, step_name, aggregator.name), aggregator.combine_fn) def get_counters(self): """Returns the current set of counters. Returns: An iterable that contains the current set of counters. To make sure that multiple threads can iterate over the set of counters, we return a new iterable here. Note that the actual set of counters may get modified after this method returns hence the returned iterable may be stale. """ with self._lock: return self.counters.values() def get_aggregator_values(self, aggregator_or_name): """Returns dict of step names to values of the aggregator.""" with self._lock: return get_aggregator_values( aggregator_or_name, self.counters, lambda counter: counter.value()) def get_aggregator_values(aggregator_or_name, counter_dict, value_extractor=None): """Extracts the named aggregator value from a set of counters. Args: aggregator_or_name: an Aggregator object or the name of one. counter_dict: a dict object of {name: value_wrapper} value_extractor: a function to convert the value_wrapper into a value. If None, no extraction is done and the value is return unchanged. Returns: dict of step names to values of the aggregator. """ name = aggregator_or_name if value_extractor is None: value_extractor = lambda x: x if not isinstance(aggregator_or_name, basestring): name = aggregator_or_name.name return {n: value_extractor(c) for n, c in counter_dict.iteritems() if n.startswith(USER_COUNTER_PREFIX) and n.endswith('-%s' % name)} DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/dependency.py000066400000000000000000000456421272340251000250450ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Support for installing custom code and required dependencies. Workflows, with the exception of very simple ones, are organized in multiple modules and packages. Typically, these modules and packages have dependencies on other standard libraries. Dataflow relies on the Python setuptools package to handle these scenarios. For further details please read: https://pythonhosted.org/setuptools/setuptools.html When a runner tries to run a pipeline it will check for a --requirements_file and a --setup_file option. If --setup_file is present then it is assumed that the folder containing the file specified by the option has the typical layout required by setuptools and it will run 'python setup.py sdist' to produce a source distribution. The resulting tarball (a file ending in .tar.gz) will be staged at the GCS staging location specified as job option. When a worker starts it will check for the presence of this file and will run 'easy_install tarball' to install the package in the worker. If --requirements_file is present then the file specified by the option will be staged in the GCS staging location. When a worker starts it will check for the presence of this file and will run 'pip install -r requirements.txt'. A requirements file can be easily generated by running 'pip freeze -r requirements.txt'. The reason a Dataflow runner does not run this automatically is because quite often only a small fraction of the dependencies present in a requirements.txt file are actually needed for remote execution and therefore a one-time manual trimming is desirable. TODO(silviuc): Staged files should have a job specific prefix. To prevent several jobs in the same project stomping on each other due to a shared staging location. TODO(silviuc): Should we allow several setup packages? TODO(silviuc): We should allow customizing the exact command for setup build. """ import glob import logging import os import shutil import tempfile from google.cloud.dataflow import utils from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.utils import names from google.cloud.dataflow.utils import processes from google.cloud.dataflow.utils.options import GoogleCloudOptions from google.cloud.dataflow.utils.options import SetupOptions from google.cloud.dataflow.version import __version__ # Standard file names used for staging files. WORKFLOW_TARBALL_FILE = 'workflow.tar.gz' REQUIREMENTS_FILE = 'requirements.txt' EXTRA_PACKAGES_FILE = 'extra_packages.txt' PACKAGES_URL_PREFIX = ( 'https://github.com/GoogleCloudPlatform/DataflowPythonSDK/archive') def _dependency_file_copy(from_path, to_path): """Copies a local file to a GCS file or vice versa.""" logging.info('file copy from %s to %s.', from_path, to_path) if from_path.startswith('gs://') or to_path.startswith('gs://'): command_args = ['gsutil', '-m', '-q', 'cp', from_path, to_path] logging.info('Executing command: %s', command_args) result = processes.call(command_args) if result != 0: raise ValueError( 'Failed to copy GCS file from %s to %s.' % (from_path, to_path)) else: # Branch used only for unit tests and integration tests. # In such environments GCS support is not available. if not os.path.isdir(os.path.dirname(to_path)): logging.info('Created folder (since we have not done yet, and any errors ' 'will follow): %s ', os.path.dirname(to_path)) os.mkdir(os.path.dirname(to_path)) shutil.copyfile(from_path, to_path) def _dependency_file_download(from_url, to_folder): """Downloads a file from a URL and returns path to the local file.""" # TODO(silviuc): We should cache downloads so we do not do it for every job. try: # We check if the file is actually there because wget returns a file # even for a 404 response (file will contain the contents of the 404 # response). response, content = __import__('httplib2').Http().request(from_url) if int(response['status']) >= 400: raise RuntimeError( 'Dataflow SDK not found at %s (response: %s)' % (from_url, response)) local_download_file = os.path.join(to_folder, 'dataflow-sdk.tar.gz') with open(local_download_file, 'w') as f: f.write(content) except Exception: logging.info('Failed to download SDK from %s', from_url) raise return local_download_file def _stage_extra_packages(extra_packages, staging_location, file_copy=_dependency_file_copy, temp_dir=None): """Stages a list of local extra packages. Args: extra_packages: Ordered list of local paths to extra packages to be staged. staging_location: Staging location for the packages. file_copy: Callable for copying files. The default version will copy from a local file to a GCS location using the gsutil tool available in the Google Cloud SDK package. temp_dir: Temporary folder where the resource building can happen. If None then a unique temp directory will be created. Used only for testing. Returns: A list of file names (no paths) for the resources staged. All the files are assumed to be staged in staging_location. Raises: RuntimeError: If files specified are not found or do not have expected name patterns. """ resources = [] tempdir = None local_packages = [] for package in extra_packages: if not os.path.basename(package).endswith('.tar.gz'): raise RuntimeError( 'The --extra_packages option expects a full path ending with ' '\'.tar.gz\' instead of %s' % package) if not os.path.isfile(package): if package.startswith('gs://'): if not tempdir: tempdir = tempfile.mkdtemp() logging.info('Downloading extra package: %s locally before staging', package) _dependency_file_copy(package, tempdir) else: raise RuntimeError( 'The file %s cannot be found. It was specified in the ' '--extra_packages command line option.' % package) else: local_packages.append(package) if tempdir: local_packages.extend( [utils.path.join(tempdir, f) for f in os.listdir(tempdir)]) for package in local_packages: basename = os.path.basename(package) staged_path = utils.path.join(staging_location, basename) file_copy(package, staged_path) resources.append(basename) # Create a file containing the list of extra packages and stage it. # The file is important so that in the worker the packages are installed # exactly in the order specified. This approach will avoid extra PyPI # requests. For example if package A depends on package B and package A # is installed first then the installer will try to satisfy the # dependency on B by downloading the package from PyPI. If package B is # installed first this is avoided. with open(os.path.join(temp_dir, EXTRA_PACKAGES_FILE), 'wt') as f: for package in local_packages: f.write('%s\n' % os.path.basename(package)) staged_path = utils.path.join(staging_location, EXTRA_PACKAGES_FILE) # Note that the caller of this function is responsible for deleting the # temporary folder where all temp files are created, including this one. file_copy(os.path.join(temp_dir, EXTRA_PACKAGES_FILE), staged_path) resources.append(EXTRA_PACKAGES_FILE) # Remove temp files created by downloading packages from GCS. if tempdir: try: temp_files = os.listdir(tempdir) for temp_file in temp_files: os.remove(utils.path.join(tempdir, temp_file)) os.rmdir(tempdir) except OSError as e: logging.info( '%s: (Ignored) Failed to delete all temporary files in %s.', e, tempdir) return resources def _populate_requirements_cache(requirements_file, cache_dir): # The 'pip download' command will not download again if it finds the # tarball with the proper version already present. # It will get the packages downloaded in the order they are presented in # the requirements file and will not download package dependencies. cmd_args = [ 'pip', 'install', '--download', cache_dir, '-r', requirements_file, # Download from PyPI source distributions. '--no-binary', ':all:'] logging.info('Executing command: %s', cmd_args) result = processes.call(cmd_args) if result != 0: raise RuntimeError( 'Failed to execute command: %s. Exit code %d', cmd_args, result) def stage_job_resources( options, file_copy=_dependency_file_copy, build_setup_args=None, temp_dir=None, populate_requirements_cache=_populate_requirements_cache): """Creates (if needed) and stages job resources to options.staging_location. Args: options: Command line options. More specifically the function will expect staging_location, requirements_file, setup_file, and save_main_session options to be present. file_copy: Callable for copying files. The default version will copy from a local file to a GCS location using the gsutil tool available in the Google Cloud SDK package. build_setup_args: A list of command line arguments used to build a setup package. Used only if options.setup_file is not None. Used only for testing. temp_dir: Temporary folder where the resource building can happen. If None then a unique temp directory will be created. Used only for testing. populate_requirements_cache: Callable for populating the requirements cache. Used only for testing. Returns: A list of file names (no paths) for the resources staged. All the files are assumed to be staged in options.staging_location. Raises: RuntimeError: If files specified are not found or error encountered while trying to create the resources (e.g., build a setup package). """ temp_dir = temp_dir or tempfile.mkdtemp() resources = [] google_cloud_options = options.view_as(GoogleCloudOptions) setup_options = options.view_as(SetupOptions) # Make sure that all required options are specified. There are a few that have # defaults to support local running scenarios. if google_cloud_options.staging_location is None: raise RuntimeError( 'The --staging_location option must be specified.') if google_cloud_options.temp_location is None: raise RuntimeError( 'The --temp_location option must be specified.') # Stage a requirements file if present. if setup_options.requirements_file is not None: if not os.path.isfile(setup_options.requirements_file): raise RuntimeError('The file %s cannot be found. It was specified in the ' '--requirements_file command line option.' % setup_options.requirements_file) staged_path = utils.path.join(google_cloud_options.staging_location, REQUIREMENTS_FILE) file_copy(setup_options.requirements_file, staged_path) resources.append(REQUIREMENTS_FILE) requirements_cache_path = ( os.path.join(tempfile.gettempdir(), 'dataflow-requirements-cache') if setup_options.requirements_cache is None else setup_options.requirements_cache) # Populate cache with packages from requirements and stage the files # in the cache. if not os.path.exists(requirements_cache_path): os.makedirs(requirements_cache_path) populate_requirements_cache( setup_options.requirements_file, requirements_cache_path) for pkg in glob.glob(os.path.join(requirements_cache_path, '*')): file_copy(pkg, utils.path.join(google_cloud_options.staging_location, os.path.basename(pkg))) resources.append(os.path.basename(pkg)) # Handle a setup file if present. # We will build the setup package locally and then copy it to the staging # location because the staging location is a GCS path and the file cannot be # created directly there. if setup_options.setup_file is not None: if not os.path.isfile(setup_options.setup_file): raise RuntimeError('The file %s cannot be found. It was specified in the ' '--setup_file command line option.' % setup_options.setup_file) if os.path.basename(setup_options.setup_file) != 'setup.py': raise RuntimeError( 'The --setup_file option expects the full path to a file named ' 'setup.py instead of %s' % setup_options.setup_file) tarball_file = _build_setup_package(setup_options.setup_file, temp_dir, build_setup_args) staged_path = utils.path.join(google_cloud_options.staging_location, WORKFLOW_TARBALL_FILE) file_copy(tarball_file, staged_path) resources.append(WORKFLOW_TARBALL_FILE) # Handle extra local packages that should be staged. if setup_options.extra_packages is not None: resources.extend( _stage_extra_packages(setup_options.extra_packages, google_cloud_options.staging_location, file_copy=file_copy, temp_dir=temp_dir)) # Pickle the main session if requested. # We will create the pickled main session locally and then copy it to the # staging location because the staging location is a GCS path and the file # cannot be created directly there. if setup_options.save_main_session: pickled_session_file = os.path.join(temp_dir, names.PICKLED_MAIN_SESSION_FILE) pickler.dump_session(pickled_session_file) staged_path = utils.path.join(google_cloud_options.staging_location, names.PICKLED_MAIN_SESSION_FILE) file_copy(pickled_session_file, staged_path) resources.append(names.PICKLED_MAIN_SESSION_FILE) if hasattr(setup_options, 'sdk_location') and setup_options.sdk_location: if setup_options.sdk_location == 'default': stage_tarball_from_remote_location = True elif (setup_options.sdk_location.startswith('gs://') or setup_options.sdk_location.startswith('http://') or setup_options.sdk_location.startswith('https://')): stage_tarball_from_remote_location = True else: stage_tarball_from_remote_location = False staged_path = utils.path.join(google_cloud_options.staging_location, names.DATAFLOW_SDK_TARBALL_FILE) if stage_tarball_from_remote_location: # If --sdk_location is not specified then the appropriate URL is built # based on the version of the currently running SDK. If the option is # present then no version matching is made and the exact URL or path # is expected. # # Unit tests running in the 'python setup.py test' context will # not have the sdk_location attribute present and therefore we # will not stage a tarball. if setup_options.sdk_location == 'default': sdk_remote_location = '%s/v%s.tar.gz' % ( PACKAGES_URL_PREFIX, __version__) else: sdk_remote_location = setup_options.sdk_location _stage_dataflow_sdk_tarball(sdk_remote_location, staged_path, temp_dir) resources.append(names.DATAFLOW_SDK_TARBALL_FILE) else: # Check if we have a local Dataflow SDK tarball present. This branch is # used by tests running with the SDK built at head. if setup_options.sdk_location == 'default': module_path = os.path.abspath(__file__) sdk_path = os.path.join( os.path.dirname(module_path), '..', names.DATAFLOW_SDK_TARBALL_FILE) elif os.path.isdir(setup_options.sdk_location): sdk_path = os.path.join( setup_options.sdk_location, names.DATAFLOW_SDK_TARBALL_FILE) else: sdk_path = setup_options.sdk_location if os.path.isfile(sdk_path): logging.info('Copying dataflow SDK "%s" to staging location.', sdk_path) file_copy(sdk_path, staged_path) resources.append(names.DATAFLOW_SDK_TARBALL_FILE) else: if setup_options.sdk_location == 'default': raise RuntimeError('Cannot find default Dataflow SDK tar file "%s"', sdk_path) else: raise RuntimeError( 'The file "%s" cannot be found. Its location was specified by ' 'the --sdk_location command-line option.' % sdk_path) # Delete all temp files created while staging job resources. shutil.rmtree(temp_dir) return resources def _build_setup_package(setup_file, temp_dir, build_setup_args=None): saved_current_directory = os.getcwd() try: os.chdir(os.path.dirname(setup_file)) if build_setup_args is None: build_setup_args = [ 'python', os.path.basename(setup_file), 'sdist', '--dist-dir', temp_dir] logging.info('Executing command: %s', build_setup_args) result = processes.call(build_setup_args) if result != 0: raise RuntimeError( 'Failed to execute command: %s. Exit code %d', build_setup_args, result) output_files = glob.glob(os.path.join(temp_dir, '*.tar.gz')) if not output_files: raise RuntimeError( 'File %s not found.' % os.path.join(temp_dir, '*.tar.gz')) return output_files[0] finally: os.chdir(saved_current_directory) def _stage_dataflow_sdk_tarball(sdk_remote_location, staged_path, temp_dir): """Stage a Dataflow SDK tarball with the appropriate version. Args: sdk_remote_location: A GCS path to a Dataflow SDK tarball or a URL from the file can be downloaded. staged_path: GCS path where the found SDK tarball should be copied. temp_dir: path to temporary location where the file should be downloaded. Raises: RuntimeError: If wget on the URL specified returs errors or the file cannot be copied from/to GCS. """ if (sdk_remote_location.startswith('http://') or sdk_remote_location.startswith('https://')): logging.info( 'Staging Dataflow SDK tarball from %s to %s', sdk_remote_location, staged_path) local_download_file = _dependency_file_download( sdk_remote_location, temp_dir) _dependency_file_copy(local_download_file, staged_path) elif sdk_remote_location.startswith('gs://'): # Stage the file to the GCS staging area. logging.info( 'Staging Dataflow SDK tarball from %s to %s', sdk_remote_location, staged_path) _dependency_file_copy(sdk_remote_location, staged_path) else: raise RuntimeError( 'The --sdk_location option was used with an unsupported ' 'type of location: %s' % sdk_remote_location) DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/dependency_test.py000066400000000000000000000365331272340251000261030ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the setup module.""" import logging import os import shutil import tempfile import unittest from google.cloud.dataflow import utils from google.cloud.dataflow.utils import dependency from google.cloud.dataflow.utils import names from google.cloud.dataflow.utils.options import GoogleCloudOptions from google.cloud.dataflow.utils.options import PipelineOptions from google.cloud.dataflow.utils.options import SetupOptions from google.cloud.dataflow.version import __version__ class SetupTest(unittest.TestCase): def update_options(self, options): setup_options = options.view_as(SetupOptions) setup_options.sdk_location = '' google_cloud_options = options.view_as(GoogleCloudOptions) if google_cloud_options.temp_location is None: google_cloud_options.temp_location = google_cloud_options.staging_location def create_temp_file(self, path, contents): with open(path, 'w') as f: f.write(contents) return f.name def populate_requirements_cache(self, requirements_file, cache_dir): _ = requirements_file self.create_temp_file(os.path.join(cache_dir, 'abc.txt'), 'nothing') self.create_temp_file(os.path.join(cache_dir, 'def.txt'), 'nothing') def test_no_staging_location(self): with self.assertRaises(RuntimeError) as cm: dependency.stage_job_resources(PipelineOptions()) self.assertEqual('The --staging_location option must be specified.', cm.exception.message) def test_no_temp_location(self): staging_dir = tempfile.mkdtemp() options = PipelineOptions() google_cloud_options = options.view_as(GoogleCloudOptions) google_cloud_options.staging_location = staging_dir self.update_options(options) google_cloud_options.temp_location = None with self.assertRaises(RuntimeError) as cm: dependency.stage_job_resources(options) self.assertEqual('The --temp_location option must be specified.', cm.exception.message) def test_no_main_session(self): staging_dir = tempfile.mkdtemp() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir options.view_as(SetupOptions).save_main_session = False self.update_options(options) self.assertEqual( [], dependency.stage_job_resources(options)) def test_default_resources(self): staging_dir = tempfile.mkdtemp() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) self.assertEqual( [names.PICKLED_MAIN_SESSION_FILE], dependency.stage_job_resources(options)) self.assertTrue( os.path.isfile( os.path.join(staging_dir, names.PICKLED_MAIN_SESSION_FILE))) def test_with_requirements_file(self): staging_dir = tempfile.mkdtemp() source_dir = tempfile.mkdtemp() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, dependency.REQUIREMENTS_FILE) self.create_temp_file( os.path.join(source_dir, dependency.REQUIREMENTS_FILE), 'nothing') self.assertEqual( sorted([dependency.REQUIREMENTS_FILE, names.PICKLED_MAIN_SESSION_FILE, 'abc.txt', 'def.txt']), sorted(dependency.stage_job_resources( options, populate_requirements_cache=self.populate_requirements_cache))) self.assertTrue( os.path.isfile( os.path.join(staging_dir, dependency.REQUIREMENTS_FILE))) self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt'))) self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt'))) def test_requirements_file_not_present(self): staging_dir = tempfile.mkdtemp() with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).requirements_file = 'nosuchfile' dependency.stage_job_resources( options, populate_requirements_cache=self.populate_requirements_cache) self.assertEqual( cm.exception.message, 'The file %s cannot be found. It was specified in the ' '--requirements_file command line option.' % 'nosuchfile') def test_with_requirements_file_and_cache(self): staging_dir = tempfile.mkdtemp() source_dir = tempfile.mkdtemp() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, dependency.REQUIREMENTS_FILE) options.view_as(SetupOptions).requirements_cache = os.path.join( tempfile.gettempdir(), 'alternative-cache-dir') self.create_temp_file( os.path.join(source_dir, dependency.REQUIREMENTS_FILE), 'nothing') self.assertEqual( sorted([dependency.REQUIREMENTS_FILE, names.PICKLED_MAIN_SESSION_FILE, 'abc.txt', 'def.txt']), sorted(dependency.stage_job_resources( options, populate_requirements_cache=self.populate_requirements_cache))) self.assertTrue( os.path.isfile( os.path.join(staging_dir, dependency.REQUIREMENTS_FILE))) self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt'))) self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt'))) def test_with_setup_file(self): staging_dir = tempfile.mkdtemp() source_dir = tempfile.mkdtemp() self.create_temp_file( os.path.join(source_dir, 'setup.py'), 'notused') options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).setup_file = os.path.join( source_dir, 'setup.py') self.assertEqual( [dependency.WORKFLOW_TARBALL_FILE, names.PICKLED_MAIN_SESSION_FILE], dependency.stage_job_resources( options, # We replace the build setup command because a realistic one would # require the setuptools package to be installed. Note that we can't # use "touch" here to create the expected output tarball file, since # touch is not available on Windows, so we invoke python to produce # equivalent behavior. build_setup_args=[ 'python', '-c', 'open(__import__("sys").argv[1], "a")', os.path.join(source_dir, dependency.WORKFLOW_TARBALL_FILE)], temp_dir=source_dir)) self.assertTrue( os.path.isfile( os.path.join(staging_dir, dependency.WORKFLOW_TARBALL_FILE))) def test_setup_file_not_present(self): staging_dir = tempfile.mkdtemp() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).setup_file = 'nosuchfile' with self.assertRaises(RuntimeError) as cm: dependency.stage_job_resources(options) self.assertEqual( cm.exception.message, 'The file %s cannot be found. It was specified in the ' '--setup_file command line option.' % 'nosuchfile') def test_setup_file_not_named_setup_dot_py(self): staging_dir = tempfile.mkdtemp() source_dir = tempfile.mkdtemp() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).setup_file = ( os.path.join(source_dir, 'xyz-setup.py')) self.create_temp_file( os.path.join(source_dir, 'xyz-setup.py'), 'notused') with self.assertRaises(RuntimeError) as cm: dependency.stage_job_resources(options) self.assertTrue( cm.exception.message.startswith( 'The --setup_file option expects the full path to a file named ' 'setup.py instead of ')) def override_file_copy(self, expected_from_path, expected_to_dir): def file_copy(from_path, to_path): if not from_path.endswith(names.PICKLED_MAIN_SESSION_FILE): self.assertEqual(expected_from_path, from_path) self.assertEqual(utils.path.join(expected_to_dir, names.DATAFLOW_SDK_TARBALL_FILE), to_path) if from_path.startswith('gs://') or to_path.startswith('gs://'): logging.info('Faking file_copy(%s, %s)', from_path, to_path) else: shutil.copyfile(from_path, to_path) dependency._dependency_file_copy = file_copy def override_file_download(self, expected_from_url, expected_to_folder): def file_download(from_url, _): self.assertEqual(expected_from_url, from_url) tarball_path = os.path.join(expected_to_folder, 'sdk-tarball') with open(tarball_path, 'w') as f: f.write('Some contents.') return tarball_path dependency._dependency_file_download = file_download return os.path.join(expected_to_folder, 'sdk-tarball') def test_sdk_location_default(self): staging_dir = tempfile.mkdtemp() expected_from_url = '%s/v%s.tar.gz' % ( dependency.PACKAGES_URL_PREFIX, __version__) expected_from_path = self.override_file_download( expected_from_url, staging_dir) self.override_file_copy(expected_from_path, staging_dir) options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).sdk_location = 'default' self.assertEqual( [names.PICKLED_MAIN_SESSION_FILE, names.DATAFLOW_SDK_TARBALL_FILE], dependency.stage_job_resources( options, file_copy=dependency._dependency_file_copy)) def test_sdk_location_local(self): staging_dir = tempfile.mkdtemp() sdk_location = tempfile.mkdtemp() self.create_temp_file( os.path.join( sdk_location, names.DATAFLOW_SDK_TARBALL_FILE), 'contents') options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location self.assertEqual( [names.PICKLED_MAIN_SESSION_FILE, names.DATAFLOW_SDK_TARBALL_FILE], dependency.stage_job_resources(options)) tarball_path = os.path.join( staging_dir, names.DATAFLOW_SDK_TARBALL_FILE) with open(tarball_path) as f: self.assertEqual(f.read(), 'contents') def test_sdk_location_local_not_present(self): staging_dir = tempfile.mkdtemp() sdk_location = 'nosuchdir' with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location dependency.stage_job_resources(options) self.assertEqual( 'The file "%s" cannot be found. Its ' 'location was specified by the --sdk_location command-line option.' % sdk_location, cm.exception.message) def test_sdk_location_gcs(self): staging_dir = tempfile.mkdtemp() sdk_location = 'gs://my-gcs-bucket/tarball.tar.gz' self.override_file_copy(sdk_location, staging_dir) options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location self.assertEqual( [names.PICKLED_MAIN_SESSION_FILE, names.DATAFLOW_SDK_TARBALL_FILE], dependency.stage_job_resources(options)) def test_with_extra_packages(self): staging_dir = tempfile.mkdtemp() source_dir = tempfile.mkdtemp() self.create_temp_file( os.path.join(source_dir, 'abc.tar.gz'), 'nothing') self.create_temp_file( os.path.join(source_dir, 'xyz.tar.gz'), 'nothing') self.create_temp_file( os.path.join(source_dir, dependency.EXTRA_PACKAGES_FILE), 'nothing') options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).extra_packages = [ os.path.join(source_dir, 'abc.tar.gz'), os.path.join(source_dir, 'xyz.tar.gz'), 'gs://my-gcs-bucket/gcs.tar.gz'] gcs_copied_files = [] def file_copy(from_path, to_path): if from_path.startswith('gs://'): gcs_copied_files.append(from_path) _, from_name = os.path.split(from_path) self.create_temp_file(os.path.join(to_path, from_name), 'nothing') logging.info('Fake copied GCS file: %s to %s', from_path, to_path) elif to_path.startswith('gs://'): logging.info('Faking file_copy(%s, %s)', from_path, to_path) else: shutil.copyfile(from_path, to_path) dependency._dependency_file_copy = file_copy self.assertEqual( ['abc.tar.gz', 'xyz.tar.gz', 'gcs.tar.gz', dependency.EXTRA_PACKAGES_FILE, names.PICKLED_MAIN_SESSION_FILE], dependency.stage_job_resources(options)) with open(os.path.join(staging_dir, dependency.EXTRA_PACKAGES_FILE)) as f: self.assertEqual(['abc.tar.gz\n', 'xyz.tar.gz\n', 'gcs.tar.gz\n'], f.readlines()) self.assertEqual(['gs://my-gcs-bucket/gcs.tar.gz'], gcs_copied_files) def test_with_extra_packages_missing_files(self): staging_dir = tempfile.mkdtemp() with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).extra_packages = ['nosuchfile.tar.gz'] dependency.stage_job_resources(options) self.assertEqual( cm.exception.message, 'The file %s cannot be found. It was specified in the ' '--extra_packages command line option.' % 'nosuchfile.tar.gz') def test_with_extra_packages_invalid_file_name(self): staging_dir = tempfile.mkdtemp() source_dir = tempfile.mkdtemp() self.create_temp_file( os.path.join(source_dir, 'abc.tgz'), 'nothing') with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).extra_packages = [ os.path.join(source_dir, 'abc.tgz')] dependency.stage_job_resources(options) self.assertEqual( cm.exception.message, 'The --extra_packages option expects a full path ending with ' '\'.tar.gz\' instead of %s' % os.path.join(source_dir, 'abc.tgz')) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/names.py000066400000000000000000000047031272340251000240230ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Various names for properties, transforms, etc.""" # Standard file names used for staging files. PICKLED_MAIN_SESSION_FILE = 'pickled_main_session' DATAFLOW_SDK_TARBALL_FILE = 'dataflow_python_sdk.tar' # String constants related to sources framework SOURCE_FORMAT = 'custom_source' SOURCE_TYPE = 'CustomSourcesType' SERIALIZED_SOURCE_KEY = 'serialized_source' class TransformNames(object): """Transform strings as they are expected in the CloudWorkflow protos.""" COLLECTION_TO_SINGLETON = 'CollectionToSingleton' COMBINE = 'CombineValues' CREATE_PCOLLECTION = 'CreateCollection' DO = 'ParallelDo' FLATTEN = 'Flatten' GROUP = 'GroupByKey' READ = 'ParallelRead' WRITE = 'ParallelWrite' class PropertyNames(object): """Property strings as they are expected in the CloudWorkflow protos.""" BIGQUERY_CREATE_DISPOSITION = 'create_disposition' BIGQUERY_DATASET = 'dataset' BIGQUERY_QUERY = 'bigquery_query' BIGQUERY_TABLE = 'table' BIGQUERY_PROJECT = 'project' BIGQUERY_SCHEMA = 'schema' BIGQUERY_WRITE_DISPOSITION = 'write_disposition' ELEMENT = 'element' ELEMENTS = 'elements' ENCODING = 'encoding' FILE_PATTERN = 'filepattern' FILE_NAME_PREFIX = 'filename_prefix' FILE_NAME_SUFFIX = 'filename_suffix' FORMAT = 'format' INPUTS = 'inputs' NON_PARALLEL_INPUTS = 'non_parallel_inputs' NUM_SHARDS = 'num_shards' OUT = 'out' OUTPUT = 'output' OUTPUT_INFO = 'output_info' OUTPUT_NAME = 'output_name' PARALLEL_INPUT = 'parallel_input' PUBSUB_TOPIC = 'pubsub_topic' PUBSUB_SUBSCRIPTION = 'pubsub_subscription' PUBSUB_ID_LABEL = 'pubsub_id_label' SERIALIZED_FN = 'serialized_fn' SHARD_NAME_TEMPLATE = 'shard_template' SOURCE_STEP_INPUT = 'custom_source_step_input' STEP_NAME = 'step_name' USER_FN = 'user_fn' USER_NAME = 'user_name' VALIDATE_SINK = 'validate_sink' VALIDATE_SOURCE = 'validate_source' VALUE = 'value' DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/options.py000066400000000000000000000430411272340251000244110ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Pipeline options obtained from command line parsing. TODO(silviuc): Should rename this module to pipeline_options. """ import argparse class PipelineOptions(object): """Pipeline options class used as container for command line options. The class is essentially a wrapper over the standard argparse Python module (see https://docs.python.org/3/library/argparse.html). To define one option or a group of options you subclass from PipelineOptions:: class XyzOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--abc', default='start') parser.add_argument('--xyz', default='end') The arguments for the add_argument() method are exactly the ones described in the argparse public documentation. Pipeline objects require an options object during initialization. This is obtained simply by initializing an options class as defined above:: p = Pipeline(options=XyzOptions()) if p.options.xyz == 'end': raise ValueError('Option xyz has an invalid value.') By default the options classes will use command line arguments to initialize the options. """ def __init__(self, flags=None, **kwargs): """Initialize an options class. The initializer will traverse all subclasses, add all their argparse arguments and then parse the command line specified by flags or by default the one obtained from sys.argv. The subclasses are not expected to require a redefinition of __init__. Args: flags: An iterable of command line arguments to be used. If not specified then sys.argv will be used as input for parsing arguments. **kwargs: Add overrides for arguments passed in flags. """ self._flags = flags self._all_options = kwargs parser = argparse.ArgumentParser() for cls in type(self).mro(): if cls == PipelineOptions: break elif '_add_argparse_args' in cls.__dict__: cls._add_argparse_args(parser) # The _visible_options attribute will contain only those options from the # flags (i.e., command line) that can be recognized. The _all_options # field contains additional overrides. self._visible_options, _ = parser.parse_known_args(flags) @classmethod def _add_argparse_args(cls, parser): # Override this in subclasses to provide options. pass @classmethod def from_dictionary(cls, options): """Returns a PipelineOptions from a dictionary of arguments. Args: options: Dictinary of argument value pairs. Returns: A PipelineOptions object representing the given arguments. """ flags = [] for k, v in options.iteritems(): if isinstance(v, bool): if v: flags.append('--%s' % k) else: flags.append('--%s=%s' % (k, v)) return cls(flags) def get_all_options(self): """Returns a dictionary of all defined arguments. Returns a dictionary of all defined arguments (arguments that are defined in any subclass of PipelineOptions) into a dictionary. Returns: Dictionary of all args and values. """ parser = argparse.ArgumentParser() for cls in PipelineOptions.__subclasses__(): cls._add_argparse_args(parser) # pylint: disable=protected-access known_args, _ = parser.parse_known_args(self._flags) result = vars(known_args) # Apply the overrides if any for k in result: if k in self._all_options: result[k] = self._all_options[k] return result def view_as(self, cls): view = cls(self._flags) view._all_options = self._all_options return view def _visible_option_list(self): return sorted(option for option in dir(self._visible_options) if option[0] != '_') def __dir__(self): return sorted(dir(type(self)) + self.__dict__.keys() + self._visible_option_list()) def __getattr__(self, name): # Special methods which may be accessed before the object is # fully constructed (e.g. in unpickling). if name[:2] == name[-2:] == '__': return object.__getattr__(self, name) elif name in self._visible_option_list(): return self._all_options.get(name, getattr(self._visible_options, name)) else: raise AttributeError("'%s' object has no attribute '%s'" % (type(self).__name__, name)) def __setattr__(self, name, value): if name in ('_flags', '_all_options', '_visible_options'): super(PipelineOptions, self).__setattr__(name, value) elif name in self._visible_option_list(): self._all_options[name] = value else: raise AttributeError("'%s' object has no attribute '%s'" % (type(self).__name__, name)) def __str__(self): return '%s(%s)' % (type(self).__name__, ', '.join('%s=%s' % (option, getattr(self, option)) for option in self._visible_option_list())) class StandardOptions(PipelineOptions): DEFAULT_RUNNER = 'DirectPipelineRunner' @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--runner', help=('Pipeline runner used to execute the workflow. Valid values are ' 'DirectPipelineRunner, DataflowPipelineRunner, ' 'and BlockingDataflowPipelineRunner.')) # Whether to enable streaming mode. parser.add_argument('--streaming', default=False, action='store_true', help='Whether to enable streaming mode.') class TypeOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): # TODO(laolu): Add a type inferencing option here once implemented. parser.add_argument('--type_check_strictness', default='DEFAULT_TO_ANY', choices=['ALL_REQUIRED', 'DEFAULT_TO_ANY'], help='The level of exhaustive manual type-hint ' 'annotation required') parser.add_argument('--no_pipeline_type_check', dest='pipeline_type_check', action='store_false', help='Disable type checking at pipeline construction ' 'time') parser.add_argument('--pipeline_type_check', action='store_true', help='Enable type checking at pipeline construction ' 'time') parser.add_argument('--runtime_type_check', default=False, action='store_true', help='Enable type checking at pipeline execution ' 'time. NOTE: only supported with the ' 'DirectPipelineRunner') class GoogleCloudOptions(PipelineOptions): """Google Cloud Dataflow service execution options.""" @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--dataflow_endpoint', default='https://dataflow.googleapis.com', help= ('The URL for the Dataflow API. If not set, the default public URL ' 'will be used.')) # Remote execution must check that this option is not None. parser.add_argument('--project', default=None, help='Name of the Cloud project owning the Dataflow ' 'job.') # Remote execution must check that this option is not None. parser.add_argument('--job_name', default=None, help='Name of the Cloud Dataflow job.') # Remote execution must check that this option is not None. parser.add_argument('--staging_location', default=None, help='GCS path for staging code packages needed by ' 'workers.') # Remote execution must check that this option is not None. parser.add_argument('--temp_location', default=None, help='GCS path for saving temporary workflow jobs.') # Options for using service account credentials. parser.add_argument('--service_account_name', default=None, help='Name of the service account for Google APIs.') parser.add_argument('--service_account_key_file', default=None, help='Path to a file containing the P12 service ' 'credentials.') parser.add_argument('--no_auth', dest='no_auth', type=bool, default=False) def validate(self, validator): errors = [] if validator.is_service_runner(): errors.extend(validator.validate_cloud_options(self)) errors.extend(validator.validate_gcs_path(self, 'staging_location')) errors.extend(validator.validate_gcs_path(self, 'temp_location')) return errors # Command line options controlling the worker pool configuration. # TODO(silviuc): Update description when autoscaling options are in. class WorkerOptions(PipelineOptions): """Worker pool configuration options.""" @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--num_workers', type=int, default=None, help= ('Number of workers to use when executing the Dataflow job. If not ' 'set, the Dataflow service will use a reasonable default.')) # TODO(silviuc): Remove --machine_type variant of the flag. parser.add_argument( '--worker_machine_type', '--machine_type', dest='machine_type', default=None, help=('Machine type to create Dataflow worker VMs as. See ' 'https://cloud.google.com/compute/docs/machine-types ' 'for a list of valid options. If not set, ' 'the Dataflow service will choose a reasonable ' 'default.')) parser.add_argument( '--disk_size_gb', type=int, default=None, help= ('Remote worker disk size, in gigabytes, or 0 to use the default size. ' 'If not set, the Dataflow service will use a reasonable default.')) # TODO(silviuc): Remove --disk_type variant of the flag. parser.add_argument( '--worker_disk_type', '--disk_type', dest='disk_type', default=None, help=('Specifies what type of persistent disk should be used.')) parser.add_argument( '--disk_source_image', default=None, help= ('Disk source image to use by VMs for jobs. See ' 'https://developers.google.com/compute/docs/images for further ' 'details. If not set, the Dataflow service will use a reasonable ' 'default.')) parser.add_argument( '--zone', default=None, help=( 'GCE availability zone for launching workers. Default is up to the ' 'Dataflow service.')) parser.add_argument( '--network', default=None, help=( 'GCE network for launching workers. Default is up to the Dataflow ' 'service.')) parser.add_argument( '--worker_harness_container_image', default=None, help=('Docker registry location of container image to use for the ' 'worker harness. Default is the container for the version of the ' 'SDK. Note: currently, only approved Google Cloud Dataflow ' 'container images may be used here.')) parser.add_argument( '--teardown_policy', choices=['TEARDOWN_ALWAYS', 'TEARDOWN_NEVER', 'TEARDOWN_ON_SUCCESS'], default=None, help= ('The teardown policy for the VMs. By default this is left unset and ' 'the service sets the default policy.')) def validate(self, validator): errors = [] if validator.is_service_runner(): errors.extend( validator.validate_optional_argument_positive(self, 'num_workers')) return errors class DebugOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--dataflow_job_file', default=None, help='Debug file to write the workflow specification.') class ProfilingOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--profile', action='store_true', help='Enable work item profiling') parser.add_argument('--profile_location', default=None, help='GCS path for saving profiler data.') class SetupOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): # Options for installing dependencies in the worker. parser.add_argument( '--requirements_file', default=None, help= ('Path to a requirements file containing package dependencies. ' 'Typically it is produced by a pip freeze command. More details: ' 'https://pip.pypa.io/en/latest/reference/pip_freeze.html. ' 'If used, all the packages specified will be downloaded, ' 'cached (use --requirements_cache to change default location), ' 'and then staged so that they can be automatically installed in ' 'workers during startup. The cache is refreshed as needed ' 'avoiding extra downloads for existing packages. Typically the ' 'file is named requirements.txt.')) parser.add_argument( '--requirements_cache', default=None, help= ('Path to a folder to cache the packages specified in ' 'the requirements file using the --requirements_file option.')) parser.add_argument( '--setup_file', default=None, help= ('Path to a setup Python file containing package dependencies. If ' 'specified, the file\'s containing folder is assumed to have the ' 'structure required for a setuptools setup package. The file must be ' 'named setup.py. More details: ' 'https://pythonhosted.org/setuptools/setuptools.html During job ' 'submission a source distribution will be built and the worker will ' 'install the resulting package before running any custom code.')) parser.add_argument( '--save_main_session', default=True, action='store_true', help= ('Save the main session state so that pickled functions and classes ' 'defined in __main__ (e.g. interactive session) can be unpickled. ' 'Some workflows do not need the session state if for instance all ' 'their functions/classes are defined in proper modules (not __main__)' ' and the modules are importable in the worker. ')) parser.add_argument('--no_save_main_session', dest='save_main_session', action='store_false') parser.add_argument( '--sdk_location', default='default', help= ('Override the default GitHub location from where Dataflow SDK is ' 'downloaded. It can be an URL, a GCS path, or a local path to an ' 'SDK tarball. Workflow submissions will download or copy an SDK ' 'tarball from here. If the string "default", ' 'a standard SDK location is used. If empty, no SDK is copied.')) parser.add_argument( '--extra_package', dest='extra_packages', action='append', default=None, help= ('Local path to a Python package file. The file is expected to be a ' 'compressed tarball with the suffix \'.tar.gz\' which can be ' 'installed using the easy_install command of the standard setuptools ' 'package. Multiple --extra_package options can be specified if more ' 'than one package is needed. During job submission the files will be ' 'staged in the staging area (--staging_location option) and the ' 'workers will install them in same order they were specified on the ' 'command line.')) # TODO(silviuc): Add autoscaling related options: # --autoscaling_algorithm, --max_num_workers. # TODO(silviuc): Add --files_to_stage option. # This could potentially replace the --requirements_file and --setup_file. # TODO(silviuc): Non-standard options. Keep them? If yes, add help too! # Remote execution must check that this option is not None. class OptionsContext(object): """Set default pipeline options for pipelines created in this block. This is particularly useful for pipelines implicitly created with the [python list] | PTransform construct. Can also be used as a decorator. """ overrides = [] def __init__(self, **options): self.options = options def __enter__(self): self.overrides.append(self.options) def __exit__(self, *exn_info): self.overrides.pop() def __call__(self, f, *args, **kwargs): def wrapper(*args, **kwargs): with self: f(*args, **kwargs) return wrapper @classmethod def augment_options(cls, options): for override in cls.overrides: for name, value in override.items(): setattr(options, name, value) return options DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/path.py000066400000000000000000000033411272340251000236510ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Utilities for dealing with file paths.""" import os def join(path, *paths): """Joins given path pieces with the appropriate separator. This function is useful for joining parts of a path that could at times refer to either a GCS path or a local path. In particular, this is useful for ensuring Windows compatibility as on Windows, the GCS path separator is different from the separator for local paths. Use os.path.join instead if a path always refers to a local path. Args: path: First part of path to join. If this part starts with 'gs:/', the GCS separator will be used in joining this path. *paths: Remaining part(s) of path to join. Returns: Pieces joined by the appropriate path separator. """ if path.startswith('gs:/'): # Note that we explicitly choose not to use posixpath.join() here, since # that function has the undesirable behavior of having, for example, # posixpath.join('gs://bucket/path', '/to/file') return '/to/file' instead # of the slightly less surprising result 'gs://bucket/path//to/file'. return '/'.join((path,) + paths) else: return os.path.join(path, *paths) DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/path_test.py000066400000000000000000000044531272340251000247150ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the path module.""" import unittest import mock from google.cloud.dataflow.utils import path def _gen_fake_join(separator): """Returns a callable that joins paths with the given separator.""" def _join(first_path, *paths): return separator.join((first_path,) + paths) return _join class Path(unittest.TestCase): def setUp(self): pass @mock.patch('google.cloud.dataflow.utils.path.os') def test_gcs_path(self, *unused_mocks): # Test joining of GCS paths when os.path.join uses Windows-style separator. path.os.path.join.side_effect = _gen_fake_join('\\') self.assertEqual('gs://bucket/path/to/file', path.join('gs://bucket/path', 'to', 'file')) self.assertEqual('gs://bucket/path/to/file', path.join('gs://bucket/path', 'to/file')) self.assertEqual('gs://bucket/path//to/file', path.join('gs://bucket/path', '/to/file')) @mock.patch('google.cloud.dataflow.utils.path.os') def test_unix_path(self, *unused_mocks): # Test joining of Unix paths. path.os.path.join.side_effect = _gen_fake_join('/') self.assertEqual('/tmp/path/to/file', path.join('/tmp/path', 'to', 'file')) self.assertEqual('/tmp/path/to/file', path.join('/tmp/path', 'to/file')) @mock.patch('google.cloud.dataflow.utils.path.os') def test_windows_path(self, *unused_mocks): # Test joining of Windows paths. path.os.path.join.side_effect = _gen_fake_join('\\') self.assertEqual(r'C:\tmp\path\to\file', path.join(r'C:\tmp\path', 'to', 'file')) self.assertEqual(r'C:\tmp\path\to\file', path.join(r'C:\tmp\path', r'to\file')) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/pipeline_options_test.py000066400000000000000000000104131272340251000273320ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the pipeline options module.""" import logging import unittest from google.cloud.dataflow.utils.options import PipelineOptions class PipelineOptionsTest(unittest.TestCase): TEST_CASES = [ {'flags': ['--num_workers', '5'], 'expected': {'num_workers': 5, 'mock_flag': False, 'mock_option': None}}, { 'flags': [ '--profile', '--profile_location', 'gs://bucket/', 'ignored'], 'expected': { 'profile': True, 'profile_location': 'gs://bucket/', 'mock_flag': False, 'mock_option': None} }, {'flags': ['--num_workers', '5', '--mock_flag'], 'expected': {'num_workers': 5, 'mock_flag': True, 'mock_option': None}}, {'flags': ['--mock_option', 'abc'], 'expected': {'mock_flag': False, 'mock_option': 'abc'}}, {'flags': ['--mock_option', ' abc def '], 'expected': {'mock_flag': False, 'mock_option': ' abc def '}}, {'flags': ['--mock_option= abc xyz '], 'expected': {'mock_flag': False, 'mock_option': ' abc xyz '}}, {'flags': ['--mock_option=gs://my bucket/my folder/my file'], 'expected': {'mock_flag': False, 'mock_option': 'gs://my bucket/my folder/my file'}}, ] # Used for testing newly added flags. class MockOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--mock_flag', action='store_true', help='mock flag') parser.add_argument('--mock_option', help='mock option') parser.add_argument('--option with space', help='mock option with space') def test_get_all_options(self): for case in PipelineOptionsTest.TEST_CASES: options = PipelineOptions(flags=case['flags']) self.assertDictContainsSubset(case['expected'], options.get_all_options()) self.assertEqual(options.view_as( PipelineOptionsTest.MockOptions).mock_flag, case['expected']['mock_flag']) self.assertEqual(options.view_as( PipelineOptionsTest.MockOptions).mock_option, case['expected']['mock_option']) def test_from_dictionary(self): for case in PipelineOptionsTest.TEST_CASES: options = PipelineOptions(flags=case['flags']) all_options_dict = options.get_all_options() options_from_dict = PipelineOptions.from_dictionary(all_options_dict) self.assertEqual(options_from_dict.view_as( PipelineOptionsTest.MockOptions).mock_flag, case['expected']['mock_flag']) self.assertEqual(options.view_as( PipelineOptionsTest.MockOptions).mock_option, case['expected']['mock_option']) def test_option_with_spcae(self): options = PipelineOptions(flags=['--option with space= value with space']) self.assertEqual( getattr(options.view_as(PipelineOptionsTest.MockOptions), 'option with space'), ' value with space') options_from_dict = PipelineOptions.from_dictionary( options.get_all_options()) self.assertEqual( getattr(options_from_dict.view_as(PipelineOptionsTest.MockOptions), 'option with space'), ' value with space') def test_override_options(self): base_flags = ['--num_workers', '5'] options = PipelineOptions(base_flags) self.assertEqual(options.get_all_options()['num_workers'], 5) self.assertEqual(options.get_all_options()['mock_flag'], False) options.view_as(PipelineOptionsTest.MockOptions).mock_flag = True self.assertEqual(options.get_all_options()['num_workers'], 5) self.assertEqual(options.get_all_options()['mock_flag'], True) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/pipeline_options_validator.py000066400000000000000000000147311272340251000303470ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Pipeline options validator. """ import re from google.cloud.dataflow.utils.options import DebugOptions from google.cloud.dataflow.utils.options import GoogleCloudOptions from google.cloud.dataflow.utils.options import SetupOptions from google.cloud.dataflow.utils.options import StandardOptions from google.cloud.dataflow.utils.options import TypeOptions from google.cloud.dataflow.utils.options import WorkerOptions class PipelineOptionsValidator(object): """Validates PipelineOptions. Goes through a list of known PipelineOption subclassess and calls:: validate(validator) if one is implemented. Aggregates a list of validation errors from all and returns an aggregated list. """ # Validator will call validate on these subclasses of PipelineOptions OPTIONS = [DebugOptions, GoogleCloudOptions, SetupOptions, StandardOptions, TypeOptions, WorkerOptions] # Possible validation errors. ERR_MISSING_OPTION = 'Missing required option: %s.' ERR_MISSING_GCS_PATH = 'Missing GCS path option: %s.' ERR_INVALID_GCS_PATH = 'Invalid GCS path (%s), given for the option: %s.' ERR_INVALID_GCS_BUCKET = ( 'Invalid GCS bucket (%s), given for the option: %s. See ' 'https://developers.google.com/storage/docs/bucketnaming ' 'for more details.') ERR_INVALID_GCS_OBJECT = 'Invalid GCS object (%s), given for the option: %s.' ERR_INVALID_JOB_NAME = ( 'Invalid job_name (%s); the name must consist of only the characters ' '[-a-z0-9], starting with a letter and ending with a letter or number') ERR_INVALID_PROJECT_NUMBER = ( 'Invalid Project ID (%s). Please make sure you specified the Project ID, ' 'not project number.') ERR_INVALID_PROJECT_ID = ( 'Invalid Project ID (%s). Please make sure you specified the Project ID, ' 'not project description.') ERR_INVALID_NOT_POSITIVE = ('Invalid value (%s) for option: %s. Value needs ' 'to be positive.') # GCS path specific patterns. GCS_URI = '(?P[^:]+)://(?P[^/]+)(/(?P.*))?' GCS_BUCKET = '^[a-z0-9][-_a-z0-9.]+[a-z0-9]$' GCS_SCHEME = 'gs' # GoogleCloudOptions specific patterns. JOB_PATTERN = '[a-z]([-a-z0-9]*[a-z0-9])?' PROJECT_ID_PATTERN = '[a-z][-a-z0-9:.]+[a-z0-9]' PROJECT_NUMBER_PATTERN = '[0-9]*' ENDPOINT_PATTERN = r'https://[\S]*googleapis\.com' def __init__(self, options, runner): self.options = options self.runner = runner def validate(self): """Calls validate on subclassess and returns a list of errors. validate will call validate method on subclasses, accumulate the returned list of errors, and returns the aggregate list. Returns: Aggregate list of errors after all calling all possible validate methods. """ errors = [] for cls in self.OPTIONS: if 'validate' in cls.__dict__: errors.extend(self.options.view_as(cls).validate(self)) return errors def is_service_runner(self): """True if pipeline will execute on the Google Cloud Dataflow service.""" is_service_runner = (self.runner is not None and type(self.runner).__name__ in [ 'DataflowPipelineRunner', 'BlockingDataflowPipelineRunner']) dataflow_endpoint = ( self.options.view_as(GoogleCloudOptions).dataflow_endpoint) is_service_endpoint = (dataflow_endpoint is not None and self.is_full_string_match( self.ENDPOINT_PATTERN, dataflow_endpoint)) return is_service_runner and is_service_endpoint def is_full_string_match(self, pattern, string): """Returns True if the pattern matches the whole string.""" pattern = '^%s$' % pattern return re.search(pattern, string) is not None def _validate_error(self, err, *args): return [err % args] def validate_gcs_path(self, view, arg_name): """Validates a GCS path against gs://bucket/object URI format.""" arg = getattr(view, arg_name, None) if arg is None: return self._validate_error(self.ERR_MISSING_GCS_PATH, arg_name) match = re.match(self.GCS_URI, arg, re.DOTALL) if match is None: return self._validate_error(self.ERR_INVALID_GCS_PATH, arg, arg_name) scheme = match.group('SCHEME') bucket = match.group('BUCKET') gcs_object = match.group('OBJECT') if ((scheme is None) or (scheme.lower() != self.GCS_SCHEME) or (bucket is None)): return self._validate_error(self.ERR_INVALID_GCS_PATH, arg, arg_name) if not self.is_full_string_match(self.GCS_BUCKET, bucket): return self._validate_error(self.ERR_INVALID_GCS_BUCKET, arg, arg_name) if gcs_object is None or '\n' in gcs_object or '\r' in gcs_object: return self._validate_error(self.ERR_INVALID_GCS_OBJECT, arg, arg_name) return [] def validate_cloud_options(self, view): """Validates job_name and project arguments.""" errors = [] job_name = view.job_name if (job_name is None or not self.is_full_string_match(self.JOB_PATTERN, job_name)): errors.extend(self._validate_error(self.ERR_INVALID_JOB_NAME, job_name)) project = view.project if project is None: errors.extend(self._validate_error(self.ERR_MISSING_OPTION, 'project')) else: if self.is_full_string_match(self.PROJECT_NUMBER_PATTERN, project): errors.extend( self._validate_error(self.ERR_INVALID_PROJECT_NUMBER, project)) elif not self.is_full_string_match(self.PROJECT_ID_PATTERN, project): errors.extend( self._validate_error(self.ERR_INVALID_PROJECT_ID, project)) return errors def validate_optional_argument_positive(self, view, arg_name): """Validates that an optional argument (if set) has a positive value.""" arg = getattr(view, arg_name, None) if arg is not None and int(arg) <= 0: return self._validate_error(self.ERR_INVALID_NOT_POSITIVE, arg, arg_name) return [] DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/pipeline_options_validator_test.py000066400000000000000000000164771272340251000314170ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the pipeline options validator module.""" import logging import unittest from google.cloud.dataflow.utils.options import PipelineOptions from google.cloud.dataflow.utils.pipeline_options_validator import PipelineOptionsValidator # Mock runners to use for validations. class MockRunners(object): class DataflowPipelineRunner(object): pass class OtherRunner(object): pass class SetupTest(unittest.TestCase): def check_errors_for_arguments(self, errors, args): """Checks that there is exactly one error for each given argument.""" missing = [] remaining = list(errors) for arg in args: found = False for error in remaining: if arg in error: remaining.remove(error) found = True break if not found: missing.append('Missing error for: ' + arg) # Return missing and remaining (not matched) errors. return missing + remaining def test_local_runner(self): runner = MockRunners.OtherRunner() options = PipelineOptions([]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) def test_missing_required_options(self): options = PipelineOptions(['']) runner = MockRunners.DataflowPipelineRunner() validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual( self.check_errors_for_arguments( errors, ['project', 'job_name', 'staging_location', 'temp_location']), []) def test_gcs_path(self): def get_validator(temp_location): options = ['--project=example:example', '--job_name=job', '--staging_location=gs://foo/bar'] if temp_location is not None: options.append('--temp_location=' + temp_location) pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowPipelineRunner() validator = PipelineOptionsValidator(pipeline_options, runner) return validator test_cases = [ {'temp_location': None, 'errors': ['temp_location']}, {'temp_location': 'gcs:/foo/bar', 'errors': ['temp_location']}, {'temp_location': 'gs:/foo/bar', 'errors': ['temp_location']}, {'temp_location': 'gs://ABC/bar', 'errors': ['temp_location']}, {'temp_location': 'gs://ABC/bar', 'errors': ['temp_location']}, {'temp_location': 'gs://foo', 'errors': ['temp_location']}, {'temp_location': 'gs://foo/', 'errors': []}, {'temp_location': 'gs://foo/bar', 'errors': []}, ] for case in test_cases: errors = get_validator(case['temp_location']).validate() self.assertEqual( self.check_errors_for_arguments(errors, case['errors']), []) def test_project(self): def get_validator(project): options = ['--job_name=job', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar'] if project is not None: options.append('--project=' + project) pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowPipelineRunner() validator = PipelineOptionsValidator(pipeline_options, runner) return validator test_cases = [ {'project': None, 'errors': ['project']}, {'project': '12345', 'errors': ['project']}, {'project': 'FOO', 'errors': ['project']}, {'project': 'foo:BAR', 'errors': ['project']}, {'project': 'fo', 'errors': ['project']}, {'project': 'foo', 'errors': []}, {'project': 'foo:bar', 'errors': []}, ] for case in test_cases: errors = get_validator(case['project']).validate() self.assertEqual( self.check_errors_for_arguments(errors, case['errors']), []) def test_job_name(self): def get_validator(job_name): options = ['--project=example:example', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar'] if job_name is not None: options.append('--job_name=' + job_name) pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowPipelineRunner() validator = PipelineOptionsValidator(pipeline_options, runner) return validator test_cases = [ {'job_name': None, 'errors': ['job_name']}, {'job_name': '12345', 'errors': ['job_name']}, {'job_name': 'FOO', 'errors': ['job_name']}, {'job_name': 'foo:bar', 'errors': ['job_name']}, {'job_name': 'fo', 'errors': []}, {'job_name': 'foo', 'errors': []}, ] for case in test_cases: errors = get_validator(case['job_name']).validate() self.assertEqual( self.check_errors_for_arguments(errors, case['errors']), []) def test_num_workers(self): def get_validator(num_workers): options = ['--project=example:example', '--job_name=job', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar'] if num_workers is not None: options.append('--num_workers=' + num_workers) pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowPipelineRunner() validator = PipelineOptionsValidator(pipeline_options, runner) return validator test_cases = [ {'num_workers': None, 'errors': []}, {'num_workers': '1', 'errors': []}, {'num_workers': '0', 'errors': ['num_workers']}, {'num_workers': '-1', 'errors': ['num_workers']}, ] for case in test_cases: errors = get_validator(case['num_workers']).validate() self.assertEqual( self.check_errors_for_arguments(errors, case['errors']), []) def test_is_service_runner(self): test_cases = [ { 'runner': MockRunners.OtherRunner(), 'options': [], 'expected': False, }, { 'runner': MockRunners.OtherRunner(), 'options': ['--dataflow_endpoint=https://dataflow.googleapis.com'], 'expected': False, }, { 'runner': MockRunners.DataflowPipelineRunner(), 'options': ['--dataflow_endpoint=https://another.service.com'], 'expected': False, }, { 'runner': MockRunners.DataflowPipelineRunner(), 'options': ['--dataflow_endpoint=https://dataflow.googleapis.com'], 'expected': True, }, { 'runner': MockRunners.DataflowPipelineRunner(), 'options': [], 'expected': True, }, ] for case in test_cases: validator = PipelineOptionsValidator( PipelineOptions(case['options']), case['runner']) self.assertEqual(validator.is_service_runner(), case['expected']) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/processes.py000066400000000000000000000027241272340251000247270ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Cross-platform utilities for creating subprocesses.""" import platform import subprocess # On Windows, we need to use shell=True when creating subprocesses for binary # paths to be resolved correctly. force_shell = platform.system() == 'Windows' # We mimic the interface of the standard Python subprocess module. PIPE = subprocess.PIPE STDOUT = subprocess.STDOUT def call(*args, **kwargs): if force_shell: kwargs['shell'] = True return subprocess.call(*args, **kwargs) def check_call(*args, **kwargs): if force_shell: kwargs['shell'] = True return subprocess.check_call(*args, **kwargs) def check_output(*args, **kwargs): if force_shell: kwargs['shell'] = True return subprocess.check_output(*args, **kwargs) def Popen(*args, **kwargs): # pylint: disable=invalid-name if force_shell: kwargs['shell'] = True return subprocess.Popen(*args, **kwargs) DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/processes_test.py000066400000000000000000000063061272340251000257660ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the processes module.""" import unittest import mock from google.cloud.dataflow.utils import processes class Exec(unittest.TestCase): def setUp(self): pass @mock.patch('google.cloud.dataflow.utils.processes.subprocess') def test_method_forwarding_not_windows(self, *unused_mocks): # Test that the correct calls are being forwarded to the subprocess module # when we are not on Windows. processes.force_shell = False processes.call(['subprocess', 'call'], shell=False, other_arg=True) processes.subprocess.call.assert_called_once_with( ['subprocess', 'call'], shell=False, other_arg=True) processes.check_call( ['subprocess', 'check_call'], shell=False, other_arg=True) processes.subprocess.check_call.assert_called_once_with( ['subprocess', 'check_call'], shell=False, other_arg=True) processes.check_output( ['subprocess', 'check_output'], shell=False, other_arg=True) processes.subprocess.check_output.assert_called_once_with( ['subprocess', 'check_output'], shell=False, other_arg=True) processes.Popen(['subprocess', 'Popen'], shell=False, other_arg=True) processes.subprocess.Popen.assert_called_once_with( ['subprocess', 'Popen'], shell=False, other_arg=True) @mock.patch('google.cloud.dataflow.utils.processes.subprocess') def test_method_forwarding_windows(self, *unused_mocks): # Test that the correct calls are being forwarded to the subprocess module # and that the shell=True flag is added when we are on Windows. processes.force_shell = True processes.call(['subprocess', 'call'], shell=False, other_arg=True) processes.subprocess.call.assert_called_once_with( ['subprocess', 'call'], shell=True, other_arg=True) processes.check_call( ['subprocess', 'check_call'], shell=False, other_arg=True) processes.subprocess.check_call.assert_called_once_with( ['subprocess', 'check_call'], shell=True, other_arg=True) processes.check_output( ['subprocess', 'check_output'], shell=False, other_arg=True) processes.subprocess.check_output.assert_called_once_with( ['subprocess', 'check_output'], shell=True, other_arg=True) processes.Popen(['subprocess', 'Popen'], shell=False, other_arg=True) processes.subprocess.Popen.assert_called_once_with( ['subprocess', 'Popen'], shell=True, other_arg=True) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/profiler.py000066400000000000000000000041631272340251000245420ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A profiler context manager based on cProfile.Profile objects.""" import cProfile import logging import os import pstats import StringIO import tempfile import time from google.cloud.dataflow.utils.dependency import _dependency_file_copy class Profile(object): """cProfile wrapper context for saving and logging profiler results.""" SORTBY = 'cumulative' def __init__(self, profile_id, profile_location=None, log_results=False): self.stats = None self.profile_id = str(profile_id) self.profile_location = profile_location self.log_results = log_results def __enter__(self): logging.info('Start profiling: %s', self.profile_id) self.profile = cProfile.Profile() self.profile.enable() return self def __exit__(self, *args): self.profile.disable() logging.info('Stop profiling: %s', self.profile_id) if self.profile_location: dump_location = os.path.join( self.profile_location, 'profile', ('%s-%s' % (time.strftime('%Y-%m-%d_%H_%M_%S'), self.profile_id))) fd, filename = tempfile.mkstemp() self.profile.dump_stats(filename) logging.info('Copying profiler data to: [%s]', dump_location) _dependency_file_copy(filename, dump_location) # pylint: disable=protected-access os.close(fd) os.remove(filename) if self.log_results: s = StringIO.StringIO() self.stats = pstats.Stats( self.profile, stream=s).sort_stats(Profile.SORTBY) self.stats.print_stats() logging.info('Profiler data: [%s]', s.getvalue()) DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/retry.py000066400000000000000000000154521272340251000240700ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Retry decorators for calls raising exceptions. This module is used mostly to decorate all integration points where the code makes calls to remote services. Searching through the code base for @retry should find all such places. For this reason even places where retry is not needed right now use a @retry.no_retries decorator. """ import logging import random import sys import time import traceback from apitools.base.py.exceptions import HttpError class PermanentException(Exception): """Base class for exceptions that should not be retried.""" pass class FuzzedExponentialIntervals(object): """Iterable for intervals that are exponentially spaced, with fuzzing. On iteration, yields retry interval lengths, in seconds. Every iteration over this iterable will yield differently fuzzed interval lengths, as long as fuzz is nonzero. Args: initial_delay_secs: The delay before the first retry, in seconds. num_retries: The total number of times to retry. factor: The exponential factor to use on subsequent retries. Default is 2 (doubling). fuzz: A value between 0 and 1, indicating the fraction of fuzz. For a given delay d, the fuzzed delay is randomly chosen between [(1 - fuzz) * d, d]. max_delay_sec: Maximum delay (in seconds). After this limit is reached, further tries use max_delay_sec instead of exponentially increasing the time. Defaults to 4 hours. """ def __init__(self, initial_delay_secs, num_retries, factor=2, fuzz=0.5, max_delay_secs=60 * 60 * 4): self._initial_delay_secs = initial_delay_secs self._num_retries = num_retries self._factor = factor if not 0 <= fuzz <= 1: raise ValueError('Fuzz parameter expected to be in [0, 1] range.') self._fuzz = fuzz self._max_delay_secs = max_delay_secs def __iter__(self): current_delay_secs = min(self._max_delay_secs, self._initial_delay_secs) for _ in xrange(self._num_retries): fuzz_multiplier = 1 - self._fuzz + random.random() * self._fuzz yield current_delay_secs * fuzz_multiplier current_delay_secs = min( self._max_delay_secs, current_delay_secs * self._factor) def retry_on_server_errors_filter(exception): """Filter allowing retries on server errors and non-HttpErrors.""" if isinstance(exception, HttpError): if exception.status_code >= 500: return True else: return False elif isinstance(exception, PermanentException): return False else: # We may get here for non HttpErrors such as socket timeouts, SSL # exceptions, etc. return True class Clock(object): """A simple clock implementing sleep().""" def sleep(self, value): time.sleep(value) def no_retries(fun): """A retry decorator for places where we do not want retries.""" return with_exponential_backoff( retry_filter=lambda _: False, clock=None)(fun) def with_exponential_backoff( num_retries=16, initial_delay_secs=5.0, logger=logging.warning, retry_filter=retry_on_server_errors_filter, clock=Clock(), fuzz=True): """Decorator with arguments that control the retry logic. Args: num_retries: The total number of times to retry. initial_delay_secs: The delay before the first retry, in seconds. logger: A callable used to report en exception. Must have the same signature as functions in the standard logging module. The default is logging.warning. retry_filter: A callable getting the exception raised and returning True if the retry should happen. For instance we do not want to retry on 404 Http errors most of the time. The default value will return true for server errors (HTTP status code >= 500) and non Http errors. clock: A clock object implementing a sleep method. The default clock will use time.sleep(). fuzz: True if the delay should be fuzzed (default). During testing False can be used so that the delays are not randomized. Returns: As per Python decorators with arguments pattern returns a decorator for the function which in turn will return the wrapped (decorated) function. The decorator is intended to be used on callables that make HTTP or RPC requests that can temporarily timeout or have transient errors. For instance the make_http_request() call below will be retried 16 times with exponential backoff and fuzzing of the delay interval (default settings). from google.cloud.dataflow.utils import retry # ... @retry.with_exponential_backoff() make_http_request(args) """ def real_decorator(fun): """The real decorator whose purpose is to return the wrapped function.""" retry_intervals = iter( FuzzedExponentialIntervals( initial_delay_secs, num_retries, fuzz=0.5 if fuzz else 0)) def wrapper(*args, **kwargs): while True: try: return fun(*args, **kwargs) except Exception as exn: # pylint: disable=broad-except if not retry_filter(exn): raise # Get the traceback object for the current exception. The # sys.exc_info() function returns a tuple with three elements: # exception type, exception value, and exception traceback. exn_traceback = sys.exc_info()[2] try: try: sleep_interval = retry_intervals.next() except StopIteration: # Re-raise the original exception since we finished the retries. raise exn, None, exn_traceback logger( 'Retry with exponential backoff: waiting for %s seconds before ' 'retrying %s because we caught exception: %s ' 'Traceback for above exception (most recent call last):\n%s', sleep_interval, getattr(fun, '__name__', str(fun)), ''.join(traceback.format_exception_only(exn.__class__, exn)), ''.join(traceback.format_tb(exn_traceback))) clock.sleep(sleep_interval) finally: # Traceback objects in locals can cause reference cycles that will # prevent garbage collection. Clear it now since we do not need # it anymore. sys.exc_clear() exn_traceback = None return wrapper return real_decorator DataflowPythonSDK-0.2.5/google/cloud/dataflow/utils/retry_test.py000066400000000000000000000131111272340251000251150ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for the retry module.""" import unittest from google.cloud.dataflow.utils import retry from apitools.base.py.exceptions import HttpError class FakeClock(object): """A fake clock object implementing sleep() and recording calls.""" def __init__(self): self.calls = [] def sleep(self, value): self.calls.append(value) class FakeLogger(object): """A fake logger object implementing log() and recording calls.""" def __init__(self): self.calls = [] def log(self, message, interval, func_name, exn_name, exn_traceback): _ = interval, exn_traceback self.calls.append((message, func_name, exn_name)) @retry.with_exponential_backoff(clock=FakeClock()) def test_function(a, b): _ = a, b raise NotImplementedError @retry.with_exponential_backoff(initial_delay_secs=1.0, num_retries=1) def test_function_with_real_clock(a, b): _ = a, b raise NotImplementedError @retry.no_retries def test_no_retry_function(a, b): _ = a, b raise NotImplementedError class RetryTest(unittest.TestCase): def setUp(self): self.clock = FakeClock() self.logger = FakeLogger() self.calls = 0 def permanent_failure(self, a, b): raise NotImplementedError def transient_failure(self, a, b): self.calls += 1 if self.calls > 8: return a + b raise NotImplementedError def http_error(self, code): raise HttpError({'status': str(code)}, '', '') def test_with_explicit_decorator(self): # We pass one argument as positional argument and one as keyword argument # so that we cover both code paths for argument handling. self.assertRaises(NotImplementedError, test_function, 10, b=20) def test_with_no_retry_decorator(self): self.assertRaises(NotImplementedError, test_no_retry_function, 1, 2) def test_with_real_clock(self): self.assertRaises(NotImplementedError, test_function_with_real_clock, 10, b=20) def test_with_default_number_of_retries(self): self.assertRaises(NotImplementedError, retry.with_exponential_backoff(clock=self.clock)( self.permanent_failure), 10, b=20) self.assertEqual(len(self.clock.calls), 16) def test_with_explicit_number_of_retries(self): self.assertRaises(NotImplementedError, retry.with_exponential_backoff( clock=self.clock, num_retries=10)( self.permanent_failure), 10, b=20) self.assertEqual(len(self.clock.calls), 10) def test_with_http_error_that_should_not_be_retried(self): self.assertRaises(HttpError, retry.with_exponential_backoff( clock=self.clock, num_retries=10)( self.http_error), 404) # Make sure just one call was made. self.assertEqual(len(self.clock.calls), 0) def test_with_http_error_that_should_be_retried(self): self.assertRaises(HttpError, retry.with_exponential_backoff( clock=self.clock, num_retries=10)( self.http_error), 500) self.assertEqual(len(self.clock.calls), 10) def test_with_explicit_initial_delay(self): self.assertRaises(NotImplementedError, retry.with_exponential_backoff( initial_delay_secs=10.0, clock=self.clock, fuzz=False)( self.permanent_failure), 10, b=20) self.assertEqual(len(self.clock.calls), 16) self.assertEqual(self.clock.calls[0], 10.0) def test_log_calls_for_permanent_failure(self): self.assertRaises(NotImplementedError, retry.with_exponential_backoff( clock=self.clock, logger=self.logger.log)( self.permanent_failure), 10, b=20) self.assertEqual(len(self.logger.calls), 16) for message, func_name, exn_name in self.logger.calls: self.assertTrue(message.startswith('Retry with exponential backoff:')) self.assertEqual(exn_name, 'NotImplementedError\n') self.assertEqual(func_name, 'permanent_failure') def test_log_calls_for_transient_failure(self): result = retry.with_exponential_backoff( clock=self.clock, logger=self.logger.log, fuzz=False)( self.transient_failure)(10, b=20) self.assertEqual(result, 30) self.assertEqual(len(self.clock.calls), 8) self.assertEqual(self.clock.calls, [5.0 * 1, 5.0 * 2, 5.0 * 4, 5.0 * 8, 5.0 * 16, 5.0 * 32, 5.0 * 64, 5.0 * 128]) self.assertEqual(len(self.logger.calls), 8) for message, func_name, exn_name in self.logger.calls: self.assertTrue(message.startswith('Retry with exponential backoff:')) self.assertEqual(exn_name, 'NotImplementedError\n') self.assertEqual(func_name, 'transient_failure') if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/version.py000066400000000000000000000012351272340251000232420ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Dataflow SDK for Python version information.""" __version__ = '0.2.5' DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/000077500000000000000000000000001272340251000225135ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/__init__.py000066400000000000000000000000001272340251000246120ustar00rootroot00000000000000DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/batchworker.py000066400000000000000000000636261272340251000254150ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Python Dataflow batch worker. A Worker implements a lease/execute work loop. Multiple processes may execute the same worker loop to get more throughput. In each worker process there are two threads running: the main worker thread and the progress updating thread. The main worker thread will lease a work item, execute it and then report completion (either success or failure). The progress updating thread will pick up the current work item executed by main thread (see the synchronized current_work_item property) and then will send periodic progress updates to the service. These progress updates are essential for ensuring the worker does not lose the lease on the worker item. This can happen if the service does not get timely updates, declares the work item lost, and reassigns to another worker. The two threads can be in contention only when work item attributes related to the progress updating protocol are accessed (reporting index, lease expiration time, duration till next report, etc.). The threads will not be in contention while the work item is getting executed. This is essential in order to make sure that long work items have progress updates sent in a timely manner and leases are not lost often. """ import BaseHTTPServer from collections import namedtuple import datetime import logging import os import random import re import resource import sys import threading import time import traceback from google.cloud.dataflow.internal import apiclient from google.cloud.dataflow.internal import auth from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.utils import names from google.cloud.dataflow.utils import options from google.cloud.dataflow.utils import profiler from google.cloud.dataflow.utils import retry from google.cloud.dataflow.worker import environment from google.cloud.dataflow.worker import executor from google.cloud.dataflow.worker import logger from google.cloud.dataflow.worker import maptask from google.cloud.dataflow.worker import workitem from apitools.base.py.exceptions import HttpError class ProgressReporter(object): """A utility class that can be used to send progress of work items to service. An instance of this should be used to send progress reports for a given work item. """ DEFAULT_MIN_REPORTING_INTERVAL_SECS = 5.0 DEFAULT_MAX_REPORTING_INTERVAL_SECS = 10 * 60.0 DEFAULT_LEASE_RENEWAL_LATENCY_SECS = 5.0 def __init__(self, work_item, work_executor, batch_worker, client): assert work_item is not None assert work_executor is not None assert batch_worker is not None assert client is not None self._work_item = work_item self._work_executor = work_executor self._batch_worker = batch_worker self._client = client self._stopped = False self._stop_reporting_progress = False self._desired_lease_duration = None # Public for testing self.dynamic_split_result_to_report = None def start_reporting_progress(self): """Starts sending progress reports.""" thread = threading.Thread(target=self.progress_reporting_thread) thread.daemon = True thread.start() def stop_reporting_progress(self): """Stops sending progress updates and shuts down the progress reporter. May fail with an exception if unable to send the last split request to the service in which case the work item should be marked as failed. """ self._stop_reporting_progress = True # Shutting down cleanly while not self._stopped: time.sleep(1) # If there is an unreported dynamic work rebalancing response, we must send # it now to guarantee correctness. This may raise an error which will # result in current WorkItem being re-tried by the service. if self.dynamic_split_result_to_report is not None: self.report_status(progress=self._work_executor.get_progress()) def progress_reporting_thread(self): """Sends progress reports for the work item till stopped.""" try: while not self._stop_reporting_progress: try: BatchWorker.log_memory_usage_if_needed(self._batch_worker.worker_id, force=False) with self._work_item.lock: # If WorkItem was marked 'done' in the main worker thread we stop # reporting progress of it. if self._work_item.done: break self.report_status(progress=self._work_executor.get_progress()) sleep_time = self.next_progress_report_interval( self._work_item.report_status_interval, self._work_item.lease_expire_time) logging.debug( 'Progress reporting thread will sleep %f secs between updates.', sleep_time) time.sleep(sleep_time) except Exception: # pylint: disable=broad-except logging.info('Progress reporting thread got error: %s', traceback.format_exc()) finally: self._stopped = True # Public for testing def next_progress_report_interval(self, suggested_interval, lease_renewal_deadline): """Returns the duration till next progress report is needed (in secs). Args: suggested_interval: Duration (as a string) until a status update for the work item should be send back to the service (e.g., '5.000s' or '5s' if zero milliseconds). lease_renewal_deadline: UTC time (a string) when the lease will expire (e.g., '2015-06-17T17:22:49.999Z' or '2015-06-17T17:22:49Z' if zero milliseconds). Returns: Seconds with fractional msecs when next report is expected. """ # Note that the calculation below will clear out a zero returned from the # cloud_time_to_timestamp() function which can happen if the service sends # cloud time strings in an unexpected format. suggested_interval = min( float(suggested_interval.rstrip('s')), self.cloud_time_to_timestamp(lease_renewal_deadline) - time.time() - self.DEFAULT_LEASE_RENEWAL_LATENCY_SECS) return min( max(self.DEFAULT_MIN_REPORTING_INTERVAL_SECS, suggested_interval), self.DEFAULT_MAX_REPORTING_INTERVAL_SECS) def cloud_time_to_timestamp(self, cloud_time_string): """Converts a cloud time string into a timestamp (seconds since EPOCH). Args: cloud_time_string: UTC time (a string) when the lease will expire (e.g., '2015-06-17T17:22:49.999Z' or '2015-06-17T17:22:49Z' if zero milliseconds). Returns: Seconds since EPOCH as a float with fractional part representing msecs. Will return 0 if the string is not in expected format. """ rgx_cloud_time = (r'^(?P\d{4})-(?P\d{2})-(?P\d{2})T' r'(?P\d{2}):(?P\d{2}):(?P\d{2})' r'(\.(?P\d{3}))?Z$') match = re.match(rgx_cloud_time, cloud_time_string) if match: dt = datetime.datetime( int(match.group('year')), int(match.group('month')), int(match.group('day')), int(match.group('hour')), int(match.group('minute')), int(match.group('second')), 0 if match.group('msec') is None else int(match.group('msec')) * 1000) return (dt - datetime.datetime.fromtimestamp(0)).total_seconds() else: # Callers will handle this unexpected error. logging.warning('Unexpected cloud time string: %s', cloud_time_string) return 0 def report_status(self, completed=False, progress=None, source_operation_response=None, exception_details=None): """Reports to the service status of a work item (completion or progress). Args: completed: True if there is no further work to be done on this work item either because it succeeded or because it failed. False if this is a progress report. progress: Progress of processing the work_item. source_operation_response: Response to a custom source operation exception_details: A string representation of the stack trace for an exception raised while executing the work item. The string is the output of the standard traceback.format_exc() function. Note. Callers of this function should acquire the work_item.lock because the function will change fields in the work item based on the response received (e.g., next_report_index, lease_expire_time, etc.). """ response = self._client.report_status( self._batch_worker.worker_info_for_client(), self.desired_lease_duration(), self._work_item, completed, progress if not completed else None, self.dynamic_split_result_to_report if not completed else None, source_operation_response=source_operation_response, exception_details=exception_details) # Resetting dynamic_split_result_to_report after reporting status # successfully. self.dynamic_split_result_to_report = None # If this a progress report (not completion report) then pick up the # new reporting parameters for the work item from the response. if not completed: self.process_report_status_response(response) # Public for testing def process_report_status_response(self, response): """Processes a response to a progress report received from the service.""" work_item_state = response.workItemServiceStates[0] self._work_item.next_report_index = work_item_state.nextReportIndex self._work_item.lease_expire_time = work_item_state.leaseExpireTime self._work_item.report_status_interval = ( work_item_state.reportStatusInterval) suggested_split_point = work_item_state.suggestedStopPoint # Along with the response to the status report, Dataflow service may # send a suggested_split_point, which basically is a request for # performing dynamic work rebalancing if possible. # # Here we pass the received suggested_split_point to current # 'SourceReader' and try to perform a dynamic split. # # If splitting is successful, the corresponding 'DynamicSplitResult' # will be sent to the Dataflow service along with the next progress # report. if suggested_split_point is not None: self.dynamic_split_result_to_report = ( self._work_executor.request_dynamic_split( apiclient.approximate_progress_to_dynamic_split_request( suggested_split_point))) def desired_lease_duration(self): """Returns the desired duration for a work item lease. This duration is send to the service in progress updates. The service may or may not honor the request. The worker has to use the progress updating timings sent by the service in the response in order to not lose the lease on the work item. Returns: The duration to request, as a string representing number of seconds. """ return (self. _desired_lease_duration or self._batch_worker.default_desired_lease_duration()) # Encapsulates information about a BatchWorker object needed when sending # requests to Dataflow service. BatchWorkerInfo = namedtuple( 'WorkerInfo', 'worker_id project_id job_id work_types capabilities ' 'formatted_current_time') class BatchWorker(object): """A worker class with all the knowledge to lease and execute work items.""" # TODO(vladum): Make this configurable via a flag. STATUS_HTTP_PORT = 0 # A value of 0 will pick a random unused port. MEMORY_USAGE_REPORTING_INTERVAL_SECS = 5 * 60 DEFAULT_LEASE_DURATION_SECS = 3 * 60.0 last_memory_usage_report_time = None def __init__(self, properties, sdk_pipeline_options): """Initializes a worker object from command line arguments.""" self.project_id = properties['project_id'] self.job_id = properties['job_id'] self.worker_id = properties['worker_id'] self.service_path = properties['service_path'] # TODO(silviuc): Make sure environment_info_path is always specified. self.environment_info_path = properties.get('environment_info_path', None) self.pipeline_options = options.PipelineOptions.from_dictionary( sdk_pipeline_options) self.capabilities = [self.worker_id, 'remote_source', 'custom_source'] self.work_types = ['map_task', 'seq_map_task', 'remote_source_task'] # The following properties are passed to the worker when its container # gets started and are not used right now. self.root_url = properties['root_url'] self.reporting_enabled = properties['reporting_enabled'] self.temp_gcs_directory = properties['temp_gcs_directory'] # Detect if the worker is running in a GCE VM. self.running_in_gce = self.temp_gcs_directory.startswith('gs://') # When running in a GCE VM the local_staging_property is always set. # For non-VM scenarios (integration tests) the local_staging_directory will # default to the temp directory. self.local_staging_directory = (properties['local_staging_directory'] if self.running_in_gce else self.temp_gcs_directory) self.client = apiclient.DataflowWorkerClient( worker=self, skip_get_credentials=(not self.running_in_gce)) self.environment = maptask.WorkerEnvironment() # If 'True' each work item will be profiled with cProfile. Results will # be logged and also saved to profile_location if set. self.work_item_profiling = sdk_pipeline_options.get('profile', False) self.profile_location = sdk_pipeline_options.get('profile_location', None) self._shutdown = False def worker_info_for_client(self): return BatchWorkerInfo(self.worker_id, self.project_id, self.job_id, self.work_types, self.capabilities, self.formatted_current_time) @property def formatted_current_time(self): # TODO(silviuc): Do we need to support milliseconds too? # The format supports also '...:5.123' (5 secs and 123 msecs). # TODO(silviuc): Consider using utcnow().strftime('%Y-%m-%dT%H:%M:%SZ') now = datetime.datetime.utcnow() return '%04d-%02d-%02dT%02d:%02d:%02d.%03dZ' % ( now.year, now.month, now.day, now.hour, now.minute, now.second, now.microsecond / 1000) def default_desired_lease_duration(self): return '%.3fs' % self.DEFAULT_LEASE_DURATION_SECS def _load_main_session(self, session_path): """Loads a pickled main session from the path specified.""" session_file = os.path.join(session_path, names.PICKLED_MAIN_SESSION_FILE) if os.path.isfile(session_file): pickler.load_session(session_file) else: logging.warning( 'No session file found: %s. Functions defined in __main__ ' '(interactive session) may fail.', session_file) @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def report_completion_status( self, current_work_item, progress_reporter, source_operation_response=None, exception_details=None): """Reports to the service a work item completion (successful or failed). Reporting completion status will do retry with exponential backoff in order to maximize the chances of getting the result to the service. An interim progress report on the other hand will not be retried since it can be sent on the next reporting cycle. The exponential backoff is done by doubling at each retry the initial delay and also introducing some fuzzing in the exact delay. Args: current_work_item: A WorkItem instance describing the work. progress_reporter: A ProgressReporter configured to process work item current_work_item. source_operation_response: Response to a custom source operation. exception_details: A string representation of the stack trace for an exception raised while executing the work item. The string is the output of the standard traceback.format_exc() function. Note. Callers of this function should acquire the work_item.lock. """ # The log message string 'Finished processing' is looked for by # internal tests. Please do not modify the prefix without checking. logging.info('Finished processing %s %s', current_work_item, 'successfully' if exception_details is None else 'with exception') progress_reporter.report_status( completed=True, source_operation_response=source_operation_response, exception_details=exception_details) @staticmethod def log_memory_usage_if_needed(worker_id, force=False): """Periodically logs memory usage of the current worker. Args: worker_id: Id of the worker. force: if True forces logging. """ if (force or BatchWorker.last_memory_usage_report_time is None or int(time.time()) - BatchWorker.last_memory_usage_report_time > BatchWorker.MEMORY_USAGE_REPORTING_INTERVAL_SECS): logging.info('Memory usage of worker %s is %d MB', worker_id, resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1000) BatchWorker.last_memory_usage_report_time = int(time.time()) def shutdown(self): self._shutdown = True def get_executor_for_work_item(self, work_item): if work_item.map_task is not None: return executor.MapTaskExecutor(work_item.map_task) elif work_item.source_operation_split_task is not None: return executor.CustomSourceSplitExecutor( work_item.source_operation_split_task) else: raise ValueError('Unknown type of work item : %s', work_item) def do_work(self, work_item, deferred_exception_details=None): """Executes worker operations and adds any failures to the report status.""" logging.info('Executing %s', work_item) BatchWorker.log_memory_usage_if_needed(self.worker_id, force=True) work_executor = self.get_executor_for_work_item(work_item) progress_reporter = ProgressReporter( work_item, work_executor, self, self.client) if deferred_exception_details: # Report (fatal) deferred exceptions that happened earlier. This # workflow will fail with the deferred exception. with work_item.lock: self.report_completion_status( work_item, progress_reporter, exception_details=deferred_exception_details) work_item.done = True logging.error('Not processing WorkItem %s since a deferred exception ' 'was found: %s', work_item, deferred_exception_details) return exception_details = None try: progress_reporter.start_reporting_progress() work_executor.execute() except Exception: # pylint: disable=broad-except exception_details = traceback.format_exc() logging.error('An exception was raised when trying to execute the ' 'work item %s : %s', work_item, exception_details, exc_info=True) finally: try: progress_reporter.stop_reporting_progress() except Exception: # pylint: disable=broad-except logging.error('An exception was raised when trying to stop the ' 'progress reporter : %s', traceback.format_exc(), exc_info=True) # If 'exception_details' was already set, we were already going to # mark this work item as failed. Hence only logging this error and # reporting the original error. if exception_details is None: # This will be reported to the service and work item will be marked as # failed. exception_details = traceback.format_exc() with work_item.lock: source_split_response = None if isinstance(work_executor, executor.CustomSourceSplitExecutor): source_split_response = work_executor.response self.report_completion_status( work_item, progress_reporter, source_operation_response=source_split_response, exception_details=exception_details) work_item.done = True def status_server(self): """Executes the serving loop for the status server.""" class StatusHttpHandler(BaseHTTPServer.BaseHTTPRequestHandler): """HTTP handler for serving stacktraces of all worker threads.""" def do_GET(self): # pylint: disable=invalid-name """Return /threadz information for any GET request.""" self.send_response(200) self.send_header('Content-Type', 'text/plain') self.end_headers() frames = sys._current_frames() # pylint: disable=protected-access for t in threading.enumerate(): self.wfile.write('--- Thread #%s name: %s ---\n' % (t.ident, t.name)) self.wfile.write(''.join(traceback.format_stack(frames[t.ident]))) def log_message(self, f, *args): """Do not log any messages.""" pass httpd = BaseHTTPServer.HTTPServer( ('localhost', self.STATUS_HTTP_PORT), StatusHttpHandler) logging.info('Status HTTP server running at %s:%s', httpd.server_name, httpd.server_port) httpd.serve_forever() def run(self): """Runs the worker loop for leasing and executing work items.""" if self.running_in_gce: auth.set_running_in_gce(worker_executing_project=self.project_id) # Deferred exceptions are used as a way to report unrecoverable errors that # happen before they could be reported to the service. If it is not None, # worker will use the first work item to report deferred exceptions and # fail eventually. # TODO(silviuc): Add the deferred exception mechanism to streaming worker deferred_exception_details = None if self.environment_info_path is not None: try: environment.check_sdk_compatibility(self.environment_info_path) except Exception: # pylint: disable=broad-except deferred_exception_details = traceback.format_exc() logging.error('SDK compatibility check failed: %s', deferred_exception_details, exc_info=True) if deferred_exception_details is None: logging.info('Loading main session from the staging area...') try: self._load_main_session(self.local_staging_directory) except Exception: # pylint: disable=broad-except deferred_exception_details = traceback.format_exc() logging.error('Could not load main session: %s', deferred_exception_details, exc_info=True) # Start status HTTP server thread. thread = threading.Thread(target=self.status_server) thread.daemon = True thread.start() # The batch execution context is currently a placeholder, so we don't yet # need to have it change between work items. execution_context = maptask.BatchExecutionContext() work_item = None # Loop forever leasing work items, executing them, and reporting status. while not self._shutdown: try: # Lease a work item. The lease_work call will retry for server errors # (e.g., 500s) however it will not retry for a 404 (no item to lease). # In such cases we introduce random sleep delays with the code below. should_sleep = False try: work = self.client.lease_work(self.worker_info_for_client(), self.default_desired_lease_duration()) work_item = workitem.get_work_items(work, self.environment, execution_context) if work_item is None: should_sleep = True except HttpError as exn: # Not found errors (404) are benign. The rest are not and must be # re-raised. if exn.status_code != 404: raise should_sleep = True if should_sleep: logging.debug('No work items. Sleeping a bit ...') # The sleeping is done with a bit of jitter to avoid having workers # requesting leases in lock step. time.sleep(1.0 * (1 - 0.5 * random.random())) continue stage_name = None if work_item.map_task: stage_name = work_item.map_task.stage_name with logger.PerThreadLoggingContext( work_item_id=work_item.proto.id, stage_name=stage_name): # TODO(silviuc): Add more detailed timing and profiling support. start_time = time.time() # Do the work. The do_work() call will mark the work completed or # failed. The progress reporting_thread will take care of sending # updates and updating in the workitem object the reporting indexes # and duration for the lease. if self.work_item_profiling: with profiler.Profile( profile_id=work_item.proto.id, profile_location=self.profile_location, log_results=True): self.do_work( work_item, deferred_exception_details=deferred_exception_details) else: self.do_work(work_item, deferred_exception_details=deferred_exception_details) logging.info('Completed work item: %s in %.9f seconds', work_item.proto.id, time.time() - start_time) except Exception: # pylint: disable=broad-except # This is an exception raised outside of executing a work item most # likely while leasing a work item. We log an error and march on. logging.error('Exception in worker loop: %s', traceback.format_exc(), exc_info=True) # sleeping a bit after Exception to prevent a busy loop. time.sleep(1) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/batchworker_test.py000066400000000000000000000217501272340251000264440ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for BatchWorker.""" import logging import threading import time import unittest import mock from mock import patch from google.cloud.dataflow.internal import apiclient from google.cloud.dataflow.worker import batchworker from google.cloud.dataflow.worker import executor from google.cloud.dataflow.worker import workitem class BatchWorkerTest(unittest.TestCase): def dummy_properties(self): return { 'project_id': 'test_project', 'job_id': 'test_job', 'worker_id': 'test_worker', 'service_path': 'test_services_path', 'root_url': 'test_root_url', 'reporting_enabled': 'test_reporting_enabled', 'temp_gcs_directory': 'test_temp_gcs_directory' } @patch('google.cloud.dataflow.worker.batchworker.workitem') @patch.object(apiclient.DataflowWorkerClient, 'lease_work') def test_worker_requests_for_work(self, mock_lease_work, mock_workitem): worker = batchworker.BatchWorker(self.dummy_properties(), {}) rand_work = object() mock_lease_work.return_value = rand_work mock_workitem.get_work_items.return_value = None thread = threading.Thread(target=worker.run) thread.start() time.sleep(5) worker.shutdown() mock_lease_work.assert_called_with( mock.ANY, worker.default_desired_lease_duration()) mock_workitem.get_work_items.assert_called_with( rand_work, mock.ANY, mock.ANY) @patch('google.cloud.dataflow.worker.batchworker.workitem') @patch.object(apiclient.DataflowWorkerClient, 'lease_work') def test_worker_requests_for_work_after_lease_error( self, mock_lease_work, mock_workitem): worker = batchworker.BatchWorker(self.dummy_properties(), {}) rand_work = object() mock_lease_work.side_effect = [Exception('test exception'), rand_work] mock_workitem.get_work_items.return_value = None thread = threading.Thread(target=worker.run) thread.start() time.sleep(5) worker.shutdown() mock_lease_work.assert_called_with( mock.ANY, worker.default_desired_lease_duration()) mock_workitem.get_work_items.assert_called_with( mock.ANY, mock.ANY, mock.ANY) @patch.object(executor.MapTaskExecutor, 'execute') @patch.object(batchworker.ProgressReporter, 'start_reporting_progress') @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress') @patch.object(batchworker.ProgressReporter, 'report_status') def test_worker_starts_and_stops_progress_reporter( self, mock_report_status, mock_stop, mock_start, mock_execute): worker = batchworker.BatchWorker(self.dummy_properties(), {}) mock_work_item = mock.MagicMock() worker.do_work(mock_work_item) mock_report_status.assert_called_with( completed=True, source_operation_response=None, exception_details=None) mock_start.assert_called_once_with() mock_execute.assert_called_once_with() mock_stop.assert_called_once_with() @patch.object(executor.MapTaskExecutor, 'execute') @patch.object(batchworker.ProgressReporter, 'start_reporting_progress') @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress') @patch.object(batchworker.ProgressReporter, 'report_status') def test_worker_fails_for_deferred_exceptions( self, mock_report_status, mock_stop, mock_start, mock_execute): worker = batchworker.BatchWorker(self.dummy_properties(), {}) mock_work_item = mock.MagicMock() worker.do_work(mock_work_item, deferred_exception_details='deferred_exc') mock_report_status.assert_called_with( completed=True, source_operation_response=None, exception_details='deferred_exc') assert not mock_stop.called assert not mock_start.called assert not mock_execute.called def _run_send_completion_test(self, mock_report_status, mock_stop, mock_start, mock_execute, expected_exception): worker = batchworker.BatchWorker(self.dummy_properties(), {}) mock_work_item = mock.MagicMock() worker.do_work(mock_work_item) class AnyStringWith(str): def __eq__(self, other): return self in other mock_report_status.assert_called_with( completed=True, source_operation_response=None, exception_details=AnyStringWith(expected_exception)) mock_start.assert_called_once_with() mock_execute.assert_called_once_with() mock_stop.assert_called_once_with() @patch.object(executor.MapTaskExecutor, 'execute') @patch.object(batchworker.ProgressReporter, 'start_reporting_progress') @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress') @patch.object(batchworker.ProgressReporter, 'report_status') def test_send_completion_execute_failure(self, mock_report_status, mock_stop, mock_start, mock_execute): mock_execute.side_effect = Exception('test_exception') self._run_send_completion_test(mock_report_status, mock_stop, mock_start, mock_execute, 'test_exception') @patch.object(executor.MapTaskExecutor, 'execute') @patch.object(batchworker.ProgressReporter, 'start_reporting_progress') @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress') @patch.object(batchworker.ProgressReporter, 'report_status') def test_send_completion_stop_progress_reporter_failure(self, mock_report_status, mock_stop, mock_start, mock_execute): mock_stop.side_effect = Exception('test_exception') self._run_send_completion_test(mock_report_status, mock_stop, mock_start, mock_execute, 'test_exception') @patch.object(executor.MapTaskExecutor, 'execute') @patch.object(batchworker.ProgressReporter, 'start_reporting_progress') @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress') @patch.object(batchworker.ProgressReporter, 'report_status') def test_send_completion_execute_and_stop_progress_reporter_failure( self, mock_report_status, mock_stop, mock_start, mock_execute): mock_execute.side_effect = Exception('test_exception_1') mock_stop.side_effect = Exception('test_exception_2') self._run_send_completion_test(mock_report_status, mock_stop, mock_start, mock_execute, 'test_exception_1') class ProgressReporterTest(unittest.TestCase): @patch.object(batchworker.ProgressReporter, 'next_progress_report_interval') @patch.object(batchworker.ProgressReporter, 'process_report_status_response') def test_progress_reporter_reports_progress( self, mock_report_response, mock_next_progress): # pylint: disable=unused-argument work_item = workitem.BatchWorkItem(proto=mock.MagicMock()) work_item.map_task = mock.MagicMock() mock_work_executor = mock.MagicMock() mock_batch_worker = mock.MagicMock() mock_client = mock.MagicMock() mock_next_progress.return_value = 1 progress_reporter = batchworker.ProgressReporter( work_item, mock_work_executor, mock_batch_worker, mock_client) progress_reporter.start_reporting_progress() time.sleep(10) progress_reporter.stop_reporting_progress() mock_client.report_status.assert_called_with( mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, exception_details=mock.ANY, source_operation_response=mock.ANY) @patch.object(batchworker.ProgressReporter, 'next_progress_report_interval') @patch.object(batchworker.ProgressReporter, 'process_report_status_response') def test_progress_reporter_sends_last_update( self, mock_report_response, mock_next_progress): # pylint: disable=unused-argument mock_work_item = mock.MagicMock() mock_work_executor = mock.MagicMock() mock_batch_worker = mock.MagicMock() mock_client = mock.MagicMock() progress_reporter = batchworker.ProgressReporter( mock_work_item, mock_work_executor, mock_batch_worker, mock_client) mock_split_result = mock.MagicMock() progress_reporter.dynamic_split_result_to_report = mock_split_result progress_reporter._stopped = True progress_reporter.stop_reporting_progress() mock_client.report_status.assert_called_with( mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock_split_result, exception_details=mock.ANY, source_operation_response=mock.ANY) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/concat_reader.py000066400000000000000000000062611272340251000256630ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """A reader that encapsulates a set of other readers. This is to be used for optimizing the execution of Dataflow jobs. Users should not use this when developing Dataflow jobs. """ from __future__ import absolute_import from google.cloud.dataflow.io import iobase class ConcatSource(iobase.NativeSource): """A wrapper source class for ConcatReader.""" def __init__(self, sub_sources): self.sub_sources = sub_sources def reader(self): return ConcatReader(self) def __eq__(self, other): return self.sub_sources == other.sub_sources class ConcatReader(iobase.NativeSourceReader): """A reader that reads elements from a given set of encoded sources. Creates readers for sources lazily, i.e. only when elements from the particular reader are about to be read. This class does does not cache readers and instead creates new set of readers evertime it is iterated on. Because of this, multiple iterators created for the same ConcatReader will not be able to share any state between each other. This design was chosen since keeping a large number of reader objects alive within a single ConcatReader could be highly resource consuming. For progress reporting ConcatReader uses a position of type iobase.ConcatPosition. """ def __init__(self, source): self.source = source self.current_reader = None self.current_reader_index = -1 def __enter__(self): return self def __iter__(self): if self.source.sub_sources is None: return for sub_source in self.source.sub_sources: with sub_source.reader() as reader: self.current_reader_index += 1 self.current_reader = reader for data in reader: yield data def __exit__(self, exception_type, exception_value, traceback): pass def get_progress(self): if self.current_reader_index < 0 or self.current_reader is None: return index = self.current_reader_index inner_position = None sub_reader_progress = self.current_reader.get_progress() if sub_reader_progress is not None: sub_reader_position = sub_reader_progress.position if sub_reader_position is not None: inner_position = sub_reader_position else: raise ValueError('A concat source should only be created with ' 'sub-sources that create readers that perform ' 'progress reporting and dynamic work rebalancing ' 'using positions') return iobase.ReaderProgress( position=iobase.ReaderPosition( concat_position=iobase.ConcatPosition(index, inner_position))) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/concat_reader_test.py000066400000000000000000000204171272340251000267210ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for ConcatReader.""" import logging import random import unittest from google.cloud.dataflow.io import iobase from google.cloud.dataflow.worker import concat_reader class TestSource(iobase.NativeSource): def __init__( self, elements, index_to_fail_reading=-1, fail_reader_at_close=False): self.elements = elements self.index_to_fail_reading = index_to_fail_reading self.fail_reader_at_close = fail_reader_at_close def reader(self): return TestReader(self.elements, self.index_to_fail_reading, self.fail_reader_at_close) class TestReader(iobase.NativeSourceReader): def __init__(self, elements, index_to_fail_reading, fail_reader_at_close): self.elements = elements self.index_to_fail_reading = index_to_fail_reading self.fail_reader_at_close = fail_reader_at_close self.current_index = 0 def __exit__(self, exception_type, exception_value, traceback): if self.fail_reader_at_close: raise ValueError def __enter__(self): return self def __iter__(self): for element in self.elements: if self.current_index == self.index_to_fail_reading: raise ValueError yield element self.current_index += 1 def get_progress(self): return iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=self.current_index)) class ConcatReaderTest(unittest.TestCase): def create_data(self, sizes): all_data = [] for size in sizes: data = [] for _ in range(size): next_data = random.randint(0, 10000) data.append(next_data) all_data.append(data) return all_data def _create_concat_source(self, sub_source_sizes, output_record, index_of_source_to_fail=-1, index_to_fail_reading=-1, fail_reader_at_close=False): sub_sources = [] all_data = self.create_data(sub_source_sizes) for data in all_data: output_record.extend(data) for index, data in enumerate(all_data): if index == index_of_source_to_fail: sub_sources.append( TestSource(data, index_to_fail_reading, fail_reader_at_close)) else: sub_sources.append(TestSource(data, -1, False)) return concat_reader.ConcatSource(sub_sources) def test_create_from_null(self): all_data = [] with concat_reader.ConcatSource(None).reader() as reader: for data in reader: all_data.append(data) self.assertEqual(0, len(all_data)) def test_read_empty_list(self): all_data = [] with concat_reader.ConcatSource([]).reader() as reader: for data in reader: all_data.append(data) self.assertEqual(0, len(all_data)) def test_read_one(self): expected_output = [] received_output = [] with self._create_concat_source([10], expected_output).reader() as reader: for data in reader: received_output.append(data) self.assertEqual(10, len(expected_output)) self.assertEqual(10, len(received_output)) self.assertEqual(expected_output.sort(), received_output.sort()) def test_read_multi_same_size(self): expected_output = [] received_output = [] source = self._create_concat_source([10, 10, 10], expected_output) with source.reader() as reader: for data in reader: received_output.append(data) self.assertEqual(30, len(expected_output)) self.assertEqual(30, len(received_output)) self.assertEqual(expected_output.sort(), received_output.sort()) def test_read_multi_different_sizes(self): expected_output = [] received_output = [] source = self._create_concat_source([10, 30, 40, 20], expected_output) with source.reader() as reader: for data in reader: received_output.append(data) self.assertEqual(100, len(expected_output)) self.assertEqual(100, len(received_output)) self.assertEqual(expected_output.sort(), received_output.sort()) def test_last_reader_empty(self): expected_output = [] received_output = [] source = self._create_concat_source([10, 30, 40, 0], expected_output) with source.reader() as reader: for data in reader: received_output.append(data) self.assertEqual(80, len(expected_output)) self.assertEqual(80, len(received_output)) self.assertEqual(expected_output.sort(), received_output.sort()) def test_empty_reader_before_non_empty_reader(self): expected_output = [] received_output = [] source = self._create_concat_source([10, 0, 40, 20], expected_output) with source.reader() as reader: for data in reader: received_output.append(data) self.assertEqual(70, len(expected_output)) self.assertEqual(70, len(received_output)) self.assertEqual(expected_output.sort(), received_output.sort()) def test_multiple_readers_are_empty(self): expected_output = [] received_output = [] source = self._create_concat_source([10, 0, 0, 20, 0, 30], expected_output) with source.reader() as reader: for data in reader: received_output.append(data) self.assertEqual(60, len(expected_output)) self.assertEqual(60, len(received_output)) self.assertEqual(expected_output.sort(), received_output.sort()) def test_a_reader_fails_at_read(self): all_data = [] received_output = [] try: source = self._create_concat_source( [10, 30, 40, 20], all_data, index_of_source_to_fail=2, index_to_fail_reading=15, fail_reader_at_close=False) with source.reader() as reader: for data in reader: received_output.append(data) except ValueError: self.assertEqual(100, len(all_data)) self.assertEqual(55, len(received_output)) self.assertEqual(all_data[:55].sort(), received_output.sort()) else: # reading should have produced a ValueError. Failing test. since it # didn't. raise ValueError def test_a_reader_fails_at_close(self): all_data = [] received_output = [] try: source = self._create_concat_source( [10, 30, 40, 20], all_data, index_of_source_to_fail=2, index_to_fail_reading=-1, fail_reader_at_close=True) with source.reader() as reader: for data in reader: received_output.append(data) except ValueError: self.assertEqual(100, len(all_data)) self.assertEqual(80, len(received_output)) self.assertEqual(all_data[:80].sort(), received_output.sort()) else: # reading should have produced a ValueError. Failing test. since it # didn't. raise ValueError def _test_progress_reporting(self, sizes): with self._create_concat_source(sizes, []).reader() as reader: reader_iter = iter(reader) for reader_index in range(0, len(sizes)): for record_index in range(0, sizes[reader_index]): reader_iter.next() progress = reader.get_progress() self.assertIsNotNone(progress) self.assertIsNotNone(progress.position) self.assertIsNotNone(progress.position.concat_position) self.assertTrue(isinstance(progress.position.concat_position, iobase.ConcatPosition)) self.assertEqual(reader_index, progress.position.concat_position.index) self.assertEqual( record_index, progress.position.concat_position.position.record_index) def test_get_progress_single(self): self._test_progress_reporting([10]) def test_get_progress_same_size(self): self._test_progress_reporting([10, 10, 10]) def test_get_progress_multiple_sizes(self): self._test_progress_reporting([20, 10, 30]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/environment.py000066400000000000000000000040271272340251000254340ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Python Dataflow worker environment compatiblity checking.""" import json import logging from google.cloud.dataflow import version def check_sdk_compatibility(environment_info_path): """Checks if the SDK is compatible with the container in which it runs. Args: environment_info_path: Path to a file in JSON format. The file is expected to contain a dictionary with at least two properties: 'language' and 'version'. Raises: RuntimeError: For version or language mismatches. Other exceptions can be raised if the environment file is not present or does not have the right contents. This can happen only if the base container was not built correctly. """ logging.info('Checking if container and SDK language and versions match ...') with open(environment_info_path) as f: info = json.loads(f.read()) if info['language'] != 'python': message = ( 'SDK language \'python\' does not match container language \'%s\'. ' 'Please rebuild the container using a matching language container.' % ( info['language'])) logging.error(message) raise RuntimeError(message) if info['version'] != version.__version__: message = ( 'SDK version %s does not match container version %s. ' 'Please rebuild the container or use a matching version ' 'of the SDK.' % ( version.__version__, info['version'])) logging.error(message) raise RuntimeError(message) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/environment_test.py000066400000000000000000000046421272340251000264760ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for worker environment compatibility checking.""" import logging import os import tempfile import unittest from google.cloud.dataflow import version from google.cloud.dataflow.worker import environment class EnvironmentTest(unittest.TestCase): def create_temp_file(self, path, contents): with open(path, 'w') as f: f.write(contents) return f.name def test_basics(self): config_path = os.path.join(tempfile.mkdtemp(), 'config') self.create_temp_file( config_path, '{"language":"python", "version": "%s"}' % version.__version__) environment.check_sdk_compatibility(config_path) # If we get here the test passes since no exception was raised. def test_language_no_match(self): config_path = os.path.join(tempfile.mkdtemp(), 'config') self.create_temp_file( config_path, '{"language":"java", "version": "%s"}' % version.__version__) with self.assertRaises(RuntimeError) as exn: environment.check_sdk_compatibility(config_path) self.assertEqual( 'SDK language \'python\' does not match container language \'java\'. ' 'Please rebuild the container using a matching language container.', exn.exception.message) def test_version_no_match(self): config_path = os.path.join(tempfile.mkdtemp(), 'config') self.create_temp_file( config_path, '{"language":"python", "version": "0.0.0"}') with self.assertRaises(RuntimeError) as exn: environment.check_sdk_compatibility(config_path) self.assertEqual( 'SDK version %s does not match container version 0.0.0. ' 'Please rebuild the container or use a matching version ' 'of the SDK.' % ( version.__version__), exn.exception.message) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/executor.pxd000066400000000000000000000042001272340251000250620ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. cimport cython cdef class Operation(object): cdef public spec cdef public counter_factory cdef public list receivers cdef public list counters cdef readonly bint debug_logging_enabled cdef public step_name # initialized lazily cpdef start(self) cpdef process(self, windowed_value) cpdef finish(self) @cython.locals(receiver=Operation) cpdef output(self, windowed_value, object coder=*, int output_index=*) cdef class ReadOperation(Operation): cdef object _current_progress cdef object _reader cdef class DoOperation(Operation): cdef object state cdef object context cdef object dofn_runner cdef class CombineOperation(Operation): cdef object phased_combine_fn cdef class ShuffleWriteOperation(Operation): cdef object shuffle_sink cdef object writer cdef object _write_coder cdef bint is_ungrouped cdef class GroupedShuffleReadOperation(Operation): cdef object shuffle_source cdef object _reader cdef class UngroupedShuffleReadOperation(Operation): cdef object shuffle_source cdef object _reader cdef class FlattenOperation(Operation): pass cdef class ReifyTimestampAndWindowsOperation(Operation): pass cdef class BatchGroupAlsoByWindowsOperation(Operation): cdef object windowing cdef object phased_combine_fn cdef class StreamingGroupAlsoByWindowsOperation(Operation): cdef object windowing cdef object phased_combine_fn cdef class PGBKCVOperation(Operation): cdef public object combine_fn cdef dict table cdef long max_keys cdef long key_count cpdef output_key(self, tuple wkey, value) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/executor.py000066400000000000000000001020571272340251000247300ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. # cython: profile=True """Worker operations executor.""" import collections import itertools import logging import random from google.cloud.dataflow import pvalue from google.cloud.dataflow.coders import BytesCoder from google.cloud.dataflow.coders import TupleCoder from google.cloud.dataflow.coders import WindowedValueCoder from google.cloud.dataflow.internal import apiclient from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.io import iobase from google.cloud.dataflow.pvalue import EmptySideInput from google.cloud.dataflow.runners import common import google.cloud.dataflow.transforms as ptransform from google.cloud.dataflow.transforms import combiners from google.cloud.dataflow.transforms import trigger from google.cloud.dataflow.transforms.combiners import curry_combine_fn from google.cloud.dataflow.transforms.combiners import PhasedCombineFnExecutor from google.cloud.dataflow.transforms.trigger import InMemoryUnmergedState from google.cloud.dataflow.transforms.window import GlobalWindows from google.cloud.dataflow.transforms.window import MIN_TIMESTAMP from google.cloud.dataflow.transforms.window import WindowedValue from google.cloud.dataflow.utils.names import PropertyNames from google.cloud.dataflow.worker import logger from google.cloud.dataflow.worker import maptask from google.cloud.dataflow.worker import opcounters from google.cloud.dataflow.worker import shuffle from google.cloud.dataflow.worker import sideinputs class ReceiverSet(object): """A ReceiverSet represents a graph edge between two Operation nodes. The ReceiverSet object collects information from the output of the Operation at one end of its edge and the input of the Operation at the other edge. ReceiverSets are attached to the outputting Operation. """ def __init__(self, counter_factory, coder, output_index=0): self.receivers = [] self.opcounter = None self.counter_factory = counter_factory self.output_index = output_index self.coder = coder def add_receiver(self, receiving_operation): self.receivers.append(receiving_operation) def start(self, step_name): self.opcounter = opcounters.OperationCounters( self.counter_factory, step_name, self.coder, self.output_index) def output(self, windowed_value, coder=None): self.update_counters_start(windowed_value, coder) for receiver in self.receivers: receiver.process(windowed_value) self.update_counters_finish() def update_counters_start(self, windowed_value, coder=None): if self.opcounter: self.opcounter.update_from(windowed_value, coder) def update_counters_finish(self): if self.opcounter: self.opcounter.update_collect() def __str__(self): return '[%s]' % ' '.join([r.str_internal(is_recursive=True) for r in self.receivers]) def __repr__(self): return '<%s %d %s [%s]>' % ( self.__class__.__name__, self.output_index, self.coder, ' '.join([r.str_internal(is_recursive=True) for r in self.receivers])) class Operation(object): """An operation representing the live version of a work item specification. An operation can have one or more outputs and for each output it can have one or more receiver operations that will take that as input. """ def __init__(self, spec, counter_factory): """Initializes a worker operation instance. Args: spec: A maptask.Worker* instance. counter_factory: the counter_factory to use for our counters. """ self.counter_factory = counter_factory self.spec = spec self.receivers = [] # Everything except WorkerSideInputSource, which is not a # top-level operation, should have output_coders if getattr(self.spec, 'output_coders', None): for i, coder in enumerate(self.spec.output_coders): self.receivers.append(ReceiverSet(counter_factory, coder, i)) def start(self): """Start operation.""" self.debug_logging_enabled = logging.getLogger().isEnabledFor( logging.DEBUG) # Start our receivers, now that we know our step name. for receiver in self.receivers: receiver.start(self.step_name) def finish(self): """Finish operation.""" pass def process(self, o): """Process element in operation.""" pass def output(self, windowed_value, coder=None, output_index=0): self.receivers[output_index].output(windowed_value, coder) def add_receiver(self, operation, output_index=0): """Adds a receiver operation for the specified output.""" self.receivers[output_index].add_receiver(operation) def __str__(self): """Generates a useful string for this object. Compactly displays interesting fields. In particular, pickled fields are not displayed. Note that we collapse the fields of the contained Worker* object into this object, since there is a 1-1 mapping between Operation and maptask.Worker*. Returns: Compact string representing this object. """ return self.str_internal() def str_internal(self, is_recursive=False): """Internal helper for __str__ that supports recursion. When recursing on receivers, keep the output short. Args: is_recursive: whether to omit some details, particularly receivers. Returns: Compact string representing this object. """ printable_name = self.__class__.__name__ if hasattr(self, 'step_name'): printable_name += ' %s' % self.step_name if is_recursive: # If we have a step name, stop here, no more detail needed. return '<%s>' % printable_name printable_fields = maptask.worker_printable_fields(self.spec) if not is_recursive and getattr(self, 'receivers', []): printable_fields.append('receivers=[%s]' % ', '.join([ str(receiver) for receiver in self.receivers])) return '<%s %s>' % (printable_name, ', '.join(printable_fields)) class ReadOperation(Operation): """A generic read operation that reads from proper input source.""" def __init__(self, spec, counter_factory): super(ReadOperation, self).__init__(spec, counter_factory) self._current_progress = None self._reader = None def start(self): # We cache reader progress to make sure that the progress reporting # thread does not get blocked due to any reader related operations. self._current_progress = None super(ReadOperation, self).start() with self.spec.source.reader() as reader: self._reader = reader for value in reader: self._current_progress = reader.get_progress() if reader.returns_windowed_values: windowed_value = value else: windowed_value = GlobalWindows.windowed_value(value) self.output(windowed_value) def request_dynamic_split(self, dynamic_split_request): if self._reader is not None: return self._reader.request_dynamic_split(dynamic_split_request) else: logging.warning('Cannot complete the dynamic split request since the ' 'reader is not set') def get_progress(self): """Returns the progress of the read operation. This method is invoked by the progress reporting thread. No need to lock the variable "current_progress" since it is updated by a simple variable assignment and we are OK with current_progress value returned here being slightly stale. Returns: Progress of the ReadOperation. """ return self._current_progress class WriteOperation(Operation): """A generic write operation that writes to a proper output sink.""" def __init__(self, spec, counter_factory): super(WriteOperation, self).__init__(spec, counter_factory) self.writer = None self.use_windowed_value = False def start(self): super(WriteOperation, self).start() self.writer = self.spec.sink.writer() self.writer.__enter__() self.use_windowed_value = self.writer.takes_windowed_values def finish(self): self.writer.__exit__(None, None, None) def process(self, o): if self.debug_logging_enabled: logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) self.receivers[0].update_counters_start(o) if self.use_windowed_value: self.writer.Write(o) else: self.writer.Write(o.value) self.receivers[0].update_counters_finish() class InMemoryWriteOperation(Operation): """A write operation that will write to an in-memory sink.""" def __init__(self, spec, counter_factory): super(InMemoryWriteOperation, self).__init__(spec, counter_factory) self.spec = spec def process(self, o): if self.debug_logging_enabled: logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) self.receivers[0].update_counters_start(o) self.spec.output_buffer.append(o.value) self.receivers[0].update_counters_finish() class GroupedShuffleReadOperation(Operation): """A shuffle read operation that will read from a grouped shuffle source.""" def __init__(self, spec, counter_factory, shuffle_source=None): super(GroupedShuffleReadOperation, self).__init__(spec, counter_factory) self.shuffle_source = shuffle_source self._reader = None def start(self): super(GroupedShuffleReadOperation, self).start() write_coder = None if self.shuffle_source is None: coders = (self.spec.coder.key_coder(), self.spec.coder.value_coder()) write_coder = WindowedValueCoder(TupleCoder(coders)) self.shuffle_source = shuffle.GroupedShuffleSource( self.spec.shuffle_reader_config, coder=coders, start_position=self.spec.start_shuffle_position, end_position=self.spec.end_shuffle_position) with self.shuffle_source.reader() as reader: for key, key_values in reader: self._reader = reader windowed_value = GlobalWindows.windowed_value((key, key_values)) self.output(windowed_value, coder=write_coder) def get_progress(self): if self._reader is not None: return self._reader.get_progress() def request_dynamic_split(self, dynamic_split_request): if self._reader is not None: return self._reader.request_dynamic_split(dynamic_split_request) class UngroupedShuffleReadOperation(Operation): """A shuffle read operation reading from an ungrouped shuffle source.""" def __init__(self, spec, counter_factory, shuffle_source=None): super(UngroupedShuffleReadOperation, self).__init__(spec, counter_factory) self.shuffle_source = shuffle_source self._reader = None def start(self): super(UngroupedShuffleReadOperation, self).start() write_coder = None if self.shuffle_source is None: coders = (BytesCoder(), self.spec.coder) write_coder = WindowedValueCoder(TupleCoder(coders)) self.shuffle_source = shuffle.UngroupedShuffleSource( self.spec.shuffle_reader_config, coder=coders, start_position=self.spec.start_shuffle_position, end_position=self.spec.end_shuffle_position) with self.shuffle_source.reader() as reader: for value in reader: self._reader = reader windowed_value = GlobalWindows.windowed_value(value) self.output(windowed_value, coder=write_coder) def get_progress(self): # 'UngroupedShuffleReader' does not support progress reporting. pass def request_dynamic_split(self, dynamic_split_request): # 'UngroupedShuffleReader' does not support dynamic work rebalancing. pass class ShuffleWriteOperation(Operation): """A shuffle write operation that will write to a shuffle sink.""" def __init__(self, spec, counter_factory, shuffle_sink=None): super(ShuffleWriteOperation, self).__init__(spec, counter_factory) self.writer = None self.shuffle_sink = shuffle_sink def start(self): super(ShuffleWriteOperation, self).start() self.is_ungrouped = self.spec.shuffle_kind == 'ungrouped' coder = self.spec.output_coders[0] if self.is_ungrouped: coders = (BytesCoder(), coder) else: coders = (coder.key_coder(), coder.value_coder()) self._write_coder = WindowedValueCoder(TupleCoder(coders)) if self.shuffle_sink is None: self.shuffle_sink = shuffle.ShuffleSink( self.spec.shuffle_writer_config, coder=coders) self.writer = self.shuffle_sink.writer() self.writer.__enter__() def finish(self): logging.debug('Finishing %s', self) self.writer.__exit__(None, None, None) def process(self, o): if self.debug_logging_enabled: logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) self.receivers[0].update_counters_start(o, coder=self._write_coder) # We typically write into shuffle key/value pairs. This is the reason why # the else branch below expects the value attribute of the WindowedValue # argument to be a KV pair. However the service may write to shuffle in # 'ungrouped' mode in which case the value written is just a plain value and # the key is randomly generated on the spot. The random keys make sure that # the resulting KV pairs are uniformly distributed. The 'ungrouped' mode is # used to reshard workflow outputs into a fixed set of files. This is # achieved by using an UngroupedShuffleSource to read back the values # written in 'ungrouped' mode. if self.is_ungrouped: # We want to spread the values uniformly to all shufflers. k, v = str(random.getrandbits(64)), o.value else: k, v = o.value # TODO(silviuc): Use timestamps for the secondary key to get values in # times-sorted order. self.writer.Write(k, '', v) self.receivers[0].update_counters_finish() class _TaggedReceivers(dict): class NullReceiver(object): def output(self, element): pass def __missing__(self, unused_key): if not getattr(self, '_null_receiver', None): self._null_receiver = _TaggedReceivers.NullReceiver() return self._null_receiver class DoOperation(Operation): """A Do operation that will execute a custom DoFn for each input element.""" def __init__(self, spec, counter_factory): super(DoOperation, self).__init__(spec, counter_factory) self.state = common.DoFnState(counter_factory) def _read_side_inputs(self, tags_and_types): """Generator reading side inputs in the order prescribed by tags_and_types. Args: tags_and_types: List of tuples (tag, type). Each side input has a string tag that is specified in the worker instruction. The type is actually a boolean which is True for singleton input (read just first value) and False for collection input (read all values). Yields: With each iteration it yields the result of reading an entire side source either in singleton or collection mode according to the tags_and_types argument. """ # We will read the side inputs in the order prescribed by the # tags_and_types argument because this is exactly the order needed to # replace the ArgumentPlaceholder objects in the args/kwargs of the DoFn # getting the side inputs. # # Note that for each tag there could be several read operations in the # specification. This can happen for instance if the source has been # sharded into several files. for side_tag, view_class, view_options in tags_and_types: sources = [] # Using the side_tag in the lambda below will trigger a pylint warning. # However in this case it is fine because the lambda is used right away # while the variable has the value assigned by the current iteration of # the for loop. # pylint: disable=cell-var-from-loop for si in itertools.ifilter( lambda o: o.tag == side_tag, self.spec.side_inputs): if not isinstance(si, maptask.WorkerSideInputSource): raise NotImplementedError('Unknown side input type: %r' % si) sources.append(si.source) iterator_fn = sideinputs.get_iterator_fn_for_sources(sources) if view_class == pvalue.SingletonPCollectionView: has_default, default = view_options has_result = False result = None for v in iterator_fn(): has_result = True result = v break if has_result: yield result elif has_default: yield default else: yield EmptySideInput() elif view_class == pvalue.IterablePCollectionView: yield sideinputs.EmulatedIterable(iterator_fn) elif view_class == pvalue.ListPCollectionView: # TODO(ccy): this is not yet suitable for lists that do not fit in # memory on a single machine. yield list(iterator_fn()) elif view_class == pvalue.DictPCollectionView: # TODO(ccy): this is not yet suitable for dictionaries that do not fit # in memory on a single machine. yield dict(iterator_fn()) else: raise NotImplementedError('Unknown PCollectionView type: %s' % view_class) def start(self): super(DoOperation, self).start() # See fn_data in dataflow_runner.py fn, args, kwargs, tags_and_types, window_fn = ( pickler.loads(self.spec.serialized_fn)) self.state.step_name = self.step_name # TODO(silviuc): What is the proper label here? PCollection being processed? self.context = ptransform.DoFnProcessContext('label', state=self.state) # Tag to output index map used to dispatch the side output values emitted # by the DoFn function to the appropriate receivers. The main output is # tagged with None and is associated with its corresponding index. tagged_receivers = _TaggedReceivers() output_tag_prefix = PropertyNames.OUT + '_' for index, tag in enumerate(self.spec.output_tags): if tag == PropertyNames.OUT: original_tag = None elif tag.startswith(output_tag_prefix): original_tag = tag[len(output_tag_prefix):] else: raise ValueError('Unexpected output name for operation: %s' % tag) tagged_receivers[original_tag] = self.receivers[index] self.dofn_runner = common.DoFnRunner( fn, args, kwargs, self._read_side_inputs(tags_and_types), window_fn, self.context, tagged_receivers, logger, self.step_name) self.dofn_runner.start() def finish(self): self.dofn_runner.finish() def process(self, o): self.dofn_runner.process(o) class CombineOperation(Operation): """A Combine operation executing a CombineFn for each input element.""" def __init__(self, spec, counter_factory): super(CombineOperation, self).__init__(spec, counter_factory) # Combiners do not accept deferred side-inputs (the ignored fourth argument) # and therefore the code to handle the extra args/kwargs is simpler than for # the DoFn's of ParDo. fn, args, kwargs = pickler.loads(self.spec.serialized_fn)[:3] self.phased_combine_fn = ( PhasedCombineFnExecutor(self.spec.phase, fn, args, kwargs)) def finish(self): logging.debug('Finishing %s', self) def process(self, o): if self.debug_logging_enabled: logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) key, values = o.value self.output( o.with_value((key, self.phased_combine_fn.apply(values)))) def create_pgbk_op(spec, counter_factory): if spec.combine_fn: return PGBKCVOperation(spec, counter_factory) else: return PGBKOperation(spec, counter_factory) class PGBKOperation(Operation): """Partial group-by-key operation. This takes (windowed) input (key, value) tuples and outputs (key, [value]) tuples, performing a best effort group-by-key for values in this bundle, memory permitting. """ def __init__(self, spec, counter_factory): super(PGBKOperation, self).__init__(spec, counter_factory) assert not self.spec.combine_fn self.table = collections.defaultdict(list) self.size = 0 # TODO(robertwb) Make this configurable. self.max_size = 10000 def process(self, o): # TODO(robertwb): Structural (hashable) values. key = o.value[0], tuple(o.windows) self.table[key].append(o) self.size += 1 if self.size > self.max_size: self.flush(9 * self.max_size // 10) def finish(self): self.flush(0) def flush(self, target): limit = self.size - target for ix, (kw, vs) in enumerate(self.table.items()): if ix >= limit: break del self.table[kw] key, windows = kw output_value = [v.value[1] for v in vs] windowed_value = WindowedValue( (key, output_value), vs[0].timestamp, windows) self.output(windowed_value) class PGBKCVOperation(Operation): def __init__(self, spec, counter_factory): super(PGBKCVOperation, self).__init__(spec, counter_factory) # Combiners do not accept deferred side-inputs (the ignored fourth # argument) and therefore the code to handle the extra args/kwargs is # simpler than for the DoFn's of ParDo. fn, args, kwargs = pickler.loads(self.spec.combine_fn)[:3] self.combine_fn = curry_combine_fn(fn, args, kwargs) # Optimization for the (known tiny accumulator, often wide keyspace) # count function. # TODO(robertwb): Bound by in-memory size rather than key count. self.max_keys = ( 1000000 if isinstance(fn, combiners.CountCombineFn) else 10000) self.key_count = 0 self.table = {} def process(self, wkv): key, value = wkv.value wkey = tuple(wkv.windows), key entry = self.table.get(wkey, None) if entry is None: if self.key_count >= self.max_keys: target = self.key_count * 9 // 10 old_wkeys = [] # TODO(robertwb): Use an LRU cache? for old_wkey, old_wvalue in self.table.iteritems(): old_wkeys.append(old_wkey) # Can't mutate while iterating. self.output_key(old_wkey, old_wvalue[0]) self.key_count -= 1 if self.key_count <= target: break for old_wkey in reversed(old_wkeys): del self.table[old_wkey] self.key_count += 1 # We save the accumulator as a one element list so we can efficiently # mutate when new values are added without searching the cache again. entry = self.table[wkey] = [self.combine_fn.create_accumulator()] entry[0] = self.combine_fn.add_inputs(entry[0], [value]) def finish(self): for wkey, value in self.table.iteritems(): self.output_key(wkey, value[0]) self.table = {} self.key_count = 0 def output_key(self, wkey, value): windows, key = wkey self.output(WindowedValue((key, value), windows[0].end, windows)) class FlattenOperation(Operation): """Flatten operation. Receives one or more producer operations, outputs just one list with all the items. """ def process(self, o): if self.debug_logging_enabled: logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) self.output(o) class ReifyTimestampAndWindowsOperation(Operation): """ReifyTimestampAndWindows operation. Maps each input KV item into a tuple of the original key and the value as the WindowedValue object of the original value and the original KV item's timestamp and windows. """ def __init__(self, spec, counter_factory): super(ReifyTimestampAndWindowsOperation, self).__init__( spec, counter_factory) def process(self, o): if self.debug_logging_enabled: logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) k, v = o.value self.output(o.with_value((k, o.with_value(v)))) class BatchGroupAlsoByWindowsOperation(Operation): """BatchGroupAlsoByWindowsOperation operation. Implements GroupAlsoByWindow for batch pipelines. """ def __init__(self, spec, counter_factory): super(BatchGroupAlsoByWindowsOperation, self).__init__( spec, counter_factory) self.windowing = pickler.loads(self.spec.window_fn) if self.spec.combine_fn: # Combiners do not accept deferred side-inputs (the ignored fourth # argument) and therefore the code to handle the extra args/kwargs is # simpler than for the DoFn's of ParDo. fn, args, kwargs = pickler.loads(self.spec.combine_fn)[:3] self.phased_combine_fn = ( PhasedCombineFnExecutor(self.spec.phase, fn, args, kwargs)) else: self.phased_combine_fn = None def process(self, o): """Process a given value.""" if self.debug_logging_enabled: logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) k, vs = o.value driver = trigger.create_trigger_driver( self.windowing, is_batch=True, phased_combine_fn=self.phased_combine_fn) state = InMemoryUnmergedState() # TODO(robertwb): Process in smaller chunks. for wvalue in driver.process_elements(state, vs, MIN_TIMESTAMP): self.output(wvalue.with_value((k, wvalue.value))) while state.timers: timers = state.get_and_clear_timers() for timer_window, (name, time_domain, timestamp) in timers: for wvalue in driver.process_timer( timer_window, name, time_domain, timestamp, state): self.output(wvalue.with_value((k, wvalue.value))) class StreamingGroupAlsoByWindowsOperation(Operation): """StreamingGroupAlsoByWindowsOperation operation. Implements GroupAlsoByWindow for streaming pipelines. """ def __init__(self, spec, counter_factory): super(StreamingGroupAlsoByWindowsOperation, self).__init__( spec, counter_factory) self.windowing = pickler.loads(self.spec.window_fn) def process(self, o): if self.debug_logging_enabled: logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) keyed_work = o.value driver = trigger.create_trigger_driver(self.windowing) state = self.spec.context.state output_watermark = self.spec.context.output_data_watermark key = keyed_work.key for wvalue in driver.process_elements( state, keyed_work.elements(), output_watermark): self.output(wvalue.with_value((key, wvalue.value))) for timer in keyed_work.timers(): timer_window = int(timer.namespace) for wvalue in driver.process_timer( timer_window, timer.name, timer.time_domain, timer.timestamp, state): self.output(wvalue.with_value((key, wvalue.value))) class Executor(object): """An abstract executor for WorkItem tasks.""" def execute(self): raise NotImplementedError def get_progress(self): pass def request_dynamic_split(self, dynamic_split_request): # pylint:disable=unused-argument pass class MapTaskExecutor(Executor): """An executor for map tasks. Stores progress of the read operation that is the first operation of a map task. """ multiple_read_instruction_error_msg = ( 'Found more than one \'read instruction\' in a single \'map task\'') def __init__( self, map_task, test_shuffle_source=None, test_shuffle_sink=None): """Initializes MapTaskExecutor. Args: map_task: The map task we are to run. test_shuffle_source: Used during tests for dependency injection into shuffle read operation objects. test_shuffle_sink: Used during tests for dependency injection into shuffle write operation objects. """ self._ops = [] self._read_operation = None self._test_shuffle_source = test_shuffle_source self._test_shuffle_sink = test_shuffle_sink self._map_task = map_task def get_progress(self): return (self._read_operation.get_progress() if self._read_operation is not None else None) def request_dynamic_split(self, dynamic_split_request): if self._read_operation is not None: return self._read_operation.request_dynamic_split(dynamic_split_request) def execute(self): """Executes all the maptask.Worker* instructions in a map task. We update the map_task with the execution status, expressed as counters. Raises: RuntimeError: if we find more than on read instruction in task spec. TypeError: if the spec parameter is not an instance of the recognized maptask.Worker* classes. """ # operations is a list of maptask.Worker* instances. The order of the # elements is important because the inputs use list indexes as references. for spec in self._map_task.operations: if isinstance(spec, maptask.WorkerRead): op = ReadOperation(spec, self._map_task.counter_factory) if self._read_operation is not None: raise RuntimeError( MapTaskExecutor.multiple_read_instruction_error_msg) else: self._read_operation = op elif isinstance(spec, maptask.WorkerWrite): op = WriteOperation(spec, self._map_task.counter_factory) elif isinstance(spec, maptask.WorkerCombineFn): op = CombineOperation(spec, self._map_task.counter_factory) elif isinstance(spec, maptask.WorkerPartialGroupByKey): op = create_pgbk_op(spec, self._map_task.counter_factory) elif isinstance(spec, maptask.WorkerDoFn): op = DoOperation(spec, self._map_task.counter_factory) elif isinstance(spec, maptask.WorkerGroupingShuffleRead): op = GroupedShuffleReadOperation( spec, self._map_task.counter_factory, shuffle_source=self._test_shuffle_source) if self._read_operation is not None: raise RuntimeError( MapTaskExecutor.multiple_read_instruction_error_msg) else: self._read_operation = op elif isinstance(spec, maptask.WorkerUngroupedShuffleRead): op = UngroupedShuffleReadOperation( spec, self._map_task.counter_factory, shuffle_source=self._test_shuffle_source) if self._read_operation is not None: raise RuntimeError( MapTaskExecutor.multiple_read_instruction_error_msg) else: self._read_operation = op elif isinstance(spec, maptask.WorkerInMemoryWrite): op = InMemoryWriteOperation(spec, self._map_task.counter_factory) elif isinstance(spec, maptask.WorkerShuffleWrite): op = ShuffleWriteOperation( spec, self._map_task.counter_factory, shuffle_sink=self._test_shuffle_sink) elif isinstance(spec, maptask.WorkerFlatten): op = FlattenOperation(spec, self._map_task.counter_factory) elif isinstance(spec, maptask.WorkerMergeWindows): if isinstance(spec.context, maptask.BatchExecutionContext): op = BatchGroupAlsoByWindowsOperation( spec, self._map_task.counter_factory) elif isinstance(spec.context, maptask.StreamingExecutionContext): op = StreamingGroupAlsoByWindowsOperation( spec, self._map_task.counter_factory) else: raise RuntimeError('Unknown execution context: %s' % spec.context) elif isinstance(spec, maptask.WorkerReifyTimestampAndWindows): op = ReifyTimestampAndWindowsOperation( spec, self._map_task.counter_factory) else: raise TypeError('Expected an instance of maptask.Worker* class ' 'instead of %s' % (spec,)) self._ops.append(op) # Every MapTask must start with a read instruction. assert self._read_operation is not None # Add receiver operations to the appropriate producers. if hasattr(op.spec, 'input'): producer, output_index = op.spec.input self._ops[producer].add_receiver(op, output_index) # Flatten has 'inputs', not 'input' if hasattr(op.spec, 'inputs'): for producer, output_index in op.spec.inputs: self._ops[producer].add_receiver(op, output_index) # Inject the step names into the operations. # This is used for logging and assigning names to counters. if self._map_task.step_names is not None: for ix, op in enumerate(self._ops): op.step_name = self._map_task.step_names[ix] ix = len(self._ops) for op in reversed(self._ops): ix -= 1 logging.debug('Starting op %d %s', ix, op) op.start() for op in self._ops: op.finish(*()) class CustomSourceSplitExecutor(Executor): """An executor for custom source split requests.""" def __init__(self, split_task): self.response = None self._split_task = split_task def execute(self): self.response = self._perform_source_split(self._split_task) def _perform_source_split(self, source_operation_split_task): """Splits a source into a set of bundles.""" source = source_operation_split_task.source splits = [] for split in source.split( source_operation_split_task.desired_bundle_size_bytes): assert isinstance(split, iobase.SourceBundle) assert isinstance(split.weight, int) or isinstance(split.weight, float) assert isinstance(split.source, iobase.BoundedSource) splits.append(split) return apiclient.splits_to_split_response(splits) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/executor_test.py000066400000000000000000000637161272340251000257770ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for work item executor functionality.""" import logging import tempfile import unittest import mock from google.cloud.dataflow import coders from google.cloud.dataflow import pvalue from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.internal import util from google.cloud.dataflow.io import bigquery from google.cloud.dataflow.io import fileio import google.cloud.dataflow.transforms as ptransform from google.cloud.dataflow.transforms import core from google.cloud.dataflow.transforms import window from google.cloud.dataflow.worker import executor from google.cloud.dataflow.worker import inmemory from google.cloud.dataflow.worker import maptask from google.cloud.dataflow.worker import workitem def pickle_with_side_inputs(fn, tag_and_type=None): tags_and_types = [] args = [] if tag_and_type is not None: args.append(util.ArgumentPlaceholder()) tags_and_types.append(tag_and_type) return pickler.dumps((fn, args, {}, tags_and_types, core.Windowing(window.GlobalWindows()))) def get_bigquery_source_coder(): return bigquery.RowAsDictJsonCoder def make_map_task(operation_list): """Make a skeletal MapTask that is good enough to pass to the executor. Args: operation_list: the work operations to embed. Returns: A maptask.MapTask object. """ return maptask.MapTask(operation_list, '', ['step-%d' % n for n in xrange(len(operation_list))]) def make_text_sink(output_path, input, coder=coders.ToStringCoder()): return maptask.WorkerWrite( fileio.NativeTextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coder), input=input, output_coders=(coder,)) class DoFnUsingStartBundle(ptransform.DoFn): """A DoFn class defining start_bundle, finish_bundle, and process methods.""" def __init__(self, finish_path): self.finish_path = finish_path def start_bundle(self, context, *args, **kwargs): context.state = 'XYZ' def process(self, context, *args, **kwargs): return ['%s: %s' % (context.state, context.element)] def finish_bundle(self, context, *args, **kwargs): """Writes a string to a file as a side effect to be checked by tests.""" with open(self.finish_path, 'w') as f: f.write('finish called.') class DoFnUsingWithUndeclaredSideOutput(ptransform.DoFn): """A DoFn class that emits an undeclared side output.""" def process(self, context, *args, **kwargs): yield pvalue.SideOutputValue('undeclared', context.element) yield context.element class ProgressRequestRecordingInMemoryReader(inmemory.InMemoryReader): def __init__(self, source): self.progress_record = [] super(ProgressRequestRecordingInMemoryReader, self).__init__(source) def get_progress(self): next_progress = super(ProgressRequestRecordingInMemoryReader, self).get_progress() self.progress_record.append(next_progress.position.record_index) return next_progress class ProgressRequestRecordingInMemorySource(inmemory.InMemorySource): def reader(self): self.last_reader = ProgressRequestRecordingInMemoryReader(self) return self.last_reader class ExecutorTest(unittest.TestCase): SHUFFLE_CODER = coders.PickleCoder() OUTPUT_CODER = coders.PickleCoder() def create_temp_file(self, content_text): """Creates a temporary file with content and returns the path to it.""" temp = tempfile.NamedTemporaryFile(delete=False) with temp.file as tmp: tmp.write(content_text) return temp.name def test_read_do_write(self): input_path = self.create_temp_file('01234567890123456789\n0123456789') output_path = '%s.out' % input_path work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( fileio.TextFileSource(file_path=input_path, start_offset=0, end_offset=15, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: ['XYZ: %s' % x])), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), make_text_sink(output_path, input=(1, 0)) ]) executor.MapTaskExecutor(work_item.map_task).execute() with open(output_path) as f: self.assertEqual('XYZ: 01234567890123456789\n', f.read()) def test_read_do_write_with_start_bundle(self): input_path = self.create_temp_file('01234567890123456789\n0123456789') output_path = '%s.out' % input_path finish_path = '%s.finish' % input_path work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( fileio.TextFileSource(file_path=input_path, start_offset=0, end_offset=15, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( DoFnUsingStartBundle(finish_path)), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), make_text_sink(output_path, input=(1, 0)) ]) executor.MapTaskExecutor(work_item.map_task).execute() with open(output_path) as f: self.assertEqual('XYZ: 01234567890123456789\n', f.read()) # Check that the finish_bundle method of the custom DoFn object left the # expected side-effect by writing a file with a specific content. with open(finish_path) as f: self.assertEqual('finish called.', f.read()) def test_read_do_write_with_undeclared_output(self): input_path = self.create_temp_file('01234567890123456789\n0123456789') output_path = '%s.out' % input_path work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( fileio.TextFileSource(file_path=input_path, start_offset=0, end_offset=15, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( DoFnUsingWithUndeclaredSideOutput()), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), make_text_sink(output_path, input=(1, 0)) ]) executor.MapTaskExecutor(work_item.map_task).execute() with open(output_path) as f: self.assertEqual('01234567890123456789\n', f.read()) def test_read_do_shuffle_write(self): input_path = self.create_temp_file('a\nb\nc\nd\n') work_spec = [ maptask.WorkerRead( fileio.TextFileSource(file_path=input_path, start_offset=0, end_offset=8, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: [(x, 1)])), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), maptask.WorkerShuffleWrite(shuffle_kind='group_keys', shuffle_writer_config='none', input=(1, 0), output_coders=(self.SHUFFLE_CODER,)) ] shuffle_sink_mock = mock.MagicMock() work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task(work_spec) executor.MapTaskExecutor(work_item.map_task, test_shuffle_sink=shuffle_sink_mock).execute() # Make sure we have seen all the (k, v) writes. shuffle_sink_mock.writer().Write.assert_has_calls( [mock.call('a', '', 1), mock.call('b', '', 1), mock.call('c', '', 1), mock.call('d', '', 1)]) def test_shuffle_read_do_write(self): output_path = self.create_temp_file('n/a') work_spec = [ maptask.WorkerGroupingShuffleRead(shuffle_reader_config='none', start_shuffle_position='aaa', end_shuffle_position='zzz', coder=self.SHUFFLE_CODER, output_coders=[self.SHUFFLE_CODER]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda (k, vs): [str((k, v)) for v in vs])), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), make_text_sink(output_path, input=(1, 0)) ] shuffle_source_mock = mock.MagicMock() shuffle_source_mock.reader().__enter__().__iter__.return_value = [ (10, [1, 2]), (20, [3])] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task(work_spec) executor.MapTaskExecutor(work_item.map_task, test_shuffle_source=shuffle_source_mock).execute() with open(output_path) as f: self.assertEqual('(10, 1)\n(10, 2)\n(20, 3)\n', f.read()) def test_ungrouped_shuffle_read_and_write(self): output_path = self.create_temp_file('n/a') work_spec = [ maptask.WorkerUngroupedShuffleRead(shuffle_reader_config='none', start_shuffle_position='aaa', end_shuffle_position='zzz', coder=self.SHUFFLE_CODER, output_coders=[self.SHUFFLE_CODER]), make_text_sink(output_path, input=(0, 0)) ] shuffle_source_mock = mock.MagicMock() shuffle_source_mock.reader().__enter__().__iter__.return_value = [1, 2, 3] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task(work_spec) executor.MapTaskExecutor(work_item.map_task, test_shuffle_source=shuffle_source_mock).execute() with open(output_path) as f: self.assertEqual('1\n2\n3\n', f.read()) def test_create_do_write(self): output_path = self.create_temp_file('n/a') elements = ['abc', 'def', 'ghi'] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], # Start at the last element. start_index=2, # Go beyond the end to test that case is handled. end_index=15), output_coders=[coders.ToStringCoder()]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: ['XYZ: %s' % x])), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), make_text_sink(output_path, input=(1, 0)) ]) executor.MapTaskExecutor(work_item.map_task).execute() with open(output_path) as f: self.assertEqual('XYZ: ghi\n', f.read()) def test_create_do_avro_write(self): output_path = self.create_temp_file('n/a') elements = ['abc', 'def', 'ghi'] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=2, # Start at the last element. end_index=3), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: ['XYZ: %s' % x])), output_tags=['out'], input=(0, 0), side_inputs=None, output_coders=[self.OUTPUT_CODER]), make_text_sink( output_path, input=(1, 0), coder=coders.Base64PickleCoder())]) executor.MapTaskExecutor(work_item.map_task).execute() with open(output_path) as f: self.assertEqual('XYZ: ghi', pickler.loads(f.read().strip())) def test_create_do_with_side_in_memory_write(self): elements = ['abc', 'def', 'ghi'] side_elements = ['x', 'y', 'z'] output_buffer = [] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=3), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda x, side: ['%s:%s' % (x, side)]), tag_and_type=('inmemory', pvalue.SingletonPCollectionView, (False, None))), output_tags=['out'], input=(0, 0), side_inputs=[ maptask.WorkerSideInputSource( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in side_elements], start_index=None, end_index=None), tag='inmemory')], output_coders=[self.OUTPUT_CODER]), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0), output_coders=(self.OUTPUT_CODER,))]) executor.MapTaskExecutor(work_item.map_task).execute() # The side source was specified as singleton therefore we should see # only the first element appended. self.assertEqual(['abc:x', 'def:x', 'ghi:x'], output_buffer) def test_in_memory_source_progress_reporting(self): elements = [101, 201, 301, 401, 501, 601, 701] output_buffer = [] source = ProgressRequestRecordingInMemorySource( elements=[pickler.dumps(e) for e in elements]) work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead(source, output_coders=[self.OUTPUT_CODER]), maptask.WorkerInMemoryWrite(output_buffer=output_buffer, input=(0, 0), output_coders=(self.OUTPUT_CODER,)) ]) executor.MapTaskExecutor(work_item.map_task).execute() self.assertEqual(elements, output_buffer) expected_progress_record = range(len(elements)) self.assertEqual(expected_progress_record, source.last_reader.progress_record) def test_create_do_with_side_text_file_write(self): input_path = self.create_temp_file('x\ny\n') elements = ['aa', 'bb'] output_buffer = [] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=2), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda x, side: ['%s:%s' % (x, s) for s in side]), tag_and_type=( 'textfile', pvalue.IterablePCollectionView, ())), output_tags=['out'], input=(0, 0), side_inputs=[ maptask.WorkerSideInputSource(fileio.TextFileSource( file_path=input_path, start_offset=None, end_offset=None, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), tag='textfile')], output_coders=[self.OUTPUT_CODER]), maptask.WorkerInMemoryWrite(output_buffer=output_buffer, input=(1, 0), output_coders=(self.OUTPUT_CODER,))]) executor.MapTaskExecutor(work_item.map_task).execute() # The side source was specified as collection therefore we should see # all elements of the side source. self.assertEqual([u'aa:x', u'aa:y', u'bb:x', u'bb:y'], sorted(output_buffer)) def test_create_do_with_singleton_side_bigquery_write(self): elements = ['abc', 'def', 'ghi'] side_elements = ['x', 'y', 'z'] output_buffer = [] patch_target = 'google.cloud.dataflow.io.bigquery.BigQueryReader' with mock.patch(target=patch_target) as mock_class: # Setup the reader so it will yield the values in 'side_elements'. reader_mock = mock_class.return_value reader_mock.__enter__.return_value = reader_mock # Use a lambda so that multiple readers can be created, each reading the # entirety of the side elements. reader_mock.__iter__.side_effect = lambda: (x for x in side_elements) pickled_elements = [pickler.dumps(e) for e in elements] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource(elements=pickled_elements, start_index=0, end_index=3), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda x, side: ['%s:%s' % (x, side)]), tag_and_type=('bigquery', pvalue.SingletonPCollectionView, (False, None))), output_tags=['out'], input=(0, 0), side_inputs=[ maptask.WorkerSideInputSource( bigquery.BigQuerySource( project='project', dataset='dataset', table='table', coder=get_bigquery_source_coder()), tag='bigquery')], output_coders=[self.OUTPUT_CODER]), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0), output_coders=(self.OUTPUT_CODER,))]) executor.MapTaskExecutor(work_item.map_task).execute() # The side source was specified as singleton therefore we should see # only the first element appended. self.assertEqual(['abc:x', 'def:x', 'ghi:x'], output_buffer) def test_create_do_with_collection_side_bigquery_write(self): elements = ['aa', 'bb'] side_elements = ['x', 'y'] output_buffer = [] patch_target = 'google.cloud.dataflow.io.bigquery.BigQueryReader' with mock.patch(target=patch_target) as mock_class: # Setup the reader so it will yield the values in 'side_elements'. reader_mock = mock_class.return_value reader_mock.__enter__.return_value = reader_mock # Use a lambda so that multiple readers can be created, each reading the # entirety of the side elements. reader_mock.__iter__.side_effect = lambda: (x for x in side_elements) work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=3), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda x, side: ['%s:%s' % (x, s) for s in side]), tag_and_type=('bigquery', pvalue.IterablePCollectionView, ())), output_tags=['out'], input=(0, 0), side_inputs=[ maptask.WorkerSideInputSource( bigquery.BigQuerySource( project='project', dataset='dataset', table='table', coder=get_bigquery_source_coder()), tag='bigquery')], output_coders=[self.OUTPUT_CODER]), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0), output_coders=(self.OUTPUT_CODER,))]) executor.MapTaskExecutor(work_item.map_task).execute() # The side source was specified as collection therefore we should see # all elements of the side source. self.assertEqual(['aa:x', 'aa:y', 'bb:x', 'bb:y'], sorted(output_buffer)) def test_create_do_with_side_avro_file_write(self): input_path1 = self.create_temp_file('%s\n' % pickler.dumps('x')) input_path2 = self.create_temp_file('%s\n' % pickler.dumps('y')) elements = ['aa', 'bb'] output_buffer = [] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=2), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda x, side: ['%s:%s' % (x, s) for s in side]), tag_and_type=( 'sometag', pvalue.IterablePCollectionView, ())), output_tags=['out'], input=(0, 0), # Note that the two side inputs have the same tag. This is quite # common for intermediary PCollections used as side inputs that # are saved as AVRO files. The files will contain the sharded # PCollection. side_inputs=[ maptask.WorkerSideInputSource( fileio.TextFileSource( file_path=input_path1, coder=coders.Base64PickleCoder()), tag='sometag'), maptask.WorkerSideInputSource( fileio.TextFileSource(file_path=input_path2, coder=coders.Base64PickleCoder()), tag='sometag')], output_coders=[self.OUTPUT_CODER]), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0), output_coders=(self.OUTPUT_CODER,))]) executor.MapTaskExecutor(work_item.map_task).execute() # The side source was specified as collection therefore we should see # all three elements of the side source. self.assertEqual([u'aa:x', u'aa:y', u'bb:x', u'bb:y'], sorted(output_buffer)) def test_combine(self): elements = [('a', [1, 2, 3]), ('b', [10])] output_buffer = [] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=100), output_coders=[self.OUTPUT_CODER]), maptask.WorkerCombineFn(serialized_fn=pickle_with_side_inputs( ptransform.CombineFn.from_callable(sum)), phase='all', input=(0, 0), output_coders=[self.OUTPUT_CODER]), maptask.WorkerInMemoryWrite(output_buffer=output_buffer, input=(1, 0), output_coders=(self.OUTPUT_CODER,)) ]) executor.MapTaskExecutor(work_item.map_task).execute() self.assertEqual([('a', 6), ('b', 10)], output_buffer) def test_pgbk(self): elements = [('a', 1), ('b', 2), ('a', 3), ('a', 4)] output_buffer = [] work_item = workitem.BatchWorkItem(None) work_item.map_task = make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=100), output_coders=[self.OUTPUT_CODER]), maptask.WorkerPartialGroupByKey( combine_fn=None, input=(0, 0), output_coders=[self.OUTPUT_CODER]), maptask.WorkerInMemoryWrite(output_buffer=output_buffer, input=(1, 0), output_coders=(self.OUTPUT_CODER,)) ]) executor.MapTaskExecutor(work_item.map_task).execute() self.assertEqual([('a', [1, 3, 4]), ('b', [2])], sorted(output_buffer)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/inmemory.py000066400000000000000000000071601272340251000247300ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """In-memory input source.""" import itertools import logging from google.cloud.dataflow import coders from google.cloud.dataflow.io import iobase from google.cloud.dataflow.io import range_trackers class InMemorySource(iobase.NativeSource): """In-memory input source.""" def __init__( self, elements, coder=coders.Base64PickleCoder(), start_index=None, end_index=None): self.elements = elements self.coder = coder if start_index is None: self.start_index = 0 else: self.start_index = start_index if end_index is None: self.end_index = len(elements) else: self.end_index = end_index def __eq__(self, other): return (self.elements == other.elements and self.coder == other.coder and self.start_index == other.start_index and self.end_index == other.end_index) def reader(self): return InMemoryReader(self) class InMemoryReader(iobase.NativeSourceReader): """A reader for in-memory source.""" def __init__(self, source): self._source = source # Index of the last item returned by InMemoryReader. # Initialized to None. self._current_index = None self._range_tracker = range_trackers.OffsetRangeTracker( self._source.start_index, self._source.end_index) def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): pass def __iter__(self): for value in itertools.islice(self._source.elements, self._source.start_index, self._source.end_index): claimed = False if self._current_index is None: claimed = self._range_tracker.try_claim( self._source.start_index) else: claimed = self._range_tracker.try_claim( self._current_index + 1) if claimed: if self._current_index is None: self._current_index = self._source.start_index else: self._current_index += 1 yield self._source.coder.decode(value) else: return def get_progress(self): if self._current_index is None: return None return iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=self._current_index)) def request_dynamic_split(self, dynamic_split_request): assert dynamic_split_request is not None progress = dynamic_split_request.progress split_position = progress.position if split_position is None: logging.debug('InMemory reader only supports split requests that are ' 'based on positions. Received : %r', dynamic_split_request) return None index_position = split_position.record_index if index_position is None: logging.debug('InMemory reader only supports split requests that are ' 'based on index positions. Received : %r', dynamic_split_request) return None if self._range_tracker.try_split(index_position): return iobase.DynamicSplitResultWithPosition(split_position) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/inmemory_test.py000066400000000000000000000126061272340251000257700ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for in-memory input source.""" import logging import unittest from google.cloud.dataflow.io import iobase from google.cloud.dataflow.worker import inmemory class FakeCoder(object): def decode(self, value): return value + 10 class InMemoryIO(unittest.TestCase): def test_inmemory(self): source = inmemory.InMemorySource([1, 2, 3, 4, 5], FakeCoder(), 1, 3) with source.reader() as reader: self.assertItemsEqual([12, 13], [i for i in reader]) def test_norange(self): source = inmemory.InMemorySource([1, 2, 3, 4, 5], coder=FakeCoder()) with source.reader() as reader: self.assertItemsEqual([11, 12, 13, 14, 15], [i for i in reader]) def test_in_memory_source_updates_progress_none(self): source = inmemory.InMemorySource([], coder=FakeCoder()) with source.reader() as reader: self.assertEqual(None, reader.get_progress()) def test_in_memory_source_updates_progress_one(self): source = inmemory.InMemorySource([1], coder=FakeCoder()) with source.reader() as reader: self.assertEqual(None, reader.get_progress()) i = 0 for item in reader: self.assertEqual(i, reader.get_progress().position.record_index) self.assertEqual(11, item) i += 1 self.assertEqual(1, i) self.assertEqual(0, reader.get_progress().position.record_index) def test_in_memory_source_updates_progress_many(self): source = inmemory.InMemorySource([1, 2, 3, 4, 5], coder=FakeCoder()) with source.reader() as reader: self.assertEqual(None, reader.get_progress()) i = 0 for item in reader: self.assertEqual(i, reader.get_progress().position.record_index) self.assertEqual(11 + i, item) i += 1 self.assertEqual(5, i) self.assertEqual(4, reader.get_progress().position.record_index) def try_splitting_reader_at(self, reader, split_request, expected_response): actual_response = reader.request_dynamic_split(split_request) if expected_response is None: self.assertIsNone(actual_response) else: self.assertIsNotNone(actual_response.stop_position) self.assertIsInstance(actual_response.stop_position, iobase.ReaderPosition) self.assertIsNotNone(actual_response.stop_position.record_index) self.assertEqual(expected_response.stop_position.record_index, actual_response.stop_position.record_index) def test_in_memory_source_dynamic_split(self): source = inmemory.InMemorySource([10, 20, 30, 40, 50, 60], coder=FakeCoder()) # Unstarted reader with source.reader() as reader: self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest( iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=2))), None) # Proposed split position out of range with source.reader() as reader: reader_iter = iter(reader) next(reader_iter) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest( iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=-1))), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest( iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=10))), None) # Already read past proposed split position with source.reader() as reader: reader_iter = iter(reader) next(reader_iter) next(reader_iter) next(reader_iter) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest( iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=1))), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest( iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=2))), None) # Successful split with source.reader() as reader: reader_iter = iter(reader) next(reader_iter) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest( iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=4))), iobase.DynamicSplitResultWithPosition( stop_position=iobase.ReaderPosition(record_index=4))) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest( iobase.ReaderProgress( position=iobase.ReaderPosition(record_index=2))), iobase.DynamicSplitResultWithPosition( stop_position=iobase.ReaderPosition(record_index=2))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/logger.py000066400000000000000000000136721272340251000243550ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Python Dataflow worker logging.""" import json import logging import threading import traceback # Per-thread worker information. This is used only for logging to set # context information that changes while work items get executed: # work_item_id, step_name, stage_name. class _PerThreadWorkerData(threading.local): def __init__(self): super(_PerThreadWorkerData, self).__init__() self.stack = [] def get_data(self): all_data = {} for datum in self.stack: all_data.update(datum) return all_data per_thread_worker_data = _PerThreadWorkerData() class PerThreadLoggingContext(object): """A context manager to add per thread attributes.""" def __init__(self, **kwargs): self.kwargs = kwargs def __enter__(self): per_thread_worker_data.stack.append(self.kwargs) def __exit__(self, exn_type, exn_value, exn_traceback): per_thread_worker_data.stack.pop() class JsonLogFormatter(logging.Formatter): """A JSON formatter class as expected by the logging standard module.""" def __init__(self, job_id, worker_id): super(JsonLogFormatter, self).__init__() self.job_id = job_id self.worker_id = worker_id def format(self, record): """Returns a JSON string based on a LogRecord instance. Args: record: A LogRecord instance. See below for details. Returns: A JSON string representing the record. A LogRecord instance has the following attributes and is used for formatting the final message. Attributes: created: A double representing the timestamp for record creation (e.g., 1438365207.624597). Note that the number contains also msecs and microsecs information. Part of this is also available in the 'msecs' attribute. msecs: A double representing the msecs part of the record creation (e.g., 624.5970726013184). msg: Logging message containing formatting instructions or an arbitrary object. This is the first argument of a log call. args: A tuple containing the positional arguments for the logging call. levelname: A string. Possible values are: INFO, WARNING, ERROR, etc. exc_info: None or a 3-tuple with exception information as it is returned by a call to sys.exc_info(). name: Logger's name. Most logging is done using the default root logger and therefore the name will be 'root'. filename: Basename of the file where logging occurred. funcName: Name of the function where logging occurred. process: The PID of the process running the worker. thread: An id for the thread where the record was logged. This is not a real TID (the one provided by OS) but rather the id (address) of a Python thread object. Nevertheless having this value can allow to filter log statement from only one specific thread. """ output = {} output['timestamp'] = { 'seconds': int(record.created), 'nanos': int(record.msecs * 1000000)} # ERROR. INFO, DEBUG log levels translate into the same for severity # property. WARNING becomes WARN. output['severity'] = ( record.levelname if record.levelname != 'WARNING' else 'WARN') # msg could be an arbitrary object, convert it to a string first. record_msg = str(record.msg) # Prepare the actual message using the message formatting string and the # positional arguments as they have been used in the log call. if record.args: try: output['message'] = record_msg % record.args except (TypeError, ValueError): output['message'] = '%s with args (%s)' % (record_msg, record.args) else: output['message'] = record_msg # The thread ID is logged as a combination of the process ID and thread ID # since workers can run in multiple processes. output['thread'] = '%s:%s' % (record.process, record.thread) # job ID and worker ID. These do not change during the lifetime of a worker. output['job'] = self.job_id output['worker'] = self.worker_id # Stage, step and work item ID come from thread local storage since they # change with every new work item leased for execution. If there is no # work item ID then we make sure the step is undefined too. data = per_thread_worker_data.get_data() if 'work_item_id' in data: output['work'] = data['work_item_id'] if 'stage_name' in data: output['stage'] = data['stage_name'] if 'step_name' in data: output['step'] = data['step_name'] # All logging happens using the root logger. We will add the basename of the # file and the function name where the logging happened to make it easier # to identify who generated the record. output['logger'] = '%s:%s:%s' % ( record.name, record.filename, record.funcName) # Add exception information if any is available. if record.exc_info: output['exception'] = ''.join( traceback.format_exception(*record.exc_info)) return json.dumps(output) def initialize(job_id, worker_id, log_path): """Initialize root logger so that we log JSON to a file and text to stdout.""" file_handler = logging.FileHandler(log_path) file_handler.setFormatter(JsonLogFormatter(job_id, worker_id)) logging.getLogger().addHandler(file_handler) # Set default level to INFO to avoid logging various DEBUG level log calls # sprinkled throughout the code. logging.getLogger().setLevel(logging.INFO) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/logger_test.py000066400000000000000000000162231272340251000254070ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for worker logging utilities.""" import json import logging import sys import threading import unittest from google.cloud.dataflow.worker import logger class PerThreadLoggingContextTest(unittest.TestCase): def thread_check_attribute(self, name): self.assertFalse(name in logger.per_thread_worker_data.get_data()) with logger.PerThreadLoggingContext(**{name: 'thread-value'}): self.assertEqual( logger.per_thread_worker_data.get_data()[name], 'thread-value') self.assertFalse(name in logger.per_thread_worker_data.get_data()) def test_no_positional_args(self): with self.assertRaises(TypeError): with logger.PerThreadLoggingContext('something'): pass def test_per_thread_attribute(self): self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') thread = threading.Thread( target=self.thread_check_attribute, args=('xyz',)) thread.start() thread.join() self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) def test_set_when_undefined(self): self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) def test_set_when_already_defined(self): self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') with logger.PerThreadLoggingContext(xyz='value2'): self.assertEqual( logger.per_thread_worker_data.get_data()['xyz'], 'value2') self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) class JsonLogFormatterTest(unittest.TestCase): SAMPLE_RECORD = { 'created': 123456.789, 'msecs': 789.654321, 'msg': '%s:%d:%.2f', 'args': ('xyz', 4, 3.14), 'levelname': 'WARNING', 'process': 'pid', 'thread': 'tid', 'name': 'name', 'filename': 'file', 'funcName': 'func', 'exc_info': None} SAMPLE_OUTPUT = { 'timestamp': {'seconds': 123456, 'nanos': 789654321}, 'severity': 'WARN', 'message': 'xyz:4:3.14', 'thread': 'pid:tid', 'job': 'jobid', 'worker': 'workerid', 'logger': 'name:file:func'} def create_log_record(self, **kwargs): class Record(object): def __init__(self, **kwargs): for k, v in kwargs.iteritems(): setattr(self, k, v) return Record(**kwargs) def test_basic_record(self): formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') record = self.create_log_record(**self.SAMPLE_RECORD) self.assertEqual(json.loads(formatter.format(record)), self.SAMPLE_OUTPUT) def execute_multiple_cases(self, test_cases): record = self.SAMPLE_RECORD output = self.SAMPLE_OUTPUT formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') for case in test_cases: record['msg'] = case['msg'] record['args'] = case['args'] output['message'] = case['expected'] self.assertEqual( json.loads(formatter.format(self.create_log_record(**record))), output) def test_record_with_format_character(self): test_cases = [ {'msg': '%A', 'args': (), 'expected': '%A'}, {'msg': '%s', 'args': (), 'expected': '%s'}, {'msg': '%A%s', 'args': ('xy'), 'expected': '%A%s with args (xy)'}, {'msg': '%s%s', 'args': (1), 'expected': '%s%s with args (1)'}, ] self.execute_multiple_cases(test_cases) def test_record_with_arbitrary_messages(self): test_cases = [ {'msg': ImportError('abc'), 'args': (), 'expected': 'abc'}, {'msg': TypeError('abc %s'), 'args': ('def'), 'expected': 'abc def'}, ] self.execute_multiple_cases(test_cases) def test_record_with_per_thread_info(self): with logger.PerThreadLoggingContext( work_item_id='workitem', stage_name='stage', step_name='step'): formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') record = self.create_log_record(**self.SAMPLE_RECORD) log_output = json.loads(formatter.format(record)) expected_output = dict(self.SAMPLE_OUTPUT) expected_output.update( {'work': 'workitem', 'stage': 'stage', 'step': 'step'}) self.assertEqual(log_output, expected_output) def test_nested_with_per_thread_info(self): formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') with logger.PerThreadLoggingContext( work_item_id='workitem', stage_name='stage', step_name='step1'): record = self.create_log_record(**self.SAMPLE_RECORD) log_output1 = json.loads(formatter.format(record)) with logger.PerThreadLoggingContext(step_name='step2'): record = self.create_log_record(**self.SAMPLE_RECORD) log_output2 = json.loads(formatter.format(record)) record = self.create_log_record(**self.SAMPLE_RECORD) log_output3 = json.loads(formatter.format(record)) record = self.create_log_record(**self.SAMPLE_RECORD) log_output4 = json.loads(formatter.format(record)) self.assertEqual(log_output1, dict( self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1')) self.assertEqual(log_output2, dict( self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step2')) self.assertEqual(log_output3, dict( self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1')) self.assertEqual(log_output4, self.SAMPLE_OUTPUT) def test_exception_record(self): formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') try: raise ValueError('Something') except ValueError: attribs = dict(self.SAMPLE_RECORD) attribs.update({'exc_info': sys.exc_info()}) record = self.create_log_record(**attribs) log_output = json.loads(formatter.format(record)) # Check if exception type, its message, and stack trace information are in. exn_output = log_output.pop('exception') self.assertNotEqual(exn_output.find('ValueError: Something'), -1) self.assertNotEqual(exn_output.find('logger_test.py'), -1) self.assertEqual(log_output, self.SAMPLE_OUTPUT) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/maptask.py000066400000000000000000001047571272340251000245430ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Worker utilities for parsing out a MapTask message. Each MapTask represents a sequence of ParallelInstruction(s): read from a source, write to a sink, parallel do, etc. """ import base64 import collections from google.cloud.dataflow import coders from google.cloud.dataflow import io from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.internal.json_value import from_json_value from google.cloud.dataflow.io import fileio from google.cloud.dataflow.io import iobase from google.cloud.dataflow.utils import names from google.cloud.dataflow.utils.counters import CounterFactory from google.cloud.dataflow.worker import concat_reader from google.cloud.dataflow.worker import inmemory from google.cloud.dataflow.worker import windmillio from google.cloud.dataflow.worker import workercustomsources def build_worker_instruction(*args): """Create an object representing a ParallelInstruction protobuf. This will be a collections.namedtuple with a custom __str__ method. Alas, this wrapper is not known to pylint, which thinks it creates constants. You may have to put a disable=invalid-name pylint annotation on any use of this, depending on your names. Args: *args: first argument is the name of the type to create. Should start with "Worker". Second arguments is alist of the attributes of this object. Returns: A new class, a subclass of tuple, that represents the protobuf. """ tuple_class = collections.namedtuple(*args) tuple_class.__str__ = worker_object_to_string tuple_class.__repr__ = worker_object_to_string return tuple_class def worker_printable_fields(workerproto): """Returns the interesting fields of a Worker* object.""" return ['%s=%s' % (name, value) # _asdict is the only way and cannot subclass this generated class # pylint: disable=protected-access for name, value in workerproto._asdict().iteritems() # want to output value 0 but not None nor [] if (value or value == 0) and name not in ('coder', 'coders', 'output_coders', 'elements', 'combine_fn', 'serialized_fn', 'window_fn', 'append_trailing_newlines', 'strip_trailing_newlines', 'compression_type', 'context', 'start_shuffle_position', 'end_shuffle_position', 'shuffle_reader_config', 'shuffle_writer_config')] def worker_object_to_string(worker_object): """Returns a string compactly representing a Worker* object.""" return '%s(%s)' % (worker_object.__class__.__name__, ', '.join(worker_printable_fields(worker_object))) # All the following Worker* definitions will have these lint problems: # pylint: disable=invalid-name # pylint: disable=pointless-string-statement WorkerRead = build_worker_instruction( 'WorkerRead', ['source', 'output_coders']) """Worker details needed to read from a source. Attributes: source: a source object. output_coders: 1-tuple of the coder for the output. """ WorkerSideInputSource = build_worker_instruction( 'WorkerSideInputSource', ['source', 'tag']) """Worker details needed to read from a side input source. Attributes: source: a source object. tag: string tag for this side input. """ WorkerGroupingShuffleRead = build_worker_instruction( 'WorkerGroupingShuffleRead', ['start_shuffle_position', 'end_shuffle_position', 'shuffle_reader_config', 'coder', 'output_coders']) """Worker details needed to read from a grouping shuffle source. Attributes: start_shuffle_position: An opaque string to be passed to the shuffle source to indicate where to start reading. end_shuffle_position: An opaque string to be passed to the shuffle source to indicate where to stop reading. shuffle_reader_config: An opaque string used to initialize the shuffle reader. Contains things like connection endpoints for the shuffle server appliance and various options. coder: The KV coder used to decode shuffle entries. output_coders: 1-tuple of the coder for the output. """ WorkerUngroupedShuffleRead = build_worker_instruction( 'WorkerUngroupedShuffleRead', ['start_shuffle_position', 'end_shuffle_position', 'shuffle_reader_config', 'coder', 'output_coders']) """Worker details needed to read from an ungrouped shuffle source. Attributes: start_shuffle_position: An opaque string to be passed to the shuffle source to indicate where to start reading. end_shuffle_position: An opaque string to be passed to the shuffle source to indicate where to stop reading. shuffle_reader_config: An opaque string used to initialize the shuffle reader. Contains things like connection endpoints for the shuffle server appliance and various options. coder: The value coder used to decode shuffle entries. """ WorkerWrite = build_worker_instruction( 'WorkerWrite', ['sink', 'input', 'output_coders']) """Worker details needed to write to a sink. Attributes: sink: a sink object. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple, coder to use to estimate bytes written. """ WorkerInMemoryWrite = build_worker_instruction( 'WorkerInMemoryWrite', ['output_buffer', 'input', 'output_coders']) """Worker details needed to write to a in-memory sink. Used only for unit testing. It makes worker tests less cluttered with code like "write to a file and then check file contents". Attributes: output_buffer: list to which output elements will be appended input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple, coder to use to estimate bytes written. """ WorkerShuffleWrite = build_worker_instruction( 'WorkerShuffleWrite', ['shuffle_kind', 'shuffle_writer_config', 'input', 'output_coders']) """Worker details needed to write to a shuffle sink. Attributes: shuffle_kind: A string describing the shuffle kind. This can control the way the worker interacts with the shuffle sink. The possible values are: 'ungrouped', 'group_keys', and 'group_keys_and_sort_values'. shuffle_writer_config: An opaque string used to initialize the shuffle write. Contains things like connection endpoints for the shuffle server appliance and various options. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple of the coder for input elements. If the shuffle_kind is grouping, this is expected to be a KV coder. """ WorkerDoFn = build_worker_instruction( 'WorkerDoFn', ['serialized_fn', 'output_tags', 'input', 'side_inputs', 'output_coders']) """Worker details needed to run a DoFn. Attributes: serialized_fn: A serialized DoFn object to be run for each input element. output_tags: The string tags used to identify the outputs of a ParDo operation. The tag is present even if the ParDo has just one output (e.g., ['out']. output_coders: array of coders, one for each output. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). side_inputs: A list of Worker...Read instances describing sources to be used for getting values. The types supported right now are WorkerInMemoryRead and WorkerTextRead. """ WorkerReifyTimestampAndWindows = build_worker_instruction( 'WorkerReifyTimestampAndWindows', ['output_tags', 'input', 'output_coders']) """Worker details needed to run a WindowInto. Attributes: output_tags: The string tags used to identify the outputs of a ParDo operation. The tag is present even if the ParDo has just one output (e.g., ['out']. output_coders: array of coders, one for each output. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). """ WorkerMergeWindows = build_worker_instruction( 'WorkerMergeWindows', ['window_fn', 'combine_fn', 'phase', 'output_tags', 'input', 'coders', 'context', 'output_coders']) """Worker details needed to run a MergeWindows (aka. GroupAlsoByWindows). Attributes: window_fn: A serialized Windowing object representing the windowing strategy. combine_fn: A serialized CombineFn object to be used after executing the GroupAlsoByWindows operation. May be None if not a combining operation. phase: Possible values are 'all', 'add', 'merge', and 'extract'. The dataflow optimizer may split the user combiner in 3 separate phases (ADD, MERGE, and EXTRACT), on separate VMs, as it sees fit. The phase attribute dictates which DoFn is actually running in the worker. May be None if not a combining operation. output_tags: The string tags used to identify the outputs of a ParDo operation. The tag is present even if the ParDo has just one output (e.g., ['out']. output_coders: array of coders, one for each output. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). coders: A 2-tuple of coders (key, value) to encode shuffle entries. context: The ExecutionContext object for the current work item. """ WorkerCombineFn = build_worker_instruction( 'WorkerCombineFn', ['serialized_fn', 'phase', 'input', 'output_coders']) """Worker details needed to run a CombineFn. Attributes: serialized_fn: A serialized CombineFn object to be used. phase: Possible values are 'all', 'add', 'merge', and 'extract'. The dataflow optimizer may split the user combiner in 3 separate phases (ADD, MERGE, and EXTRACT), on separate VMs, as it sees fit. The phase attribute dictates which DoFn is actually running in the worker. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple of the coder for the output. """ WorkerPartialGroupByKey = build_worker_instruction( 'WorkerPartialGroupByKey', ['combine_fn', 'input', 'output_coders']) """Worker details needed to run a partial group-by-key. Attributes: combine_fn: A serialized CombineFn object to be used. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple of the coder for the output. """ WorkerFlatten = build_worker_instruction( 'WorkerFlatten', ['inputs', 'output_coders']) """Worker details needed to run a Flatten. Attributes: inputs: A list of tuples, each (producer index, output index), representing the ParallelInstruction operations whose output feeds into this operation. The output index is 0 unless the input is from a multi-output operation (such as ParDo). output_coders: 1-tuple of the coder for the output. """ class ExecutionContext(object): """Context for executing the operations within a MapTask.""" pass class BatchExecutionContext(ExecutionContext): """ExecutionContext used in the batch worker.""" pass class StreamingExecutionContext(ExecutionContext): def start(self, computation_id, work_item, input_data_watermark, output_data_watermark, workitem_commit_request, windmill, state): self.computation_id = computation_id self.work_item = work_item self.input_data_watermark = input_data_watermark self.output_data_watermark = output_data_watermark self.workitem_commit_request = workitem_commit_request self.windmill = windmill self.state = state def get_coder_from_spec(coder_spec, kv_pair=False): """Return a coder instance from a coder spec. Args: coder_spec: A dict where the value of the '@type' key is a pickled instance of a Coder instance. kv_pair: True if a 2-tuple of coders (key and value) must be returned. Returns: A coder instance (has encode/decode methods). It is possible to return a 2-tuple of (key coder, value coder) if the spec is for a shuffle source or sink. Such shuffle source and sinks can take a 2-tuple of coders as parameter. Raises: ValueError: if KV coder requested but coder spec is not of a KV coder. """ assert coder_spec is not None # Ignore the wrappers in these encodings. ignored_wrappers = ( 'kind:stream', 'com.google.cloud.dataflow.sdk.util.TimerOrElement$TimerOrElementCoder') if coder_spec['@type'] in ignored_wrappers: assert len(coder_spec['component_encodings']) == 1 coder_spec = coder_spec['component_encodings'][0] return get_coder_from_spec(coder_spec, kv_pair=kv_pair) # We pass coders in the form "$" to make the job # description JSON more readable. coder = coders.deserialize_coder(coder_spec['@type']) # If this is a coder with components potentially modified by the service, # use these components. # # TODO(ccy): This is necessary since the service may move around the # wrapped types of WindowedValueCoders and TupleCoders. We should refactor # coder serialization so these special cases is not necessary. if isinstance(coder, coders.WindowedValueCoder): value_coder, timestamp_coder, window_coder = [ get_coder_from_spec(c) for c in coder_spec['component_encodings']] coder = coders.WindowedValueCoder(value_coder, timestamp_coder, window_coder) elif isinstance(coder, coders.TupleCoder): component_coders = [ get_coder_from_spec(c) for c in coder_spec['component_encodings']] coder = coders.TupleCoder(component_coders) if kv_pair: if not coder.is_kv_coder(): raise ValueError('Coder is not a KV coder: %s.' % coder) return coder.key_coder(), coder.value_coder() else: return coder def get_output_coders(work): """Return a list of coder instances for the output(s) of this work item. Args: work: a ParallelInstruction protobuf Returns: A list of coders. """ return [get_coder_from_spec({p.key: from_json_value(p.value) for p in output.codec.additionalProperties}) for output in work.outputs] def get_custom_source_read_spec(source_spec): source_info = pickler.loads(source_spec['serialized_source']['value']) assert isinstance(source_info, tuple) assert len(source_info) == 3 assert isinstance(source_info[0], iobase.BoundedSource) return WorkerRead( workercustomsources.NativeBoundedSource(source_info[0], source_info[1], source_info[2]), [source_info[0].default_output_coder()]) def get_read_work_item(work, env, context): """Parses a read parallel instruction into the appropriate Worker* object.""" source_spec = {p.key: from_json_value(p.value) for p in work.read.source.spec.additionalProperties} # Only sources for which a custom coder can be specified have the # codec property (e.g. TextSource). codec_specs = None if work.read.source.codec: codec_specs = { p.key: from_json_value(p.value) for p in work.read.source.codec.additionalProperties} source = env.parse_source(source_spec, codec_specs, context) if source: return WorkerRead(source, output_coders=get_output_coders(work)) coder = get_coder_from_spec(codec_specs) # TODO(ccy): Reconcile WindowedValueCoder wrappings for sources with custom # coders so this special case won't be necessary. if isinstance(coder, coders.WindowedValueCoder): coder = coder.wrapped_value_coder if source_spec['@type'] == 'GroupingShuffleSource': return WorkerGroupingShuffleRead( start_shuffle_position=source_spec['start_shuffle_position']['value'], end_shuffle_position=source_spec['end_shuffle_position']['value'], shuffle_reader_config=source_spec['shuffle_reader_config']['value'], coder=coder, output_coders=get_output_coders(work)) elif source_spec['@type'] == 'UngroupedShuffleSource': return WorkerUngroupedShuffleRead( start_shuffle_position=source_spec['start_shuffle_position']['value'], end_shuffle_position=source_spec['end_shuffle_position']['value'], shuffle_reader_config=source_spec['shuffle_reader_config']['value'], coder=coder, output_coders=get_output_coders(work)) elif source_spec['@type'] == names.SOURCE_TYPE: return get_custom_source_read_spec(source_spec) else: raise NotImplementedError('Unknown source type: %r' % source_spec) # pylint: enable=invalid-name # pylint: enable=pointless-string-statement def get_input_spec(instruction_input): """Returns a (producer, output) indexes tuple based on input specified. Args: instruction_input: an InstructionInput protobuf. Returns: A (producer, output) indexes tuple. """ producer_index = instruction_input.producerInstructionIndex if producer_index is None: producer_index = 0 output_index = instruction_input.outputNum return producer_index, 0 if output_index is None else output_index def get_side_input_sources(side_inputs_spec, env, context): """Returns a list of Worker...Read objects for the side sources specified.""" side_inputs = [] for side_spec in side_inputs_spec: assert side_spec.tag # All side input sources have tags. # Make sure we got a side input type we understand. specs = {p.key: from_json_value(p.value) for p in side_spec.kind.additionalProperties} assert specs['@type'] == 'collection' for source in side_spec.sources: source_spec = { p.key: from_json_value(p.value) for p in source.spec.additionalProperties} # Only sources for which a custom coder can be specified have the # codec property (e.g. TextSource). if source.codec: source_codec_spec = { p.key: from_json_value(p.value) for p in source.codec.additionalProperties} parsed_source = env.parse_source(source_spec, source_codec_spec, context) if parsed_source: side_inputs.append(WorkerSideInputSource(parsed_source, side_spec.tag)) else: raise NotImplementedError( 'Unknown side input source type: %r' % source_spec) return side_inputs def get_write_work_item(work, env, context): """Parses a write parallel instruction into the appropriate Worker* object.""" specs = {p.key: from_json_value(p.value) for p in work.write.sink.spec.additionalProperties} # Only sinks for which a custom coder can be specified have the # codec property (e.g. TextSink. codec_specs = None if work.write.sink.codec: codec_specs = { p.key: from_json_value(p.value) for p in work.write.sink.codec.additionalProperties} sink = env.parse_sink(specs, codec_specs, context) if sink: write_coder = get_coder_from_spec(codec_specs) # All Worker items have an "output_coders", even if they have no # output, so that the executor can estimate bytes in a uniform way. return WorkerWrite(sink, input=get_input_spec(work.write.input), output_coders=(write_coder,)) if specs['@type'] == 'ShuffleSink': coder = get_coder_from_spec(codec_specs) # TODO(ccy): Reconcile WindowedValueCoder wrappings for sources with custom # coders so this special case won't be necessary. if isinstance(coder, coders.WindowedValueCoder): coder = coder.wrapped_value_coder return WorkerShuffleWrite( shuffle_kind=specs['shuffle_kind']['value'], shuffle_writer_config=specs['shuffle_writer_config']['value'], input=get_input_spec(work.write.input), output_coders=(coder,)) else: raise NotImplementedError('Unknown sink type: %r' % specs) def get_do_work_item(work, env, context): """Parses a do parallel instruction into the appropriate Worker* object.""" # Get side inputs if any. side_inputs = [] if hasattr(work.parDo, 'sideInputs'): side_inputs = get_side_input_sources(work.parDo.sideInputs, env, context) specs = {p.key: from_json_value(p.value) for p in work.parDo.userFn.additionalProperties} if specs['@type'] == 'DoFn': return WorkerDoFn( serialized_fn=specs['serialized_fn']['value'], output_tags=[o.tag for o in work.parDo.multiOutputInfos], output_coders=get_output_coders(work), input=get_input_spec(work.parDo.input), side_inputs=side_inputs) elif specs['@type'] == 'CombineValuesFn': # Note: CombineFn's do not take side inputs like DoFn's so far. return WorkerCombineFn( serialized_fn=specs['serialized_fn']['value'], phase=specs['phase']['value'], # 'add' is one possible value. input=get_input_spec(work.parDo.input), output_coders=get_output_coders(work)) elif specs['@type'] == 'ReifyTimestampAndWindowsDoFn': return WorkerReifyTimestampAndWindows( output_tags=[o.tag for o in work.parDo.multiOutputInfos], output_coders=get_output_coders(work), input=get_input_spec(work.parDo.input)) elif specs['@type'] == 'MergeBucketsDoFn': return WorkerMergeWindows( window_fn=specs['serialized_fn']['value'], combine_fn=specs.get('combine_fn', {}).get('value', None), phase=specs.get('phase', {}).get('value', None), output_tags=[o.tag for o in work.parDo.multiOutputInfos], output_coders=get_output_coders(work), input=get_input_spec(work.parDo.input), coders=None, context=context) # AssignBucketsDoFn is intentionally unimplemented. The implementation of # WindowInto in transforms/core.py does not use a service primitive. else: raise NotImplementedError('Unknown ParDo type: %r' % specs) def get_flatten_work_item(instruction, unused_env, unused_context): """Parses a flatten instruction into the appropriate Worker* object. Args: instruction: a ParallelInstruction protobuf with a FlattenInstruction in it. Returns: A WorkerFlatten object. """ return WorkerFlatten( inputs=[get_input_spec(inp) for inp in instruction.flatten.inputs], output_coders=get_output_coders(instruction)) def get_partial_gbk_work_item(instruction, unused_env, unused_context): """Parses a partial GBK instruction into the appropriate Worker* object. Args: instruction: a ParallelInstruction protobuf with a PartialGroupByKeyInstruction in it. Returns: A WorkerPartialGroupByKey object. """ combine_fn = None if instruction.partialGroupByKey.valueCombiningFn: combine_fn_specs = { p.key: from_json_value(p.value) for p in (instruction.partialGroupByKey.valueCombiningFn .additionalProperties)} combine_fn = combine_fn_specs.get('serialized_fn', {}).get('value', None) return WorkerPartialGroupByKey( combine_fn=combine_fn, input=get_input_spec(instruction.partialGroupByKey.input), output_coders=get_output_coders(instruction)) class MapTask(object): """A map task decoded into operations and ready to be executed. Attributes: operations: A list of Worker* object created by parsing the instructions within the map task. stage_name: The name of this map task execution stage. step_names: The names of the step corresponding to each map task operation. """ def __init__(self, operations, stage_name, step_names): self.operations = operations self.stage_name = stage_name self.step_names = step_names self.counter_factory = CounterFactory() def itercounters(self): return self.counter_factory.get_counters() def __str__(self): return '<%s %s steps=%s>' % (self.__class__.__name__, self.stage_name, '+'.join(self.step_names)) class WorkerEnvironment(object): """A worker execution environment.""" def __init__(self): self._sources = [] self._sinks = [] self._register_predefined() def _register_predefined(self): """Register predefined sources and sinks.""" self.register_source_parser(WorkerEnvironment._parse_text_source) self.register_source_parser(WorkerEnvironment._parse_inmemory_source) self.register_source_parser(WorkerEnvironment._parse_avro_source) self.register_source_parser(WorkerEnvironment._parse_big_query_source) self.register_source_parser(WorkerEnvironment._parse_pubsub_source) self.register_source_parser(WorkerEnvironment._parse_concat_source) self.register_source_parser(WorkerEnvironment._parse_windmill_source) # TODO(silviuc): Implement support for PartitioningShuffleSource # TODO(silviuc): Implement support for AvroSource # TODO(silviuc): Implement support for custom sources self.register_sink_parser(WorkerEnvironment._parse_text_sink) self.register_sink_parser(WorkerEnvironment._parse_avro_sink) self.register_sink_parser(WorkerEnvironment._parse_pubsub_sink) self.register_sink_parser(WorkerEnvironment._parse_windmill_sink) def register_source_parser(self, fn): self._sources.append(fn) def register_sink_parser(self, fn): self._sinks.append(fn) def parse_source(self, specs, codec_specs, context): for fn in self._sources: result = fn(specs, codec_specs, context) if result: return result return None def parse_sink(self, specs, codec_specs, context): for fn in self._sinks: result = fn(specs, codec_specs, context) if result: return result return None @staticmethod def _parse_text_source(specs, codec_specs, unused_context): if specs['@type'] == 'TextSource': coder = get_coder_from_spec(codec_specs) start_offset = None if 'start_offset' in specs: start_offset = int(specs['start_offset']['value']) end_offset = None if 'end_offset' in specs: end_offset = int(specs['end_offset']['value']) return io.TextFileSource( file_path=specs['filename']['value'], start_offset=start_offset, end_offset=end_offset, compression_type=specs['compression_type']['value'], strip_trailing_newlines=specs['strip_trailing_newlines']['value'], coder=coder, ) @staticmethod def _parse_concat_source(specs, _, unused_context): if specs['@type'] == 'ConcatSource': assert unused_context.worker_environment is not None sub_sources = [] for sub_source_dict in specs['sources']: sub_source_specs = sub_source_dict['spec'] sub_source_codec_specs = None if 'encoding' in sub_source_dict: sub_source_codec_specs = sub_source_dict['encoding'] sub_sources.append(unused_context.worker_environment.parse_source( sub_source_specs, sub_source_codec_specs, unused_context)) return concat_reader.ConcatSource(sub_sources) @staticmethod def _parse_inmemory_source(specs, codec_specs, unused_context): if specs['@type'] == 'InMemorySource': # We do not wrap values sent to the service in a Create transform and # received here in a WindowedValue wrapper, but the service needs to be # sent the wrapped encoding so subsequent GroupByKey operations work # correctly. # # Note: The service may create a dummy empty InMemorySource that is a # windowed value when processing a BigQuerySource. In that case, we do # not unwrap this coder. # TODO(ccy): investigate if we can make these semantics cleaner. coder = get_coder_from_spec(codec_specs) if isinstance(coder, coders.WindowedValueCoder): coder = coder.wrapped_value_coder # Handle the case where 'elements' for an InMemory source is empty # list. if specs['elements']: # start_index/end_index could be missing if default behavior should be # used. For instance a list with one element will have start_index=0 and # end_index=1 by default. start_index = ( None if 'start_index' not in specs else int( specs['start_index']['value'])) end_index = ( None if 'end_index' not in specs else int(specs['end_index']['value'])) return inmemory.InMemorySource( elements=[base64.b64decode(v['value']) for v in specs['elements']], coder=coder, start_index=start_index, end_index=end_index) else: return inmemory.InMemorySource(elements=[], coder=coder) @staticmethod def _parse_avro_source(specs, unused_codec_specs, unused_context): if specs['@type'] == 'AvroSource': # Note that the worker does not really implement AVRO yet.It takes # advantage that both reading and writing is done through the worker to # choose a supported format (text files with one pickled object per line). start_offset = None if 'start_offset' in specs: start_offset = int(specs['start_offset']['value']) end_offset = None if 'end_offset' in specs: end_offset = int(specs['end_offset']['value']) return io.TextFileSource( file_path=specs['filename']['value'], start_offset=start_offset, end_offset=end_offset, strip_trailing_newlines=True, coder=coders.Base64PickleCoder()) @staticmethod def _parse_big_query_source(specs, codec_specs, unused_context): if specs['@type'] == 'BigQuerySource': coder = get_coder_from_spec(codec_specs) if 'table' in specs: return io.BigQuerySource( project=specs['project']['value'], dataset=specs['dataset']['value'], table=specs['table']['value'], coder=coder) elif 'bigquery_query' in specs: return io.BigQuerySource( query=specs['bigquery_query']['value'], coder=coder) else: raise ValueError('BigQuery source spec must specify either a table' ' or a query') @staticmethod def _parse_pubsub_source(specs, codec_specs, context): if specs['@type'] == 'PubsubReader': topic = specs['pubsub_topic']['value'] subscription = specs['pubsub_subscription_name']['value'] coder = coders.deserialize_coder(codec_specs['@type']) return windmillio.PubSubWindmillSource(context, topic, subscription, coder) @staticmethod def _parse_windmill_source(specs, codec_specs, context): if specs['@type'] == 'WindowingWindmillReader': stream_id = specs['stream_id']['value'] coder = get_coder_from_spec(codec_specs) return windmillio.WindowingWindmillSource(context, stream_id, coder) @staticmethod def _parse_text_sink(specs, codec_specs, unused_context): if specs['@type'] == 'TextSink': coder = get_coder_from_spec(codec_specs) return fileio.NativeTextFileSink( file_path_prefix=specs['filename']['value'], append_trailing_newlines=specs['append_trailing_newlines']['value'], coder=coder) @staticmethod def _parse_avro_sink(specs, unused_codec_specs, unused_context): # Note that the worker does not really implement AVRO yet.It takes # advantage that both reading and writing is done through the worker to # choose a supported format (text files with one pickled object per line). if specs['@type'] == 'AvroSink': return fileio.NativeTextFileSink( specs['filename']['value'], append_trailing_newlines=True, coder=coders.Base64PickleCoder()) @staticmethod def _parse_pubsub_sink(specs, codec_specs, context): if specs['@type'] == 'PubsubSink': coder = get_coder_from_spec(codec_specs) topic = specs['pubsub_topic']['value'] timestamp_label = specs['pubsub_timestamp_label']['value'] id_label = specs['pubsub_id_label']['value'] return windmillio.PubSubWindmillSink(context, coder, topic, timestamp_label, id_label) @staticmethod def _parse_windmill_sink(specs, codec_specs, context): if specs['@type'] == 'WindmillSink': coder = get_coder_from_spec(codec_specs) stream_id = specs['stream_id']['value'] return windmillio.WindmillSink(context, stream_id, coder) def decode_map_task(map_task_proto, env=WorkerEnvironment(), context=ExecutionContext()): """Parses a map task proto into operations within a MapTask object. The response is received by the worker as a result of a LeaseWorkItem request to the Dataflow service. Args: map_task_proto: A MapTask protobuf object returned by the service. env: An environment object with worker configuration. context: An ExecutionContext object providing context for operations to be executed. Returns: A tuple of work item id and the list of Worker* objects (see definitions above) representing the list of operations to be executed as part of the work item. """ operations = [] stage_name = map_task_proto.stageName step_names = [] context.worker_environment = env # Parse the MapTask instructions. for work in map_task_proto.instructions: step_names.append(work.name) if work.read is not None: operations.append(get_read_work_item(work, env, context)) elif work.write is not None: operations.append(get_write_work_item(work, env, context)) elif work.parDo is not None: operations.append(get_do_work_item(work, env, context)) elif work.flatten is not None: operations.append(get_flatten_work_item(work, env, context)) elif work.partialGroupByKey is not None: operations.append(get_partial_gbk_work_item(work, env, context)) else: raise NotImplementedError('Unknown instruction: %r' % work) return MapTask(operations, stage_name, step_names) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/maptask_test.py000066400000000000000000000024171272340251000255700ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for MapTask behavior. """ import logging import unittest from google.cloud.dataflow.utils.counters import Counter from google.cloud.dataflow.worker import maptask class MapTaskTest(unittest.TestCase): def test_itercounters(self): map_task = maptask.MapTask([], 'test-name', []) map_task.counter_factory.get_counter('counter-abc', Counter.SUM) counters_found = 0 for counter in map_task.itercounters(): logging.info('iterator yielded %s', counter) self.assertEqual('counter-abc', counter.name) counters_found += 1 self.assertEqual(1, counters_found) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/opcounters.py000066400000000000000000000044021272340251000252660ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Counters collect the progress of the Worker for reporting to the service.""" from __future__ import absolute_import from google.cloud.dataflow.utils.counters import Counter class OperationCounters(object): """The set of basic counters to attach to an Operation.""" def __init__(self, counter_factory, step_name, coder, output_index): self.element_counter = counter_factory.get_counter( '%s-out%d-ElementCount' % (step_name, output_index), Counter.SUM) self.mean_byte_counter = counter_factory.get_counter( '%s-out%d-MeanByteCount' % (step_name, output_index), Counter.MEAN) self.coder = coder def update_from(self, windowed_value, coder=None): # pylint: disable=unused-argument """Add one value to this counter.""" self.element_counter.update(1) # TODO(silviuc): Implement estimated size sampling. # TODO(gildea): # Actually compute the encoded size of this value. # In spirit, something like this: # if coder is None: # coder = self.coder # coder.store_estimated_size(windowed_value, byte_size_accumulator) # but will need to do sampling. def update_collect(self): """Collects the accumulated size estimates. Now that the element has been processed, we ask our accumulator for the total and store the result in a counter. """ # TODO(silviuc): Implement estimated size sampling. pass def __str__(self): return '<%s [%s]>' % (self.__class__.__name__, ', '.join([str(x) for x in self.__iter__()])) def __repr__(self): return '<%s %s at %s>' % (self.__class__.__name__, [x for x in self.__iter__()], hex(id(self))) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/opcounters_test.py000066400000000000000000000064771272340251000263430ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for worker counters.""" import logging import unittest from google.cloud.dataflow import coders from google.cloud.dataflow.transforms.window import GlobalWindows from google.cloud.dataflow.utils.counters import CounterFactory from google.cloud.dataflow.worker.opcounters import OperationCounters # Classes to test that we can handle a variety of objects. # These have to be at top level so the pickler can find them. class OldClassThatDoesNotImplementLen: # pylint: disable=old-style-class def __init__(self): pass class ObjectThatDoesNotImplementLen(object): def __init__(self): pass class OperationCountersTest(unittest.TestCase): def verify_counters(self, opcounts, expected_elements): self.assertEqual(expected_elements, opcounts.element_counter.value()) def test_update_int(self): opcounts = OperationCounters(CounterFactory(), 'some-name', coders.PickleCoder(), 0) self.verify_counters(opcounts, 0) opcounts.update_from(GlobalWindows.windowed_value(1)) opcounts.update_collect() self.verify_counters(opcounts, 1) def test_update_str(self): opcounts = OperationCounters(CounterFactory(), 'some-name', coders.PickleCoder(), 0) self.verify_counters(opcounts, 0) opcounts.update_from(GlobalWindows.windowed_value('abcde')) opcounts.update_collect() self.verify_counters(opcounts, 1) def test_update_old_object(self): opcounts = OperationCounters(CounterFactory(), 'some-name', coders.PickleCoder(), 0) self.verify_counters(opcounts, 0) obj = OldClassThatDoesNotImplementLen() opcounts.update_from(GlobalWindows.windowed_value(obj)) opcounts.update_collect() self.verify_counters(opcounts, 1) def test_update_new_object(self): opcounts = OperationCounters(CounterFactory(), 'some-name', coders.PickleCoder(), 0) self.verify_counters(opcounts, 0) obj = ObjectThatDoesNotImplementLen() opcounts.update_from(GlobalWindows.windowed_value(obj)) opcounts.update_collect() self.verify_counters(opcounts, 1) def test_update_multiple(self): opcounts = OperationCounters(CounterFactory(), 'some-name', coders.PickleCoder(), 0) self.verify_counters(opcounts, 0) opcounts.update_from(GlobalWindows.windowed_value('abcde')) opcounts.update_from(GlobalWindows.windowed_value('defghij')) opcounts.update_collect() self.verify_counters(opcounts, 2) opcounts.update_from(GlobalWindows.windowed_value('klmnop')) opcounts.update_collect() self.verify_counters(opcounts, 3) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/shuffle.py000066400000000000000000000452001272340251000245220ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Shuffle sources and sinks. The worker code communicates with the shuffler using a shuffle client library (see shuffle_client below). The shuffle operates with entries consisting of a 4-tuple: position, key, secondary key (a.k.a. 2nd-key), value. All values are just raw bytes. On the wire a shuffle entry is represented as a sequence of length and bytes tuples in the order mentioned above. The length is represented as a 4 byte big endian integer. The semantics when reading from shuffle is that values are grouped by key and the values associated with a key are sorted by the secondary key. The opaque position information returned for each shuffle entry can be used to reiterate over values several times and in general to read in a non-sequential manner. The shuffle source supports reiterating over values and values returned have indefinite lifetimes, are stateless and immutable. """ from __future__ import absolute_import import base64 import cStringIO as StringIO import logging import struct from google.cloud.dataflow.coders import observable from google.cloud.dataflow.io import iobase from google.cloud.dataflow.io import range_trackers # The following import works perfectly fine for the Dataflow SDK properly # installed. However in the testing environment the module is not available # since it is built elsewhere. The tests rely on the test_reader/test_writer # arguments for shuffle readers and writers respectively to inject alternative # implementations. try: from google.cloud.dataflow.worker import shuffle_client # pylint: disable=g-import-not-at-top except ImportError: pass def _shuffle_decode(parameter): """Decodes a shuffle parameter. The parameters used to initialize a shuffle source or shuffle sink are sent by the service as urlsafe_base64 Unicode strings. In addition, the encoding does not contain the '=' padding expected by the base64 library. The parameters using this encoding are: shuffle reader positions (start/end), and shuffle reader/writer configuration protobufs. Args: parameter: A Unicode string encoded using urlsafe_base64. Returns: Decoded string. """ # Convert to str and compensate for the potential lack of padding. parameter = str(parameter) if len(parameter) % 4 != 0: parameter += '=' * (4 - len(parameter) % 4) return base64.urlsafe_b64decode(parameter) class ShuffleEntry(object): """A (position, key, 2nd-key, value) tuple as used by the shuffle library.""" def __init__(self, key, secondary_key, value, position): self.key = key self.secondary_key = secondary_key self.value = value self.position = position def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return 'ShuffleEntry %s%s' % (self.key, '/%s' % self.secondary_key if self.secondary_key != self.key else '') def __eq__(self, other): return (self.key == other.key and self.secondary_key == other.secondary_key and self.value == other.value and self.position == other.position) @property def size(self): """Returns the size in bytes of the serialized entry.""" return (16 + len(self.key) + len(self.secondary_key) + len(self.value) + (len(self.position) if self.position else 0)) def to_bytes(self, stream, with_position=True): """Writes the serialized shuffle entry to the stream. Args: stream: A StringIO where the bytes are written to. with_position: True whenever reading from shuffle. False when we write an entry to the shuffle. """ # The struct.pack '>I' specifier means 32 bit big endian integer. if with_position: stream.write(struct.pack('>I', len(self.position))) stream.write(self.position) stream.write(struct.pack('>I', len(self.key))) stream.write(self.key) stream.write(struct.pack('>I', len(self.secondary_key))) stream.write(self.secondary_key) stream.write(struct.pack('>I', len(self.value))) stream.write(self.value) @staticmethod def from_stream(stream, with_position=True): """Returns a shuffle entry read from a StringIO stream. Args: stream: StringIO stream to read the bytes from. with_position: False only for tests when we want to read something that was written to the shuffle without a position. During normal execution when reading the position is always there. Returns: A fully initialized shuffle entry read from the StringIO stream. """ if with_position: position_length = struct.unpack('>I', stream.read(4)) position = stream.read(position_length[0]) else: position = None key_length = struct.unpack('>I', stream.read(4)) key = stream.read(key_length[0]) secondary_key_length = struct.unpack('>I', stream.read(4)) secondary_key = stream.read(secondary_key_length[0]) value_length = struct.unpack('>I', stream.read(4)) value = stream.read(value_length[0]) return ShuffleEntry(key, secondary_key, value, position) class ShuffleEntriesIterable(object): """An iterable over all entries between two positions filtered by key. The method can be used to iterate over all values in the shuffle if key is None and start and nd positions are ''. """ def __init__(self, reader, start_position='', end_position='', key=None): """Constructs an iterable for reading sequentially entries in a range. The iterable object can be used to get all the shuffle entries associated with a key (repeatedly) or simply iterating over all entries (if key is None). Args: reader: A shuffle reader object. These are shared among all iterables since there are networking costs associated to setting one up. start_position: The first shuffle position to read from. end_position: The shuffle position where reading will stop. key: The key to match for all shuffle entries if not None. The iteration stops when a record with a different key is encountered. """ self.reader = reader self.start_position = start_position self.end_position = end_position self.key = key self._pushed_back_entry = None def push_back(self, entry): """Pushes back one entry to support simple look ahead scenarios.""" if self._pushed_back_entry is not None: raise RuntimeError('There is already an entry pushed back.') self._pushed_back_entry = entry def __iter__(self): last_chunk_seen = False start_position = self.start_position end_position = self.end_position while not last_chunk_seen: chunk, next_position = self.reader.Read(start_position, end_position) if not next_position: # An empty string signals the last chunk. last_chunk_seen = True # Yield records inside the chunk just read. read_bytes, total_bytes = 0, len(chunk) stream = StringIO.StringIO(chunk) while read_bytes < total_bytes: entry = ShuffleEntry.from_stream(stream) if self.key is not None and self.key != entry.key: return read_bytes += entry.size yield entry # Check if anything was pushed back. We do this until there is no # value pushed back since it is quite possible to have values pushed # back multiple times by the upper callers. while self._pushed_back_entry is not None: to_return, self._pushed_back_entry = self._pushed_back_entry, None yield to_return # Move on to the next chunk. start_position = next_position class ShuffleEntriesIterator(object): """An iterator object for a ShuffleEntryIterable with push back support. The class supports also the iterable protocol (__iter__) and it is careful to not create a new iterator from the underlying iterable when iter() is called. This is important because shuffle entries iterators are passed around and we want to keep reading sequentially while the passing happens. More specifically they are kept as the underlying iterators for the key values iterables returned for each key. """ def __init__(self, iterable): self.iterable = iterable self.iterator = iter(self.iterable) def __iter__(self): return self def push_back(self, entry): self.iterable.push_back(entry) def next(self): return next(self.iterator) def clone(self, start_position, end_position, key): """Clones the current iterator with a new key, start, and end position.""" return ShuffleEntriesIterator( ShuffleEntriesIterable( self.iterable.reader, start_position, end_position, key)) class ShuffleKeyValuesIterable(observable.ObservableMixin): """An iterable over all values associated with a key. The class supports reiteration over the values by cloning the underlying iterables every time __iter__ gets called. This way the values can be reiterated. The first time __iter__ is called no cloning happens. This supports the very common case of going once over all values for all keys. """ def __init__(self, entries_iterator, key, value_coder, start_position, end_position=''): super(ShuffleKeyValuesIterable, self).__init__() self.key = key self.value_coder = value_coder self.start_position = start_position self.end_position = end_position self.entries_iterator = entries_iterator self.first_values_iterator = None def __iter__(self): if self.first_values_iterator is None: # We safe the first values iterator returned because upper layers # can use it to drain the values in it. This is an optimization needed # to make efficient the very common case of iterating over all key values # available. self.first_values_iterator = self.values_iterator() return self.first_values_iterator else: # If this is not the first time __iter__ is called we will clone the # underlying iterables so that we can reiterate as many times as we # want over the key's values. return ShuffleKeyValuesIterable( self.entries_iterator.clone( self.start_position, self.end_position, self.key), self.key, self.value_coder, self.start_position, self.end_position).values_iterator() def values_iterator(self): for entry in self.entries_iterator: if self.key != entry.key: # Remember the end_position so that if we reiterate over the values # we can do that without reading too much beyond the key. self.end_position = entry.position self.entries_iterator.push_back(entry) break decoded_value = self.value_coder.decode(entry.value) self.notify_observers(entry.value, is_encoded=True) yield decoded_value def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s on %s' % (self.__class__.__name__, self.key) class ShuffleReaderBase(iobase.NativeSourceReader): """A base class for grouped and ungrouped shuffle readers.""" def __init__(self, shuffle_source, reader=None): self.source = shuffle_source self.reader = reader self.entries_iterable = None self.key_coder = self.source.key_coder.get_impl() self.value_coder = self.source.value_coder.get_impl() def __enter__(self): if self.reader is None: self.reader = shuffle_client.PyShuffleReader( _shuffle_decode(self.source.config_bytes)) # Initialize the shuffle entries iterable. For now we read from start to # end which is enough for plain GroupByKey operations. if self.entries_iterable is None: self.entries_iterable = ShuffleEntriesIterable( self.reader, self.source.start_position, self.source.end_position) return self def __exit__(self, exception_type, exception_value, traceback): pass class GroupedShuffleReader(ShuffleReaderBase): """A shuffle reader providing grouped reading.""" def __init__(self, shuffle_source, reader=None): super(GroupedShuffleReader, self).__init__(shuffle_source, reader) self._range_tracker = range_trackers.GroupedShuffleRangeTracker( decoded_start_pos=shuffle_source.start_position, decoded_stop_pos=shuffle_source.end_position) def __iter__(self): entries_iterator = ShuffleEntriesIterator(self.entries_iterable) for entry in entries_iterator: entries_iterator.push_back(entry) key_values = ShuffleKeyValuesIterable( entries_iterator, entry.key, self.value_coder, entry.position) group_start = entry.position last_group_start = self._range_tracker.last_group_start() is_at_split_point = ( last_group_start is None or group_start != last_group_start) if is_at_split_point: if not self._range_tracker.try_claim(group_start): # If an end position is defined, reader has read all records up to the # defined end position, otherwise, reader has read all records of the # source. return else: self._range_tracker.set_current_position(group_start) yield (self.key_coder.decode(entry.key), key_values) # We need to drain the iterator returned just in case this # was not done by the caller. Otherwise we will not properly advance # to the next key but rather return the next entry for the current # key (if there are multiple values). drain_iterator = key_values.first_values_iterator if drain_iterator is None: drain_iterator = iter(key_values) for _ in drain_iterator: pass def get_progress(self): last_group_start = self._range_tracker.last_group_start() if last_group_start is None: return None reader_position = iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode(last_group_start)) return iobase.ReaderProgress(position=reader_position) def request_dynamic_split(self, dynamic_split_request): assert dynamic_split_request is not None split_request_progress = dynamic_split_request.progress if split_request_progress.position is None: logging.warning('GroupingShuffleReader only supports split at a Position.' ' Requested: %r', dynamic_split_request) return encoded_shuffle_position = split_request_progress.position.shuffle_position if encoded_shuffle_position is None: logging.warning('GroupingShuffleReader only supports split at a shuffle' ' position. Requested: %r' , split_request_progress.position) return if self._range_tracker.try_split(_shuffle_decode(encoded_shuffle_position)): logging.info('Split GroupedShuffleReader at %s', encoded_shuffle_position) split_position = iobase.ReaderPosition( shuffle_position=encoded_shuffle_position) return iobase.DynamicSplitResultWithPosition(split_position) else: logging.info('Refusing to split GroupedShuffleReader %r at %s' , self, encoded_shuffle_position) class UngroupedShuffleReader(ShuffleReaderBase): """A shuffle reader providing ungrouped reading.""" def __init__(self, shuffle_source, reader=None): super(UngroupedShuffleReader, self).__init__(shuffle_source, reader) def __iter__(self): for entry in self.entries_iterable: yield self.value_coder.decode(entry.value) class ShuffleSourceBase(iobase.NativeSource): """A base class for grouped and ungrouped shuffle sources.""" def __init__(self, config_bytes, coder, start_position='', end_position=''): self.config_bytes = config_bytes self.key_coder, self.value_coder = ( coder if isinstance(coder, tuple) else (coder, coder)) self.start_position = (start_position if not start_position else _shuffle_decode(start_position)) self.end_position = (end_position if not end_position else _shuffle_decode(end_position)) class GroupedShuffleSource(ShuffleSourceBase): """A source that reads from a shuffled dataset and yields key-grouped data. The value for each key will be an iterable object that will yield values. """ def reader(self, test_reader=None): return GroupedShuffleReader(self, reader=test_reader) class UngroupedShuffleSource(ShuffleSourceBase): """A source that reads from a shuffled dataset and yields values. This source will drop the keys of the key-value pairs and yield just the values. This source is used in resharding operations. """ def reader(self, test_reader=None): return UngroupedShuffleReader(self, reader=test_reader) class ShuffleSinkWriter(iobase.NativeSinkWriter): """A sink writer for ShuffleSink.""" def __init__(self, shuffle_sink, writer=None): self.sink = shuffle_sink self.writer = writer self.stream = StringIO.StringIO() self.bytes_buffered = 0 self.key_coder = self.sink.key_coder.get_impl() self.value_coder = self.sink.value_coder.get_impl() def __enter__(self): if self.writer is None: self.writer = shuffle_client.PyShuffleWriter( _shuffle_decode(self.sink.config_bytes)) return self def __exit__(self, exception_type, exception_value, traceback): value = self.stream.getvalue() if value: self.writer.Write(value) self.bytes_buffered = 0 self.stream.close() self.writer.Close() def Write(self, key, secondary_key, value): entry = ShuffleEntry( self.key_coder.encode(key), secondary_key, self.value_coder.encode(value), position=None) entry.to_bytes(self.stream, with_position=False) self.bytes_buffered += entry.size if self.bytes_buffered > 10 << 20: self.writer.Write(self.stream.getvalue()) self.stream.close() self.stream = StringIO.StringIO() self.bytes_buffered = 0 class ShuffleSink(iobase.NativeSink): """A sink that writes to a shuffled dataset.""" def __init__(self, config_bytes, coder): self.config_bytes = config_bytes self.key_coder, self.value_coder = ( coder if isinstance(coder, tuple) else (coder, coder)) def writer(self, test_writer=None): return ShuffleSinkWriter(self, writer=test_writer) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/shuffle_test.py000066400000000000000000000317241272340251000255670ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for shuffle sources and sinks.""" import base64 import cStringIO as StringIO import logging import unittest from google.cloud.dataflow import coders from google.cloud.dataflow.io import iobase from google.cloud.dataflow.worker.shuffle import GroupedShuffleSource from google.cloud.dataflow.worker.shuffle import ShuffleEntry from google.cloud.dataflow.worker.shuffle import ShuffleSink from google.cloud.dataflow.worker.shuffle import UngroupedShuffleSource class Base64Coder(coders.Coder): """Simple base64 coder used throughout the tests.""" def decode(self, o): return base64.b64decode(o) def encode(self, o): return base64.b64encode(o) class FakeShuffleReader(object): """A fake shuffle reader returning a known set of shuffle chunks. This is not an implementation of a shuffler. The class simply returns the shuffle chunks as specified and is injected as a dependency when a shuffle reader gets created. There are several properties of the sequence of shuffle entries in a chunk that are expected to be true. Tests for the real shuffler make sure that this is the case. For instance: - keys appear in lexicographic order """ def __init__(self, chunk_descriptors): """Initializes the fake shuffle from a list of lists of (k,v) pairs.""" self.all_vals = [] self.chunk_starts = [] last_index = 0 for chunk in chunk_descriptors: self.all_vals.extend(chunk) last_index += len(chunk) self.chunk_starts.append(last_index) def _make_chunk(self, descriptor, start_index): """Returns an encoded shuffle chunk from a list of (k,v) pairs.""" stream = StringIO.StringIO() coder = Base64Coder() position = start_index for key, value in descriptor: ShuffleEntry( coder.encode(key), '', coder.encode(value), position=str(position)).to_bytes(stream) position += 1 value = stream.getvalue() stream.close() return value def _get_first_chunk_key_values(self, start, end): result = [self.all_vals[start]] index = start + 1 while index < end: if index in self.chunk_starts: return result, index result.append(self.all_vals[index]) index += 1 return result, end def Read(self, first, last): # pylint: disable=invalid-name first = 0 if not first else int(first) last = len(self.all_vals) if not last else int(last) key_values, next_position = self._get_first_chunk_key_values(first, last) return ( self._make_chunk(key_values, next_position - len(key_values)), str(next_position) if next_position < last else '') class FakeShuffleWriter(object): """A fake shuffle writter recording what entries were written.""" def __init__(self): # The list of (key, 2nd-key, value) tuples written. The attribute will # get its real value only when close() is called. self.values = [] self._entries = [] def Write(self, entries): # pylint: disable=invalid-name stream = StringIO.StringIO(entries) # TODO(silviuc): Find a better way to detect EOF for a string stream. while stream.tell() < len(stream.getvalue()): self._entries.append( ShuffleEntry.from_stream(stream, with_position=False)) def Close(self): # pylint: disable=invalid-name coder = Base64Coder() for entry in self._entries: self.values.append(( coder.decode(entry.key), coder.decode(entry.secondary_key), coder.decode(entry.value))) class TestShuffleEntry(unittest.TestCase): def test_basics(self): entry = ShuffleEntry('abc', 'xyz123', '0123456789', position='zyx') stream = StringIO.StringIO() entry.to_bytes(stream) self.assertEqual( ShuffleEntry.from_stream(StringIO.StringIO(stream.getvalue())), entry) def test_size(self): """Test that the computed size property returns expected values.""" params = ['abc', 'xyz123', '0123456789', 'zyx'] entry = ShuffleEntry(params[0], params[1], params[2], position=params[3]) expected_size = 4 * len(params) + sum(len(x) for x in params) stream = StringIO.StringIO() entry.to_bytes(stream) self.assertEqual(entry.size, expected_size) self.assertEqual( ShuffleEntry.from_stream(StringIO.StringIO(stream.getvalue())).size, expected_size) def test_big_endian(self): """Tests that lengths are written as big endian ints.""" entry = ShuffleEntry('abc', 'xyz123', '0123456789', position='zyx') stream = StringIO.StringIO() entry.to_bytes(stream) entry_bytes = stream.getvalue() # The length of 'abc' is represented in big endian form. self.assertEqual(entry_bytes[0], '\x00') self.assertEqual(entry_bytes[1], '\x00') self.assertEqual(entry_bytes[2], '\x00') self.assertEqual(entry_bytes[3], '\x03') TEST_CHUNK1 = [('a', '1'), ('b', '0'), ('b', '1'), ('c', '0')] TEST_CHUNK2 = [('c', '1'), ('c', '2'), ('c', '3'), ('c', '4')] class TestGroupedShuffleSource(unittest.TestCase): def test_basics(self): result = [] source = GroupedShuffleSource( config_bytes='not used', coder=Base64Coder()) chunks = [TEST_CHUNK1, TEST_CHUNK2] with source.reader(test_reader=FakeShuffleReader(chunks)) as reader: for key, key_values in reader: for value in key_values: result.append((key, value)) self.assertEqual(TEST_CHUNK1 + TEST_CHUNK2, result) def test_progress_reporting(self): result = [] progress_record = [] source = GroupedShuffleSource( config_bytes='not used', coder=Base64Coder()) chunks = [TEST_CHUNK1, TEST_CHUNK2] expected_progress_record = [] expected_progress_record += 1 * [base64.urlsafe_b64encode('0')] expected_progress_record += 2 * [base64.urlsafe_b64encode('1')] expected_progress_record += 5 * [base64.urlsafe_b64encode('3')] with source.reader(test_reader=FakeShuffleReader(chunks)) as reader: for key, key_values in reader: for value in key_values: result.append((key, value)) progress_record.append( reader.get_progress().position.shuffle_position) self.assertEqual(TEST_CHUNK1 + TEST_CHUNK2, result) self.assertEqual(expected_progress_record, progress_record) def try_splitting_reader_at(self, reader, split_request, expected_response): actual_response = reader.request_dynamic_split(split_request) if expected_response is None: self.assertIsNone(actual_response) else: self.assertIsNotNone(actual_response.stop_position) self.assertIsInstance(actual_response.stop_position, iobase.ReaderPosition) self.assertIsNotNone(actual_response.stop_position.shuffle_position) self.assertEqual(expected_response.stop_position.shuffle_position, actual_response.stop_position.shuffle_position) def test_dynamic_splitting(self): source = GroupedShuffleSource( config_bytes='not used', coder=Base64Coder()) chunks = [TEST_CHUNK1, TEST_CHUNK2] with source.reader(test_reader=FakeShuffleReader(chunks)) as reader: # Cannot split an unstarted reader self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('1')))), None) reader_iter = iter(reader) next(reader_iter) next(reader_iter) # Cannot split since the provided split position is smaller than or equal # to the current position '1'. self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('0')))), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('1')))), None) # Successful split. self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('3')))), iobase.DynamicSplitResultWithPosition(iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('3')))) def test_dynamic_splitting_with_range(self): source = GroupedShuffleSource( config_bytes='not used', coder=Base64Coder(), start_position=base64.urlsafe_b64encode('0'), end_position=base64.urlsafe_b64encode('3')) chunks = [TEST_CHUNK1, TEST_CHUNK2] with source.reader(test_reader=FakeShuffleReader(chunks)) as reader: reader_iter = iter(reader) next(reader_iter) # Cannot split if split request is out of range self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('0')))), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('3')))), None) self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('4')))), None) # Successful split. self.try_splitting_reader_at( reader, iobase.DynamicSplitRequest(iobase.ReaderProgress( position=iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('2')))), iobase.DynamicSplitResultWithPosition(iobase.ReaderPosition( shuffle_position=base64.urlsafe_b64encode('2')))) def test_reiteration(self): """Tests that key values iterators can be iterated repeatedly.""" source = GroupedShuffleSource( config_bytes='not used', coder=Base64Coder()) chunks = [TEST_CHUNK1, TEST_CHUNK2] saved_iterators = {} with source.reader(test_reader=FakeShuffleReader(chunks)) as reader: for key, key_values in reader: saved_iterators[key] = key_values # Iterate once. self.assertEqual(list(saved_iterators['a']), ['1']) self.assertEqual(list(saved_iterators['b']), ['0', '1']) self.assertEqual(list(saved_iterators['c']), ['0', '1', '2', '3', '4']) # Iterate twice. self.assertEqual(list(saved_iterators['a']), ['1']) self.assertEqual(list(saved_iterators['b']), ['0', '1']) self.assertEqual(list(saved_iterators['c']), ['0', '1', '2', '3', '4']) def test_iterator_drained(self): result = [] source = GroupedShuffleSource( config_bytes='not used', coder=Base64Coder()) chunks = [TEST_CHUNK1, TEST_CHUNK2] with source.reader(test_reader=FakeShuffleReader(chunks)) as reader: for key, key_values in reader: for value in key_values: result.append((key, value)) # We stop after getting the first shuffle entry for eack key. # We need to check that the iterator is properly drained and the # rest of the entries for the same key are discarded. break # We expect only the first entry for each key to show up. self.assertEqual([('a', '1'), ('b', '0'), ('c', '0')], result) class TestUngroupedShuffleSource(unittest.TestCase): def test_basics(self): result = [] source = UngroupedShuffleSource( config_bytes='not used', coder=Base64Coder()) chunks = [TEST_CHUNK1, TEST_CHUNK2] with source.reader(test_reader=FakeShuffleReader(chunks)) as reader: for v in reader: result.append(v) # We get only the values from the (k, 2nd-k, v) tuples. self.assertEqual([e[1] for e in TEST_CHUNK1 + TEST_CHUNK2], result) class TestShuffleSink(unittest.TestCase): def test_basics(self): source = ShuffleSink(config_bytes='not used', coder=Base64Coder()) entries = [('a', '', '1'), ('b', '', '0'), ('b', '', '1')] fake_writer = FakeShuffleWriter() with source.writer(test_writer=fake_writer) as writer: for entry in entries: writer.Write(*entry) self.assertEqual(entries, fake_writer.values) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/sideinputs.py000066400000000000000000000022001272340251000252460ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Utilities for handling side inputs.""" import collections def get_iterator_fn_for_sources(sources): """Returns callable that returns iterator over elements for given sources.""" def _inner(): for source in sources: with source.reader() as reader: for value in reader: yield value return _inner class EmulatedIterable(collections.Iterable): """Emulates an iterable for a side input.""" def __init__(self, iterator_fn): self.iterator_fn = iterator_fn def __iter__(self): return self.iterator_fn() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/sideinputs_test.py000066400000000000000000000030141272340251000263110ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for side input utilities.""" import logging import unittest from google.cloud.dataflow.worker import sideinputs class EmulatedCollectionsTest(unittest.TestCase): def test_emulated_iterable(self): def _iterable_fn(): for i in range(10): yield i iterable = sideinputs.EmulatedIterable(_iterable_fn) # Check that multiple iterations are supported. for _ in range(0, 5): for i, j in enumerate(iterable): self.assertEqual(i, j) def test_large_iterable_values(self): def _iterable_fn(): for i in range(10): yield ('%d' % i) * (200 * 1024 * 1024) iterable = sideinputs.EmulatedIterable(_iterable_fn) # Check that multiple iterations are supported. for _ in range(0, 3): for i, j in enumerate(iterable): self.assertEqual(('%d' % i) * (200 * 1024 * 1024), j) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/start.py000066400000000000000000000051501272340251000242230ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Python Dataflow started script for the worker main loop.""" # Import _strptime to work around http://bugs.python.org/issue7980. There is a # thread-safety issue with datetime.datetime.strptime if this module is not # already imported. import _strptime # pylint: disable=unused-import import json import logging import random import re import sys from google.cloud.dataflow.worker import batchworker from google.cloud.dataflow.worker import logger def parse_properties(args): properties = {} unused_args = [] for arg in args: match = re.search(r'-D(.+)=(.+)', arg) if match: properties[match.group(1)] = match.group(2) else: unused_args.append(arg) return properties, unused_args def main(): properties, unused_args = parse_properties(sys.argv[1:]) # Initialize the logging machinery. job_id = properties['job_id'] worker_id = properties['worker_id'] log_path = properties['dataflow.worker.logging.location'] logger.initialize(job_id, worker_id, log_path) logging.info('Worker started with properties: %s', properties) sdk_pipeline_options = json.loads( properties.get('sdk_pipeline_options', '{}')) logging.info('Worker started with sdk_pipeline_options: %s', sdk_pipeline_options) if unused_args: logging.warning('Unrecognized arguments %s', unused_args) if properties.get('is_streaming', False): # TODO(ccy): right now, if we pull this in when not in the worker # environment, this will fail on not being able to pull in the correct gRPC # C dependencies. I am investigating a fix. from google.cloud.dataflow.worker import streamingworker # pylint: disable=g-import-not-at-top # Initialize the random number generator, which is used to generate Windmill # client IDs. random.seed() logging.info('Starting streaming worker.') streamingworker.StreamingWorker(properties).run() else: logging.info('Starting batch worker.') batchworker.BatchWorker(properties, sdk_pipeline_options).run() if __name__ == '__main__': main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/streamingworker.py000066400000000000000000000222221272340251000263100ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Python Dataflow streaming worker.""" from __future__ import absolute_import import logging import random import sys import time import traceback from grpc.beta import implementations from google.cloud.dataflow.internal import windmill_pb2 from google.cloud.dataflow.internal import windmill_service_pb2 from google.cloud.dataflow.utils import retry from google.cloud.dataflow.worker import executor from google.cloud.dataflow.worker import maptask from google.cloud.dataflow.worker import windmillio from google.cloud.dataflow.worker import windmillstate import apitools.base.py as apitools_base import google.cloud.dataflow.internal.clients.dataflow as dataflow # pylint: disable=invalid-name class WindmillClient(object): """Client for communication with Windmill.""" def __init__(self, host, port, request_timeout=10): self.host = host self.port = port self.request_timeout = request_timeout channel = implementations.insecure_channel(host, port) self.stub = ( windmill_service_pb2.beta_create_CloudWindmillServiceV1Alpha1_stub( channel)) @retry.with_exponential_backoff() def GetWork(self, request): return self.stub.GetWork(request, self.request_timeout) @retry.with_exponential_backoff() def GetData(self, request): return self.stub.GetData(request, self.request_timeout) @retry.with_exponential_backoff() def CommitWork(self, request): return self.stub.CommitWork(request, self.request_timeout) @retry.with_exponential_backoff() def GetConfig(self, request): return self.stub.GetConfig(request, self.request_timeout) @retry.with_exponential_backoff() def ReportStats(self, request): return self.stub.ReportStats(request, self.request_timeout) # pylint: enable=invalid-name class StreamingWorker(object): """A streaming worker that communicates with Windmill.""" # Maximum size of the result of a GetWork request. MAX_GET_WORK_FETCH_BYTES = 64 << 20 # 64m # Maximum number of items to return in a GetWork request. MAX_GET_WORK_ITEMS = 100 # Delay to use before retrying work items locally, in seconds. RETRY_LOCALLY_DELAY = 10.0 def __init__(self, properties): self.project_id = properties['project_id'] self.job_id = properties['job_id'] self.worker_id = properties['worker_id'] self.client_id = random.getrandbits(63) windmill_host = properties['windmill.host'] windmill_port = int(properties['windmill.grpc_port']) logging.info('Using gRPC to connect to Windmill at %s:%d.', windmill_host, windmill_port) self.windmill = WindmillClient(windmill_host, windmill_port) self.instruction_map = {} self.system_name_to_computation_id_map = {} def run(self): self.running = True # TODO(ccy): support multi-threaded or multi-process execution. self.dispatch_loop() def get_work(self): request = windmill_pb2.GetWorkRequest( client_id=self.client_id, max_items=StreamingWorker.MAX_GET_WORK_ITEMS, max_bytes=StreamingWorker.MAX_GET_WORK_FETCH_BYTES) return self.windmill.GetWork(request) def add_computation(self, map_task): computation_id = self.system_name_to_computation_id_map.get( map_task.systemName, map_task.systemName) if computation_id not in self.instruction_map: self.instruction_map[computation_id] = map_task def parse_map_task(self, serialized_map_task): return apitools_base.JsonToMessage(dataflow.MapTask, serialized_map_task) def get_config(self, computation_id): """Load the config for a given computation from Windmill.""" request = windmill_pb2.GetConfigRequest(computations=[computation_id]) response = self.windmill.GetConfig(request) for map_entry in response.system_name_to_computation_id_map: self.system_name_to_computation_id_map[ map_entry.system_name] = map_entry.computation_id for serialized_map_task in response.cloud_works: # Print the serialized version here as it's more readable. logging.info('Adding config for computation %s: %r', computation_id, serialized_map_task) self.add_computation(self.parse_map_task(serialized_map_task)) return response def dispatch_loop(self): while self.running: backoff_seconds = 0.001 while self.running: work_response = self.get_work() if work_response.work: break time.sleep(backoff_seconds) backoff_seconds = min(1.0, backoff_seconds * 2) for computation_work in work_response.work: self.process_computation(computation_work) def process_computation(self, computation_work): computation_id = computation_work.computation_id input_data_watermark = windmillio.windmill_to_harness_timestamp( computation_work.input_data_watermark) if computation_id not in self.instruction_map: self.get_config(computation_id) map_task_proto = self.instruction_map[computation_id] for work_item in computation_work.work: retry_locally = True while retry_locally: try: self.process_work_item(computation_id, map_task_proto, input_data_watermark, work_item) break except: # pylint: disable=bare-except logging.error( 'Exception while processing work item for computation %r: ' '%s, %s', computation_id, work_item, traceback.format_exc()) # Send exception details to Windmill, retry locally if possible. retry_locally = self.report_failure(computation_id, work_item, sys.exc_info()) # TODO(ccy): handle token expiration in retry logic. # TODO(ccy): handle out-of-memory error in retry logic. if retry_locally: logging.error('Execution of work in computation %s for key %r ' 'failed; will retry locally.', computation_id, work_item.key) time.sleep(StreamingWorker.RETRY_LOCALLY_DELAY) else: logging.error('Execution of work in computation %s for key %r ' 'failed; Windmill indicated to not retry ' 'locally.', computation_id, work_item.key) def report_failure(self, computation_id, work_item, exc_info): """Send exception details to Windmill; returns whether to retry locally.""" exc_type, exc_value, exc_traceback = exc_info messages = list(line.strip() for line in (traceback.format_exception_only(exc_type, exc_value) + traceback.format_tb(exc_traceback))) wm_exception = windmill_pb2.Exception(stack_frames=messages) report_stats_request = windmill_pb2.ReportStatsRequest( computation_id=computation_id, key=work_item.key, sharding_key=work_item.sharding_key, work_token=work_item.work_token, exceptions=[wm_exception]) response = self.windmill.ReportStats(report_stats_request) return not response.failed def process_work_item(self, computation_id, map_task_proto, input_data_watermark, work_item): """Process a work item.""" workitem_commit_request = windmill_pb2.WorkItemCommitRequest( key=work_item.key, work_token=work_item.work_token) env = maptask.WorkerEnvironment() context = maptask.StreamingExecutionContext() reader = windmillstate.WindmillStateReader( computation_id, work_item.key, work_item.work_token, self.windmill) state_internals = windmillstate.WindmillStateInternals(reader) state = windmillstate.WindmillUnmergedState(state_internals) output_data_watermark = windmillio.windmill_to_harness_timestamp( work_item.output_data_watermark) context.start(computation_id, work_item, input_data_watermark, output_data_watermark, workitem_commit_request, self.windmill, state) map_task = maptask.decode_map_task(map_task_proto, env, context) map_task_executor = executor.MapTaskExecutor(map_task) map_task_executor.execute() state_internals.persist_to(workitem_commit_request) # Send result to Windmill. # TODO(ccy): in the future, this will not be done serially with respect to # work execution. commit_request = windmill_pb2.CommitWorkRequest() computation_commit_request = windmill_pb2.ComputationCommitWorkRequest( computation_id=computation_id, requests=[workitem_commit_request]) commit_request.requests.extend([computation_commit_request]) self.windmill.CommitWork(commit_request) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/streamingworker_test.py000066400000000000000000000077201272340251000273550ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for the streaming worker. These tests check that the streaming worker harness runs properly, with mocked interactions with Windmill. """ import logging import unittest import mock from google.cloud.dataflow.internal import windmill_pb2 from google.cloud.dataflow.worker.streamingworker import StreamingWorker class StreamingWorkerTest(unittest.TestCase): @mock.patch( 'google.cloud.dataflow.worker.streamingworker.WindmillClient') def _get_worker(self, *unused_mocks): fake_properties = { 'project_id': 'fake_project', 'job_id': 'fake_job', 'worker_id': 'fake_worker', 'windmill.host': 'fake_host', 'windmill.grpc_port': '12345', } return StreamingWorker(fake_properties) def _get_worker_and_single_computation(self): worker = self._get_worker() computation_work = windmill_pb2.ComputationWorkItems( computation_id='A1', work=[windmill_pb2.WorkItem( key='k', work_token=12345)]) worker.instruction_map['A1'] = mock.Mock() return worker, computation_work @mock.patch('google.cloud.dataflow.worker.streamingworker.StreamingWorker.' 'process_work_item') def test_successful_work_item(self, *unused_mocks): worker, computation_work = self._get_worker_and_single_computation() worker.process_computation(computation_work) self.assertEqual(0, len(worker.windmill.ReportStats.call_args_list)) self.assertEqual(1, len(worker.process_work_item.call_args_list)) @mock.patch('google.cloud.dataflow.worker.streamingworker.StreamingWorker.' 'process_work_item') @mock.patch('logging.error') def test_failed_work_item(self, *unused_mocks): worker, computation_work = self._get_worker_and_single_computation() worker.windmill.ReportStats.return_value = ( windmill_pb2.ReportStatsResponse(failed=True)) worker.process_work_item.side_effect = Exception worker.process_computation(computation_work) # Verify number of attempts and that failed work was reported. self.assertEqual(1, len(worker.windmill.ReportStats.call_args_list)) self.assertEqual(1, len(worker.process_work_item.call_args_list)) logging.error.assert_called_with( 'Execution of work in computation %s for key %r failed; Windmill ' 'indicated to not retry locally.', u'A1', 'k') @mock.patch('google.cloud.dataflow.worker.streamingworker.StreamingWorker.' 'process_work_item') @mock.patch('logging.error') @mock.patch('time.sleep') def test_retrying_failed_work_item(self, *unused_mocks): worker, computation_work = self._get_worker_and_single_computation() retries = 5 worker.windmill.ReportStats.side_effect = ( [windmill_pb2.ReportStatsResponse(failed=False)] * retries) worker.process_work_item.side_effect = ( [Exception] * retries + [None]) worker.process_computation(computation_work) # Verify number of attempts and that failed work was reported the correct # number of times. self.assertEqual(retries, len(worker.windmill.ReportStats.call_args_list)) self.assertEqual(retries + 1, len(worker.process_work_item.call_args_list)) logging.error.assert_called_with( 'Execution of work in computation %s for key %r failed; will retry ' 'locally.', u'A1', 'k') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/version.py000066400000000000000000000012341272340251000245520ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Dataflow Python worker version information.""" __version__ = '0.2.5' DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/windmillio.py000066400000000000000000000224651272340251000252450ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Windmill sources and sinks. Windmill sources and sinks are used internally in streaming pipelines. """ from __future__ import absolute_import from google.cloud.dataflow.coders import observable from google.cloud.dataflow.io import coders from google.cloud.dataflow.io import iobase from google.cloud.dataflow.io import pubsub from google.cloud.dataflow.transforms.timeutil import TimeDomain from google.cloud.dataflow.transforms.timeutil import Timestamp from google.cloud.dataflow.transforms.window import GlobalWindows from google.cloud.dataflow.transforms.window import WindowedValue def harness_to_windmill_timestamp(timestamp): # The timestamp taken by Windmill is in microseconds. return timestamp.micros def windmill_to_harness_timestamp(windmill_timestamp): # The timestamp given by Windmill is in microseconds. return Timestamp(micros=windmill_timestamp) class PubSubWindmillSource(pubsub.PubSubSource): """Internal worker PubSubSource which reads from Windmill.""" def __init__(self, context, topic, subscription, coder): super(PubSubWindmillSource, self).__init__(topic, subscription, coder) self.context = context def reader(self): return PubSubWindmillReader(self) class PubSubWindmillReader(iobase.NativeSourceReader): """Internal worker Windmill PubSub reader.""" def __init__(self, source): self.source = source def __iter__(self): for bundle in self.source.context.work_item.message_bundles: for message in bundle.messages: yield GlobalWindows.windowed_value( self.source.coder.decode(message.data), timestamp=windmill_to_harness_timestamp(message.timestamp)) def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): pass @property def returns_windowed_values(self): """Returns whether this reader returns windowed values.""" return True class PubSubWindmillSink(pubsub.PubSubSink): """Internal worker PubSubSink which writes to Windmill.""" def __init__(self, context, coder, topic, timestamp_label, id_label): super(PubSubWindmillSink, self).__init__(topic, coder=coder) self.context = context self.timestamp_label = timestamp_label self.id_label = id_label def writer(self): return PubSubWindmillWriter(self) class PubSubWindmillWriter(iobase.NativeSinkWriter): """Internal worker Windmill PubSub writer.""" def __init__(self, sink): self.sink = sink # Avoid dependency on gRPC during testing. # pylint: disable=g-import-not-at-top from google.cloud.dataflow.internal import windmill_pb2 # pylint: enable=g-import-not-at-top self.windmill_pb2 = windmill_pb2 def __enter__(self): self.output_message_bundle = self.windmill_pb2.PubSubMessageBundle( topic=self.sink.topic, timestamp_label=self.sink.timestamp_label, id_label=self.sink.id_label) return self @property def takes_windowed_values(self): """Returns whether this writer takes windowed values.""" return True def Write(self, windowed_value): data = self.sink.coder.encode(windowed_value.value) timestamp = harness_to_windmill_timestamp(windowed_value.timestamp) self.output_message_bundle.messages.add(data=data, timestamp=timestamp) def __exit__(self, exception_type, exception_value, traceback): if self.output_message_bundle and self.output_message_bundle.messages: self.sink.context.workitem_commit_request.pubsub_messages.extend( [self.output_message_bundle]) self.output_message_bundle = None class WindmillSink(iobase.NativeSink): """Sink for writing to a given Cloud Pubsub topic.""" def __init__(self, context, stream_id, coder): self.context = context self.stream_id = stream_id self.coder = coder @property def format(self): """Sink format name required for remote execution.""" return 'windmill' def writer(self): return WindmillWriter(self) class WindmillWriter(iobase.NativeSinkWriter): """Internal worker Windmill writer.""" def __init__(self, sink): self.sink = sink self.key_coder = self.sink.coder.key_coder() value_coder = self.sink.coder.value_coder() self.wv_coder = coders.WindowedValueCoder(value_coder) # Avoid dependency on gRPC during testing. # pylint: disable=g-import-not-at-top from google.cloud.dataflow.internal import windmill_pb2 # pylint: enable=g-import-not-at-top self.windmill_pb2 = windmill_pb2 def __enter__(self): self.keyed_output = {} return self @property def takes_windowed_values(self): """Returns whether this writer takes windowed values.""" return True def Write(self, windowed_kv): # WindmillWriter takes windowed values, reifies the windows and writes the # resulting windowed value to Windmill. Note that in this streaming case, # the service does not add a ReifyWindows step, so we do that here. key, value = windowed_kv.value timestamp = windowed_kv.timestamp wm_timestamp = harness_to_windmill_timestamp(timestamp) windows = windowed_kv.windows windowed_value = WindowedValue(value, timestamp, windows) encoded_key = self.key_coder.encode(key) encoded_value = self.wv_coder.encode(windowed_value) # TODO(ccy): In the future, we will populate metadata with PaneInfo # details. metadata = '' # Add to output for key. if encoded_key not in self.keyed_output: self.keyed_output[encoded_key] = ( self.windmill_pb2.KeyedMessageBundle(key=encoded_key)) self.keyed_output[encoded_key].messages.add( timestamp=wm_timestamp, data=encoded_value, metadata=metadata) def __exit__(self, exception_type, exception_value, traceback): self.sink.context.workitem_commit_request.output_messages.add( destination_stream_id=self.sink.stream_id, bundles=self.keyed_output.values()) del self.keyed_output class WindmillTimer(object): """Timer sent by Windmill.""" def __init__(self, key, namespace, name, time_domain, timestamp, state_family): self.key = key self.namespace = namespace self.name = name self.time_domain = time_domain self.timestamp = timestamp self.state_family = state_family def __repr__(self): return ('WindmillTimer(key=%s, namespace=%s, name=%s, time_domain=%s, ' 'timestamp=%s, state_family=%s)') % (self.key, self.namespace, self.name, self.time_domain, self.timestamp, self.state_family) class KeyedWorkItem(observable.ObservableMixin): """Keyed work item used by a StreamingGroupAlsoByWindowsOperation.""" def __init__(self, work_item, coder): super(KeyedWorkItem, self).__init__() self.work_item = work_item self.coder = coder self.key_coder = coder.key_coder() value_coder = coder.value_coder() self.wv_coder = coders.WindowedValueCoder(value_coder) self.key = self.key_coder.decode(work_item.key) # Avoid dependency on gRPC during testing. # pylint: disable=g-import-not-at-top from google.cloud.dataflow.internal import windmill_pb2 # pylint: enable=g-import-not-at-top self.windmill_pb2 = windmill_pb2 def elements(self): for bundle in self.work_item.message_bundles: for message in bundle.messages: element = self.wv_coder.decode(message.data) self.notify_observers(message.data, is_encoded=True) yield element def timers(self): if self.work_item.timers: for timer_item in self.work_item.timers.timers: (namespace, name, unused_time_domain) = timer_item.tag.split('|') yield WindmillTimer( key=self.key, namespace=namespace, name=name, time_domain=TimeDomain.from_string( self.windmill_pb2.Timer.Type.Name(timer_item.type)), timestamp=windmill_to_harness_timestamp(timer_item.timestamp), state_family=timer_item.state_family) def __repr__(self): return '<%s %s>' % (self.__class__.__name__, self.key) class WindowingWindmillSource(iobase.NativeSource): """Internal worker PubSubSource which reads from Windmill.""" def __init__(self, context, stream_id, coder): self.context = context self.coder = coder def reader(self): return WindowingWindmillReader(self) class WindowingWindmillReader(iobase.NativeSourceReader): """Internal worker Windmill PubSub reader.""" def __init__(self, source): self.source = source def __iter__(self): return iter([KeyedWorkItem(self.source.context.work_item, self.source.coder)]) def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): pass DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/windmillstate.py000066400000000000000000000436301272340251000257530ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Implementation of UnmergedState, backed by Windmill.""" from __future__ import absolute_import from abc import ABCMeta from abc import abstractmethod import cPickle as pickle import logging from google.cloud.dataflow.internal import windmill_pb2 from google.cloud.dataflow.transforms import trigger from google.cloud.dataflow.worker import windmillio # Max timestamp value used in Windmill requests. MAX_TIMESTAMP = 0x7fffffffffffffff class WindmillUnmergedState(trigger.UnmergedState): """UnmergedState implementation, backed by Windmill.""" def __init__(self, state_internals): self.internals = state_internals def set_global_state(self, tag, value): self.internals.access('_global_', tag).add(value) def get_global_state(self, tag, default=None): return self.internals.access('_global_', tag).get() or default def set_timer(self, window, name, time_domain, timestamp): namespace = self._encode_window(window) self.internals.add_output_timer(namespace, name, time_domain, timestamp) def clear_timer(self, window, name, time_domain): namespace = self._encode_window(window) self.internals.clear_output_timer(namespace, name, time_domain) def get_window(self, timer_id): return timer_id def _encode_window(self, window): # TODO(robertwb): This is only true for merging windows (but we currently # consider all windows to be merging and pay the costs). assert isinstance(window, int) return str(window) def add_state(self, window, tag, value): namespace = self._encode_window(window) self.internals.access(namespace, tag).add(value) def get_state(self, window, tag): namespace = self._encode_window(window) return self.internals.access(namespace, tag).get() def clear_state(self, window, tag): namespace = self._encode_window(window) self.internals.access(namespace, tag).clear() class WindmillStateInternals(object): """Internal interface to access data in Windmill via state tags.""" def __init__(self, reader): self.reader = reader self.accessed = {} self.output_timers = {} def access(self, namespace, state_tag): """Returns accessor for given namespace and state tag.""" # Note: namespace currently is either a numeric string or "_global_", and so # cannot contain "/". If this changes, we need to be careful in our # construction of the state_key below. state_key = '%s/%s' % (namespace, state_tag.tag) if state_key not in self.accessed: if isinstance(state_tag, trigger.ListStateTag): # List state. self.accessed[state_key] = WindmillBagAccessor(self.reader, state_key) elif isinstance(state_tag, trigger.ValueStateTag): # Value state without combiner. self.accessed[state_key] = WindmillValueAccessor(self.reader, state_key) elif isinstance(state_tag, trigger.CombiningValueStateTag): # Value state with combiner. self.accessed[state_key] = WindmillCombiningValueAccessor( self.reader, state_key, state_tag.combine_fn) elif isinstance(state_tag, trigger.WatermarkHoldStateTag): # Watermark hold state. self.accessed[state_key] = WindmillWatermarkHoldAccessor( self.reader, state_key, state_tag.output_time_fn_impl) else: raise ValueError('Invalid state tag.') return self.accessed[state_key] def add_output_timer(self, namespace, name, time_domain, timestamp): windmill_ts = windmillio.harness_to_windmill_timestamp(timestamp) # Note: The character "|" must not be in the given namespace or name # since we use it as the delimiter in the combined tag string. assert '|' not in namespace assert '|' not in name self.output_timers[(namespace, name, time_domain)] = windmill_pb2.Timer( tag='%s|%s|%s' % (namespace, name, time_domain), timestamp=windmill_ts, type=time_domain, state_family='') def clear_output_timer(self, namespace, name, time_domain): self.output_timers[(namespace, name, time_domain)] = windmill_pb2.Timer( tag='%s|%s|%s' % (namespace, name, time_domain), type=time_domain, state_family='') def persist_to(self, commit_request): for unused_key, accessor in self.accessed.iteritems(): accessor.persist_to(commit_request) commit_request.output_timers.extend(self.output_timers.values()) class WindmillStateReader(object): """Reader of raw state from Windmill.""" # The size of Windmill list request responses is capped at this size (or at # least one list element, if a single such element would exceed this size). MAX_LIST_BYTES = 8 << 20 # 8MB def __init__(self, computation_id, key, work_token, windmill): self.computation_id = computation_id self.key = key self.work_token = work_token self.windmill = windmill def fetch_value(self, state_key): """Get the value at given state tag.""" request = windmill_pb2.GetDataRequest() computation_request = windmill_pb2.ComputationGetDataRequest( computation_id=self.computation_id) keyed_request = windmill_pb2.KeyedGetDataRequest( key=self.key, work_token=self.work_token) keyed_request.values_to_fetch.add( tag=state_key, state_family='') computation_request.requests.extend([keyed_request]) request.requests.extend([computation_request]) return self.windmill.GetData(request) def fetch_list(self, state_key, request_token=None): """Get the list at given state tag.""" request = windmill_pb2.GetDataRequest() computation_request = windmill_pb2.ComputationGetDataRequest( computation_id=self.computation_id) keyed_request = windmill_pb2.KeyedGetDataRequest( key=self.key, work_token=self.work_token) keyed_request.lists_to_fetch.add( tag=state_key, state_family='', end_timestamp=MAX_TIMESTAMP, request_token=request_token or '', fetch_max_bytes=WindmillStateReader.MAX_LIST_BYTES) computation_request.requests.extend([keyed_request]) request.requests.extend([computation_request]) return self.windmill.GetData(request) def fetch_watermark_hold(self, state_key): """Get the watermark hold at given state tag.""" request = windmill_pb2.GetDataRequest() computation_request = windmill_pb2.ComputationGetDataRequest( computation_id=self.computation_id) keyed_request = windmill_pb2.KeyedGetDataRequest( key=self.key, work_token=self.work_token) keyed_request.watermark_holds_to_fetch.add( tag=state_key, state_family='') computation_request.requests.extend([keyed_request]) request.requests.extend([computation_request]) return self.windmill.GetData(request) # TODO(ccy): investigate use of coders for Windmill state data. def encode_value(value): return pickle.dumps(value) def decode_value(encoded): return pickle.loads(encoded) class StateAccessor(object): """Interface for accessing state bound to a given tag.""" __metaclass__ = ABCMeta @abstractmethod def get(self): """Get the state at the bound tag. Returns: the current value (or accumulated value) for a ValueTag; an interable of current values for a ListTag. """ pass @abstractmethod def add(self, value): """Add the given value to the state at the bound tag. For a ValueTag with a combiner, this adds the given value through the combiner's accumulator. For a ListTag, this inserts the given value at the end of the list state. For a ValueTag without a combiner, this replaces the single value stored in the value state. Args: value: the value to add. """ pass @abstractmethod def clear(self): """Clears the state at the bound tag.""" pass @abstractmethod def persist_to(self, commit_request): """Writes state changes to the given WorkItemCommitRequest message.""" pass class WindmillValueAccessor(StateAccessor): """Accessor for value state in Windmill.""" def __init__(self, reader, state_key): self.reader = reader self.state_key = state_key self.value = None self.fetched = False self.modified = False self.cleared = False def get(self): if not self.fetched: self._fetch() return self.value def add(self, value): self.modified = True self.cleared = False # Note: we don't do a deep copy of the added value; it is the caller's # responsibility to make sure the value doesn't change until the value # is committed to Windmill. self.value = value def clear(self): self.modified = True self.cleared = True self.value = None def _fetch(self): """Fetch state from Windmill.""" result = self.reader.fetch_value(self.state_key) for wrapper in result.data: for item in wrapper.data: for value in item.values: if value.value.data == '': # pylint: disable=g-explicit-bool-comparison # When uninitialized, Windmill returns the empty string as the # initial value. self.value = None else: try: self.value = decode_value(value.value.data) except Exception: # pylint: disable=broad-except logging.error( 'Error: could not decode value for key %r; ' 'setting to None: %r.', self.state_key, value.value.data) self.value = None self.fetched = True def persist_to(self, commit_request): if not self.modified: return if self.cleared: encoded_value = '' else: encoded_value = encode_value(self.value) commit_request.value_updates.add( tag=self.state_key, state_family='', value=windmill_pb2.Value( data=encoded_value, timestamp=MAX_TIMESTAMP)) class WindmillCombiningValueAccessor(StateAccessor): """Accessor for combining value state in Windmill.""" def __init__(self, reader, state_key, combine_fn): self.reader = reader self.state_key = state_key self.combine_fn = combine_fn self.accum = None self.fetched = False self.modified = False self.cleared = False def get(self): if not self.fetched: self._fetch() if self.cleared: return ( self.combine_fn.extract_output(self.combine_fn.create_accumulator())) return self.combine_fn.extract_output(self.accum) def add(self, value): # TODO(ccy): once WindmillStateReader supports asynchronous I/O, we won't # have to do this synchronously, i.e. we can fire off the fetch here and # return, queuing up (possibly eagerly-combined) values to be accumulated # for until we have the response. We also want to do blind writes, combine # new values in persist, and combine all values in fetch. if not self.fetched: self._fetch() if self.cleared: self.accum = self.combine_fn.create_accumulator() self.cleared = False self.modified = True self.accum = self.combine_fn.add_inputs(self.accum, [value]) def clear(self): self.modified = True self.cleared = True def _fetch(self): """Fetch state from Windmill.""" result = self.reader.fetch_value(self.state_key) for wrapper in result.data: for item in wrapper.data: for value in item.values: if value.value.data == '': # pylint: disable=g-explicit-bool-comparison # When uninitialized, Windmill returns the empty string as the # initial value. self.accum = self.combine_fn.create_accumulator() else: try: self.accum = decode_value(value.value.data) except Exception: # pylint: disable=broad-except logging.error( 'Error: could not decode value; resetting accumulator: %r.', value.value.data) self.accum = self.combine_fn.create_accumulator() self.fetched = True def persist_to(self, commit_request): if not self.modified: return if self.cleared: encoded_value = '' else: encoded_value = encode_value(self.accum) commit_request.value_updates.add( tag=self.state_key, state_family='', value=windmill_pb2.Value( data=encoded_value, timestamp=MAX_TIMESTAMP)) class WindmillBagAccessor(StateAccessor): """Accessor for list state in Windmill.""" class WindmillBagIterable(object): def __init__(self, accessor): self.accessor = accessor def __iter__(self): return self.accessor._get_iter() # pylint: disable=protected-access def __init__(self, reader, state_key): self.reader = reader self.state_key = state_key self.cleared = False self.encoded_new_values = [] def get(self): # Don't directly iterate here; we want to return an iterable object so that # the user may restart iteration if desired. return WindmillBagAccessor.WindmillBagIterable(self) def _get_iter(self): if not self.cleared: pass # Fetch values from Windmill, followed by values added in this sesison. for value in self._fetch(): yield value for value in self.encoded_new_values: yield decode_value(value) def _fetch(self): """Fetch state from Windmill.""" # TODO(ccy): the Java SDK caches the first page and at the start of each # page of values, fires off an asynchronous read for the next page. We # should do this too once we have asynchronous Windmill state reading. should_fetch_more = True next_request_token = None while should_fetch_more: result = self.reader.fetch_list(self.state_key, request_token=next_request_token) next_request_token = None for wrapper in result.data: for datum in wrapper.data: for item in datum.lists: next_request_token = item.continuation_token for value in item.values: try: yield decode_value(value.data) except Exception: # pylint: disable=broad-except logging.error('Could not decode value: %r.', value.data) yield None should_fetch_more = next_request_token != '' # pylint: disable=g-explicit-bool-comparison def add(self, value): # Encode the value here to ensure further mutations of the value don't # affect the value eventually committed to Windmill. self.encoded_new_values.append(encode_value(value)) def clear(self): self.cleared = True self.encoded_new_values = [] def persist_to(self, commit_request): if self.cleared: commit_request.list_updates.add( tag=self.state_key, state_family='', end_timestamp=MAX_TIMESTAMP) if self.encoded_new_values: list_updates = commit_request.list_updates.add( tag=self.state_key, state_family='') for encoded_value in self.encoded_new_values: list_updates.values.add(data=encoded_value, timestamp=MAX_TIMESTAMP) class WindmillWatermarkHoldAccessor(StateAccessor): """Accessor for watermark hold state in Windmill.""" def __init__(self, reader, state_key, output_time_fn_impl): self.reader = reader self.state_key = state_key self.output_time_fn_impl = output_time_fn_impl self.hold_time = None self.fetched = False self.modified = False self.cleared = False def get(self): if not self.fetched: self._fetch() if self.cleared: return None return self.hold_time def add(self, value): # TODO(ccy): once WindmillStateReader supports asynchronous I/O, we won't # have to do this synchronously, i.e. we can fire off the fetch here and # return, queuing up (possibly eagerly-combined) values to be accumulated # into the hold time for until we have the response. if not self.fetched: self._fetch() self.cleared = False self.modified = True if self.hold_time is None: self.hold_time = value else: self.hold_time = self.output_time_fn_impl.combine(self.hold_time, value) def clear(self): self.modified = True self.cleared = True def _fetch(self): """Fetch state from Windmill.""" result = self.reader.fetch_watermark_hold(self.state_key) for wrapper in result.data: for item in wrapper.data: for value in item.watermark_holds: if (len(value.timestamps) == 1 and value.timestamps[0] == MAX_TIMESTAMP): # When uninitialized, Windmill returns MAX_TIMESTAMP self.hold_time = None else: for wm_timestamp in value.timestamps: timestamp = windmillio.windmill_to_harness_timestamp( wm_timestamp) if self.hold_time is None: self.hold_time = timestamp else: self.hold_time = self.output_time_fn_impl.combine( self.hold_time, timestamp) self.fetched = True def persist_to(self, commit_request): # TODO(ccy): Apparently sending reset=True below is expensive for Windmill # if we haven't done a read. We will need to optimize this if we ever do # blind writes here. if not self.modified: return if self.cleared: value_to_persist = None else: value_to_persist = [ windmillio.harness_to_windmill_timestamp(self.hold_time)] commit_request.watermark_holds.add( tag=self.state_key, state_family='', timestamps=value_to_persist, reset=True) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/workercustomsources.py000066400000000000000000000051611272340251000272400ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Worker utilities related to custom sources.""" from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.internal.json_value import from_json_value from google.cloud.dataflow.io import iobase from google.cloud.dataflow.utils import names DEFAULT_DESIRED_BUNDLE_SIZE = 64 * (1 << 20) # 64MB class NativeBoundedSource(iobase.NativeSource): def __init__(self, bounded_source, start_position, stop_position): self.bounded_source = bounded_source self.start_position = start_position self.stop_position = stop_position def reader(self): return NativeBoundedSourceReader(self) class NativeBoundedSourceReader(iobase.NativeSourceReader): """A native source reader for reading custom sources.""" def __init__(self, source): self._source = source def __exit__(self, exception_type, exception_value, traceback): pass def __enter__(self): return self def __iter__(self): range_tracker = self._source.bounded_source.get_range_tracker( self._source.start_position, self._source.stop_position) return self._source.bounded_source.read(range_tracker) or iter([]) class SourceOperationSplitTask(object): def __init__(self, source_operation_split_proto): source_spec = {p.key: from_json_value(p.value) for p in source_operation_split_proto.source.spec .additionalProperties} if not source_spec.has_key(names.SERIALIZED_SOURCE_KEY): raise ValueError( 'Source split spec must contain a serialized source. Received: %r', source_operation_split_proto) self.source = pickler.loads( source_spec[names.SERIALIZED_SOURCE_KEY]['value']) assert self.source is not None assert isinstance(self.source, iobase.BoundedSource) desired_bundle_size_bytes = ( source_operation_split_proto.options.desiredBundleSizeBytes) if not desired_bundle_size_bytes: self.desired_bundle_size_bytes = DEFAULT_DESIRED_BUNDLE_SIZE else: self.desired_bundle_size_bytes = long(desired_bundle_size_bytes) DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/workercustomsources_test.py000066400000000000000000000176761272340251000303150ustar00rootroot00000000000000 # Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Unit tests for workercustomsources module.""" import logging import unittest from google.cloud.dataflow.internal import pickler from google.cloud.dataflow.internal.json_value import to_json_value from google.cloud.dataflow.io import iobase from google.cloud.dataflow.io import range_trackers from google.cloud.dataflow.utils import names from google.cloud.dataflow.worker import executor from google.cloud.dataflow.worker import workercustomsources from google.cloud.dataflow.worker import workitem import google.cloud.dataflow.internal.clients.dataflow as dataflow class TestSource(iobase.BoundedSource): def __init__(self, start_position, stop_position, test_range_tracker_fn): self._start_position = start_position self._stop_position = stop_position self._test_range_tracker_fn = test_range_tracker_fn def split(self, desired_bundle_size, start=None, stop=None): if not start: start = self._start_position if not stop: stop = self._stop_position range_start = start while range_start < stop: range_stop = min(range_start + desired_bundle_size, stop) yield iobase.SourceBundle(1, self, range_start, range_stop) range_start = range_stop def read(self, range_tracker): if not range_tracker: range_tracker = self.get_range_tracker(self._start_position, self._stop_position) val = range_tracker.start_position() while range_tracker.try_claim(val): yield val val += 1 def get_range_tracker(self, start_position, stop_position): if self._test_range_tracker_fn: return self._test_range_tracker_fn() else: return range_trackers.OffsetRangeTracker(start_position, stop_position) class WorkerCustomSourcesTest(unittest.TestCase): def test_native_bounded_source_read_all(self): source = TestSource(10, 24, None) splits = [split for split in source.split(5)] self.assertEquals(3, len(splits)) read_data = [] for split in splits: _, bundle, start, stop = split native_source = workercustomsources.NativeBoundedSource( bundle, start, stop) read_data.extend([val for val in native_source.reader()]) self.assertEquals(14, len(read_data)) self.assertItemsEqual(range(10, 24), read_data) def test_native_bounded_source_gets_range_tracker(self): def create_dummy_tracker(): return range_trackers.OffsetRangeTracker(0, 3) source = TestSource(0, 20, create_dummy_tracker) read_data = [] for split in source.split(5): _, bundle, start, stop = split native_source = workercustomsources.NativeBoundedSource( bundle, start, stop) read_data.extend([val for val in native_source.reader()]) self.assertEquals(12, len(read_data)) self.assertItemsEqual(range(0, 3) * 4, read_data) def build_split_proto(self, bounded_source, desired_bundle_size): split_proto = dataflow.SourceSplitRequest() split_proto.options = dataflow.SourceSplitOptions() split_proto.options.desiredBundleSizeBytes = desired_bundle_size source = dataflow.Source() spec = dataflow.Source.SpecValue() if bounded_source: spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty( key=names.SERIALIZED_SOURCE_KEY, value=to_json_value({'value': pickler.dumps(bounded_source), '@type': 'http://schema.org/Text'}))) spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty( key='@type', value=to_json_value('CustomSourcesType'))) source.spec = spec split_proto.source = source return split_proto def build_split_work_item(self, split_proto): lease_work_item_response_proto = dataflow.LeaseWorkItemResponse() work_item_proto = dataflow.WorkItem() lease_work_item_response_proto.workItems = [work_item_proto] source_operation_task = dataflow.SourceOperationRequest() work_item_proto.sourceOperationTask = source_operation_task source_operation_task.split = split_proto return workitem.get_work_items(lease_work_item_response_proto) def test_split_task_with_source_and_desired_size(self): test_source = TestSource(start_position=123, stop_position=456, test_range_tracker_fn=None) split_proto = self.build_split_proto(test_source, 1234) split_task = workercustomsources.SourceOperationSplitTask(split_proto) self.assertEquals(1234, split_task.desired_bundle_size_bytes) self.assertIsNotNone(split_task.source) self.assertEquals(123, split_task.source._start_position) self.assertEquals(456, split_task.source._stop_position) def test_split_task_finds_source_no_desired_bundle_size(self): test_source = TestSource(start_position=123, stop_position=456, test_range_tracker_fn=None) split_proto = self.build_split_proto(test_source, None) split_task = workercustomsources.SourceOperationSplitTask(split_proto) self.assertEquals(workercustomsources.DEFAULT_DESIRED_BUNDLE_SIZE, split_task.desired_bundle_size_bytes) self.assertIsNotNone(split_task.source) self.assertEquals(123, split_task.source._start_position) self.assertEquals(456, split_task.source._stop_position) def test_split_task_finds_source_no_source_fails(self): split_proto = self.build_split_proto(None, 1234) with self.assertRaisesRegexp( ValueError, 'Source split spec must contain a serialized source'): workercustomsources.SourceOperationSplitTask(split_proto) def test_get_split_work_item(self): test_source = TestSource(start_position=123, stop_position=456, test_range_tracker_fn=None) split_proto = self.build_split_proto(test_source, 1234) batch_work_item = self.build_split_work_item(split_proto) self.assertIsNotNone(batch_work_item.source_operation_split_task) self.assertEquals( 1234, batch_work_item.source_operation_split_task.desired_bundle_size_bytes) self.assertIsNotNone(batch_work_item.source_operation_split_task.source) self.assertEquals( 123, batch_work_item.source_operation_split_task.source._start_position) self.assertEquals( 456, batch_work_item.source_operation_split_task.source._stop_position) def test_split_executor_generates_splits(self): test_source = TestSource(start_position=10, stop_position=30, test_range_tracker_fn=None) split_proto = self.build_split_proto(test_source, 5) work_item = self.build_split_work_item(split_proto) split_executor = executor.CustomSourceSplitExecutor( work_item.source_operation_split_task) split_executor.execute() self.assertIsNotNone(split_executor.response) self.assertEquals( (dataflow.SourceSplitResponse.OutcomeValueValuesEnum. SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED), split_executor.response.split.outcome) bundles = split_executor.response.split.bundles self.assertEquals(4, len(bundles)) for bundle in bundles: self.assertEquals( (dataflow.DerivedSource.DerivationModeValueValuesEnum. SOURCE_DERIVATION_MODE_INDEPENDENT), bundle.derivationMode) self.assertTrue(bundle.source.doesNotNeedSplitting) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/workitem.py000066400000000000000000000123131272340251000247260ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Worker utilities for parsing out a LeaseWorkItemResponse message. The worker requests work items in a loop. Every response is a description of a complex operation to be executed. For now only MapTask(s) are supported. These tasks represent a sequence of ParallelInstruction(s): read from a source, write to a sink, parallel do, etc. """ import threading from google.cloud.dataflow.worker import maptask from google.cloud.dataflow.worker import workercustomsources class BatchWorkItem(object): """A work item wrapper over the work item proto returned by the service. Attributes: proto: The proto returned by the service for this work item. Some of the fields in the proto are surfaced as attributes of the wrapper class for convenience. map_task: The parsed MapTask object describing the work to perform. next_report_index: The reporting index (an int64) to be used when reporting status. This is returned in the response proto. If there are several status updates for the work item then each update response will contain the next reporting index to be used. This protocol is very important for the service to be able to handle update errors (missed, duplicated, etc.). lease_expire_time: UTC time (a string) when the lease will expire (e.g., '2015-06-17T17:22:49.999Z' or '2015-06-17T17:22:49Z' if zero milliseconds). report_status_interval: Duration (as a string) until a status update for the work item should be send back to the service (e.g., '5.000s' or '5s' if zero milliseconds). """ def __init__(self, proto): self.proto = proto self._map_task = None self._source_operation_split_task = None # Lock to be acquired when reporting status (either reporting progress or # reporting completion). The attributes following the lock attribute (e.g., # 'done', 'next_report_index', etc.) must be accessed using the lock because # the main worker thread executing a work item and the progress reporting # thread handling progress reports will modify them in parallel. self.lock = threading.Lock() self.done = False if self.proto is not None: self.next_report_index = self.proto.initialReportIndex self.lease_expire_time = self.proto.leaseExpireTime self.report_status_interval = self.proto.reportStatusInterval @property def map_task(self): return self._map_task @map_task.setter def map_task(self, map_task): self._map_task = map_task @property def source_operation_split_task(self): return self._source_operation_split_task @source_operation_split_task.setter def source_operation_split_task(self, source_operation_split_task): self._source_operation_split_task = source_operation_split_task def __str__(self): stage_name = self.map_task.stage_name if self.map_task else '' step_names = '+'.join(self.map_task.step_names) if self.map_task else '' return '<%s %s steps=%s %s>' % ( self.__class__.__name__, stage_name, step_names, self.proto.id) def get_work_items(response, env=maptask.WorkerEnvironment(), context=maptask.ExecutionContext()): """Parses a lease work item response into a list of Worker* objects. The response is received by the worker as a result of a LeaseWorkItem request to the Dataflow service. Args: response: A LeaseWorkItemResponse protobuf object returned by the service. env: An environment object with worker configuration. context: A maptask.ExecutionContext object providing context for operations to be executed. Returns: A tuple of work item id and the list of Worker* objects (see definitions above) representing the list of operations to be executed as part of the work item. Raises: ValueError: if type of WorkItem cannot be determined. """ # Check if the request for work did not return anything. if not response.workItems: return None # For now service always sends one work item only. assert len(response.workItems) == 1 work_item_proto = response.workItems[0] work_item = BatchWorkItem(work_item_proto) if work_item_proto.mapTask is not None: map_task = maptask.decode_map_task(work_item_proto.mapTask, env, context) work_item.map_task = map_task elif ( work_item_proto.sourceOperationTask and work_item_proto.sourceOperationTask.split): source_operation_split_task = workercustomsources.SourceOperationSplitTask( work_item_proto.sourceOperationTask.split) work_item.source_operation_split_task = source_operation_split_task else: raise ValueError('Unknown type of work item: %s', work_item_proto) return work_item DataflowPythonSDK-0.2.5/google/cloud/dataflow/worker/workitem_test.py000066400000000000000000000446501272340251000257760ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Tests for LeaseWorkItemResponse parsing utilities. The tests create synthetic LeaseWorkItemResponse messages and then check if the utility routines get the expected maptask.Worker* objects and attributes. The messages built are not necessarily correct in all respects but are good enough to exercise the functionality. """ import base64 import logging import unittest from google.cloud.dataflow import coders from google.cloud.dataflow import io from google.cloud.dataflow.internal.json_value import to_json_value from google.cloud.dataflow.io import fileio from google.cloud.dataflow.worker import concat_reader from google.cloud.dataflow.worker import inmemory from google.cloud.dataflow.worker import maptask from google.cloud.dataflow.worker import workitem import google.cloud.dataflow.internal.clients.dataflow as dataflow # Sample specifications for various worker operations. Note that these samples # are used just to test that the specifications are parsed correctly by the # worker code even though they may not be logically correct. For instance # offsets and indeces may be out of range. TEXT_SOURCE_SPEC = { '@type': 'TextSource', 'start_offset': {'value': '123', '@type': 'http://int'}, 'end_offset': {'value': '123123', '@type': 'http://int'}, 'filename': {'value': 'gs://somefile', '@type': 'http://text'}, 'compression_type': {'value': 'AUTO', '@type': 'http://text'}, 'strip_trailing_newlines': {'value': True, '@type': 'http://bool'}, } IN_MEMORY_ELEMENTS = [ {'value': base64.b64encode('1'), '@type': 'http://int'}, {'value': base64.b64encode('2'), '@type': 'http://int'}, {'value': base64.b64encode('3'), '@type': 'http://int'}] IN_MEMORY_SOURCE_SPEC = { '@type': 'InMemorySource', 'start_index': {'value': '1', '@type': 'http://int'}, 'end_index': {'value': '3', '@type': 'http://int'}, 'elements': IN_MEMORY_ELEMENTS, } GROUPING_SHUFFLE_SOURCE_SPEC = { '@type': 'GroupingShuffleSource', 'start_shuffle_position': {'value': 'opaque', '@type': 'xyz'}, 'end_shuffle_position': {'value': 'opaque', '@type': 'xyz'}, 'shuffle_reader_config': {'value': 'opaque', '@type': 'xyz'}, } UNGROUPED_SHUFFLE_SOURCE_SPEC = { '@type': 'UngroupedShuffleSource', 'start_shuffle_position': {'value': 'opaque', '@type': 'xyz'}, 'end_shuffle_position': {'value': 'opaque', '@type': 'xyz'}, 'shuffle_reader_config': {'value': 'opaque', '@type': 'xyz'}, } TEXT_SINK_SPEC = { '@type': 'TextSink', 'filename': {'value': 'gs://somefile', '@type': 'http://text'}, 'append_trailing_newlines': {'value': True, '@type': 'http://bool'}, } SHUFFLE_SINK_SPEC = { '@type': 'ShuffleSink', 'shuffle_kind': {'value': 'group_keys', '@type': 'xyz'}, 'shuffle_writer_config': {'value': 'opaque', '@type': 'xyz'}, } PARDO_DOFN_SPEC = { '@type': 'DoFn', 'serialized_fn': {'value': 'code', '@type': 'xyz'}, } PARDO_COMBINEFN_SPEC = { '@type': 'DoFn', 'serialized_fn': {'value': 'code', '@type': 'xyz'}, 'phase': {'value': 'add', '@type': 'xyz'}, } CONCAT_SOURCE_SPEC = { 'sources': [ { 'spec': { 'strip_trailing_newlines': {'@type': 'http://schema.org/Boolean', 'value': True}, 'start_offset': {'@type': 'http://schema.org/Integer', 'value': '0'}, 'compression_type': {'@type': 'http://schema.org/Text', 'value': 'AUTO'}, 'end_offset': {'@type': 'http://schema.org/Integer', 'value': '1000000'}, 'filename': {'@type': 'http://schema.org/Text', 'value': 'gs://sort_g/input_small_files/' 'ascii_sort_1MB_input.0000006'}, '@type': 'TextSource' }, 'encoding': { 'component_encodings': [{'@type': 'notused'}, {'@type': 'notused'}], '@type': coders.serialize_coder(coders.PickleCoder()) } }, { 'spec': { 'strip_trailing_newlines': {'@type': 'http://schema.org/Boolean', 'value': True}, 'start_offset': {'@type': 'http://schema.org/Integer', 'value': '0'}, 'compression_type': {'@type': 'http://schema.org/Text', 'value': 'AUTO'}, 'end_offset': {'@type': 'http://schema.org/Integer', 'value': '1000000'}, 'filename': {'@type': 'http://schema.org/Text', 'value': 'gs://sort_g/input_small_files/' 'ascii_sort_1MB_input.0000007'}, '@type': 'TextSource' }, 'encoding': { 'component_encodings': [{'@type': 'notused'}, {'@type': 'notused'}], '@type': coders.serialize_coder(coders.PickleCoder()) } } ], '@type': 'ConcatSource' } CODER = coders.PickleCoder() WINDOWED_CODER = coders.WindowedValueCoder(CODER) CODER_SPEC = CODER.as_cloud_object() WINDOWED_CODER_SPEC = WINDOWED_CODER.as_cloud_object() def add_source_codec_spec(target): target.source.codec = dataflow.Source.CodecValue() for k, v in CODER_SPEC.iteritems(): target.source.codec.additionalProperties.append( dataflow.Source.CodecValue.AdditionalProperty( key=k, value=to_json_value(v))) def add_source_windowed_codec_spec(target): target.source.codec = dataflow.Source.CodecValue() for k, v in WINDOWED_CODER_SPEC.iteritems(): target.source.codec.additionalProperties.append( dataflow.Source.CodecValue.AdditionalProperty( key=k, value=to_json_value(v))) def add_sink_codec_spec(target): target.sink.codec = dataflow.Sink.CodecValue() for k, v in CODER_SPEC.iteritems(): target.sink.codec.additionalProperties.append( dataflow.Sink.CodecValue.AdditionalProperty( key=k, value=to_json_value(v))) def get_instruction_with_outputs(num_outputs=1, **kwargs): pi = dataflow.ParallelInstruction(**kwargs) for _ in xrange(num_outputs): output = dataflow.InstructionOutput() output.codec = dataflow.InstructionOutput.CodecValue() for k, v in CODER_SPEC.iteritems(): output.codec.additionalProperties.append( dataflow.InstructionOutput.CodecValue.AdditionalProperty( key=k, value=to_json_value(v))) pi.outputs.append(output) return pi def get_concat_source_to_shuffle_sink_message(): ri = dataflow.ReadInstruction() ri.source = dataflow.Source() ri.source.spec = dataflow.Source.SpecValue() for k, v in CONCAT_SOURCE_SPEC.iteritems(): ri.source.spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) di = dataflow.ParDoInstruction() di.input = dataflow.InstructionInput() di.input.producerInstructionIndex = 1 di.multiOutputInfos = [dataflow.MultiOutputInfo(tag='out')] di.userFn = dataflow.ParDoInstruction.UserFnValue() for k, v in PARDO_DOFN_SPEC.iteritems(): di.userFn.additionalProperties.append( dataflow.ParDoInstruction.UserFnValue.AdditionalProperty( key=k, value=to_json_value(v))) wsi = dataflow.WriteInstruction() wsi.input = dataflow.InstructionInput() wsi.input.producerInstructionIndex = 1 di.input.outputNum = 0 wsi.sink = dataflow.Sink() wsi.sink.spec = dataflow.Sink.SpecValue() for k, v in SHUFFLE_SINK_SPEC.iteritems(): wsi.sink.spec.additionalProperties.append( dataflow.Sink.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) add_sink_codec_spec(wsi) mt = dataflow.MapTask() mt.instructions.append(get_instruction_with_outputs(read=ri)) mt.instructions.append(get_instruction_with_outputs(parDo=di)) mt.instructions.append(dataflow.ParallelInstruction(write=wsi)) wi = dataflow.WorkItem() wi.id = 1234 wi.projectId = 'project' wi.jobId = 'job' wi.mapTask = mt m = dataflow.LeaseWorkItemResponse() m.workItems.append(wi) return m def get_text_source_to_shuffle_sink_message(): ri = dataflow.ReadInstruction() ri.source = dataflow.Source() ri.source.spec = dataflow.Source.SpecValue() for k, v in TEXT_SOURCE_SPEC.iteritems(): ri.source.spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) add_source_codec_spec(ri) di = dataflow.ParDoInstruction() di.input = dataflow.InstructionInput() di.input.producerInstructionIndex = 1 di.multiOutputInfos = [dataflow.MultiOutputInfo(tag='out')] di.userFn = dataflow.ParDoInstruction.UserFnValue() for k, v in PARDO_DOFN_SPEC.iteritems(): di.userFn.additionalProperties.append( dataflow.ParDoInstruction.UserFnValue.AdditionalProperty( key=k, value=to_json_value(v))) wsi = dataflow.WriteInstruction() wsi.input = dataflow.InstructionInput() wsi.input.producerInstructionIndex = 1 di.input.outputNum = 0 wsi.sink = dataflow.Sink() wsi.sink.spec = dataflow.Sink.SpecValue() for k, v in SHUFFLE_SINK_SPEC.iteritems(): wsi.sink.spec.additionalProperties.append( dataflow.Sink.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) add_sink_codec_spec(wsi) mt = dataflow.MapTask() mt.instructions.append(get_instruction_with_outputs(read=ri)) mt.instructions.append(get_instruction_with_outputs(parDo=di)) mt.instructions.append(dataflow.ParallelInstruction(write=wsi)) wi = dataflow.WorkItem() wi.id = 1234 wi.projectId = 'project' wi.jobId = 'job' wi.mapTask = mt m = dataflow.LeaseWorkItemResponse() m.workItems.append(wi) return m def get_shuffle_source_to_text_sink_message(shuffle_source_spec): rsi = dataflow.ReadInstruction() rsi.source = dataflow.Source() rsi.source.spec = dataflow.Source.SpecValue() for k, v in shuffle_source_spec.iteritems(): rsi.source.spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) add_source_codec_spec(rsi) wi = dataflow.WriteInstruction() wi.input = dataflow.InstructionInput() wi.sink = dataflow.Sink() wi.sink.spec = dataflow.Sink.SpecValue() for k, v in TEXT_SINK_SPEC.iteritems(): wi.sink.spec.additionalProperties.append( dataflow.Sink.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) add_sink_codec_spec(wi) mt = dataflow.MapTask() mt.instructions.append(get_instruction_with_outputs(read=rsi)) mt.instructions.append(dataflow.ParallelInstruction(write=wi)) wi = dataflow.WorkItem() wi.id = 1234 wi.projectId = 'project' wi.jobId = 'job' wi.mapTask = mt m = dataflow.LeaseWorkItemResponse() m.workItems.append(wi) return m def get_in_memory_source_to_text_sink_message(): rsi = dataflow.ReadInstruction() rsi.source = dataflow.Source() rsi.source.spec = dataflow.Source.SpecValue() for k, v in IN_MEMORY_SOURCE_SPEC.iteritems(): rsi.source.spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) # Note that the in-memory source spec requires a windowed coder. add_source_windowed_codec_spec(rsi) wi = dataflow.WriteInstruction() wi.input = dataflow.InstructionInput() wi.sink = dataflow.Sink() wi.sink.spec = dataflow.Sink.SpecValue() for k, v in TEXT_SINK_SPEC.iteritems(): wi.sink.spec.additionalProperties.append( dataflow.Sink.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) add_sink_codec_spec(wi) mt = dataflow.MapTask() mt.instructions.append(get_instruction_with_outputs(read=rsi)) mt.instructions.append(dataflow.ParallelInstruction(write=wi)) wi = dataflow.WorkItem() wi.id = 1234 wi.projectId = 'project' wi.jobId = 'job' wi.mapTask = mt m = dataflow.LeaseWorkItemResponse() m.workItems.append(wi) return m def get_in_memory_source_to_flatten_message(): rsi = dataflow.ReadInstruction() rsi.source = dataflow.Source() add_source_codec_spec(rsi) rsi.source.spec = dataflow.Source.SpecValue() for k, v in IN_MEMORY_SOURCE_SPEC.iteritems(): rsi.source.spec.additionalProperties.append( dataflow.Source.SpecValue.AdditionalProperty( key=k, value=to_json_value(v))) # Note that the in-memory source spec requires a windowed coder. add_source_windowed_codec_spec(rsi) fi = dataflow.FlattenInstruction() fi.inputs = [dataflow.InstructionInput()] mt = dataflow.MapTask() mt.instructions.append(get_instruction_with_outputs(read=rsi)) mt.instructions.append(get_instruction_with_outputs(flatten=fi)) wi = dataflow.WorkItem() wi.id = 1234 wi.projectId = 'project' wi.jobId = 'job' wi.mapTask = mt m = dataflow.LeaseWorkItemResponse() m.workItems.append(wi) return m class WorkItemTest(unittest.TestCase): def test_concat_source_to_shuffle_sink(self): work = workitem.get_work_items(get_concat_source_to_shuffle_sink_message()) self.assertIsNotNone(work) expected_sub_sources = [] expected_sub_sources.append( io.TextFileSource( file_path='gs://sort_g/input_small_files/' 'ascii_sort_1MB_input.0000006', start_offset=0, end_offset=1000000, strip_trailing_newlines=True, coder=CODER)) expected_sub_sources.append( io.TextFileSource( file_path='gs://sort_g/input_small_files/' 'ascii_sort_1MB_input.0000007', start_offset=0, end_offset=1000000, strip_trailing_newlines=True, coder=CODER)) expected_concat_source = concat_reader.ConcatSource(expected_sub_sources) self.assertEqual( (work.proto.id, work.map_task.operations), (1234, [ maptask.WorkerRead( expected_concat_source, output_coders=[CODER]), maptask.WorkerDoFn( serialized_fn='code', output_tags=['out'], input=(1, 0), side_inputs=[], output_coders=[CODER]), maptask.WorkerShuffleWrite( shuffle_kind='group_keys', shuffle_writer_config='opaque', input=(1, 0), output_coders=(CODER,))])) def test_text_source_to_shuffle_sink(self): work = workitem.get_work_items(get_text_source_to_shuffle_sink_message()) self.assertEqual( (work.proto.id, work.map_task.operations), (1234, [ maptask.WorkerRead(io.TextFileSource( file_path='gs://somefile', start_offset=123, end_offset=123123, strip_trailing_newlines=True, coder=CODER), output_coders=[CODER]), maptask.WorkerDoFn( serialized_fn='code', output_tags=['out'], input=(1, 0), side_inputs=[], output_coders=[CODER]), maptask.WorkerShuffleWrite( shuffle_kind='group_keys', shuffle_writer_config='opaque', input=(1, 0), output_coders=(CODER,))])) def test_shuffle_source_to_text_sink(self): work = workitem.get_work_items( get_shuffle_source_to_text_sink_message(GROUPING_SHUFFLE_SOURCE_SPEC)) self.assertEqual( (work.proto.id, work.map_task.operations), (1234, [ maptask.WorkerGroupingShuffleRead( start_shuffle_position='opaque', end_shuffle_position='opaque', shuffle_reader_config='opaque', coder=CODER, output_coders=[CODER]), maptask.WorkerWrite(fileio.NativeTextFileSink( file_path_prefix='gs://somefile', append_trailing_newlines=True, coder=CODER), input=(0, 0), output_coders=(CODER,))])) def test_ungrouped_shuffle_source_to_text_sink(self): work = workitem.get_work_items( get_shuffle_source_to_text_sink_message(UNGROUPED_SHUFFLE_SOURCE_SPEC)) self.assertEqual( (work.proto.id, work.map_task.operations), (1234, [ maptask.WorkerUngroupedShuffleRead( start_shuffle_position='opaque', end_shuffle_position='opaque', shuffle_reader_config='opaque', coder=CODER, output_coders=[CODER]), maptask.WorkerWrite(fileio.NativeTextFileSink( file_path_prefix='gs://somefile', append_trailing_newlines=True, coder=CODER), input=(0, 0), output_coders=(CODER,))])) def test_in_memory_source_to_text_sink(self): work = workitem.get_work_items(get_in_memory_source_to_text_sink_message()) self.assertEqual( (work.proto.id, work.map_task.operations), (1234, [ maptask.WorkerRead( inmemory.InMemorySource( start_index=1, end_index=3, elements=[base64.b64decode(v['value']) for v in IN_MEMORY_ELEMENTS], coder=CODER), output_coders=[CODER]), maptask.WorkerWrite(fileio.NativeTextFileSink( file_path_prefix='gs://somefile', append_trailing_newlines=True, coder=CODER), input=(0, 0), output_coders=(CODER,))])) def test_in_memory_source_to_flatten(self): work = workitem.get_work_items(get_in_memory_source_to_flatten_message()) self.assertEqual( (work.proto.id, work.map_task.operations), (1234, [ maptask.WorkerRead( inmemory.InMemorySource( start_index=1, end_index=3, elements=[base64.b64decode(v['value']) for v in IN_MEMORY_ELEMENTS], coder=CODER), output_coders=[CODER]), maptask.WorkerFlatten( inputs=[(0, 0)], output_coders=[CODER])])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.5/setup.cfg000066400000000000000000000017431272340251000166450ustar00rootroot00000000000000[nosetests] # Allow discovery of Python test files marked executable. exe=True verbosity=2 # Exclude some unit tests because they define command line options. Nose runs # tests by loading the corresponding modules in the same process and the # side-effect of defining command line options makes other tests fail. # # TODO(silviuc): Find a way to run the remaining tests excluded here. # # The following tests are excluded because they try to load the Cython-based # fast_coders module which is not available when running unit tests: # fast_coders_test, typecoders_test, workitem_test, and executor_test. # # The streamingworker_test test is excluded because it depends on protobuf and # gRPC. # # The batchworker_test is excluded because it imports batchworker.py which # in-turn tries to import module 'resource' which does not work for Windows. exclude=examples|bigquery_test|ptransform_test|fast_coders_test|typecoders_test|workitem_test|executor_test|streamingworker_test|batchworker_test DataflowPythonSDK-0.2.5/setup.py000066400000000000000000000110301272340251000165240ustar00rootroot00000000000000# Copyright 2016 Google Inc. All Rights Reserved. # # 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 # # http://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. """Python Dataflow SDK and Worker setup configuration.""" import os import platform import re import setuptools # Currently all compiled modules are optional (for performance only). # Cython is available on the workers, but we don't require it for development. if platform.system() == 'Windows': # Windows doesn't always provide int64_t. cythonize = lambda *args, **kwargs: [] else: try: # pylint: disable=g-statement-before-imports,g-import-not-at-top from Cython.Build import cythonize except ImportError: cythonize = lambda *args, **kwargs: [] # Configure the required packages and scripts to install. REQUIRED_PACKAGES = [ 'dill>=0.2.5', 'google-apitools>=0.5.2', # TODO(silviuc): Reenable api client package dependencies when we can # update the packages to the latest version without affecting previous # SDK releases. # 'google-apitools-bigquery-v2', # 'google-apitools-dataflow-v1b3>=0.4.20160217', # 'google-apitools-storage-v1', 'httplib2>=0.8', 'mock>=1.0.1', 'oauth2client>=2.0.1', 'protorpc>=0.9.1', 'python-gflags>=2.0', 'pyyaml>=3.10', ] CONSOLE_SCRIPTS = [ ] def get_dataflow_version(): global_names = {} execfile(os.path.normpath('./google/cloud/dataflow/version.py'), global_names) return global_names['__version__'] def get_dataflow_docstring(): """Get docstring for Dataflow module and give it an rST title.""" init_file_path = os.path.normpath('./google/cloud/dataflow/__init__.py') try: with open(init_file_path, 'r') as init_file: init_file_contents = init_file.read() except IOError: return None doc_match = re.search(r'"""(.*)"""', init_file_contents, flags=re.DOTALL) if not doc_match: return None docstring = doc_match.group(1).rstrip() title_match = re.match(r'(.*)\.\n\n', docstring) if title_match: # A module docstring has a first line that ends with a period and has a # blank line after it. reStructuredText, the format used by setuptools # (and other Python API documentation tools), wants no trailing period # and a highlighting line of equal signs under the title line. # Convert by removing the period and adding a highlighting line. equalsigns_fill_format = '\n{:=^%d}\n' % title_match.end(1) title_underline = equalsigns_fill_format.format('=') docstring = re.sub(r'\.\n', title_underline, docstring, count=1) return docstring _PYTHON_DATAFLOW_VERSION = get_dataflow_version() def get_download_url(): """Calculate a permanent download URL for this version.""" return ('https://github.com/GoogleCloudPlatform/DataflowPythonSDK/' 'archive/v%s.tar.gz' % _PYTHON_DATAFLOW_VERSION) setuptools.setup( name='python_dataflow', version=_PYTHON_DATAFLOW_VERSION, description='Google Cloud Dataflow SDK for Python', long_description=get_dataflow_docstring(), url='https://cloud.google.com/dataflow/', download_url=get_download_url(), author='Google, Inc.', packages=setuptools.find_packages(), namespace_packages=['google', 'google.cloud'], entry_points={ 'console_scripts': CONSOLE_SCRIPTS, }, ext_modules=cythonize([ '**/*.pyx', 'google/cloud/dataflow/coders/coder_impl.py', 'google/cloud/dataflow/runners/common.py', 'google/cloud/dataflow/worker/executor.py', 'google/cloud/dataflow/utils/counters.py', ]), setup_requires=['nose>=1.0'], install_requires=REQUIRED_PACKAGES, test_suite='nose.collector', zip_safe=False, # PyPI package information. classifiers=[ 'Intended Audience :: End Users/Desktop', 'License :: OSI Approved :: Apache Software License', 'Operating System :: POSIX :: Linux', 'Programming Language :: Python :: 2.7', 'Topic :: Software Development :: Libraries', 'Topic :: Software Development :: Libraries :: Python Modules', ], license='Apache 2.0', keywords='google cloud dataflow', ) DataflowPythonSDK-0.2.5/tox.ini000066400000000000000000000001311272340251000163250ustar00rootroot00000000000000[tox] envlist = py27 [testenv:py27] commands = python setup.py test passenv = TRAVIS*