Initial commit

Author: msukmanowsky

.gitignore

@@ -0,0 +1,3 @@
+.python-version
+drpyspark.egg-info
+*.pyc

README.md

@@ -0,0 +1,67 @@
+# drpyspark
+
+The doctor is in.
+
+drpyspark provides handy utilities for debugging and tuning pyspark programs.
+A work in progress.
+
+
+## Better debugging
+
+Just add
+
+```python
+import drpyspark
+drpyspark.enable_debug_output()
+```
+
+To your Spark script and then you'll get wonderful output that shows you
+exactly what is happening at every stage in your pyspark script without having
+to add 500 print statements in between things.
+
+
+```python
+def is_even(val):
+    if val % 2 == 0:
+        return [val]
+    return []
+
+
+def square(val):
+    return val ** 2
+
+
+with pyspark.SparkContext() as sc:
+    numbers = sc.parallelize([str(x) for x in xrange(10)])
+    even_squares = (numbers
+                    .map(lambda l: int(l))
+                    .map(square)
+                    .flatMap(is_even))
+    div_100 = even_squares.map(lambda l: l / 100.0)
+    print(div_100.collect())
+
+```
+
+When run with `enable_debug_output`, provides:
+```
+/Users/mikesukmanowsky/code/drpyspark/examples/rdd_1.py:18: numbers = sc.parallelize([str(x) for x in xrange(10)])
+['0', '1', '2', '3', '4']
+/Users/mikesukmanowsky/code/drpyspark/examples/rdd_1.py:20: .map(lambda l: int(l))
+[0, 1, 2, 3, 4]
+/Users/mikesukmanowsky/code/drpyspark/examples/rdd_1.py:21: .map(square)
+[0, 1, 4, 9, 16]
+/Users/mikesukmanowsky/code/drpyspark/examples/rdd_1.py:22: .flatMap(is_even))
+[0, 4, 16, 36, 64]
+/Users/mikesukmanowsky/code/drpyspark/examples/rdd_1.py:23: div_100 = even_squares.map(lambda l: l / 100.0)
+[0.0, 0.04, 0.16, 0.36, 0.64]
+```
+
+## Running examples
+
+You'll need to [download a release of Apache Spark](http://spark.apache.org/)
+With a virtualenv built, install `drpyspark` with `pip install drpyspark` (or
+`python setup.py develop`) then run:
+
+```
+PYSPARK_PYTHON=$(which python) $SPARK_HOME/bin/spark-submit examples/<filename>
+```

drpyspark/__init__.py

@@ -0,0 +1,62 @@
+from __future__ import print_function
+import functools
+import inspect
+import logging
+import re
+import pprint
+import sys
+
+try:
+    import pyspark
+except ImportError:
+    pyspark = None
+
+from .version import VERSION
+log = logging.getLogger(__package__)
+
+
+def print_output(f):
+    @functools.wraps(f)
+    def _debug_pyspark_call(*args, **kwargs):
+        log.debug('%s called', f.__name__)
+        stack = inspect.stack()
+        caller = stack[1]
+        caller_package = inspect.getmodule(caller[0]).__package__
+
+        result = f(*args, **kwargs)
+        if caller_package in ('pyspark', __package__):
+            log.debug('%s internal call to %s, returning to avoid infinite '
+                      'recursion', caller_package, f.__name__)
+            return result
+        if not isinstance(result, (pyspark.RDD, pyspark.sql.DataFrame)):
+            log.debug('%s returned non RDD/DataFrame value, returning',
+                      f.__name__)
+            return result
+
+        log.debug('printing 5 from %s', f.__name__)
+        sample = result.take(5)
+        file, line_no, code = stack[-1][1], stack[-1][2], ''.join(stack[-1][4]).strip()
+        print('{}:{}: {}'.format(file, line_no, code))
+        pprint.pprint(sample)
+        return result
+
+    return _debug_pyspark_call
+
+
+def enable_debug_output(num_elements=5):
+    if pyspark is None:
+        print('pyspark not found in PYTHONPATH, did you run via spark-submit?',
+              file=sys.stderr)
+        sys.exit(1)
+
+    classes_to_patch = (pyspark.SparkContext, pyspark.sql.SQLContext,
+                        pyspark.sql.HiveContext, pyspark.RDD,
+                        pyspark.sql.DataFrame)
+    for klass in classes_to_patch:
+        members = inspect.getmembers(klass)
+        # get all public methods not starting with _ or save
+        methods = [(name, member) for (name, member) in members
+                   if not name.startswith('_') and inspect.ismethod(member)]
+        for name, method in methods:
+            setattr(klass, name, print_output(method))
+            log.debug('Patched %s.%s.', klass.__name__, name)

drpyspark/version.py

@@ -0,0 +1,11 @@
+def _safe_int(string):
+    """Simple function to convert strings into ints without dying.
+    Helps when we define versions like 0.1.0dev"""
+    try:
+        return int(string)
+    except ValueError:
+        return string
+
+
+__version__ = '0.1.0'
+VERSION = tuple(_safe_int(x) for x in __version__.split('.'))

examples/rdd_1.py

@@ -0,0 +1,25 @@
+from __future__ import print_function
+import pyspark
+import drpyspark
+drpyspark.enable_debug_output()
+
+
+def is_even(val):
+    if val % 2 == 0:
+        return [val]
+    return []
+
+
+def square(val):
+    return val ** 2
+
+
+with pyspark.SparkContext() as sc:
+    numbers = sc.parallelize([str(x) for x in xrange(10)])
+    even_squares = (numbers
+                    .map(lambda l: int(l))
+                    .map(square)
+                    .flatMap(is_even))
+    div_100 = even_squares.map(lambda l: l / 100.0)
+    print(div_100.collect())
+