pax_global_header00006660000000000000000000000064126636760360014531gustar00rootroot0000000000000052 comment=3a56ce7bee50714da06f2dd16e471a328cb4a655 DataflowPythonSDK-0.2.0/000077500000000000000000000000001266367603600150355ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/.gitignore000066400000000000000000000001161266367603600170230ustar00rootroot00000000000000*~ *.py[cod] *.egg-info/ .eggs/ build/ dist/ distribute-* # Test files .tox/ DataflowPythonSDK-0.2.0/.travis.yml000066400000000000000000000001501266367603600171420ustar00rootroot00000000000000language: python sudo: false env: - TOX_ENV=py27 install: - pip install tox script: tox -e $TOX_ENV DataflowPythonSDK-0.2.0/CONTRIBUTING.md000066400000000000000000000025701266367603600172720ustar00rootroot00000000000000Want 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.0/LICENSE000066400000000000000000000261361266367603600160520ustar00rootroot00000000000000 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.0/README.md000066400000000000000000000333331266367603600163210ustar00rootroot00000000000000# 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. ## Status of this Release This is the Google Cloud Dataflow SDK for Python version 0.2.0. It is still early in its development, and significant changes should be expected before the first stable version. Google recently [announced its desire](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. ## 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), then you will need to ensure `setuptools` version 17.1 or newer is installed on your system (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 Clone the SDK from GitHub: git clone https://github.com/GoogleCloudPlatform/DataflowPythonSDK #### Install With a virtual environment active, install the Dataflow package: cd DataflowPythonSDK python setup.py install #### Test After install, run the tests to make sure everything is okay. python setup.py test ## Local execution of a pipeline 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 google/cloud/dataflow/examples/wordcount.py --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 = '%s:DATASET.TABLENAME' % project 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 = '%s:DATASET.TABLENAME' % project 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.0/google/000077500000000000000000000000001266367603600163115ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/__init__.py000066400000000000000000000002261266367603600204220ustar00rootroot00000000000000try: __import__('pkg_resources').declare_namespace(__name__) except ImportError: __path__ = __import__('pkgutil').extend_path(__path__, __name__) DataflowPythonSDK-0.2.0/google/cloud/000077500000000000000000000000001266367603600174175ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/__init__.py000066400000000000000000000002261266367603600215300ustar00rootroot00000000000000try: __import__('pkg_resources').declare_namespace(__name__) except ImportError: __path__ = __import__('pkgutil').extend_path(__path__, __name__) DataflowPythonSDK-0.2.0/google/cloud/dataflow/000077500000000000000000000000001266367603600212205ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/__init__.py000066400000000000000000000042251266367603600233340ustar00rootroot00000000000000# 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. """ 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.0/google/cloud/dataflow/coders/000077500000000000000000000000001266367603600224775ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/coders/__init__.py000066400000000000000000000013051266367603600246070ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/coders/coder_impl.pxd000066400000000000000000000041411266367603600253310ustar00rootroot00000000000000# 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 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 BytesCoderImpl(CoderImpl): pass cdef list small_ints cdef class VarIntCoderImpl(StreamCoderImpl): @cython.locals(ivalue=libc.stdint.int64_t) cpdef bytes encode(self, 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): """A coder for tuple objects.""" pass cdef class WindowedValueCoderImpl(AbstractComponentCoderImpl): """A coder for windowed values.""" cdef object wrapped_value_coder cdef object timestamp_coder cdef object window_coder DataflowPythonSDK-0.2.0/google/cloud/dataflow/coders/coder_impl.py000066400000000000000000000142511266367603600251710ustar00rootroot00000000000000# 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 minimial overhead. """ 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 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 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 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 WindowedValueCoderImpl(AbstractComponentCoderImpl): """A coder for windowed values.""" def __init__(self, wrapped_value_coder, timestamp_coder, window_coder): super(WindowedValueCoderImpl, self).__init__( (wrapped_value_coder, timestamp_coder, window_coder)) self.wrapped_value_coder = wrapped_value_coder self.timestamp_coder = timestamp_coder self.window_coder = window_coder def _extract_components(self, value): return [value.value, value.timestamp, value.windows] def _construct_from_components(self, components): return WindowedValue(components[0], # value components[1], # timestamp components[2]) # windows DataflowPythonSDK-0.2.0/google/cloud/dataflow/coders/coders.py000066400000000000000000000312521266367603600243330ustar00rootroot00000000000000# 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 cPickle as pickle import struct 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 except ImportError: # We fall back to using the stock dill library in tests that don't use the # full Python SDK. import dill 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 # TODO(ccy): Write a Cython implementation of FloatCoder. class FloatCoder(Coder): """A coder used for floating-point values.""" def encode(self, value): return struct.pack('>= 7 if v: bits |= 0x80 self.write_byte(bits) if not v: break def get(self): return ''.join(self.data) 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 DataflowPythonSDK-0.2.0/google/cloud/dataflow/coders/stream.pxd000066400000000000000000000023021266367603600245040ustar00rootroot00000000000000# 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 bytes get(self) cdef extend(self, size_t missing) 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 bytes read_all(self, bint nested=*) DataflowPythonSDK-0.2.0/google/cloud/dataflow/coders/stream.pyx000066400000000000000000000065221266367603600245410ustar00rootroot00000000000000# 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): assert 0 <= val <= 0xFF 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 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 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 DataflowPythonSDK-0.2.0/google/cloud/dataflow/coders/stream_test.py000066400000000000000000000057311266367603600254110ustar00rootroot00000000000000# 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 math import unittest from google.cloud.dataflow.coders import slow_stream class StreamTest(unittest.TestCase): InputStream = slow_stream.InputStream OutputStream = slow_stream.OutputStream 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]) 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 class SlowFastStreamTest(StreamTest): """Runs the test with compiled and uncompiled stream classes.""" InputStream = stream.InputStream OutputStream = slow_stream.OutputStream class FastSlowStreamTest(StreamTest): """Runs the test with uncompiled and compiled stream classes.""" InputStream = slow_stream.InputStream OutputStream = stream.OutputStream except ImportError: pass if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/coders/typecoders.py000066400000000000000000000140561266367603600252400ustar00rootroot00000000000000# 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(long, 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 else: logging.warning('Using fallback coder for typehint: %r.', typehint) coder = self._fallback_coder return coder.from_type_hint(typehint, self) def get_windowed_coder(self, typehint, timestamp_coder=None, window_coder=None): # Values passed between steps are implicitly WindowedValue objects # with a value as well as windows and so should be coded with a # WindowedValueCoder. return coders.WindowedValueCoder( self.get_coder(typehint), timestamp_coder, window_coder) 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] registry = CoderRegistry() DataflowPythonSDK-0.2.0/google/cloud/dataflow/coders/typecoders_test.py000066400000000000000000000102421266367603600262700ustar00rootroot00000000000000# 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))) real_coder = typecoders.registry.get_coder(long) expected_coder = coders.VarIntCoder() 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.0/google/cloud/dataflow/dataflow_test.py000066400000000000000000000326671266367603600244500ustar00rootroot00000000000000# 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('odd', 'even', main='main') # 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.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_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('odd', 'even', main='main') assert_that(results.main, 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_list_without_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) with self.assertRaises(RuntimeError) as e: _ = main_input | FlatMap( 'test', lambda x, ls1, ls2: [[x, ls1, ls2]], AsList(side_list), AsList(side_list)) self.assertTrue( e.exception.message.startswith( 'Transform with label AsList already applied.')) 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, a_list)) pipeline.run() def test_as_dict_without_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) with self.assertRaises(RuntimeError) as e: _ = main_input | FlatMap( 'test', lambda x, dct1, dct2: [[x, dct1, dct2]], AsDict(side_kvs), AsDict(side_kvs)) self.assertTrue( e.exception.message.startswith( 'Transform with label AsDict already applied.')) 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.0/google/cloud/dataflow/error.py000066400000000000000000000023371266367603600227300ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/000077500000000000000000000000001266367603600230365ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/__init__.py000066400000000000000000000000001266367603600251350ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/complete/000077500000000000000000000000001266367603600246465ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/complete/autocomplete.py000066400000000000000000000045651266367603600277330ustar00rootroot00000000000000# 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 sys import google.cloud.dataflow as df def run(argv=sys.argv[1:]): 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.0/google/cloud/dataflow/examples/complete/autocomplete_test.py000066400000000000000000000045521266367603600307660ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/complete/estimate_pi.py000066400000000000000000000070321266367603600275250ustar00rootroot00000000000000# -*- 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 sys 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=sys.argv[1:]): 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.0/google/cloud/dataflow/examples/complete/estimate_pi_test.py000066400000000000000000000030011266367603600305540ustar00rootroot00000000000000# 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('not used') estimate_pi.run([ '--output=%s' % temp_path]) # Parse result file and compare. with open(temp_path) 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.0/google/cloud/dataflow/examples/complete/juliaset/000077500000000000000000000000001266367603600264665ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/complete/juliaset/juliaset/000077500000000000000000000000001266367603600303065ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/complete/juliaset/juliaset/__init__.py000066400000000000000000000000001266367603600324050ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/complete/juliaset/juliaset/juliaset.py000066400000000000000000000102501266367603600324760ustar00rootroot00000000000000# 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 sys 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=sys.argv[1:]): # 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.0/google/cloud/dataflow/examples/complete/juliaset/juliaset/juliaset_test.py000066400000000000000000000051571266367603600335470ustar00rootroot00000000000000# 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']) 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.0/google/cloud/dataflow/examples/complete/juliaset/juliaset_main.py000066400000000000000000000040131266367603600316620ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/complete/juliaset/setup.py000066400000000000000000000100471266367603600302020ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/complete/tfidf.py000066400000000000000000000167531266367603600263300ustar00rootroot00000000000000# 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 sys 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=sys.argv[1:]): """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.0/google/cloud/dataflow/examples/complete/tfidf_test.py000066400000000000000000000055621266367603600273630ustar00rootroot00000000000000# 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')) 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.0/google/cloud/dataflow/examples/complete/top_wikipedia_sessions.py000066400000000000000000000126341266367603600320040ustar00rootroot00000000000000# 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=sys.argv[1:]): """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.0/google/cloud/dataflow/examples/complete/top_wikipedia_sessions_test.py000066400000000000000000000040751266367603600330430ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/cookbook/000077500000000000000000000000001266367603600246445ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/cookbook/bigquery_tornadoes.py000066400000000000000000000065221266367603600311300ustar00rootroot00000000000000# 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 sys 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=sys.argv[1:]): 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.0/google/cloud/dataflow/examples/cookbook/bigquery_tornadoes_test.py000066400000000000000000000026541266367603600321710ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/cookbook/bigshuffle.py000066400000000000000000000054331266367603600273410ustar00rootroot00000000000000# 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 sys import google.cloud.dataflow as df def crc32line(line): return binascii.crc32(line) & 0xffffffff def run(argv=sys.argv[1:]): # 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.0/google/cloud/dataflow/examples/cookbook/bigshuffle_test.py000066400000000000000000000040301266367603600303700ustar00rootroot00000000000000# 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') 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') as input_csum_file: input_csum = input_csum_file.read().strip() output_csum = '' with open(temp_path + '.checksum-output') 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.0/google/cloud/dataflow/examples/cookbook/coders.py000066400000000000000000000054001266367603600264740ustar00rootroot00000000000000# 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 sys 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=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) (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.0/google/cloud/dataflow/examples/cookbook/coders_test.py000066400000000000000000000033551266367603600275420ustar00rootroot00000000000000# 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') 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.0/google/cloud/dataflow/examples/cookbook/combiners_test.py000066400000000000000000000045301266367603600302400ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/cookbook/custom_ptransform.py000066400000000000000000000074161266367603600310130ustar00rootroot00000000000000# 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 sys 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=sys.argv[1:]): 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.0/google/cloud/dataflow/examples/cookbook/custom_ptransform_test.py000066400000000000000000000040751266367603600320500ustar00rootroot00000000000000# 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)) 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.0/google/cloud/dataflow/examples/cookbook/filters.py000066400000000000000000000075031266367603600266730ustar00rootroot00000000000000# 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 sys 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=sys.argv[1:]): """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.0/google/cloud/dataflow/examples/cookbook/filters_test.py000066400000000000000000000045321266367603600277310ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/cookbook/group_with_coder.py000066400000000000000000000077431266367603600305740ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/examples/cookbook/group_with_coder_test.py000066400000000000000000000063001266367603600316170ustar00rootroot00000000000000# 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') 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') 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.0/google/cloud/dataflow/examples/cookbook/mergecontacts.py000066400000000000000000000120031266367603600300500ustar00rootroot00000000000000# 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 sys import google.cloud.dataflow as df def run(argv=sys.argv[1:], 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.0/google/cloud/dataflow/examples/cookbook/mergecontacts_test.py000066400000000000000000000117461266367603600311240ustar00rootroot00000000000000# 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' % 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.0/google/cloud/dataflow/examples/cookbook/multiple_output_pardo.py000066400000000000000000000141271266367603600316630ustar00rootroot00000000000000# 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 --stagingLocation 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 sys 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=sys.argv[1:]): """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.0/google/cloud/dataflow/examples/cookbook/multiple_output_pardo_test.py000066400000000000000000000044331266367603600327210ustar00rootroot00000000000000# 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') as f: contents = f.read() self.assertEqual(expected_char_count, int(contents)) short_words = self.get_wordcount_results(result_prefix + '-short-words') self.assertEqual(sorted(short_words), sorted(self.EXPECTED_SHORT_WORDS)) words = self.get_wordcount_results(result_prefix + '-words') self.assertEqual(sorted(words), sorted(self.EXPECTED_WORDS)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/snippets/000077500000000000000000000000001266367603600247035ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/snippets/bigquery_side_input.py000066400000000000000000000032251266367603600313310ustar00rootroot00000000000000# 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 using a BigQuery source as a side input to a ParDo operation. The workflow will read a table that has a 'month' field (among others) and will write the values for the field to a text sink. """ import logging import google.cloud.dataflow as df from google.cloud.dataflow.utils.options import add_option from google.cloud.dataflow.utils.options import get_options def run(options=None): p = df.Pipeline(options=get_options(options)) (p # pylint: disable=expression-not-assigned | df.Create('one element', ['ignored']) | df.FlatMap( 'filter month', lambda _, rows: [row['month'] for row in rows], df.pvalue.AsIter( p | df.io.Read('read table', df.io.BigQuerySource(p.options.input)))) | df.io.Write('write file', df.io.TextFileSink(p.options.output))) # Actually run the pipeline (all operations above are deferred). p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) add_option('--input', dest='input', required=True) add_option('--output', dest='output', required=True) run() DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/snippets/snippets.py000066400000000000000000000371361266367603600271340ustar00rootroot00000000000000# 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 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.0/google/cloud/dataflow/examples/snippets/snippets_test.py000066400000000000000000000411711266367603600301650ustar00rootroot00000000000000# 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 pvalue from google.cloud.dataflow import typehints from google.cloud.dataflow.examples.snippets import snippets from google.cloud.dataflow.utils.options import PipelineOptions from google.cloud.dataflow.utils.options import TypeOptions 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', options=PipelineOptions(sys.argv)) # [START type_hints_missing_define_numbers] numbers = p | df.Create(['1', '2', '3']) # [START 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]) 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): lines = [] with open(path) 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)) 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 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.0/google/cloud/dataflow/examples/streaming_wordcap.py000066400000000000000000000035741266367603600271310ustar00rootroot00000000000000# 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 sys import google.cloud.dataflow as df def run(argv=sys.argv[1:]): """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.0/google/cloud/dataflow/examples/streaming_wordcount.py000066400000000000000000000050171266367603600275100ustar00rootroot00000000000000# 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 sys import google.cloud.dataflow as df from google.cloud.dataflow.transforms.trigger import AccumulationMode from google.cloud.dataflow.transforms.trigger import AfterCount import google.cloud.dataflow.transforms.window as window def run(argv=sys.argv[1:]): """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(60, 0), trigger=AfterCount(3), accumulation_mode=AccumulationMode.DISCARDING) | 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.0/google/cloud/dataflow/examples/wordcount.py000066400000000000000000000055061266367603600254420ustar00rootroot00000000000000# 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 sys import google.cloud.dataflow as df empty_line_aggregator = df.Aggregator('emptyLines') 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) return re.findall(r'[A-Za-z\']+', text_line) def run(argv=sys.argv[1:]): """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). p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() DataflowPythonSDK-0.2.0/google/cloud/dataflow/examples/wordcount_debugging.py000066400000000000000000000134511266367603600274530ustar00rootroot00000000000000# 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 --stagingLocation 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 sys 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=sys.argv[1:]): """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.0/google/cloud/dataflow/examples/wordcount_debugging_test.py000066400000000000000000000034341266367603600305120ustar00rootroot00000000000000# 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') 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.0/google/cloud/dataflow/examples/wordcount_minimal.py000066400000000000000000000044231266367603600271450ustar00rootroot00000000000000# 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.""" from __future__ import absolute_import import argparse import logging import re import sys import google.cloud.dataflow as df def run(argv=sys.argv[1:]): """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.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.0/google/cloud/dataflow/examples/wordcount_minimal_test.py000066400000000000000000000035161266367603600302060ustar00rootroot00000000000000# 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') 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.0/google/cloud/dataflow/examples/wordcount_test.py000066400000000000000000000033551266367603600265010ustar00rootroot00000000000000# 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') 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.0/google/cloud/dataflow/internal/000077500000000000000000000000001266367603600230345ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/internal/__init__.py000066400000000000000000000000001266367603600251330ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/internal/apiclient.py000066400000000000000000001043271266367603600253650ustar00rootroot00000000000000# 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.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.utils import dependency 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 apitools.clients import storage import apitools.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=False): """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) append_metric( status_object, counter.name, counter.total, counter.elements if counter.aggregation_kind == counter.MEAN else None, tentative=tentative) def append_metric(status_object, metric_name, value1, value2=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 value1: scalar for a Sum or mean_sum for a Mean value2: mean_count for a Mean aggregation (do not provide for a Sum). 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 value2 is None: if is_counter: # Counters are distinguished by having a kind; metrics do not. metric_update.kind = 'Sum' metric_update.scalar = to_json_value(value1, with_type=True) elif value2 > 0: metric_update.kind = 'Mean' metric_update.meanSum = to_json_value(value1, with_type=True) metric_update.meanCount = to_json_value(value2, 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. pass 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.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.is_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.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.is_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) 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).is_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.""" # Checks the whitelisting status of this account. This is just an early # courtesy check to show a warning in case of potential whitelisting errors. # It will not block job submission. Jobs submitted from non-whitelisted # projects will fail to download required files, make no progress and fail # eventually. # # This check will provide a false warning if a project is whitelisted but # not the current user. In that case job will still execute successfully # in the service. # # TODO(altay): Remove once the whitelisting requirements are lifted. try: request = storage.StorageObjectsListRequest( bucket='dataflow-python-docker') self._storage_client.objects.List(request) except exceptions.HttpError as e: if e.status_code == 403: logging.error( '\n*************************************************************\n' 'This account is not whitelisted to run Python-based pipelines ' 'using the Google Cloud Dataflow service. ' 'Make sure that your project is whitelisted before submitting your ' 'job. \nPlease see documentation for getting more information on ' 'getting your project whitelisted.' '\n*************************************************************\n') else: logging.warning('Could not verify whitelisting status.') # 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) 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): """Leases a work item from the service.""" work_request = dataflow.LeaseWorkItemRequest() work_request.workerId = worker.worker_id work_request.requestedLeaseDuration = worker.desired_lease_duration() work_request.currentWorkerTime = worker.current_time work_request.workerCapabilities.append(worker.worker_id) for value in worker.capabilities: work_request.workerCapabilities.append(value) for value in worker.work_types: work_request.workItemTypes.append(value) request = dataflow.DataflowProjectsJobsWorkItemsLeaseRequest() request.jobId = worker.job_id request.projectId = worker.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, work_item, completed, progress, dynamic_split_result_to_report=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: The Worker instance executing the work item. 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. 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 = worker.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) # Look through the work item for metrics to send. for op in work_item.map_task.executed_operations: for counter in op.itercounters(): append_counter(work_item_status, counter, tentative=not completed) report_request = dataflow.ReportWorkItemStatusRequest() report_request.currentWorkerTime = worker.current_time report_request.workerId = worker.worker_id report_request.workItemStatuses.append(work_item_status) request = dataflow.DataflowProjectsJobsWorkItemsReportStatusRequest() request.jobId = worker.job_id request.projectId = worker.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)) DataflowPythonSDK-0.2.0/google/cloud/dataflow/internal/apiclient_test.py000066400000000000000000000107171266367603600264230ustar00rootroot00000000000000# 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 apitools.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.0/google/cloud/dataflow/internal/auth.py000066400000000000000000000122341266367603600243510ustar00rootroot00000000000000# 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 oauth2client.client import SignedJwtAssertionCredentials from google.cloud.dataflow.utils import processes 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 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.') raise 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 Exception('key file not provided for service account.') if not os.path.exists(google_cloud_options.service_account_key_file): raise Exception('Specified service account key file does not exist.') with file(google_cloud_options.service_account_key_file) as f: service_account_key = f.read() 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' ] 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.0/google/cloud/dataflow/internal/json_value.py000066400000000000000000000107651266367603600255640ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/internal/json_value_test.py000066400000000000000000000041741266367603600266200ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/internal/module_test.py000066400000000000000000000023111266367603600257270ustar00rootroot00000000000000# 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) XYZ_OBJECT = Xyz() DataflowPythonSDK-0.2.0/google/cloud/dataflow/internal/pickler.py000066400000000000000000000101671266367603600250440ustar00rootroot00000000000000# 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 types import dill # Monkey patch dill to handle DictionaryTypes correctly. # TODO(silviuc): Make sure we submit the fix upstream to GitHub dill project. # pylint: disable=protected-access dill.dill._reverse_typemap['DictionaryType'] = type({}) dill.dill._reverse_typemap['DictType'] = type({}) # pylint: enable=protected-access 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 v in outer.__dict__.values(): if v is nested_class: return outer 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 = find_containing_class(obj) if containing_class is not None: return pickler.save_reduce( getattr, (containing_class, obj.__name__), obj=obj) return fun(pickler, obj) 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]) # 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): return base64.b64encode(dill.dumps(o)) def loads(s): return dill.loads(base64.b64decode(s)) 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.0/google/cloud/dataflow/internal/pickler_test.py000066400000000000000000000045101266367603600260760ustar00rootroot00000000000000# 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) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/internal/util.py000066400000000000000000000065051266367603600243710ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/internal/util_test.py000066400000000000000000000041111266367603600254170ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/internal/windmill_pb2.py000066400000000000000000002551661266367603600260070ustar00rootroot00000000000000# 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\"\\\n\x12KeyedMessageBundle\x12\x0b\n\x03key\x18\x01 \x02(\x0c\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\"\xbe\x02\n\x08WorkItem\x12\x0b\n\x03key\x18\x01 \x02(\x0c\x12\x12\n\nwork_token\x18\x02 \x02(\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\"\xc9\x01\n\x13KeyedGetDataRequest\x12\x0b\n\x03key\x18\x01 \x02(\x0c\x12\x12\n\nwork_token\x18\x02 \x02(\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\"\x8e\x05\n\x15WorkItemCommitRequest\x12\x0b\n\x03key\x18\x01 \x02(\x0c\x12\x12\n\nwork_token\x18\x02 \x02(\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\"\xb2\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\'\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=3238, serialized_end=3281, ) _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='messages', full_name='windmill.KeyedMessageBundle.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='messages_ids', full_name='windmill.KeyedMessageBundle.messages_ids', index=2, 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=488, ) _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=491, serialized_end=626, ) _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=628, serialized_end=744, ) _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=746, serialized_end=792, ) _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=794, serialized_end=856, ) _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=858, serialized_end=935, ) _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=938, serialized_end=1157, ) _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=1159, serialized_end=1203, ) _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=1205, serialized_end=1312, ) _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=1314, serialized_end=1364, ) _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=1366, serialized_end=1455, ) _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='cache_token', full_name='windmill.WorkItem.cache_token', index=2, 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=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='timers', full_name='windmill.WorkItem.timers', index=4, 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=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), _descriptor.FieldDescriptor( name='source_state', full_name='windmill.WorkItem.source_state', index=6, 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=7, 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=1458, serialized_end=1776, ) _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=1779, serialized_end=1977, ) _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=1980, serialized_end=2121, ) _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=2123, serialized_end=2186, ) _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='values_to_fetch', full_name='windmill.KeyedGetDataRequest.values_to_fetch', 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_to_fetch', full_name='windmill.KeyedGetDataRequest.lists_to_fetch', 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_to_fetch', full_name='windmill.KeyedGetDataRequest.watermark_holds_to_fetch', 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=2189, serialized_end=2390, ) _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=2392, serialized_end=2492, ) _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=2495, serialized_end=2701, ) _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=2704, serialized_end=2875, ) _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=2877, serialized_end=2975, ) _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=2977, serialized_end=3089, ) _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=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=3092, serialized_end=3281, ) _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=3284, serialized_end=3425, ) _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='output_messages', full_name='windmill.WorkItemCommitRequest.output_messages', 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='pubsub_messages', full_name='windmill.WorkItemCommitRequest.pubsub_messages', index=3, 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=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='value_updates', full_name='windmill.WorkItemCommitRequest.value_updates', 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), _descriptor.FieldDescriptor( name='list_updates', full_name='windmill.WorkItemCommitRequest.list_updates', index=6, 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=7, 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=8, 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=9, 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=10, 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=11, 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=12, 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=13, 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=3428, serialized_end=4082, ) _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=4084, serialized_end=4189, ) _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=4191, serialized_end=4284, ) _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=4286, serialized_end=4306, ) _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=4308, serialized_end=4364, ) _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=4575, serialized_end=4629, ) _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=4631, serialized_end=4711, ) _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=4367, serialized_end=4711, ) _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=4713, serialized_end=4782, ) _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='exceptions', full_name='windmill.ReportStatsRequest.exceptions', 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='counter_updates', full_name='windmill.ReportStatsRequest.counter_updates', 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=4785, serialized_end=4963, ) _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=4965, serialized_end=5002, ) _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')) import abc from grpc.beta import implementations as beta_implementations from grpc.framework.common import cardinality from grpc.framework.interfaces.face import utilities as face_utilities # @@protoc_insertion_point(module_scope) DataflowPythonSDK-0.2.0/google/cloud/dataflow/internal/windmill_service_pb2.py000066400000000000000000000213531266367603600275140ustar00rootroot00000000000000# 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.ReportStatsResponse') , dependencies=[windmill__pb2.DESCRIPTOR,]) _sym_db.RegisterFileDescriptor(DESCRIPTOR) import abc from grpc.beta import implementations as beta_implementations from grpc.framework.common import cardinality from grpc.framework.interfaces.face import utilities as face_utilities class BetaCloudWindmillServiceV1Alpha1Servicer(object): """""" __metaclass__ = abc.ABCMeta @abc.abstractmethod def GetWork(self, request, context): raise NotImplementedError() @abc.abstractmethod def GetData(self, request, context): raise NotImplementedError() @abc.abstractmethod def CommitWork(self, request, context): raise NotImplementedError() @abc.abstractmethod def GetConfig(self, request, context): raise NotImplementedError() @abc.abstractmethod def ReportStats(self, request, context): raise NotImplementedError() class BetaCloudWindmillServiceV1Alpha1Stub(object): """The interface to which stubs will conform.""" __metaclass__ = abc.ABCMeta @abc.abstractmethod def GetWork(self, request, timeout): raise NotImplementedError() GetWork.future = None @abc.abstractmethod def GetData(self, request, timeout): raise NotImplementedError() GetData.future = None @abc.abstractmethod def CommitWork(self, request, timeout): raise NotImplementedError() CommitWork.future = None @abc.abstractmethod def GetConfig(self, request, timeout): raise NotImplementedError() GetConfig.future = None @abc.abstractmethod def ReportStats(self, request, timeout): raise NotImplementedError() ReportStats.future = None def beta_create_CloudWindmillServiceV1Alpha1_server(servicer, pool=None, pool_size=None, default_timeout=None, maximum_timeout=None): import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 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): import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 import windmill_pb2 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.0/google/cloud/dataflow/io/000077500000000000000000000000001266367603600216275ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/io/__init__.py000066400000000000000000000021031266367603600237340ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/io/bigquery.py000066400000000000000000001033161266367603600240340ustar00rootroot00000000000000# 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 apitools.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.Source): """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: A bigquery.TableSchema instance or a dictionary associating field names with types. Possible types are: STRING, INTEGER, FLOAT, BOOLEAN, TIMESTAMP, RECORD (e.g., {'month': 'STRING', 'count': 'INTEGER'}). 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 == '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.SourceReader): """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: self.query = 'SELECT * FROM [%s:%s.%s];' % ( self.source.table_reference.projectId, 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 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) 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.0/google/cloud/dataflow/io/bigquery_test.py000066400000000000000000000430601266367603600250720ustar00rootroot00000000000000# 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 apitools.base.py.exceptions import HttpError from apitools.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) 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))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/io/fileio.py000066400000000000000000000370551266367603600234620ustar00rootroot00000000000000# 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 logging import os import re import tempfile 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.Source): """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) class TextFileSink(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()): """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. """ if not isinstance(file_path_prefix, basestring): raise TypeError( '%s: file_path_prefix must be a string; got %r instead' % (self.__class__.__name__, file_path_prefix)) if not isinstance(file_name_suffix, basestring): raise TypeError( '%s: file_name_suffix must be a string; got %r instead' % (self.__class__.__name__, file_name_suffix)) # 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. 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) self.shard_name_template = ( shard_name_template if shard_name_template is not None else '-SSSSS-of-NNNNN') # 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.SourceReader): """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_return_record_at( is_at_split_point=True, 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.current_offset)) 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.get_position_for_fraction_consumed( 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_at_position(split_position.byte_offset): return iobase.DynamicSplitResultWithPosition(split_position) else: return class TextMultiFileReader(iobase.SourceReader): """A reader for a multi-file text source.""" def __init__(self, source): self.source = source if source.is_gcs_source: # pylint: disable=g-import-not-at-top from google.cloud.dataflow.io import gcsio self.file_paths = gcsio.GcsIO().glob(self.source.file_path) else: self.file_paths = glob.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.0/google/cloud/dataflow/io/fileio_test.py000066400000000000000000000317001266367603600245100ustar00rootroot00000000000000# -*- 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 logging import tempfile import unittest 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: # Starting value of percent_complete might be larger than zero since we # will ignore the first line of a record if the starting position of the # line is smaller than the start offset of the source. self.assertGreaterEqual(reader.get_progress().position.byte_offset, 0) 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(0, reader.get_progress().position.byte_offset) for line in reader: self.assertIsNotNone(line) progress_record.append(reader.get_progress().position.byte_offset) self.assertEqual(18, reader.get_progress().position.byte_offset) self.assertEqual(len(progress_record), 3) self.assertEqual(progress_record, [6, 13, 18]) 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 TestTextFileSink(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.TextFileSink(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) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/io/gcsio.py000066400000000000000000000372621266367603600233170ustar00rootroot00000000000000# 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 apitools.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 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.0/google/cloud/dataflow/io/gcsio_test.py000066400000000000000000000307241266367603600243520ustar00rootroot00000000000000# 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 apitools.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.0/google/cloud/dataflow/io/iobase.py000066400000000000000000001007561266367603600234540ustar00rootroot00000000000000# 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. """ import logging import uuid from google.cloud.dataflow import pvalue 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=%s' % (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', 'elements', 'append_trailing_newlines', 'strip_trailing_newlines', 'compression_type'] class Source(object): """Generic source.""" def reader(self): """Returns a SourceReader 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 SourceReader(object): """A generic reader for a source.""" 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 'SourceReader' 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 'SourceReader' 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 'SourceReader' 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 SourceReader has read through the source.""" 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 through a 'SourceReader'.""" 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 'SourceReader.request_dynamic_split' should split the input. """ 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 class RangeTracker(object): """A thread-safe helper object for implementing dynamic work rebalancing. 'RangeTracker' can be used for implementing dynamic work rebalancing in position-based subclasses of 'iobase.SourceReader'. Usage of the RangeTracker class hierarchy ----------------------------------------- The 'RangeTracker' class should not be used per se---all users should use its subclasses directly. We declare it here because all subclasses have roughly the same interface and the same properties, to centralize the documentation. Currently we provide one implementation: 'iobase.OffsetRangeTracker'. 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_return_record_at(self, is_at_split_point, record_start): """Atomically determines if a record at the given position can be returned. Additionally, Updates the internal state of the 'RangeTracker'. In particular: * If 'is_at_split_point' is 'True', and 'record_start' is outside the current range, returns 'False'; * Otherwise, updates the last-consumed position to 'record_start' and returns 'True'. This method MUST be called on all split point records. It may be called on every record. Method 'try_return_record_at' and method 'try_split_at_position' will be accessed by different threads and implementor must ensure that only one of these methods is executed at a given time. Args: is_at_split_point: 'True' if record is at a split point, 'False' otherwise. record_start: starting position of the record. """ raise NotImplementedError def try_split_at_position(self, split_position): """Atomically splits the current range. Splits the current range '[get_start_position(), get_stop_position())' into a "primary" part '[get_start_position(), split_position())' and a "residual" part '[split_position(), get_stop_position())', assuming the current last-consumed position is within '[get_start_position(), split_position())' (i.e., 'split_position()' has not been consumed yet). Updates the current range to be the primary and returns 'True'. This means that all further calls on the current object will interpret their arguments relative to the primary range. If the split position has already been consumed, or if no 'try_return_record_at' call was made yet, returns 'False'. The second condition is to prevent dynamic splitting during reader start-up. Method 'try_return_record_at' and method 'try_split_at_position' will be accessed by different threads and implementor must ensure that only one of these methods is executed at a given time. Args: split_position: an instance of ReaderPosition that gives the position where the current range should be split at. """ raise NotImplementedError def fraction_consumed(self): """Returns the approximate fraction of consumed positions in the source. Returns the approximate fraction of positions that have been consumed by successful 'try_return_record_at()' 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. 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's. 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's 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. """ def __init__(self, *args, **kwargs): """Initializes a Write transform. Args: *args: A tuple of position arguments. **kwargs: A dictionary of keyword arguments. Returns: a 'pvalue.PDone'. 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.PValue(pipeline=pcoll.pipeline, transform=self) 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(pipeline=self.pipeline, transform=self) 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: sink.finalize_write(init_result, write_results), 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 self.writer.close() DataflowPythonSDK-0.2.0/google/cloud/dataflow/io/pubsub.py000066400000000000000000000030341266367603600235010ustar00rootroot00000000000000# 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.Source): """Source for reading from a given Cloud Pub/Sub topic.""" def __init__(self, topic, subscription=None, coder=coders.StrUtf8Coder()): self.topic = topic self.subscription = subscription self.coder = coder @property def format(self): """Source format name required for remote execution.""" return 'pubsub' 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' DataflowPythonSDK-0.2.0/google/cloud/dataflow/io/range_trackers.py000066400000000000000000000226141266367603600252000ustar00rootroot00000000000000# 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): self._start_offset = start self._stop_offset = end self.last_record_start = -1 self.offset_of_last_split_point = -1 self.lock = threading.Lock() @property def start_position(self): return self._start_offset @property def stop_position(self): return self._stop_offset def try_return_record_at(self, is_at_split_point, record_start): with self.lock: if self.last_record_start == -1 and not is_at_split_point: raise Exception( 'The first record [starting at %d] must be at a split point' % record_start) if record_start < self.last_record_start: raise Exception( '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 is_at_split_point: if (self.offset_of_last_split_point is not -1 and record_start is self.offset_of_last_split_point): raise Exception( 'Record at a split point has same offset as the previous split ' 'point: %d' % record_start) 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 try_split_at_position(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 False if self.last_record_start == -1: logging.debug('Refusing to split %r at %d: unstarted', self, split_offset) return False if split_offset <= self.last_record_start: logging.debug( 'Refusing to split %r at %d: already past proposed stop offset', self, split_offset) return False 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 False logging.debug('Agreeing to split %r at %d', self, split_offset) self._stop_offset = split_offset return True @property def fraction_consumed(self): with self.lock: fraction = ((1.0 * (self.last_record_start - self.start_position + 1) / (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 get_position_for_fraction_consumed(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): 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() @property def start_position(self): return self._decoded_start_pos @property def stop_position(self): return self._decoded_stop_pos @property def last_group_start(self): return self._decoded_last_group_start def try_return_record_at(self, is_at_split_point, decoded_group_start): with self.lock: if self.last_group_start is None and not is_at_split_point: raise ValueError('The first group [at %r] must be at a split point' % decoded_group_start) 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 is not None 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 is_at_split_point: if (self.last_group_start is not None 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 (self.stop_position and decoded_group_start >= self.stop_position): return False elif (self.last_group_start is not None and 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-'))) self._decoded_last_group_start = decoded_group_start self.last_group_was_at_a_split_point = is_at_split_point return True def try_split_at_position(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 False 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 False 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 False logging.info('Agreeing to split %r at %r' , self, decoded_split_position) self._decoded_stop_pos = decoded_split_position return True @property 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.0/google/cloud/dataflow/io/range_trackers_test.py000066400000000000000000000351261266367603600262410ustar00rootroot00000000000000# 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_return_record_at(True, 110)) self.assertTrue(tracker.try_return_record_at(True, 140)) self.assertTrue(tracker.try_return_record_at(True, 183)) self.assertFalse(tracker.try_return_record_at(True, 210)) def test_try_return_record_simple_dense(self): tracker = range_trackers.OffsetRangeTracker(3, 6) self.assertTrue(tracker.try_return_record_at(True, 3)) self.assertTrue(tracker.try_return_record_at(True, 4)) self.assertTrue(tracker.try_return_record_at(True, 5)) self.assertFalse(tracker.try_return_record_at(True, 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_return_record_at(True, 10)) self.assertTrue(tracker.try_return_record_at(False, 12)) self.assertTrue(tracker.try_return_record_at(False, 14)) self.assertTrue(tracker.try_return_record_at(True, 16)) # Out of range, but not a split point... self.assertTrue(tracker.try_return_record_at(False, 18)) self.assertTrue(tracker.try_return_record_at(False, 20)) # Out of range AND a split point. self.assertFalse(tracker.try_return_record_at(True, 22)) def test_split_at_offset_fails_if_unstarted(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertFalse(tracker.try_split_at_position(150)) def test_split_at_offset(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertTrue(tracker.try_return_record_at(True, 110)) # Example positions we shouldn't split at, when last record starts at 110: self.assertFalse(tracker.try_split_at_position(109)) self.assertFalse(tracker.try_split_at_position(110)) self.assertFalse(tracker.try_split_at_position(200)) self.assertFalse(tracker.try_split_at_position(210)) # Example positions we *should* split at: self.assertTrue(copy.copy(tracker).try_split_at_position(111)) self.assertTrue(copy.copy(tracker).try_split_at_position(129)) self.assertTrue(copy.copy(tracker).try_split_at_position(130)) self.assertTrue(copy.copy(tracker).try_split_at_position(131)) self.assertTrue(copy.copy(tracker).try_split_at_position(150)) self.assertTrue(copy.copy(tracker).try_split_at_position(199)) # If we split at 170 and then at 150: self.assertTrue(tracker.try_split_at_position(170)) self.assertTrue(tracker.try_split_at_position(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_return_record_at(True, 135)) self.assertTrue(copy.copy(tracker).try_return_record_at(True, 135)) # Should be able to return a record starting right before the new stop # offset. self.assertTrue(copy.copy(tracker).try_return_record_at(True, 149)) # Should not be able to return a record starting at or after the new stop # offset. self.assertFalse(tracker.try_return_record_at(True, 150)) self.assertFalse(tracker.try_return_record_at(True, 151)) # Should accept non-splitpoint records starting after stop offset. self.assertTrue(tracker.try_return_record_at(False, 135)) self.assertTrue(tracker.try_return_record_at(False, 152)) self.assertTrue(tracker.try_return_record_at(False, 160)) self.assertTrue(tracker.try_return_record_at(False, 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.get_position_for_fraction_consumed(0.0)) # [3, 4) represents up to 1/3 of [3, 6) self.assertEqual(4, tracker.get_position_for_fraction_consumed(1.0 / 6)) self.assertEqual(4, tracker.get_position_for_fraction_consumed(0.333)) # [3, 5) represents up to 2/3 of [3, 6) self.assertEqual(5, tracker.get_position_for_fraction_consumed(0.334)) self.assertEqual(5, tracker.get_position_for_fraction_consumed(0.666)) # Any fraction consumed over 2/3 means the whole [3, 6) has been consumed. self.assertEqual(6, tracker.get_position_for_fraction_consumed(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_return_record_at(True, 3)) self.assertEqual(1.0 / 3, tracker.fraction_consumed) self.assertTrue(tracker.try_return_record_at(True, 4)) self.assertEqual(2.0 / 3, tracker.fraction_consumed) self.assertTrue(tracker.try_return_record_at(True, 5)) self.assertEqual(1.0, tracker.fraction_consumed) self.assertTrue(tracker.try_return_record_at(False, 6)) # non-split-point self.assertTrue(tracker.try_return_record_at(False, 7)) # non-split-point self.assertFalse(tracker.try_return_record_at(True, 7)) def test_get_fraction_consumed_sparse(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertEqual(0, tracker.fraction_consumed) self.assertTrue(tracker.try_return_record_at(True, 110)) # Consumed positions through 110 = total 11 positions of 100. self.assertEqual(0.11, tracker.fraction_consumed) self.assertTrue(tracker.try_return_record_at(True, 150)) self.assertEqual(0.51, tracker.fraction_consumed) self.assertTrue(tracker.try_return_record_at(True, 195)) self.assertEqual(0.96, tracker.fraction_consumed) def test_everything_with_unbounded_range(self): tracker = range_trackers.OffsetRangeTracker( 100, range_trackers.OffsetRangeTracker.OFFSET_INFINITY) self.assertTrue(tracker.try_return_record_at(True, 150)) self.assertTrue(tracker.try_return_record_at(True, 250)) # get_position_for_fraction_consumed should fail for an unbounded range with self.assertRaises(Exception): tracker.get_position_for_fraction_consumed(0.5) def test_try_return_first_record_not_split_point(self): with self.assertRaises(Exception): range_trackers.OffsetRangeTracker(100, 200).try_return_record_at( False, 120) def test_try_return_record_non_monotonic(self): tracker = range_trackers.OffsetRangeTracker(100, 200) tracker.try_return_record_at(True, 120) with self.assertRaises(Exception): tracker.try_return_record_at(True, 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_return_record_at( True, self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([1, 2, 5]))) self.assertTrue(tracker.try_return_record_at( True, 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_return_record_at( True, self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([1, 2, 5]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([3, 6, 8, 10]))) self.assertTrue(tracker.try_return_record_at( True, 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_return_record_at( True, self.bytes_to_position([5, 0, 0]))) self.assertFalse(copy.copy(tracker).try_return_record_at( True, self.bytes_to_position([5, 0, 1]))) self.assertFalse(copy.copy(tracker).try_return_record_at( True, 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_return_record_at( True, self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_return_record_at( False, self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_return_record_at( False, self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([1, 2, 5]))) self.assertTrue(tracker.try_return_record_at( False, self.bytes_to_position([1, 2, 5]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([3, 6, 8, 10]))) self.assertTrue(tracker.try_return_record_at( True, 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.try_return_record_at(False, 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])) tracker.try_return_record_at(True, self.bytes_to_position([3, 4, 5])) with self.assertRaises(ValueError): tracker.try_return_record_at(False, 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_return_record_at(True, 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])) tracker.try_return_record_at(True, self.bytes_to_position([3, 4, 5])) tracker.try_return_record_at(True, self.bytes_to_position([3, 4, 6])) with self.assertRaises(ValueError): tracker.try_return_record_at(True, 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])) tracker.try_return_record_at(True, self.bytes_to_position([3, 4, 5])) with self.assertRaises(ValueError): tracker.try_return_record_at(True, 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_at_position( self.bytes_to_position([3, 4, 5, 6]))) tracker.try_return_record_at(True, self.bytes_to_position([1, 2, 3])) # Should now succeed. self.assertTrue(tracker.try_split_at_position( self.bytes_to_position([3, 4, 5, 6]))) # Should not split at same or larger position. self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([3, 4, 5, 6]))) self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([3, 4, 5, 6, 7]))) self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([4, 5, 6, 7]))) # Should split at smaller position. self.assertTrue(tracker.try_split_at_position( self.bytes_to_position([3, 2, 1]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([2, 3, 4]))) # Should not split at a position we're already past. self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([2, 3, 4]))) self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([2, 3, 3]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([3, 2, 0]))) self.assertFalse(tracker.try_return_record_at( True, 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_at_position( self.bytes_to_position([0, 0, 0]))) self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([3, 4, 5, 6]))) tracker.try_return_record_at(True, self.bytes_to_position([1, 2, 3])) # Should now succeed. self.assertTrue(tracker.try_split_at_position( self.bytes_to_position([3, 4, 5, 6]))) # Should not split at same or larger position. self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([3, 4, 5, 6]))) self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([3, 4, 5, 6, 7]))) self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([4, 5, 6, 7]))) # Should split at smaller position. self.assertTrue(tracker.try_split_at_position( self.bytes_to_position([3, 2, 1]))) # But not at a position at or before last returned record. self.assertFalse(tracker.try_split_at_position( self.bytes_to_position([1, 2, 3]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([2, 3, 4]))) self.assertTrue(tracker.try_return_record_at( True, self.bytes_to_position([3, 2, 0]))) self.assertFalse(tracker.try_return_record_at( True, self.bytes_to_position([3, 2, 1]))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/pipeline.py000066400000000000000000000367511266367603600234130ustar00rootroot00000000000000# 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 from google.cloud.dataflow import error from google.cloud.dataflow import pvalue from google.cloud.dataflow import typehints 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 StandardOptions from google.cloud.dataflow.utils.options import TypeOptions 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 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.') # List of PValue objects representing a DAG of transformations. self._nodes = [] # 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() def _add_pvalue(self, pval): """Adds a PValue to the pipeline's node list.""" if pval not in self._nodes: self._nodes.append(pval) 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.""" return self.runner.run(self) def visit(self, visitor, node=None): """Visits depth-first every node of a pipeline's DAG. If node is specified then only that node's predecessors (inputs and recursively their creating transforms) and outputs will be visited. Args: visitor: PipelineVisitor object whose callbacks will be called for each node visited. See PipelineVisitor comments. node: if specified it is expected to be a PValue and only the nodes of the DAG reachable from this node will be visited. 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. """ # Make sure the specified node has its transform registered as an output # producer. We can have this situation for PCollections created as results # of accessing a tag of a FlatMap().with_outputs() result. if node is not None: if not isinstance(node, pvalue.PValue): raise TypeError( 'Expected a PValue for the node argument instead of: %r' % node) if node not in self._nodes: raise error.PipelineError('PValue not in pipeline: %r' % node) assert node.producer is not None visited = set() start_transform = self._root_transform() if node is None else node.producer start_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): class CallableTransform(ptransform.PTransform): def __init__(self, callee): super(CallableTransform, self).__init__( label=getattr(callee, '__name__', 'Callable')) self._callee = callee def apply(self, *args, **kwargs): return self._callee(*args, **kwargs) assert callable(transform) transform = CallableTransform(transform) full_label = format_full_label(self._current_transform(), transform) if full_label in self.applied_labels: raise RuntimeError( 'Transform with label %s already applied. Please clone the current ' 'instance using a new label or alternatively create a new instance. ' 'To clone a transform 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)) child = AppliedPTransform( self._current_transform(), transform, full_label, inputs) self._current_transform().add_part(child) self.transforms_stack.append(child) 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 = child self._current_transform().add_output(result) # TODO(robertwb): Multi-input, multi-output inference. if (type_options is not None and type_options.pipeline_type_check and isinstance(result, pvalue.PCollection) 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) 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) self.transforms_stack.pop() return pvalueish_result 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 = [] def add_output(self, output): assert (isinstance(output, pvalue.PValue) or isinstance(output, pvalue.DoOutputsTuple)) self.outputs.append(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 exmaple of a transform that is not a producer is one that returns its inputs instead.) """ return self.parts or all(pval.producer is not self for pval in self.outputs) 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 # Visit side inputs. for pval in self.side_inputs: if isinstance(pval, pvalue.AsSideInput) and pval.pvalue not in visited: pval = pval.pvalue # Unpack marker-object-wrapped pvalue. 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: 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.0/google/cloud/dataflow/pipeline_test.py000066400000000000000000000206071266367603600244430ustar00rootroot00000000000000# 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 unittest from google.cloud.dataflow.io.iobase import Source 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.runners import DirectPipelineRunner from google.cloud.dataflow.transforms import Create from google.cloud.dataflow.transforms import FlatMap 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(Source): """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): @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('DirectPipelineRunner') 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(DirectPipelineRunner()) pcoll = pipeline | Create('label', [[1, 2, 3]]) assert_that(pcoll, equal_to([[1, 2, 3]])) pipeline.run() def test_read(self): pipeline = Pipeline('DirectPipelineRunner') 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(DirectPipelineRunner()) 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_visit_node_sub_graph(self): pipeline = Pipeline('DirectPipelineRunner') 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]) visitor = PipelineTest.Visitor(visited=[]) pipeline.visit(visitor, node=pcoll3) self.assertFalse(pcoll4 in visitor.visited) self.assertEqual(set([pcoll1, pcoll2, pcoll3]), set(visitor.visited)) def test_apply_custom_transform(self): pipeline = Pipeline(DirectPipelineRunner()) 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(DirectPipelineRunner()) 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 with label %s already applied. Please clone the current ' 'instance using a new label or alternatively create a new instance. ' 'To clone a transform use: transform.clone(\'NEW LABEL\').' % 'CustomTransform') def test_reuse_cloned_custom_transform_instance(self): pipeline = Pipeline(DirectPipelineRunner()) 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('DirectPipelineRunner') pcoll = pipeline | Create('pcoll', [1, 2, 3]) result = pipeline.apply(PipelineTest.custom_callable, pcoll) assert_that(result, equal_to([2, 3, 4])) pipeline.run() 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( ['slices', 'style', 'view_as'], [attr for attr in dir(options) if not attr.startswith('_')]) self.assertEquals( ['style', 'view_as'], [attr for attr in dir(options.view_as(Eggs)) if not attr.startswith('_')]) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/pvalue.py000066400000000000000000000330571266367603600230760ustar00rootroot00000000000000# 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 from google.cloud.dataflow import error from google.cloud.dataflow import typehints class PValue(object): """Base class for PCollection. 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, **kwargs): """Initializes a PValue with all arguments hidden behind keyword arguments. Args: **kwargs: keyword arguments. Raises: ValueError: if the expected keyword arguments (pipeline, transform, and optionally tag) are not present. The method expects a pipeline and a transform keyword argument. However in order to give a signal to users that they should not create these PValues directly we obfuscate the arguments. """ if 'pipeline' not in kwargs or 'transform' not in kwargs: raise ValueError( 'Missing required arguments (pipeline and transform): %s' % kwargs.keys) self.pipeline = kwargs.pop('pipeline') # TODO(silviuc): Remove usage of the transform argument from all call sites. # It is not used anymore and has been replaced with the producer attribute. kwargs.pop('transform') self.tag = kwargs.pop('tag', None) self.element_type = kwargs.pop('element_type', None) if kwargs: raise ValueError('Unexpected keyword arguments: %s' % kwargs.keys()) self.pipeline._add_pvalue(self) # 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.""" def __init__(self, **kwargs): """Initializes a PCollection. Do not call directly.""" super(PCollection, self).__init__(**kwargs) @property def windowing(self): if not hasattr(self, '_windowing'): self._windowing = self.producer.transform.get_windowing( self.producer.inputs) return self._windowing def _get_un_windowed_value(self, values): """Converts an iterable of WindowedValue(s) into underlying values.""" for v in values: yield v.value # TODO(silviuc): Remove uses of this method and delete it. def _get_values(self): def _get_internal(self, runner=None): """Materializes a PValue by executing its subtree of transforms.""" runner = runner or self.pipeline.runner if not runner: raise error.RunnerError( 'get() cannot find a runner to execute pipeline.') runner.run(self.pipeline, node=self) # Internally all values are WindowedValue(s) and we want to return only # the underlying value or values depending on the type of the PValue. return self._get_un_windowed_value(runner.get_pvalue(self)) return list(_get_internal(self)) 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 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): 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( pipeline=self._pipeline, transform=self._transform, tag=tag) # Transfer the producer from the DoOutputsTuple to the resulting # PCollection. pcoll.producer = self.producer self._pcolls[tag] = pcoll return pcoll def _get_values(self): return ExplicitDoOutputsTuple(self) class ExplicitDoOutputsTuple(DoOutputsTuple): def __init__(self, deferred): super(ExplicitDoOutputsTuple, self).__init__( None, None, deferred._tags, deferred._main_tag) self._deferred = deferred def __getitem__(self, tag): pcoll = self._deferred[tag] return pcoll._get_un_windowed_value( self._deferred._pipeline.runner.get_pvalue(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 AsSideInput(object): """Marker specifying that a PCollection will be used as a side input.""" @property def element_type(self): return typehints.Any class AsSingleton(AsSideInput): """Marker specifying that an entire PCollection is to be used as a side input. When a PCollection is supplied as a side input to a PTransform, it is necessary to indicate whether the entire PCollection should be made available as a PTransform side argument (in the form of an iterable), or whether just one value should be pulled from the PCollection and supplied as the side argument (as an ordinary value). Wrapping a PCollection side input argument to a PTransform in this container (e.g., data.apply('label', MyPTransform(), AsSingleton(my_side_input) ) selects the latter behavor. (Note: This marker is agnostic to whether the PValue it wraps is a PCollection. Although PCollections are the only PValues available now, there may be additional PValue types for which AsIter and AsSingleton are useful markers. """ _NO_DEFAULT = object() def __init__(self, pvalue, default_value = _NO_DEFAULT): self.pvalue = pvalue self.default_value = default_value def __repr__(self): return 'AsSingleton(%s)' % self.pvalue @property def element_type(self): return self.pvalue.element_type class AsIter(AsSideInput): """Marker specifying that an entire PCollection is to be used as a side input. When a PCollection is supplied as a side input to a PTransform, it is necessary to indicate whether the entire PCollection should be made available as a PTransform side argument (in the form of an iterable), or whether just one value should be pulled from the PCollection and supplied as the side argument (as an ordinary value). Wrapping a PCollection side input argument to a PTransform in this container (e.g., data.apply('label', MyPTransform(), AsIter(my_side_input) ) selects the former behavor. (Note: This marker is agnostic to whether the PValue it wraps is a PCollection. Although PCollection is the only PValue available now, there may be additional PValue types for which AsIter and AsSingleton are useful markers. """ def __init__(self, pvalue): self.pvalue = pvalue def __repr__(self): return 'AsIter(%s)' % self.pvalue @property def element_type(self): return typehints.Iterable[self.pvalue.element_type] def AsList(pcoll, label='AsList'): # pylint: disable=invalid-name """Convenience function packaging an entire PCollection as a side input list. Intended for use in side-argument specification---the same places where AsSingleton and AsIter are used. Unlike those wrapper classes, AsList (as implemented) is a function that schedules a Combiner to condense pcoll into a single list, then wraps the resulting one-element PCollection in AsSingleton. Args: pcoll: Input pcollection. label: Label to be specified if several AsList's are used in the pipeline at same depth level. Returns: An AsList-wrapper around a PCollection whose one element is a list containing all elements in pcoll. """ # 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. # TODO(silviuc): read directly on the worker from google.cloud.dataflow.transforms import combiners # pylint: disable=g-import-not-at-top return AsSingleton(pcoll | combiners.ToList(label)) def AsDict(pcoll, label='AsDict'): # pylint: disable=invalid-name """Convenience function packaging an entire PCollection as a side input dict. Intended for use in side-argument specification---the same places where AsSingleton and AsIter are used. Unlike those wrapper classes, AsDict (as implemented) is a function that schedules a Combiner to condense pcoll into a single dict, then wraps the resulting one-element PCollection in AsSingleton. Args: pcoll: Input pcollection. All elements should be key-value pairs (i.e. 2-tuples) with unique keys. label: Label to be specified if several AsDict's are used in the pipeline at same depth level. Returns: An AsDict-wrapper around a PCollection whose one element is a dict with entries for uniquely-keyed pairs in pcoll. """ # 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 combiners # pylint: disable=g-import-not-at-top return AsSingleton(pcoll | combiners.ToDict(label)) 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.0/google/cloud/dataflow/pvalue_test.py000066400000000000000000000030201266367603600241200ustar00rootroot00000000000000# 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 PValue from google.cloud.dataflow.transforms import PTransform 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') transform = PTransform() value = PValue(pipeline=pipeline, transform=transform) self.assertEqual(pipeline, value.pipeline) def test_pvalue_missing_arguments(self): self.assertRaises(ValueError, PValue, pipeline=Pipeline('DirectPipelineRunner')) self.assertRaises(ValueError, PValue, transform=PTransform()) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/python_sdk_releases.py000066400000000000000000000031761266367603600256460ustar00rootroot00000000000000# 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.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.0/google/cloud/dataflow/runners/000077500000000000000000000000001266367603600227145ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/runners/__init__.py000066400000000000000000000020421266367603600250230ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/runners/common.py000066400000000000000000000132141266367603600245570ustar00rootroot00000000000000# 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 operations executor.""" import logging 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 from google.cloud.dataflow.utils import counters 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, tagged_counters, 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.tagged_counters = tagged_counters self.logger = logger or FakeLogger() self.step_name = step_name def start(self): self.context.set_element(None) self._process_outputs(None, self.dofn.start_bundle(self.context)) def finish(self): logging.debug('Finishing %s', self) self.context.set_element(None) self._process_outputs(None, self.dofn.finish_bundle(self.context)) def process(self, element): with self.logger.PerThreadLoggingContext(step_name=self.step_name): assert isinstance(element, WindowedValue) self.context.set_element(element) self._process_outputs(element, self.dofn.process(self.context)) 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 = WindowedValue( result, element.timestamp, element.windows) for receiver in self.tagged_receivers[tag]: # TODO(robertwb): Should the counters be on the context? self.tagged_counters[tag].update(windowed_value) receiver.process(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 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): self.step_name = '' self._user_counters = {} def counter_for(self, aggregator): """Looks up the counter for this aggregator, creating one if necessary.""" if aggregator not in self._user_counters: self._user_counters[aggregator] = counters.AggregatorCounter( self.step_name, aggregator) return self._user_counters[aggregator] def itercounters(self): """Returns an iterable of Counters (to be sent to the service).""" return self._user_counters.values() DataflowPythonSDK-0.2.0/google/cloud/dataflow/runners/dataflow_runner.py000066400000000000000000000643321266367603600264700ustar00rootroot00000000000000# 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.pvalue import AsSideInput from google.cloud.dataflow.pvalue import AsSingleton 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.names import PropertyNames from google.cloud.dataflow.utils.names import TransformNames from apitools.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. 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 # 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: logging.info('Job %s is in state %s.', job_id, str(response.currentState)) if str(response.currentState) != 'JOB_STATE_RUNNING': break time.sleep(5.0) # 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, node=None): """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 # starting from the "node" argument (or entire graph if node is None). super(DataflowPipelineRunner, self).run(pipeline, node) # Get a Dataflow API client and submit the job. self.dataflow_client = apiclient.DataflowApplicationClient( pipeline.options, DataflowPipelineRunner.ENVIRONMENT_MAJOR_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_value=True): """Returns an encoding based on a typehint onject.""" return self._get_cloud_encoding(self._get_coder(typehint, window_value=window_value)) def _get_coder(self, typehint, window_value=True): """Returns a coder based on a typehint onject.""" if window_value: coder = coders.registry.get_windowed_coder(typehint) else: coder = coders.registry.get_coder(typehint) return coder 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_transform_type_hint(self, transform_node): """Returns the typehint for a applied transform node or a default.""" if transform_node.outputs[0].element_type is not None: # TODO(robertwb): Handle type hints for multi-output transforms. return 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). return typehints.Any 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 _add_singleton_step(self, label, full_label, tag, input_step): """Creates a CollectionToSingleton step used to handle ParDo side inputs.""" # 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(TransformNames.COLLECTION_TO_SINGLETON, label) self.job.proto.steps.append(step.proto) step.add_property(PropertyNames.USER_NAME, full_label) step.add_property( PropertyNames.PARALLEL_INPUT, {'@type': 'OutputReference', PropertyNames.STEP_NAME: input_step.proto.name, PropertyNames.OUTPUT_NAME: input_step.get_output(tag)}) step.encoding = self._get_side_input_encoding(input_step.encoding) step.add_property( PropertyNames.OUTPUT_INFO, [{PropertyNames.USER_NAME: ( '%s.%s' % (full_label, PropertyNames.OUTPUT)), PropertyNames.ENCODING: step.encoding, PropertyNames.OUTPUT_NAME: PropertyNames.OUTPUT}]) 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_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_typehint_based_encoding( self._get_transform_type_hint(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) if not coder.is_kv_coder(): raise ValueError(('Coder for the GroupByKey operation "%s" is not a ' 'key-value coder: %s.') % (transform.label, coder)) key_coder = coder.key_coder() if not key_coder.is_deterministic(): logging.warning('The key coder "%s" for the GroupByKey operation "%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, transform.label) return pvalue.PCollection(pipeline=pcoll.pipeline, transform=transform) 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_typehint_based_encoding( self._get_transform_type_hint(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]) # Must generate any required steps for side inputs here before creating the # step for the ParDo transform. si_dict = {} si_tags_and_types = [] for side_pval in transform_node.side_inputs: assert isinstance(side_pval, AsSideInput) si_label = self._get_unique_step_name() si_full_label = '%s/%s' % (transform_node.full_label, si_label) self._add_singleton_step( si_label, si_full_label, side_pval.pvalue.tag, self._cache.get_pvalue(side_pval.pvalue)) si_dict[si_label] = { '@type': 'OutputReference', PropertyNames.STEP_NAME: si_label, PropertyNames.OUTPUT_NAME: PropertyNames.OUTPUT} # 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, isinstance(side_pval, AsSingleton))) # 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_typehint_based_encoding( self._get_transform_type_hint(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(pipeline=pcoll.pipeline, transform=transform) 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_typehint_based_encoding( self._get_transform_type_hint(transform_node), window_value=False) output_encoding = self._get_typehint_based_encoding( self._get_transform_type_hint(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 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': step.add_property(PropertyNames.PUBSUB_TOPIC, transform.source.topic) if transform.source.subscription: step.add_property(PropertyNames.PUBSUB_SUBSCRIPTION, transform.source.topic) elif transform.source.format == 'custom': # TODO(silviuc): Implement custom sources. raise NotImplementedError else: raise ValueError( 'Source %r has unexpected format %s.' % ( transform.source, transform.source.format)) step.add_property(PropertyNames.FORMAT, transform.source.format) step.encoding = self._get_cloud_encoding(transform.source.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': 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 occured in running this pipeline.""" def __init__(self, msg, result): super(DataflowRuntimeException, self).__init__(msg) self.result = result DataflowPythonSDK-0.2.0/google/cloud/dataflow/runners/direct_runner.py000066400000000000000000000201071266367603600261310ustar00rootroot00000000000000# 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 from google.cloud.dataflow import coders from google.cloud.dataflow import error from google.cloud.dataflow.pvalue import AsIter from google.cloud.dataflow.pvalue import AsSingleton from google.cloud.dataflow.pvalue import EmptySideInput from google.cloud.dataflow.runners.common import DoFnRunner from google.cloud.dataflow.runners.common import DoFnState from google.cloud.dataflow.runners.runner import PipelineRunner 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.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() 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): if self._cache.is_cached(pvalue): # pylint: disable=protected-access return else: func(self, pvalue, *args, **kwargs) return func_wrapper @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()) # Construct the list of values from side-input PCollections that we'll # substitute into the arguments for DoFn methods. def get_side_input_value(si): if isinstance(si, AsSingleton): # User wants one item from the PCollection as side input, or an # EmptySideInput if no value exists. pcoll_vals = self._cache.get_pvalue(si.pvalue) if len(pcoll_vals) == 1: return pcoll_vals[0].value elif len(pcoll_vals) > 1: raise ValueError("PCollection with more than one element " "accessed as a singleton view.") elif si.default_value != si._NO_DEFAULT: return si.default_value else: # TODO(robertwb): Should be an error like Java. return EmptySideInput() if isinstance(si, AsIter): # User wants the entire PCollection as side input. List permits # repeatable iteration. return [v.value for v in self._cache.get_pvalue(si.pvalue)] side_inputs = [get_side_input_value(e) for e 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 NoOpCounters(object): def update(self, element): pass class RecordingReciever(object): def __init__(self, tag): self.tag = tag def process(self, element): results[self.tag].append(element) class TaggedRecievers(dict): def __missing__(self, key): return [RecordingReciever(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, TaggedRecievers(), collections.defaultdict(NoOpCounters)) 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._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) self._cache.cache_output( transform_node, map(GlobalWindows.WindowedValue, ((key_coder.decode(k), v) for k, v in result_dict.iteritems()))) @skip_if_cached def run_Create(self, transform_node): transform = transform_node.transform self._cache.cache_output( transform_node, [GlobalWindows.WindowedValue(v) for v in transform.value]) @skip_if_cached def run_Flatten(self, transform_node): self._cache.cache_output( transform_node, list(itertools.chain.from_iterable( self._cache.get_pvalue(pc) for pc in transform_node.inputs))) @skip_if_cached def run_Read(self, transform_node): # TODO(chamikara) Implement a more generic way for passing PipelineOption # to sources when using DirectRunner. source = transform_node.transform.source source.pipeline_options = transform_node.inputs[0].pipeline.options with source.reader() as reader: self._cache.cache_output( transform_node, [GlobalWindows.WindowedValue(e) for e in reader]) @skip_if_cached def run__NativeWrite(self, transform_node): transform = transform_node.transform with transform.sink.writer() as writer: for v in self._cache.get_pvalue(transform_node.inputs[0]): writer.Write(v.value) DataflowPythonSDK-0.2.0/google/cloud/dataflow/runners/runner.py000066400000000000000000000224111266367603600245770ustar00rootroot00000000000000# 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 copy 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() 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, 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, node=None): """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): self.runner.run_transform(transform_node) pipeline.visit(RunVisitor(self), node=node) return PipelineResult(state=PipelineState.DONE) 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): # 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 PValue instance addresses obtained using id(). self._cache = {} def __len__(self): return len(self._cache) def _get_pvalue_with_real_producer(self, pvalue): """Returns a pvalue with the real producer for the passed in pvalue. Args: pvalue: A PValue instance whose cached value is requested. Returns: A pvalue containing the real producer. 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. """ real_producer = pvalue.producer while real_producer.parts: real_producer = real_producer.parts[-1] if real_producer != pvalue.producer: pvalue = copy.copy(pvalue) pvalue.producer = real_producer return pvalue def is_cached(self, pobj): # Import here to avoid circular dependencies. from google.cloud.dataflow.pipeline import AppliedPTransform if isinstance(pobj, AppliedPTransform): transform = pobj else: pobj = self._get_pvalue_with_real_producer(pobj) transform = pobj.producer return (id(transform), None) 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[id(transform), tag] = value def get_pvalue(self, pvalue): """Gets the value associated with a PValue from the cache.""" pvalue = self._get_pvalue_with_real_producer(pvalue) try: return self._cache[self.key(pvalue)] except KeyError: if pvalue.tag is not None and (id(pvalue.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 clear_pvalue(self, pvalue): """Removes a PValue from the cache.""" pvalue = self._get_pvalue_with_real_producer(pvalue) if self.is_cached(pvalue): del self._cache[self.key(pvalue)] def key(self, pobj): return id(pobj.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 DataflowPythonSDK-0.2.0/google/cloud/dataflow/runners/runner_test.py000066400000000000000000000050261266367603600256410ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/transforms/000077500000000000000000000000001266367603600234165ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/transforms/__init__.py000066400000000000000000000016561266367603600255370ustar00rootroot00000000000000# 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.util import * DataflowPythonSDK-0.2.0/google/cloud/dataflow/transforms/aggregator.py000066400000000000000000000077061266367603600261240ustar00rootroot00000000000000# 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 combiners from google.cloud.dataflow.utils.counters import Counter class Aggregator(object): """A user-specified aggregator of statistics about pipeline data. Arg: combine_fn: how to combine values input to the aggregation. It must be one of these arithmetic functions: - Python's built-in sum - Python's built-in min - Python's built-in max - Python's built-in df.Mean() The default is sum. type: describes the numeric 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=None): # inferred from combine_fn self.name = name self.combine_fn = combine_fn self.aggregation_kind = self._aggregator_counter_kind(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) @staticmethod def _aggregator_counter_kind(combine_fn): """Returns the counter aggregation kind for the combine_fn passed in. Args: combine_fn: The combining function used in an Aggregator. Returns: The aggregation_kind (to use in a Counter) that matches combine_fn. Raises: ValueError if the combine_fn doesn't map to any supported aggregation kind. """ # We don't have combiner types that implement AND or OR. combine_kind_map = {sum: Counter.SUM, max: Counter.MAX, min: Counter.MIN, combiners.Mean: Counter.MEAN} try: return combine_kind_map[combine_fn] except KeyError: try: return combine_kind_map[combine_fn.__class__] except KeyError: raise ValueError( 'combine_fn %r (class %r) ' 'does not map to a supported aggregation kind' % (combine_fn, combine_fn.__class__)) DataflowPythonSDK-0.2.0/google/cloud/dataflow/transforms/aggregator_test.py000066400000000000000000000024741266367603600271600ustar00rootroot00000000000000# 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 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.Mean(), float) self.assertEqual('', str(for_mean)) if __name__ == '__main__': unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/transforms/combiners.py000066400000000000000000000334001266367603600257510ustar00rootroot00000000000000# 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 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) 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. """ def __init__(self, n, compare): self._n = n self._compare = compare 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, *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 = compare self.args = args self.kwargs = kwargs def __lt__(self, other): return self.compare(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, *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) class Smallest(TopCombineFn): def __init__(self, n): super(Smallest, self).__init__(n, lambda a, b: b < a) 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 DataflowPythonSDK-0.2.0/google/cloud/dataflow/transforms/combiners_test.py000066400000000000000000000212621266367603600270130ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/transforms/core.py000066400000000000000000001323361266367603600247300ustar00rootroot00000000000000# 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 logging import uuid import warnings 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 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 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 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 __init__(self): super(CombineFn, self).__init__() 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) @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): return self.add_inputs([element], *args, **kwargs) def add_inputs(self, accumulator, elements, *args, **kwargs): if accumulator is not self._EMPTY: def union(): yield accumulator for e in elements: yield e else: union = lambda: elements 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 class PartitionFn(object): """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 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(pipeline=pcoll.pipeline, transform=self) 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 **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')) 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 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 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 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(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.AsSingleton._NO_DEFAULT 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))) def Combine(*args, **kwargs): # pylint: disable=invalid-name warnings.warn('Combine is deprecated, use CombineGlobally instead', DeprecationWarning) return CombineGlobally(*args, **kwargs) @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 out_window, values in driver.process_elements(vs, state): yield window.WindowedValue((k, values), out_window.end, [out_window]) while state.timers: for timer_window, (tag, timestamp) in state.get_and_clear_timers(): for out_window, values in driver.process_timer(timer_window, timestamp, tag, state): yield window.WindowedValue( (k, values), out_window.end, [out_window]) 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) key_coder = typecoders.registry.get_coder(key_type) if not key_coder.is_deterministic(): logging.warning('The key coder "%s" for the GroupByKey operation "%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, 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(pipeline=pcoll.pipeline, transform=self) 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 arguments: 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): 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._is_default = ( self.windowfn == window.GlobalWindows() and self.triggerfn == DefaultTrigger() and self.accumulation_mode == AccumulationMode.DISCARDING) def __repr__(self): return "Windowing(%s, %s, %s)" % (self.windowfn, self.triggerfn, self.accumulation_mode) 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) label, windowfn = self.parse_label_and_arg(args, kwargs, 'windowfn') self.windowing = Windowing(windowfn, triggerfn, accumulation_mode) 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(pipeline=self.pipeline, transform=self) 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(pipeline=self.pipeline, transform=self) 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.0/google/cloud/dataflow/transforms/ptransform.py000066400000000000000000000551321266367603600261710ustar00rootroot00000000000000# 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 sys from google.cloud.dataflow import error from google.cloud.dataflow import pvalue from google.cloud.dataflow import typehints 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 _MaterializePValues(_PValueishTransform): def visit(self, node): if isinstance(node, (pvalue.PValue, pvalue.DoOutputsTuple)): return node._get_values() 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 def __init__(self, label=None): super(PTransform, self).__init__() self.has_user_label = label is not None # This field may get changed using the result returned by default_label() # when the transform gets applied. This happens only if the label was not # user supplied. self.label = label or 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: deferred = True 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.') # 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: p.run() return _MaterializePValues().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. 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(pipeline=pvalueish, transform=None) return pvalueish, (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 if p.has_user_label else p.__class__.__name__ 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: fn_or_label = str(args[0]) if isinstance(fn_or_label, str): label = fn_or_label self.fn, args = self.make_fn(args[0]), args[1:] else: label = None self.fn = self.make_fn(fn_or_label) # 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.AsSingleton, pvalue.AsIter)) self.raw_side_inputs = args, kwargs 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.AsSideInput): 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 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 = self.fn.__name__, args[1:] elif args and isinstance(args[0], str): label, self._args = args[0], args[1:] else: label, self._args = self.fn.__name__, 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 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 (e.g., |) 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:] DataflowPythonSDK-0.2.0/google/cloud/dataflow/transforms/ptransform_test.py000066400000000000000000002117421266367603600272310ustar00rootroot00000000000000# 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 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') pa | df.Create('a_label', [1, 2]) self.assertEqual('', str(pa._nodes[1].producer.transform)) pc = Pipeline('DirectPipelineRunner') pc | df.Create('with_inputs', [1, 2]) inputs_tr = pc._nodes[1].producer.transform inputs_tr.inputs = ('ci',) self.assertEqual( """""", str(inputs_tr)) pd = Pipeline('DirectPipelineRunner') pd | df.Create('with_sidei', [1, 2]) side_tr = pd._nodes[1].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_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.assertEqual('Runtime type violation detected within ' 'ParDo(D/reify_windows): Input to GroupByKey must be ' 'a PCollection with elements compatible with KV[A, B]', e.exception.message) 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_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' 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() class PTransformTypeCheckTestCase(TypeHintTestCase): 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.assertEqual("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 .", e.exception.message) 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.assertEqual( "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.", e.exception.message) 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.assertEqual("Runtime type violation detected within ParDo(to int): " "Type-hint for argument: 'x' violated. " "Expected an instance of , " "instead found 1, an instance of .", e.exception.message) 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.assertEqual("Runtime type violation detected within ParDo(add): " "Type-hint for argument: 'y' violated. " "Expected an instance of , " "instead found 3.0, an instance of .", e.exception.message) 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.assertEqual("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 .", e.exception.message) 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.assertEqual("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.", e.exception.message) 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.assertEqual("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 .", e.exception.message) 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.assertEqual("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 .", e.exception.message) 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.assertEqual( "Runtime type violation detected within " "ParDo(mul/CombinePerKey/Combine/ParDo): " "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.", e.exception.message) 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.assertEqual("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 .", e.exception.message) 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)', 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': " "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': " "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.assertEqual("Runtime type violation detected within " "ParDo(odd mean/CombinePerKey/Combine/ParDo): " "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.", e.exception.message) 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': " "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()", 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.0/google/cloud/dataflow/transforms/trigger.py000066400000000000000000000644771266367603600254550ustar00rootroot00000000000000# 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.transforms import combiners from google.cloud.dataflow.transforms import core from google.cloud.dataflow.transforms.window import GlobalWindow 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) # 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('', 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('') 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('') 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('', 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('') 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 # TODO(robertwb): Time domains. def set_timer(self, tag, timestamp): self._outer.set_timer(self._window, tag, timestamp) def clear_timer(self, timer): self._outer.clear_timer(self._window, timer) 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, tag, timestamp): self._outer.set_timer(self._prefix + tag, timestamp) def clear_timer(self, tag): self._outer.clear_timer(self._prefix + tag) 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, tag, timestamp): pass @abstractmethod def get_window(self, timer_id): pass @abstractmethod def clear_timer(self, window, timer): 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 a 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, tag, timestamp): self.raw_state.set_timer(self._get_id(window), tag, timestamp) def clear_timer(self, window, timer): for window_id in self._get_ids(window): self.raw_state.clear_timer(window_id, timer) 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] 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, timer_id): for window, ids in self.window_ids.items(): if timer_id in ids: return window raise ValueError('No window for %s' % timer_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): # TODO(robertwb): We can do more if we know elements are in timestamp # sorted order. if windowing.is_default() and is_batch: return DefaultGlobalBatchTriggerDriver() else: return GeneralTriggerDriver(windowing) class TriggerDriver(object): """Breaks a series of bundle and timer firings into window (pane)s.""" __metaclass__ = ABCMeta @abstractmethod def process_elements(self, windowed_values, state): pass @abstractmethod def process_timer(self, timer_id, timestamp, unused_tag, state): pass class DefaultGlobalBatchTriggerDriver(TriggerDriver): """Breaks a bundles into window (pane)s according to the default triggering. """ def __init__(self): pass def process_elements(self, windowed_values, state): if isinstance(windowed_values, list): unwindowed = [wv.value for wv in windowed_values] else: class UnwindowedValues(object): def __iter__(self): return (wv.value for wv in windowed_values) unwindowed = UnwindowedValues() yield GlobalWindow(), unwindowed def process_timer(self, timer_id, timestamp, unused_tag, state): raise TypeError('Triggers never set or called for batch default windowing.') 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.trigger_fn = windowing.triggerfn self.accumulation_mode = windowing.accumulation_mode self.is_merging = True def process_elements(self, windowed_values, state): 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) # 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 # Next handle element adding. for window, values in windows_to_elements.items(): if state.get_state(window, self.TOMBSTONE): continue context = state.at(window) for value in values: state.add_state(window, self.ELEMENTS, value) self.trigger_fn.on_element(value, window, context) # Maybe fire this window. watermark = float('-inf') 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, timer_id, timestamp, unused_tag, state): if self.is_merging: state = MergeableStateAdapter(state) window = state.get_window(timer_id) if state.get_state(window, self.TOMBSTONE): return 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) def _output(self, window, finished, state): values = state.get_state(window, self.ELEMENTS) if finished: # TODO(robertwb): allowed lateness state.clear_state(window, None) state.add_state(window, self.TOMBSTONE, 1) elif self.accumulation_mode == AccumulationMode.DISCARDING: state.clear_state(window, self.ELEMENTS) return window, values 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, tag, timestamp): self.timers[window][tag] = timestamp def clear_timer(self, window, tag): self.timers[window].pop(tag, None) def get_window(self, timer_id): return timer_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) 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 else: raise ValueError('Invalid tag.', tag) def clear_state(self, window, tag): if tag is None: self.state.pop(window, None) else: self.state[window].pop(tag.tag, None) def get_and_clear_timers(self, watermark=float('inf')): expired = [] for window, timers in list(self.timers.items()): for tag, timestamp in list(timers.items()): if timestamp <= watermark: expired.append((window, (tag, timestamp))) del timers[tag] 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.0/google/cloud/dataflow/transforms/trigger_test.py000066400000000000000000000443041266367603600264770ustar00rootroot00000000000000# 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 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 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 out_window, values in driver.process_elements(bundle, state): actual_panes[out_window].append(set(values)) while state.timers: for timer_window, (tag, timestamp) in state.get_and_clear_timers(): for out_window, values in driver.process_timer( timer_window, timestamp, tag, state): actual_panes[out_window].append(set(values)) for bundle in late_bundles: for out_window, values in driver.process_elements(bundle, state): actual_panes[out_window].append(set(values)) while state.timers: for timer_window, (tag, timestamp) in state.get_and_clear_timers(): for out_window, values in driver.process_timer( timer_window, timestamp, tag, state): actual_panes[out_window].append(set(values)) 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_names = dict(window_module.__dict__) trigger_names = {'Default': DefaultTrigger} trigger_names.update(trigger_module.__dict__) window_fn = parse_fn(spec.get('window_fn', 'GlobalWindows'), window_names) trigger_fn = parse_fn(spec.get('trigger_fn', 'Default'), trigger_names) accumulation_mode = getattr( AccumulationMode, spec.get('accumulation_mode', 'ACCUMULATING').upper()) allowed_lateness = float(spec.get('allowed_lateness', '-inf')) driver = GeneralTriggerDriver( Windowing(window_fn, trigger_fn, accumulation_mode)) state = InMemoryUnmergedState() output = [] watermark = float('-inf') def fire_timers(): to_fire = state.get_and_clear_timers(watermark) while to_fire: for timer_window, (tag, timestamp) in to_fire: for window, values in driver.process_timer( timer_window, timestamp, tag, state): output.append({'window': [window.start, window.end - 1], 'values': sorted(values)}) 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': [window.start, window.end - 1], 'values': sorted(values)} for window, values in driver.process_elements(bundle, state)] 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.0/google/cloud/dataflow/transforms/trigger_transcripts.yaml000066400000000000000000000054161266367603600304070ustar00rootroot00000000000000name: 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 transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], final: false} - {window: [10, 19], values: [10, 11]} - {window: [20, 29], values: [25], timestamp: 29, late: false} - input: [7] - expect: - {window: [0, 9], values: [1, 2, 3, 7], late: true} --- name: early_late_sessions window_fn: Sessions(10) trigger_fn: AfterWatermark(early=AfterCount(2), late=AfterCount(3)) transcript: - input: [1, 2, 3] - expect: - {window: [1, 12], values: [1, 2, 3]} - input: [4] # no output - input: [5] - expect: - {window: [1, 14], values: [1, 2, 3, 4, 5]} - input: [6] - watermark: 100 - expect: - {window: [1, 15], values:[1, 2, 3, 4, 5, 6]} - input: [1] - input: [3, 4] - expect: - {window: [1, 15], values: [1, 1, 2, 3, 3, 4, 4, 5, 6]} --- name: garbage_collection window_fn: FixedWindows(10) trigger_fn: AfterCount(2) allowed_lateness: 10 accumulation_mode: discarding transcript: - input: [1, 2, 3, 10, 11, 25] - expect: - {window: [0, 9]} - {window: [10, 19]} - state: present: [[20, 29]] absent: [[0, 9]] tombstone: [[10, 19]] # 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.0/google/cloud/dataflow/transforms/util.py000066400000000000000000000173511266367603600247540ustar00rootroot00000000000000# 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 an input of (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_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): def _equal(actual): if sorted(expected) != sorted(actual): raise DataflowAssertException( 'Failed assert: %r == %r' % (expected, 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) actual.pipeline | Create(label, [None]) | Map(match, AllOf(actual)) DataflowPythonSDK-0.2.0/google/cloud/dataflow/transforms/window.py000066400000000000000000000237251266367603600253100ustar00rootroot00000000000000# 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. Floating-point seconds are used as the time unit for built-in windowing strategies available here. Custom windowing function classes can be created, by subclassing from WindowFn. """ from __future__ import absolute_import MIN_TIMESTAMP = float('-Inf') MAX_TIMESTAMP = float('Inf') 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 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 class BoundedWindow(object): """A window for timestamps in range (-infinity, end). Attributes: end: End of window as floating-point seconds since Unix epoch. """ def __init__(self, end): self.end = 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 NotImplemented def __hash__(self): return hash(self.end) def __repr__(self): return '[?, %s)' % self.end class IntervalWindow(BoundedWindow): """A window for timestamps in range [start, end). Attributes: start: Start of window as floating-point seconds since Unix epoch. end: End of window as floating-point seconds since Unix epoch. """ def __init__(self, start, end): super(IntervalWindow, self).__init__(end) self.start = 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)' % (self.start, 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 floating-point 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 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 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) class TimestampedValue(object): """A timestamped value having a value and a timestamp. Attributes: value: The underlying value. timestamp: Timestamp associated with the value as floating-point seconds since Unix epoch. """ def __init__(self, value, timestamp): self.value = value self.timestamp = 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) 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 WindowedValue(cls, value, timestamp=MIN_TIMESTAMP): # pylint: disable=invalid-name return WindowedValue(value, timestamp, [GlobalWindow()]) def assign(self, assign_context): return [GlobalWindow()] def merge(self, merge_context): pass # No merging. 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 floating-point seconds. offset: Offset of this window as floating-point 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 = size self.offset = offset % 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 floating-point seconds. period: Period of the windows as floating-point seconds. offset: Offset of this window as floating-point 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 = size self.period = period self.offset = offset % size def assign(self, context): timestamp = context.timestamp start = timestamp - (timestamp - self.offset) % self.period return [IntervalWindow(s, 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 = 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.0/google/cloud/dataflow/transforms/window_test.py000066400000000000000000000201021266367603600263310ustar00rootroot00000000000000# 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, 2)', [1]), ('key @ [0, 4)', [1, 2, 3]), ('key @ [2, 6)', [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, 13)', [1, 2, 3]), ('key @ [20, 45)', [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.0/google/cloud/dataflow/transforms/write_ptransform_test.py000066400000000000000000000123211266367603600304330ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/typehints/000077500000000000000000000000001266367603600232475ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/typehints/__init__.py000066400000000000000000000014661266367603600253670ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/typehints/decorators.py000066400000000000000000000440271266367603600257750ustar00rootroot00000000000000# 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 import warnings 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 takes(*args, **kwargs): warnings.warn('takes is deprectaed, use with_input_types instead', warnings.DeprecationWarning) return with_input_types(*args, **kwargs) 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 returns(*args, **kwargs): warnings.warn('returns is deprecated, use with_output_types instead', warnings.DeprecationWarning) return with_output_types(*args, **kwargs) 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.0/google/cloud/dataflow/typehints/opcodes.py000066400000000000000000000170151266367603600252610ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/typehints/trivial_inference.py000066400000000000000000000267511266367603600273240ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/typehints/trivial_inference_test.py000066400000000000000000000107661266367603600303620ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/typehints/typecheck.py000066400000000000000000000144721266367603600256100ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/typehints/typed_pipeline_test.py000066400000000000000000000214221266367603600276730ustar00rootroot00000000000000# 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(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.0/google/cloud/dataflow/typehints/typehints.py000066400000000000000000001054441266367603600256600ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/typehints/typehints_test.py000066400000000000000000001140341266367603600267120ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/utils/000077500000000000000000000000001266367603600223605ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/utils/__init__.py000066400000000000000000000014361266367603600244750ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/utils/counters.py000066400000000000000000000053771266367603600246100ustar00rootroot00000000000000# 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.""" 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.) 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 """ # Aggregation kinds. The protocol uses string names, so the values # assigned here are not externally visible. # Numeric: SUM = 1 MAX = 2 MIN = 3 MEAN = 4 # arithmetic mean # Boolean AND = 5 OR = 6 _KIND_NAME_MAP = {SUM: 'SUM', MAX: 'MAX', MIN: 'MIN', MEAN: 'MEAN', AND: 'AND', OR: 'OR'} def aggregation_kind_str(self): return self._KIND_NAME_MAP.get(self.aggregation_kind, 'kind%d' % self.aggregation_kind) def __init__(self, name, aggregation_kind): """Creates a Counter object. Args: name: the name of this counter. Typically has three parts: "step-output-counter". aggregation_kind: one of the kinds defined by this class. """ self.name = name self.aggregation_kind = aggregation_kind self.total = 0 self.elements = 0 def update(self, count): self.total += count self.elements += 1 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/%s' % (self.name, self.aggregation_kind_str(), self.total, self.elements) class AggregatorCounter(Counter): """A Counter that represents a step-specific instance of an Aggregator.""" def __init__(self, step_name, aggregator): super(AggregatorCounter, self).__init__( 'user-%s-%s' % (step_name, aggregator.name), aggregator.aggregation_kind) DataflowPythonSDK-0.2.0/google/cloud/dataflow/utils/dependency.py000066400000000000000000000351101266367603600250500ustar00rootroot00000000000000# 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' 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 _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 = [] for package in extra_packages: if not os.path.isfile(package): raise RuntimeError( 'The file %s cannot be found. It was specified in the ' '--extra_packages command line option.' % package) 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) 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 extra_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) return resources def stage_job_resources(options, file_copy=_dependency_file_copy, build_setup_args=None, temp_dir=None): """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. 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) # 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_gcs = True elif setup_options.sdk_location.startswith('gs://'): stage_tarball_from_gcs = True else: stage_tarball_from_gcs = False staged_path = utils.path.join(google_cloud_options.staging_location, names.DATAFLOW_SDK_TARBALL_FILE) if stage_tarball_from_gcs: # 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_folder = 'gs://dataflow-sdk-for-python' else: sdk_folder = setup_options.sdk_location _stage_dataflow_sdk_tarball(sdk_folder, staged_path) 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_dir = os.path.join(os.path.dirname(module_path), '..') else: sdk_dir = setup_options.sdk_location sdk_path = os.path.join(sdk_dir, names.DATAFLOW_SDK_TARBALL_FILE) 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 directory 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_release_folder, staged_path): """Stage a Dataflow SDK tarball with the appropriate version. This function expects to find in the SDK release folder a file with the name google-cloud-dataflow-python-sdk-VERSION.tgz where VERSION is the version string for the containing SDK from google.cloud.dataflow.version.__version__. Args: sdk_release_folder: A local path or GCS path containing a Dataflow SDK tarball with the appropriate version. staged_path: GCS path where the found SDK tarball should be copied. """ # Build the expected tarball file path. tarball_file_name = 'google-cloud-dataflow-python-sdk-%s.tgz' % __version__ tarball_file_path = utils.path.join(sdk_release_folder, tarball_file_name) # Stage the file to the GCS staging area. logging.info( 'Staging Dataflow SDK tarball from %s to %s', tarball_file_path, staged_path) _dependency_file_copy(tarball_file_path, staged_path) DataflowPythonSDK-0.2.0/google/cloud/dataflow/utils/dependency_test.py000066400000000000000000000311141266367603600261070ustar00rootroot00000000000000# 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 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( [dependency.REQUIREMENTS_FILE, names.PICKLED_MAIN_SESSION_FILE], dependency.stage_job_resources(options)) self.assertTrue( os.path.isfile( os.path.join(staging_dir, dependency.REQUIREMENTS_FILE))) 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) 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_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 test_sdk_location_default(self): staging_dir = tempfile.mkdtemp() expected_from_path = utils.path.join( 'gs://dataflow-sdk-for-python', 'google-cloud-dataflow-python-sdk-%s.tgz' % __version__) 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 ' 'directory was specified by the --sdk_location command-line option.' % os.path.join(sdk_location, names.DATAFLOW_SDK_TARBALL_FILE), cm.exception.message) def test_sdk_location_gcs(self): staging_dir = tempfile.mkdtemp() sdk_location = 'gs://my-gcs-bucket' expected_from_path = utils.path.join( sdk_location, 'google-cloud-dataflow-python-sdk-%s.tgz' % __version__) 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 = 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')] self.assertEqual( ['abc.tar.gz', 'xyz.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'], f.readlines()) 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'] 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') 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.0/google/cloud/dataflow/utils/names.py000066400000000000000000000043151266367603600240400ustar00rootroot00000000000000# 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' 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' SERIALIZED_FN = 'serialized_fn' SHARD_NAME_TEMPLATE = 'shard_template' 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.0/google/cloud/dataflow/utils/options.py000066400000000000000000000350011266367603600244240ustar00rootroot00000000000000# 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 # Raw (unparsed) options. They are also added by other modules that want to # contribute modules other than the ones defined in this file. See add_option(), # below. OPTIONS = [] 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 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): if 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): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--runner', default='DirectPipelineRunner', help=('Pipeline runner used to execute the workflow. Valid values are ' 'DirectPipelineRunner, DataflowPipelineRunner, ' 'and BlockingDataflowPipelineRunner.')) # Whether this is a streaming job. # TODO(ccy): This should be an option on the pipeline runner. parser.add_argument('--is_streaming', default=False, help='Whether the job is a streaming job.') 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): @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) # Command line options controlling the worker pool configuration. # TODO(silviuc): Update description when autoscaling options are in. class WorkerOptions(PipelineOptions): @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.')) parser.add_argument( '--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.')) parser.add_argument( '--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( '--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.')) 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 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 ' 'specified, the worker will install the required dependenciesi before' ' running any custom code. Typically the file is named ' 'requirements.txt.')) 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= ('GCS folder or local directory containing the Dataflow SDK for Python ' 'tar file. Workflow submissions will copy an SDK tarball from here. ' 'If this is a GCS folder, the corresponding SDK version will be added ' 'to the tar file name. Otherwise (used only for testing the SDK), ' 'this is the name of a local directory. 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.0/google/cloud/dataflow/utils/path.py000066400000000000000000000033411266367603600236670ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/utils/path_test.py000066400000000000000000000044531266367603600247330ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/utils/processes.py000066400000000000000000000027241266367603600247450ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/utils/processes_test.py000066400000000000000000000063061266367603600260040ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/utils/retry.py000066400000000000000000000151731266367603600241060ustar00rootroot00000000000000# 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 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 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.0/google/cloud/dataflow/utils/retry_test.py000066400000000000000000000131111266367603600251330ustar00rootroot00000000000000# 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.0/google/cloud/dataflow/version.py000066400000000000000000000012351266367603600232600ustar00rootroot00000000000000# 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.0' DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/000077500000000000000000000000001266367603600225315ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/__init__.py000066400000000000000000000000001266367603600246300ustar00rootroot00000000000000DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/batchworker.py000066400000000000000000000546071266367603600254320ustar00rootroot00000000000000# 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 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 retry 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 BatchWorker(object): """A worker class with all the knowledge to lease and execute work items.""" DEFAULT_MIN_REPORTING_INTERVAL_SECS = 5.0 DEFAULT_MAX_REPORTING_INTERVAL_SECS = 10 * 60.0 DEFAULT_LEASE_RENEWAL_LATENCY_SECS = 5.0 DEFAULT_LEASE_DURATION_SECS = 3 * 60.0 # 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 def __init__(self, properties): """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'] 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'] # TODO(silviuc): Figure out how to use this logging location. # Should we log something to this specific location? self.log_path = properties['dataflow.worker.logging.location'] 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) # Initialize the logging machinery. logger.initialize(job_id=self.job_id, worker_id=self.worker_id, log_path=self.log_path) logging.info('Batch worker started with: %s', properties) self.client = apiclient.DataflowWorkerClient( worker=self, skip_get_credentials=(not self.running_in_gce)) self.dynamic_split_result_to_report = None self.last_memory_usage_report_time = None # Lock used to access/modify the _current_work_item and _current_executor # attributes which is used to coordinate work between main thread and # progress reporting thread. self.lock = threading.Lock() self._current_work_item = None self._current_executor = None self.environment = maptask.WorkerEnvironment() @property def current_work_item(self): with self.lock: return self._current_work_item @property def current_executor(self): with self.lock: return self._current_executor def set_current_work_item_and_executor(self, work_item, work_executor): with self.lock: self._current_work_item = work_item self._current_executor = work_executor @property def 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 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 '%.3fs' % self.DEFAULT_LEASE_DURATION_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 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 _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, work_item, 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: work_item: A WorkItem instance describing the work. 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', work_item, 'successfully' if exception_details is None else 'with exception') self.report_status(work_item, completed=True, exception_details=exception_details) @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def report_progress_status(self, work_item, exception_details=None): """Reports to the service a work item progress status. Args: work_item: A WorkItem instance describing the work. 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. """ self.report_status(work_item, completed=False, progress=self._current_executor.get_progress(), exception_details=exception_details) def report_status(self, work_item, completed=False, progress=None, exception_details=None): """Reports to the service status of a work item (completion or progress). Args: work_item: A WorkItem instance describing the work. 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. 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.). """ # If there is an unsent dynamic_split_result_to_report we must send it to # the service before reporting completion. # This will not result in an infinite loop since a finished reader will # not be creating a new dynamic_split_result_to_report. if completed and self.dynamic_split_result_to_report is not None: self.report_status(work_item, False, progress, exception_details) response = self.client.report_status( self, work_item, completed, progress if not completed else None, self.dynamic_split_result_to_report if not completed else None, 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: work_item_state = response.workItemServiceStates[0] work_item.next_report_index = work_item_state.nextReportIndex work_item.lease_expire_time = work_item_state.leaseExpireTime 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.current_executor.request_dynamic_split( apiclient.approximate_progress_to_dynamic_split_request( suggested_split_point))) def log_memory_usage_if_needed(self, force=False): """Periodically logs memory usage of the current worker. Args: force: if True forces logging. """ if (force or self.last_memory_usage_report_time is None or int(time.time()) - self.last_memory_usage_report_time > self.MEMORY_USAGE_REPORTING_INTERVAL_SECS): logging.info('Memory usage of worker %s is %d MB', self.worker_id, resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1000) self.last_memory_usage_report_time = int(time.time()) def progress_reporting_thread(self): """Sends progress reports in an infinite loop.""" work_item = None while True: try: self.log_memory_usage_if_needed(force=False) # If thread does not work on something then try to get the current work # item from the worker object. This in turn can be None if there where # no work items to lease from the service. if work_item is None: work_item = self.current_work_item if work_item is None: logging.debug('Progress reporting thread has no work items. ' 'Sleeping a bit ...') time.sleep(1.0) continue with work_item.lock: # Make sure we drop a work item that was marked 'done' in the main # worker thread. if work_item.done: work_item = None continue sleep_time = self.next_progress_report_interval( work_item.report_status_interval, work_item.lease_expire_time) logging.debug( 'Progress reporting thread will sleep %f secs between updates.', sleep_time) time.sleep(sleep_time) with work_item.lock: # Make sure we drop a work item that was marked done. if work_item.done: work_item = None else: self.report_progress_status(work_item) except Exception: # pylint: disable=broad-except logging.info('Progress reporting thread got error: %s', traceback.format_exc()) def do_work(self, work_item): """Executes worker operations and adds any failures to the report status.""" logging.info('Executing %s', work_item) self.log_memory_usage_if_needed(force=True) try: with work_item.lock: self.set_current_work_item_and_executor(work_item, executor.MapTaskExecutor()) self.current_executor.execute(work_item.map_task) except Exception: # pylint: disable=broad-except exception_details = traceback.format_exc() logging.error('Exception: %s', exception_details, exc_info=True) # Completed with errors means failed. with work_item.lock: self.report_completion_status(work_item, exception_details=exception_details) else: with work_item.lock: self.report_completion_status(work_item) with work_item.lock: 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. deferred_exception_details = 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() # Start the progress reporting thread. thread = threading.Thread(target=self.progress_reporting_thread) 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 True: # TODO(silviuc): Do we still need the outer try/except? 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) 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 with logger.PerThreadLoggingContext( work_item_id=work_item.proto.id, stage_name=work_item.map_task.stage_name): # TODO(silviuc): Add more detailed timing and profiling support. start_time = time.time() if deferred_exception_details: # Report (fatal) deferred exceptions that happened earlier. This # workflow will fail with the deferred exception. with work_item.lock: self.set_current_work_item_and_executor( work_item, executor.MapTaskExecutor()) work_item.map_task.executed_operations = [] self.report_completion_status( work_item, exception_details=deferred_exception_details) work_item.done = True else: # 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. self.do_work(work_item) 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) DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/concat_reader.py000066400000000000000000000062451266367603600257030ustar00rootroot00000000000000# 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.Source): """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.SourceReader): """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.0/google/cloud/dataflow/worker/concat_reader_test.py000066400000000000000000000204031266367603600267320ustar00rootroot00000000000000# 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.Source): 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.SourceReader): 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.0/google/cloud/dataflow/worker/executor.py000066400000000000000000000673321266367603600247540ustar00rootroot00000000000000# 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 operations executor.""" import collections import itertools import logging import random from google.cloud.dataflow.internal import pickler 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 trigger from google.cloud.dataflow.transforms import window from google.cloud.dataflow.transforms.trigger import InMemoryUnmergedState from google.cloud.dataflow.transforms.window import GlobalWindows 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 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. TODO(gildea): Refactor "receivers[OUTPUT][RECEIVER]" as "outputs[INDEX][RECEIVER]" """ def __init__(self, spec): """Initializes a worker operation instance. Args: spec: A maptask.Worker* instance. """ self.spec = spec self.receivers = collections.defaultdict(list) # Initially we have no counters. Initializing this here makes it # safe to call itercounters() at any time, even if start() has # not been called yet. self.counters = collections.defaultdict(self.new_operation_counters) def new_operation_counters(self, output_index=0): return opcounters.OperationCounters(self.step_name, output_index) def start(self): # If the operation has receivers, create one counter set per receiver. for output_index in self.receivers: self.counters[output_index] = self.new_operation_counters(output_index) def itercounters(self): for opcounter in self.counters.values(): for counter in opcounter: yield counter def finish(self): pass def add_receiver(self, operation, output_index=0): """Adds a receiver operation for the specified output.""" self.receivers[output_index].append(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([ rop.str_internal(is_recursive=True) for oplist in self.receivers.values() for rop in oplist])) 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): super(ReadOperation, self).__init__(spec) 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.WindowedValue(value) for receiver in self.receivers[0]: self.counters[0].update(windowed_value) receiver.process(windowed_value) def side_read_all(self, singleton=False): # TODO(mairbek): Should we return WindowedValue here? with self.spec.source.reader() as reader: for value in reader: yield value if singleton: return 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): super(WriteOperation, self).__init__(spec) 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): logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) self.counters[0].update(o) if self.use_windowed_value: self.writer.Write(o) else: self.writer.Write(o.value) class InMemoryWriteOperation(Operation): """A write operation that will write to an in-memory sink.""" def __init__(self, spec): super(InMemoryWriteOperation, self).__init__(spec) self.spec = spec def process(self, o): logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) self.counters[0].update(o) self.spec.output_buffer.append(o.value) class GroupedShuffleReadOperation(Operation): """A shuffle read operation that will read from a grouped shuffle source.""" def __init__(self, spec, shuffle_source=None): super(GroupedShuffleReadOperation, self).__init__(spec) self.shuffle_source = shuffle_source self._reader = None def start(self): super(GroupedShuffleReadOperation, self).start() if self.shuffle_source is None: self.shuffle_source = shuffle.GroupedShuffleSource( self.spec.shuffle_reader_config, coder=self.spec.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 for receiver in self.receivers[0]: windowed_value = GlobalWindows.WindowedValue((key, key_values)) self.counters[0].update(windowed_value) receiver.process(windowed_value) 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, shuffle_source=None): super(UngroupedShuffleReadOperation, self).__init__(spec) self.shuffle_source = shuffle_source self._reader = None def start(self): super(UngroupedShuffleReadOperation, self).start() if self.shuffle_source is None: self.shuffle_source = shuffle.UngroupedShuffleSource( self.spec.shuffle_reader_config, coder=self.spec.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 for receiver in self.receivers[0]: windowed_value = GlobalWindows.WindowedValue(value) self.counters[0].update(windowed_value) receiver.process(windowed_value) 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, shuffle_sink=None): super(ShuffleWriteOperation, self).__init__(spec) self.writer = None self.shuffle_sink = shuffle_sink def start(self): super(ShuffleWriteOperation, self).start() # TODO(silviuc): Shuffle 'kind' is ignored! if self.shuffle_sink is None: self.shuffle_sink = shuffle.ShuffleSink( self.spec.shuffle_writer_config, coder=self.spec.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): logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) self.counters[0].update(o) # 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.spec.shuffle_kind == '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): Figure out what is the proper value for the secondary key. # For now the secondary key is a duplicate of the primary key just because # they both use the same coder. self.writer.Write(k, k, v) class DoOperation(Operation): """A Do operation that will execute a custom DoFn for each input element.""" def __init__(self, spec): super(DoOperation, self).__init__(spec) self.state = common.DoFnState() 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, side_type in tags_and_types: # 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 results = [] for si in itertools.ifilter( lambda o: o.tag == side_tag, self.spec.side_inputs): if isinstance(si, maptask.WorkerRead): op = ReadOperation(si) else: raise NotImplementedError('Unknown side input type: %r' % si) for v in op.side_read_all(singleton=side_type): results.append(v) if side_type: break if side_type: yield results[0] if results else EmptySideInput() else: yield results def itercounters(self): """Return an iterator over all our counters. Yields: Counters associated with this operation. """ for counter in super(DoOperation, self).itercounters(): yield counter for custom_counter in self.state.itercounters(): yield custom_counter 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 = {} tagged_counters = {} self._tag_map = {} 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) # There may be no receiver for this output, in which case the # lookup will create one, and this value will be processed # for any side effect. This is desirable. There are two (known) # cases where there is no receiver for an output: # 1. ParDo without anything following it, executed for side effect. # 2. Partition (shows up here in the worker as Flatten), which # generates a default output that isn't used. tagged_receivers[original_tag] = self.receivers[index] tagged_counters[original_tag] = self.counters[index] self.dofn_runner = common.DoFnRunner( fn, args, kwargs, self._read_side_inputs(tags_and_types), window_fn, self.context, tagged_receivers, tagged_counters, logger, self.step_name) self.dofn_runner.start() def finish(self): self.dofn_runner.finish() def process(self, o): with logger.PerThreadLoggingContext(step_name=self.step_name): self.dofn_runner.process(o) class CombineOperation(Operation): """A Combine operation executing a CombineFn for each input element.""" def __init__(self, spec): super(CombineOperation, self).__init__(spec) # 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] if not args and not kwargs: self.combine_fn = fn else: class CurriedFn(ptransform.CombineFn): 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) self.combine_fn = CurriedFn() if self.spec.phase == 'all': self.apply = self.full_combine elif self.spec.phase == 'add': self.apply = self.add_only elif self.spec.phase == 'merge': self.apply = self.merge_only elif self.spec.phase == 'extract': self.apply = self.extract_only else: raise ValueError('Unexpected phase: %s' % self.spec.phase) def finish(self): logging.debug('Finishing %s', self) def process(self, o): logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) key, values = o.value windowed_result = WindowedValue( (key, self.apply(values)), o.timestamp, o.windows) for receiver in self.receivers[0]: self.counters[0].update(windowed_result) receiver.process(windowed_result) def full_combine(self, elements): return self.combine_fn.apply(elements) def add_only(self, elements): return self.combine_fn.add_inputs( self.combine_fn.create_accumulator(), elements) def merge_only(self, accumulators): return self.combine_fn.merge_accumulators(accumulators) def extract_only(self, accumulator): return self.combine_fn.extract_output(accumulator) 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): super(PGBKOperation, self).__init__(spec) 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 windowed_value = WindowedValue( (key, [v.value[1] for v in vs]), vs[0].timestamp, windows) for receiver in self.receivers[0]: self.counters[0].update(windowed_value) receiver.process(windowed_value) class FlattenOperation(Operation): """Flatten operation. Receives one or more producer operations, outputs just one list with all the items. """ def process(self, o): logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) windowed_result = WindowedValue(o.value, o.timestamp, o.windows) for receiver in self.receivers[0]: self.counters[0].update(windowed_result) receiver.process(windowed_result) 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): super(ReifyTimestampAndWindowsOperation, self).__init__(spec) def process(self, o): logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) k, v = o.value self.output( window.WindowedValue( (k, window.WindowedValue(v, o.timestamp, o.windows)), o.timestamp, o.windows)) def output(self, windowed_result): for receiver in self.receivers[0]: self.counters[0].update(windowed_result) receiver.process(windowed_result) class BatchGroupAlsoByWindowsOperation(Operation): """BatchGroupAlsoByWindowsOperation operation. Implements GroupAlsoByWindow for batch pipelines. """ def __init__(self, spec): super(BatchGroupAlsoByWindowsOperation, self).__init__(spec) self.windowing = pickler.loads(self.spec.window_fn) def process(self, o): """Process a given value.""" logging.debug('Processing [%s] in %s', o, self) assert isinstance(o, WindowedValue) k, vs = o.value driver = trigger.create_trigger_driver(self.windowing, True) state = InMemoryUnmergedState() # TODO(robertwb): Process in smaller chunks. for out_window, values in driver.process_elements(vs, state): self.output( window.WindowedValue((k, values), out_window.end, [out_window])) while state.timers: for timer_window, (tag, timestamp) in state.get_and_clear_timers(): for out_window, values in driver.process_timer(timer_window, timestamp, tag, state): self.output(window.WindowedValue( (k, values), out_window.end, [out_window])) def output(self, windowed_result): for receiver in self.receivers[0]: self.counters[0].update(windowed_result) receiver.process(windowed_result) class StreamingGroupAlsoByWindowsOperation(Operation): """BatchGroupAlsoByWindowsOperation operation. Implements GroupAlsoByWindow for streaming pipelines. """ def __init__(self, spec): super(StreamingGroupAlsoByWindowsOperation, self).__init__(spec) self.windowing = pickler.loads(self.spec.window_fn) def process(self, o): 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 for out_window, values in driver.process_elements(keyed_work.elements(), state): self.output(window.WindowedValue((keyed_work.key, values), out_window.end, [out_window])) def output(self, windowed_result): for receiver in self.receivers[0]: self.counters[0].update(windowed_result) receiver.process(windowed_result) class MapTaskExecutor(object): """A class for executing 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): self._ops = [] self._read_operation = None 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, map_task, test_shuffle_source=None, test_shuffle_sink=None): """Executes all the maptask.Worker* instructions in a map task. We update the map_task with the execution status, expressed as counters. 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. 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 map_task.operations: if isinstance(spec, maptask.WorkerRead): op = ReadOperation(spec) 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) elif isinstance(spec, maptask.WorkerCombineFn): op = CombineOperation(spec) elif isinstance(spec, maptask.WorkerPartialGroupByKey): op = PGBKOperation(spec) elif isinstance(spec, maptask.WorkerDoFn): op = DoOperation(spec) elif isinstance(spec, maptask.WorkerGroupingShuffleRead): op = GroupedShuffleReadOperation( spec, shuffle_source=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, shuffle_source=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) elif isinstance(spec, maptask.WorkerShuffleWrite): op = ShuffleWriteOperation(spec, shuffle_sink=test_shuffle_sink) elif isinstance(spec, maptask.WorkerFlatten): op = FlattenOperation(spec) elif isinstance(spec, maptask.WorkerMergeWindows): if isinstance(spec.context, maptask.BatchExecutionContext): op = BatchGroupAlsoByWindowsOperation(spec) elif isinstance(spec.context, maptask.StreamingExecutionContext): op = StreamingGroupAlsoByWindowsOperation(spec) else: raise RuntimeError('Unknown execution context: %s' % spec.context) elif isinstance(spec, maptask.WorkerReifyTimestampAndWindows): op = ReifyTimestampAndWindowsOperation(spec) 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, index = op.spec.input self._ops[producer].add_receiver(op, index) # Flatten has 'inputs', not 'input' if hasattr(op.spec, 'inputs'): for producer, index in op.spec.inputs: self._ops[producer].add_receiver(op, index) # Inject the step names into the operations. # This is used for logging and assigning names to counters. if map_task.step_names is not None: for ix, op in enumerate(self._ops): op.step_name = map_task.step_names[ix] # Attach the ops back to the map_task, so we can report their counters. map_task.executed_operations = self._ops 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() DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/executor_test.py000066400000000000000000000542161266367603600260100ustar00rootroot00000000000000# 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 from google.cloud.dataflow import coders 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 import mock 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))]) 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 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.percent_complete) return next_progress class ProgressRequestRecordingInMemorySource(inmemory.InMemorySource): def reader(self): self.last_reader = ProgressRequestRecordingInMemoryReader(self) return self.last_reader class ExecutorTest(unittest.TestCase): SHUFFLE_CODERS = (coders.PickleCoder(), 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 executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( fileio.TextFileSource(file_path=input_path, start_offset=0, end_offset=15, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), tag=None), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: ['XYZ: %s' % x])), output_tags=['out'], input=(0, 0), side_inputs=None), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(1, 0)) ])) 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 executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( fileio.TextFileSource(file_path=input_path, start_offset=0, end_offset=15, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), tag=None), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( DoFnUsingStartBundle(finish_path)), output_tags=['out'], input=(0, 0), side_inputs=None), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(1, 0)) ])) 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_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()), tag=None), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: [(x, 1)])), output_tags=['out'], input=(0, 0), side_inputs=None), maptask.WorkerShuffleWrite(shuffle_kind='group_keys', shuffle_writer_config='none', input=(1, 0), coders=self.SHUFFLE_CODERS) ] shuffle_sink_mock = mock.MagicMock() executor.MapTaskExecutor().execute( make_map_task(work_spec), test_shuffle_sink=shuffle_sink_mock) # Make sure we have seen all the (k, v) writes. shuffle_sink_mock.writer().Write.assert_has_calls( [mock.call('a', 'a', 1), mock.call('b', 'b', 1), mock.call('c', 'c', 1), mock.call('d', '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', coders=self.SHUFFLE_CODERS), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda (k, vs): [str((k, v)) for v in vs])), output_tags=['out'], input=(0, 0), side_inputs=None), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(1, 0)) ] shuffle_source_mock = mock.MagicMock() shuffle_source_mock.reader().__enter__().__iter__.return_value = [ (10, [1, 2]), (20, [3])] executor.MapTaskExecutor().execute( make_map_task(work_spec), test_shuffle_source=shuffle_source_mock) 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', coders=self.SHUFFLE_CODERS), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(0, 0)) ] shuffle_source_mock = mock.MagicMock() shuffle_source_mock.reader().__enter__().__iter__.return_value = [1, 2, 3] executor.MapTaskExecutor().execute( make_map_task(work_spec), test_shuffle_source=shuffle_source_mock) 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'] executor.MapTaskExecutor().execute(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), tag=None), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: ['XYZ: %s' % x])), output_tags=['out'], input=(0, 0), side_inputs=None), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(1, 0)) ])) 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'] executor.MapTaskExecutor().execute(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), tag=None), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: ['XYZ: %s' % x])), output_tags=['out'], input=(0, 0), side_inputs=None), maptask.WorkerWrite(fileio.TextFileSink( file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.Base64PickleCoder()), input=(1, 0))])) 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 = [] executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=3), tag=None), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda x, side: ['%s:%s' % (x, side)]), tag_and_type=('inmemory', True)), # True => type is singleton. output_tags=['out'], input=(0, 0), side_inputs=[ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in side_elements], start_index=None, end_index=None), tag='inmemory')]), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0))])) # 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]) executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead(source, tag=None), maptask.WorkerInMemoryWrite(output_buffer= output_buffer, input=(0, 0)) ])) self.assertEqual(elements, output_buffer) expected_progress_record = [] len_elements = len(elements) for i in range(len_elements): expected_progress_record.append(float(i + 1) / 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 = [] executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=2), tag=None), 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', False)), # False => type is collection. output_tags=['out'], input=(0, 0), side_inputs=[ maptask.WorkerRead(fileio.TextFileSource( file_path=input_path, start_offset=None, end_offset=None, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), tag='textfile')]), maptask.WorkerInMemoryWrite(output_buffer=output_buffer, input=(1, 0))])) # 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 reader_mock.__iter__.return_value = (x for x in side_elements) pickled_elements = [pickler.dumps(e) for e in elements] executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( inmemory.InMemorySource(elements=pickled_elements, start_index=0, end_index=3), tag=None), maptask.WorkerDoFn( serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda x, side: ['%s:%s' % (x, side)]), tag_and_type=( 'bigquery', True)), # True => type is singleton. output_tags=['out'], input=(0, 0), side_inputs=[ maptask.WorkerRead( bigquery.BigQuerySource( project='project', dataset='dataset', table='table', coder=get_bigquery_source_coder()), tag='bigquery')]), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0))])) # 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 reader_mock.__iter__.return_value = (x for x in side_elements) executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=3), tag=None), 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', False)), # False => type is collection. output_tags=['out'], input=(0, 0), side_inputs=[ maptask.WorkerRead( bigquery.BigQuerySource( project='project', dataset='dataset', table='table', coder=get_bigquery_source_coder()), tag='bigquery')]), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0))])) # 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 = [] executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=2), tag=None), 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', False)), # False => type is collection. 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.WorkerRead( fileio.TextFileSource( file_path=input_path1, coder=coders.Base64PickleCoder()), tag='sometag'), maptask.WorkerRead( fileio.TextFileSource(file_path=input_path2, coder=coders.Base64PickleCoder()), tag='sometag')]), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0))])) # 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 = [] executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], start_index=0, end_index=100), tag=None), maptask.WorkerCombineFn(serialized_fn=pickle_with_side_inputs( ptransform.CombineFn.from_callable(sum)), phase='all', input=(0, 0)), maptask.WorkerInMemoryWrite( output_buffer=output_buffer, input=(1, 0)) ])) self.assertEqual([('a', 6), ('b', 10)], output_buffer) def test_pgbk(self): elements = [('a', 1), ('b', 2), ('a', 3), ('a', 4)] output_buffer = [] executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( inmemory.InMemorySource(elements=[pickler.dumps(e) for e in elements ], start_index=0, end_index=100), tag=None), maptask.WorkerPartialGroupByKey(input=( 0, 0)), maptask.WorkerInMemoryWrite(output_buffer=output_buffer, input=(1, 0)) ])) self.assertEqual([('a', [1, 3, 4]), ('b', [2])], sorted(output_buffer)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/inmemory.py000066400000000000000000000050601266367603600247430ustar00rootroot00000000000000# 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 from google.cloud.dataflow import coders from google.cloud.dataflow.io import iobase class InMemorySource(iobase.Source): """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.SourceReader): """A reader for in-memory source.""" def __init__(self, source): self.source = source # Index of the next item to be read by the InMemoryReader. # Starts at source.start_index. self.current_index = source.start_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): self.current_index += 1 yield self.source.coder.decode(value) def get_progress(self): if (self.current_index >= self.source.end_index or self.source.start_index >= self.source.end_index): percent_complete = 1 elif self.current_index == self.source.start_index: percent_complete = 0 else: percent_complete = ( float(self.current_index - self.source.start_index) / ( self.source.end_index - self.source.start_index)) return iobase.ReaderProgress(percent_complete=percent_complete) DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/inmemory_test.py000066400000000000000000000047311266367603600260060ustar00rootroot00000000000000# 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.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(1, reader.get_progress().percent_complete) def test_in_memory_source_updates_progress_one(self): source = inmemory.InMemorySource([1], coder=FakeCoder()) with source.reader() as reader: self.assertEqual(0, reader.get_progress().percent_complete) i = 0 for item in reader: i += 1 self.assertEqual(11, item) self.assertEqual(1, reader.get_progress().percent_complete) self.assertEqual(1, i) self.assertEqual(1, reader.get_progress().percent_complete) 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(0, reader.get_progress().percent_complete) i = 0 for item in reader: i += 1 self.assertEqual(i + 10, item) self.assertEqual(float(i) / 5, reader.get_progress().percent_complete) self.assertEqual(5, i) self.assertEqual(1, reader.get_progress().percent_complete) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/logger.py000066400000000000000000000136311266367603600243660ustar00rootroot00000000000000# 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. per_thread_worker_data = threading.local() class PerThreadLoggingContext(object): """A context manager to add per thread attributes.""" def __init__(self, *args, **kwargs): if args: raise ValueError( 'PerThreadLoggingContext expects only keyword arguments.') self.kwargs = kwargs self.previous = {} def __enter__(self): for key in self.kwargs: if hasattr(per_thread_worker_data, key): self.previous[key] = getattr(per_thread_worker_data, key) setattr(per_thread_worker_data, key, self.kwargs[key]) return self def __exit__(self, exn_type, exn_value, exn_traceback): for key in self.kwargs: if key in self.previous: setattr(per_thread_worker_data, key, self.previous[key]) else: delattr(per_thread_worker_data, key) 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. 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') # Prepare the actual message using the message formatting string and the # positional arguments as they have been used in the log call. output['message'] = record.msg % record.args # 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. if hasattr(per_thread_worker_data, 'work_item_id'): output['work'] = getattr(per_thread_worker_data, 'work_item_id') if hasattr(per_thread_worker_data, 'stage_name'): output['stage'] = getattr(per_thread_worker_data, 'stage_name') if hasattr(per_thread_worker_data, 'step_name'): output['step'] = getattr(per_thread_worker_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.0/google/cloud/dataflow/worker/logger_test.py000066400000000000000000000116661266367603600254330ustar00rootroot00000000000000# 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(hasattr(logger.per_thread_worker_data, name)) with logger.PerThreadLoggingContext(xyz='thread-value'): self.assertEqual( getattr(logger.per_thread_worker_data, name), 'thread-value') self.assertFalse(hasattr(logger.per_thread_worker_data, name)) def test_no_positional_args(self): with self.assertRaises(ValueError) as exn: with logger.PerThreadLoggingContext('something'): pass self.assertEqual( exn.exception.message, 'PerThreadLoggingContext expects only keyword arguments.') def test_per_thread_attribute(self): self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz')) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.xyz, 'value') thread = threading.Thread( target=self.thread_check_attribute, args=('xyz',)) thread.start() thread.join() self.assertEqual(logger.per_thread_worker_data.xyz, 'value') self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz')) def test_set_when_undefined(self): self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz')) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.xyz, 'value') self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz')) def test_set_when_already_defined(self): self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz')) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.xyz, 'value') with logger.PerThreadLoggingContext(xyz='value2'): self.assertEqual(logger.per_thread_worker_data.xyz, 'value2') self.assertEqual(logger.per_thread_worker_data.xyz, 'value') self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz')) 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 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_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.0/google/cloud/dataflow/worker/maptask.py000066400000000000000000000740271266367603600245550ustar00rootroot00000000000000# 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.json_value import from_json_value from google.cloud.dataflow.worker import concat_reader from google.cloud.dataflow.worker import inmemory from google.cloud.dataflow.worker import windmillio 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', 'elements', 'serialized_fn', 'append_trailing_newlines', 'strip_trailing_newlines', 'compression_type', '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))) # TODO(silviuc): Figure out what to do with the codec properties. # They appear in many places (e.g., InstructionOutput, Source, Sink, etc.). # For now they are ignored and not even parsed out of the response message. # 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', 'tag'] ) """Worker details needed to read from a source. Attributes: source: a source object. tag: string tag used if source is used on a side input. None otherwise. """ WorkerGroupingShuffleRead = build_worker_instruction( 'WorkerGroupingShuffleRead', ['start_shuffle_position', 'end_shuffle_position', 'shuffle_reader_config', '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. coders: A 2-tuple of coders (key, value) to decode shuffle entries. """ WorkerUngroupedShuffleRead = build_worker_instruction( 'WorkerUngroupedShuffleRead', ['start_shuffle_position', 'end_shuffle_position', 'shuffle_reader_config', '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. coders: A 2-tuple of coders (key, value) to decode shuffle entries. """ WorkerWrite = build_worker_instruction( 'WorkerWrite', ['sink', 'input']) """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). """ WorkerInMemoryWrite = build_worker_instruction( 'WorkerInMemoryWrite', ['output_buffer', 'input']) """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). """ WorkerShuffleWrite = build_worker_instruction( 'WorkerShuffleWrite', ['shuffle_kind', 'shuffle_writer_config', 'input', '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). coders: A 2-tuple of coders (key, value) to encode shuffle entries. """ WorkerDoFn = build_worker_instruction( 'WorkerDoFn', ['serialized_fn', 'output_tags', 'input', 'side_inputs']) """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']. 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']) """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']. 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', 'output_tags', 'input', 'coders', 'context']) """Worker details needed to run a WindowInto. Attributes: window_fn: A serialized Windowing object representing the windowing strategy. 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']. 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']) """Worker details needed to run a CombineFn. Attributes: serialized_fn: A serialized CombineFn object to be run. 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). """ WorkerPartialGroupByKey = build_worker_instruction( 'WorkerPartialGroupByKey', ['input']) """Worker details needed to run a partial group-by-key. Attributes: 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). """ WorkerFlatten = build_worker_instruction( 'WorkerFlatten', ['inputs']) """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). """ 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, workitem_commit_request, windmill, state): self.computation_id = computation_id self.work_item = work_item self.input_data_watermark = input_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_read_work_item(work, env, context): """Parses a read parallel instruction into the appropriate Worker* object.""" specs = {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(specs, codec_specs, context) if source: return WorkerRead(source, tag=None) # TODO(mairbek) create Shuffler Source/Reader kv_coders = get_coder_from_spec(codec_specs, kv_pair=True) if specs['@type'] == 'GroupingShuffleSource': return WorkerGroupingShuffleRead( start_shuffle_position=specs['start_shuffle_position']['value'], end_shuffle_position=specs['end_shuffle_position']['value'], shuffle_reader_config=specs['shuffle_reader_config']['value'], coders=kv_coders) elif specs['@type'] == 'UngroupedShuffleSource': return WorkerUngroupedShuffleRead( start_shuffle_position=specs['start_shuffle_position']['value'], end_shuffle_position=specs['end_shuffle_position']['value'], shuffle_reader_config=specs['shuffle_reader_config']['value'], coders=kv_coders) else: raise NotImplementedError('Unknown source type: %r' % specs) # 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(WorkerRead(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: return WorkerWrite(sink, input=get_input_spec(work.write.input)) # TODO(mairbek) create Shuffler Sink/Writer if specs['@type'] == 'ShuffleSink': kv_coders = get_coder_from_spec(codec_specs, kv_pair=True) return WorkerShuffleWrite( shuffle_kind=specs['shuffle_kind']['value'], shuffle_writer_config=specs['shuffle_writer_config']['value'], input=get_input_spec(work.write.input), coders=kv_coders) 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], 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)) elif specs['@type'] == 'ReifyTimestampAndWindowsDoFn': return WorkerReifyTimestampAndWindows( output_tags=[o.tag for o in work.parDo.multiOutputInfos], input=get_input_spec(work.parDo.input)) elif specs['@type'] == 'MergeBucketsDoFn': return WorkerMergeWindows( window_fn=specs['serialized_fn']['value'], output_tags=[o.tag for o in work.parDo.multiOutputInfos], input=get_input_spec(work.parDo.input), coders=None, context=context) # TODO(silviuc): Implement support for MergeBucketsDoFn. # TODO(silviuc): Implement support for ReifytimestampsAndWindowsDoFn. # 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]) 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. """ return WorkerPartialGroupByKey( input=get_input_spec(instruction.partialGroupByKey.input)) 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 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). return io.TextFileSource( file_path=specs['filename']['value'], start_offset=None, end_offset=None, 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 io.TextFileSink( 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 io.TextFileSink( 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.0/google/cloud/dataflow/worker/opcounters.py000066400000000000000000000040511266367603600253040ustar00rootroot00000000000000# 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 numbers import Number from google.cloud.dataflow.utils.counters import Counter class OperationCounters(object): """The set of basic counters to attach to an Operation.""" def __init__(self, step_name, output_index=0): self.element_counter = Counter( '%s-out%d-ElementCount' % (step_name, output_index), Counter.SUM) self.mean_byte_counter = Counter( '%s-out%d-MeanByteCount' % (step_name, output_index), Counter.MEAN) def update(self, windowed_value): """Add one value to this counter.""" if isinstance(windowed_value.value, Number): size = 4 # numbers take 4 bytes else: try: # len() gives the right answer for at least strings size = len(windowed_value.value) except (AttributeError, TypeError): # it's an object, not data, and there's nothing to count. size = 0 self.mean_byte_counter.update(size) self.element_counter.update(1) def __iter__(self): """Iterator over all our counters.""" yield self.element_counter yield self.mean_byte_counter 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.0/google/cloud/dataflow/worker/opcounters_test.py000066400000000000000000000055501266367603600263500ustar00rootroot00000000000000# 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.transforms.window import GlobalWindows from google.cloud.dataflow.worker.opcounters import OperationCounters class OperationCountersTest(unittest.TestCase): def verify_counters(self, opcounts, expected_elements, expected_total_bytes): self.assertEqual(expected_elements, opcounts.element_counter.total) self.assertEqual(expected_elements, opcounts.element_counter.elements) self.assertEqual(expected_elements, opcounts.mean_byte_counter.elements) self.assertEqual(expected_total_bytes, opcounts.mean_byte_counter.total) def test_update_int(self): opcounts = OperationCounters('some-name') self.verify_counters(opcounts, 0, 0) opcounts.update(GlobalWindows.WindowedValue(1)) self.verify_counters(opcounts, 1, 4) # an int is 4 bytes def test_update_str(self): opcounts = OperationCounters('some-name') self.verify_counters(opcounts, 0, 0) opcounts.update(GlobalWindows.WindowedValue('abcde')) self.verify_counters(opcounts, 1, 5) # the string is 5 bytes long def test_update_old_object(self): class OldClassThatDoesNotImplementLen: # pylint: disable=old-style-class def __init__(self): pass opcounts = OperationCounters('some-name') self.verify_counters(opcounts, 0, 0) obj = OldClassThatDoesNotImplementLen() opcounts.update(GlobalWindows.WindowedValue(obj)) self.verify_counters(opcounts, 1, 0) # objects contribute no length def test_update_new_object(self): class ObjectThatDoesNotImplementLen(object): def __init__(self): pass opcounts = OperationCounters('some-name') self.verify_counters(opcounts, 0, 0) obj = ObjectThatDoesNotImplementLen() opcounts.update(GlobalWindows.WindowedValue(obj)) self.verify_counters(opcounts, 1, 0) # objects contribute no length def test_update_multiple(self): opcounts = OperationCounters('some-name') self.verify_counters(opcounts, 0, 0) opcounts.update(GlobalWindows.WindowedValue('abcde')) opcounts.update(GlobalWindows.WindowedValue('defghij')) self.verify_counters(opcounts, 2, 12) # the strings add up to 12 characters if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/shuffle.py000066400000000000000000000434771266367603600245560ustar00rootroot00000000000000# 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.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 '' % ( self.key, self.secondary_key, self.value, self.position) 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(object): """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=''): 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 yield self.value_coder.decode(entry.value) class ShuffleReaderBase(iobase.SourceReader): """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 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.source.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 not self._range_tracker.try_return_record_at(is_at_split_point, 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 yield (self.source.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_at_position( _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.source.value_coder.decode(entry.value) class ShuffleSourceBase(iobase.Source): """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 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.sink.key_coder.encode(key), self.sink.key_coder.encode(secondary_key), self.sink.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.0/google/cloud/dataflow/worker/shuffle_test.py000066400000000000000000000317161266367603600256060ustar00rootroot00000000000000# 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.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(object): """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('2nd-%s' % 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', '2nd-a', '1'), ('b', '2nd-b', '0'), ('b', '2nd-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.0/google/cloud/dataflow/worker/start.py000066400000000000000000000041471266367603600242460ustar00rootroot00000000000000# 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 logging import random import re import sys from google.cloud.dataflow.worker import batchworker 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:]) 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. # pylint: disable=g-import-not-at-top from google.cloud.dataflow.worker import streamingworker # pylint: enable=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).run() if __name__ == '__main__': main() DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/streamingworker.py000066400000000000000000000167471266367603600263450ustar00rootroot00000000000000# 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 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 logger from google.cloud.dataflow.worker import maptask from google.cloud.dataflow.worker import windmillstate import apitools.base.py as apitools_base import apitools.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 # TODO(altay): Remove windmill default port and host. WINDMILL_DEFAULT_PORT = 12355 def __init__(self, properties): self.project_id = properties['project_id'] self.job_id = properties['job_id'] self.worker_id = properties['worker_id'] self.log_path = properties['dataflow.worker.logging.location'] # Initialize the logging machinery. logger.initialize(job_id=self.job_id, worker_id=self.worker_id, log_path=self.log_path) logging.info('Streaming worker started with: %s', properties) self.client_id = random.getrandbits(63) windmill_host = properties.get('windmill.host', self.worker_id) windmill_port = int(properties.get('windmill.hostport', StreamingWorker.WINDMILL_DEFAULT_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: computation_id = computation_work.computation_id input_data_watermark = 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: try: self.process(computation_id, map_task_proto, input_data_watermark, work_item) except: logging.error( 'Exception while processing work item for computation %r: ' '%s, %s', computation_id, work_item, traceback.format_exc()) raise def process(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) context.start(computation_id, work_item, input_data_watermark, workitem_commit_request, self.windmill, state) map_task_executor = executor.MapTaskExecutor() map_task = maptask.decode_map_task(map_task_proto, env, context) map_task_executor.execute(map_task) 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.0/google/cloud/dataflow/worker/windmillio.py000066400000000000000000000165261266367603600252640ustar00rootroot00000000000000# 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 import time 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.window import GlobalWindows from google.cloud.dataflow.transforms.window import WindowedValue def harness_to_windmill_timestamp(float_timestamp): # The timestamp taken by Windmill is in microseconds. return int(float_timestamp * 1000000) def windmill_to_harness_timestamp(windmill_timestamp): # The timestamp given by Windmill is in microseconds. return windmill_timestamp / 1000000.0 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.SourceReader): """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.WindowedValue( 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 def Write(self, line): data = self.sink.coder.encode(line) timestamp = harness_to_windmill_timestamp(time.time()) 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 = harness_to_windmill_timestamp(windowed_kv.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=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 KeyedWorkItem(object): """Keyed work item used by a StreamingGroupAlsoByWindowsOperation.""" def __init__(self, work_item, coder): 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) def elements(self): for bundle in self.work_item.message_bundles: for message in bundle.messages: yield self.wv_coder.decode(message.data) def __repr__(self): return 'KeyedWorkItem(%r)' % self.key class WindowingWindmillSource(iobase.Source): """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.SourceReader): """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.0/google/cloud/dataflow/worker/windmillstate.py000066400000000000000000000327061266367603600257730ustar00rootroot00000000000000# 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 # 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, tag, timestamp): # TODO(ccy): implement this. logging.info('Ignoring set_timer(%s).', (window, tag, timestamp)) def clear_timer(self, window, tag): # TODO(ccy): implement this. logging.info('Ignoring clear_timer(%s).', (window, tag)) 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): if tag is None: # TODO(ccy): either implement this, or if this primitive is not supported # by Windmill, modify the upstream caller in the triggering code. logging.info('Ignoring clear all state request for window %s.', window) return 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 = {} 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) else: raise ValueError('Invalid state tag.') return self.accessed[state_key] def persist_to(self, commit_request): for unused_key, accessor in self.accessed.iteritems(): accessor.persist_to(commit_request) class WindmillStateReader(object): """Reader of raw state from Windmill.""" 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): """Get the list at given state tag.""" # TODO(ccy): refactor to support continuation tokens for paginated reading. # The current implementation returns up to one page of values from the # list from windmill. 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, fetch_max_bytes=WindmillStateReader.MAX_LIST_BYTES) 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 def get(self): if not self.fetched: self._fetch() return self.value def add(self, value): self.modified = True # 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.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 commit_request.value_updates.add( tag=self.state_key, state_family='', value=windmill_pb2.Value( data=encode_value(self.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): currently, we only look at the first page of the result # since we do not support pagination. We should support pagination. # 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. result = self.reader.fetch_list(self.state_key) for wrapper in result.data: for datum in wrapper.data: for item in datum.lists: 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 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) DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/workitem.py000066400000000000000000000100041266367603600247370ustar00rootroot00000000000000# 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 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, map_task): self.proto = proto self.map_task = map_task # 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 self.next_report_index = self.proto.initialReportIndex self.lease_expire_time = self.proto.leaseExpireTime self.report_status_interval = self.proto.reportStatusInterval def __str__(self): return '<%s %s steps=%s %s>' % ( self.__class__.__name__, self.map_task.stage_name, '+'.join(self.map_task.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. """ # 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 = response.workItems[0] map_task = maptask.decode_map_task(work_item.mapTask, env, context) return BatchWorkItem(work_item, map_task) DataflowPythonSDK-0.2.0/google/cloud/dataflow/worker/workitem_test.py000066400000000000000000000432621266367603600260120ustar00rootroot00000000000000# 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.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 apitools.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_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(dataflow.ParallelInstruction(read=ri)) mt.instructions.append(dataflow.ParallelInstruction(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(dataflow.ParallelInstruction(read=ri)) mt.instructions.append(dataflow.ParallelInstruction(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(dataflow.ParallelInstruction(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(dataflow.ParallelInstruction(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() 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(dataflow.ParallelInstruction(read=rsi)) mt.instructions.append(dataflow.ParallelInstruction(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, tag=None), maptask.WorkerDoFn( serialized_fn='code', output_tags=['out'], input=(1, 0), side_inputs=[]), maptask.WorkerShuffleWrite( shuffle_kind='group_keys', shuffle_writer_config='opaque', input=(1, 0), coders=(CODER.key_coder(), CODER.value_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), tag=None), maptask.WorkerDoFn( serialized_fn='code', output_tags=['out'], input=(1, 0), side_inputs=[]), maptask.WorkerShuffleWrite( shuffle_kind='group_keys', shuffle_writer_config='opaque', input=(1, 0), coders=(CODER.key_coder(), CODER.value_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', coders=(CODER.key_coder(), CODER.value_coder())), maptask.WorkerWrite(io.TextFileSink( file_path_prefix='gs://somefile', append_trailing_newlines=True, coder=CODER), input=(0, 0))])) 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', coders=(CODER.key_coder(), CODER.value_coder())), maptask.WorkerWrite(io.TextFileSink( file_path_prefix='gs://somefile', append_trailing_newlines=True, coder=CODER), input=(0, 0))])) 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), tag=None), maptask.WorkerWrite(io.TextFileSink( file_path_prefix='gs://somefile', append_trailing_newlines=True, coder=CODER), input=(0, 0))])) 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), tag=None), maptask.WorkerFlatten( inputs=[(0, 0)])])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() DataflowPythonSDK-0.2.0/setup.cfg000066400000000000000000000013061266367603600166560ustar00rootroot00000000000000[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. exclude=examples|bigquery_test|ptransform_test|fast_coders_test|typecoders_test|workitem_test|executor_test DataflowPythonSDK-0.2.0/setup.py000066400000000000000000000101241266367603600165450ustar00rootroot00000000000000# 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 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. 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.2', # Pin the version of APItools since 0.4.12 is broken and 0.4.11 is the # last known good. # TODO(silviuc): Redesign requirements to have pinned versions. 'google-apitools==0.4.11', 'google-apitools-bigquery-v2', 'google-apitools-dataflow-v1b3', 'google-apitools-storage-v1', 'httplib2>=0.8', 'mock>=1.0.1', # Pin the version of oauth2client since 2.0.0 does not have # oauth2client.client.SignedJwtAssertionCredentials and 1.5.2 is the last # known good version. 'oauth2client==1.5.2', '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() setuptools.setup( name='python_dataflow', version=_PYTHON_DATAFLOW_VERSION, description='Google Cloud Dataflow SDK for Python', long_description=get_dataflow_docstring(), url='http://cloud.google.com/dataflow/', author='Google, Inc.', author_email='dataflow-feedback@google.com', packages=setuptools.find_packages(), entry_points={ 'console_scripts': CONSOLE_SCRIPTS, }, ext_modules=cythonize( ['**/*.pyx', 'google/cloud/dataflow/coders/coder_impl.py']), setup_requires=['nose>=1.0'], install_requires=REQUIRED_PACKAGES, test_suite='nose.collector', # 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.0/tox.ini000066400000000000000000000001311266367603600163430ustar00rootroot00000000000000[tox] envlist = py27 [testenv:py27] commands = python setup.py test passenv = TRAVIS*