#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

import copy

import sys

import os

import re

import operator

import shlex

import warnings

import heapq

import bisect

import random

import socket

from subprocess import Popen, PIPE

from tempfile import NamedTemporaryFile

from threading import Thread

from collections import defaultdict

from itertools import chain

from functools import reduce

from math import sqrt, log, isinf, isnan, pow, ceil

if sys.version > '3':

basestring = unicode = str

else:

from itertools import imap as map, ifilter as filter

from pyspark.serializers import NoOpSerializer, CartesianDeserializer, \

BatchedSerializer, CloudPickleSerializer, PairDeserializer, \

PickleSerializer, pack_long, AutoBatchedSerializer

from pyspark.join import python_join, python_left_outer_join, \

python_right_outer_join, python_full_outer_join, python_cogroup

from pyspark.statcounter import StatCounter

from pyspark.rddsampler import RDDSampler, RDDRangeSampler, RDDStratifiedSampler

from pyspark.storagelevel import StorageLevel

from pyspark.resultiterable import ResultIterable

from pyspark.shuffle import Aggregator, ExternalMerger, \

get_used_memory, ExternalSorter, ExternalGroupBy

from pyspark.traceback_utils import SCCallSiteSync

__all__ = ["RDD"]

def portable_hash(x):

"""

This function returns consistent hash code for builtin types, especially

for None and tuple with None.

The algorithm is similar to that one used by CPython 2.7

>>> portable_hash(None)

0

>>> portable_hash((None, 1)) & 0xffffffff

219750521

"""

if sys.version >= '3.3' and 'PYTHONHASHSEED' not in os.environ:

raise Exception("Randomness of hash of string should be disabled via PYTHONHASHSEED")

if x is None:

return 0

if isinstance(x, tuple):

h = 0x345678

for i in x:

h ^= portable_hash(i)

h *= 1000003

h &= sys.maxsize

h ^= len(x)

if h == -1:

h = -2

return int(h)

return hash(x)

class BoundedFloat(float):

"""

Bounded value is generated by approximate job, with confidence and low

bound and high bound.

>>> BoundedFloat(100.0, 0.95, 95.0, 105.0)

100.0

"""

def __new__(cls, mean, confidence, low, high):

obj = float.__new__(cls, mean)

obj.confidence = confidence

obj.low = low

obj.high = high

return obj

def _parse_memory(s):

"""

Parse a memory string in the format supported by Java (e.g. 1g, 200m) and

return the value in MB

>>> _parse_memory("256m")

256

>>> _parse_memory("2g")

2048

"""

units = {'g': 1024, 'm': 1, 't': 1 << 20, 'k': 1.0 / 1024}

if s[-1].lower() not in units:

raise ValueError("invalid format: " + s)

return int(float(s[:-1]) * units[s[-1].lower()])

def _load_from_socket(port, serializer):

sock = None

# Support for both IPv4 and IPv6.

# On most of IPv6-ready systems, IPv6 will take precedence.

for res in socket.getaddrinfo("localhost", port, socket.AF_UNSPEC, socket.SOCK_STREAM):

af, socktype, proto, canonname, sa = res

sock = socket.socket(af, socktype, proto)

try:

sock.settimeout(3)

sock.connect(sa)

except socket.error:

sock.close()

sock = None

continue

break

if not sock:

raise Exception("could not open socket")

try:

rf = sock.makefile("rb", 65536)

for item in serializer.load_stream(rf):

yield item

finally:

sock.close()

def ignore_unicode_prefix(f):

"""

Ignore the 'u' prefix of string in doc tests, to make it works

in both python 2 and 3

"""

if sys.version >= '3':

# the representation of unicode string in Python 3 does not have prefix 'u',

# so remove the prefix 'u' for doc tests

literal_re = re.compile(r"(\W|^)[uU](['])", re.UNICODE)

f.__doc__ = literal_re.sub(r'\1\2', f.__doc__)

return f

class Partitioner(object):

def __init__(self, numPartitions, partitionFunc):

self.numPartitions = numPartitions

self.partitionFunc = partitionFunc

def __eq__(self, other):

return (isinstance(other, Partitioner) and self.numPartitions == other.numPartitions

and self.partitionFunc == other.partitionFunc)

def __call__(self, k):

return self.partitionFunc(k) % self.numPartitions

class RDD(object):

"""

A Resilient Distributed Dataset (RDD), the basic abstraction in Spark.

Represents an immutable, partitioned collection of elements that can be

operated on in parallel.

"""

def __init__(self, jrdd, ctx, jrdd_deserializer=AutoBatchedSerializer(PickleSerializer())):

self._jrdd = jrdd

self.is_cached = False

self.is_checkpointed = False

self.ctx = ctx

self._jrdd_deserializer = jrdd_deserializer

self._id = jrdd.id()

self.partitioner = None

def _pickled(self):

return self._reserialize(AutoBatchedSerializer(PickleSerializer()))

def id(self):

"""

A unique ID for this RDD (within its SparkContext).

"""

return self._id

def __repr__(self):

return self._jrdd.toString()

def __getnewargs__(self):

# This method is called when attempting to pickle an RDD, which is always an error:

raise Exception(

"It appears that you are attempting to broadcast an RDD or reference an RDD from an "

"action or transformation. RDD transformations and actions can only be invoked by the "

"driver, not inside of other transformations; for example, "

"rdd1.map(lambda x: rdd2.values.count() * x) is invalid because the values "

"transformation and count action cannot be performed inside of the rdd1.map "

"transformation. For more information, see SPARK-5063."

)

@property

def context(self):

"""

The L{SparkContext} that this RDD was created on.

"""

return self.ctx

def cache(self):

"""

Persist this RDD with the default storage level (C{MEMORY_ONLY}).

"""

self.is_cached = True

self.persist(StorageLevel.MEMORY_ONLY)

return self

def persist(self, storageLevel=StorageLevel.MEMORY_ONLY):

"""

Set this RDD's storage level to persist its values across operations

after the first time it is computed. This can only be used to assign

a new storage level if the RDD does not have a storage level set yet.

If no storage level is specified defaults to (C{MEMORY_ONLY}).

>>> rdd = sc.parallelize(["b", "a", "c"])

>>> rdd.persist().is_cached

True

"""

self.is_cached = True

javaStorageLevel = self.ctx._getJavaStorageLevel(storageLevel)

self._jrdd.persist(javaStorageLevel)

return self

def unpersist(self):

"""

Mark the RDD as non-persistent, and remove all blocks for it from

memory and disk.

"""

self.is_cached = False

self._jrdd.unpersist()

return self

def checkpoint(self):

"""

Mark this RDD for checkpointing. It will be saved to a file inside the

checkpoint directory set with L{SparkContext.setCheckpointDir()} and

all references to its parent RDDs will be removed. This function must

be called before any job has been executed on this RDD. It is strongly

recommended that this RDD is persisted in memory, otherwise saving it

on a file will require recomputation.

"""

self.is_checkpointed = True

self._jrdd.rdd().checkpoint()

def isCheckpointed(self):

"""

Return whether this RDD is checkpointed and materialized, either reliably or locally.

"""

return self._jrdd.rdd().isCheckpointed()

def localCheckpoint(self):

"""

Mark this RDD for local checkpointing using Spark's existing caching layer.

This method is for users who wish to truncate RDD lineages while skipping the expensive

step of replicating the materialized data in a reliable distributed file system. This is

useful for RDDs with long lineages that need to be truncated periodically (e.g. GraphX).

Local checkpointing sacrifices fault-tolerance for performance. In particular, checkpointed

data is written to ephemeral local storage in the executors instead of to a reliable,

fault-tolerant storage. The effect is that if an executor fails during the computation,

the checkpointed data may no longer be accessible, causing an irrecoverable job failure.

This is NOT safe to use with dynamic allocation, which removes executors along

with their cached blocks. If you must use both features, you are advised to set

L{spark.dynamicAllocation.cachedExecutorIdleTimeout} to a high value.

The checkpoint directory set through L{SparkContext.setCheckpointDir()} is not used.

"""

self._jrdd.rdd().localCheckpoint()

def isLocallyCheckpointed(self):

"""

Return whether this RDD is marked for local checkpointing.

Exposed for testing.

"""

return self._jrdd.rdd().isLocallyCheckpointed()

def getCheckpointFile(self):

"""

Gets the name of the file to which this RDD was checkpointed

Not defined if RDD is checkpointed locally.

"""

checkpointFile = self._jrdd.rdd().getCheckpointFile()

if checkpointFile.isDefined():

return checkpointFile.get()

def map(self, f, preservesPartitioning=False):

"""

Return a new RDD by applying a function to each element of this RDD.

>>> rdd = sc.parallelize(["b", "a", "c"])

>>> sorted(rdd.map(lambda x: (x, 1)).collect())

[('a', 1), ('b', 1), ('c', 1)]

"""

def func(_, iterator):

return map(f, iterator)

return self.mapPartitionsWithIndex(func, preservesPartitioning)

def flatMap(self, f, preservesPartitioning=False):

"""

Return a new RDD by first applying a function to all elements of this

RDD, and then flattening the results.

>>> rdd = sc.parallelize([2, 3, 4])

>>> sorted(rdd.flatMap(lambda x: range(1, x)).collect())

[1, 1, 1, 2, 2, 3]

>>> sorted(rdd.flatMap(lambda x: [(x, x), (x, x)]).collect())

[(2, 2), (2, 2), (3, 3), (3, 3), (4, 4), (4, 4)]

"""

def func(s, iterator):

return chain.from_iterable(map(f, iterator))

return self.mapPartitionsWithIndex(func, preservesPartitioning)

def mapPartitions(self, f, preservesPartitioning=False):

"""

Return a new RDD by applying a function to each partition of this RDD.

>>> rdd = sc.parallelize([1, 2, 3, 4], 2)

>>> def f(iterator): yield sum(iterator)

>>> rdd.mapPartitions(f).collect()

[3, 7]

"""

def func(s, iterator):

return f(iterator)

return self.mapPartitionsWithIndex(func, preservesPartitioning)

def mapPartitionsWithIndex(self, f, preservesPartitioning=False):

"""

Return a new RDD by applying a function to each partition of this RDD,

while tracking the index of the original partition.

>>> rdd = sc.parallelize([1, 2, 3, 4], 4)

>>> def f(splitIndex, iterator): yield splitIndex

>>> rdd.mapPartitionsWithIndex(f).sum()

6

"""

return PipelinedRDD(self, f, preservesPartitioning)

def mapPartitionsWithSplit(self, f, preservesPartitioning=False):

"""

Deprecated: use mapPartitionsWithIndex instead.

Return a new RDD by applying a function to each partition of this RDD,

while tracking the index of the original partition.

>>> rdd = sc.parallelize([1, 2, 3, 4], 4)

>>> def f(splitIndex, iterator): yield splitIndex

>>> rdd.mapPartitionsWithSplit(f).sum()

6

"""

warnings.warn("mapPartitionsWithSplit is deprecated; "

"use mapPartitionsWithIndex instead", DeprecationWarning, stacklevel=2)

return self.mapPartitionsWithIndex(f, preservesPartitioning)

def getNumPartitions(self):

"""

Returns the number of partitions in RDD

>>> rdd = sc.parallelize([1, 2, 3, 4], 2)

>>> rdd.getNumPartitions()

2

"""

return self._jrdd.partitions().size()

def filter(self, f):

"""

Return a new RDD containing only the elements that satisfy a predicate.

>>> rdd = sc.parallelize([1, 2, 3, 4, 5])

>>> rdd.filter(lambda x: x % 2 == 0).collect()

[2, 4]

"""

def func(iterator):

return filter(f, iterator)

return self.mapPartitions(func, True)

def distinct(self, numPartitions=None):

"""

Return a new RDD containing the distinct elements in this RDD.

>>> sorted(sc.parallelize([1, 1, 2, 3]).distinct().collect())

[1, 2, 3]

"""

return self.map(lambda x: (x, None)) \

.reduceByKey(lambda x, _: x, numPartitions) \

.map(lambda x: x[0])

def sample(self, withReplacement, fraction, seed=None):

"""

Return a sampled subset of this RDD.

:param withReplacement: can elements be sampled multiple times (replaced when sampled out)

:param fraction: expected size of the sample as a fraction of this RDD's size

without replacement: probability that each element is chosen; fraction must be [0, 1]

with replacement: expected number of times each element is chosen; fraction must be >= 0

:param seed: seed for the random number generator

.. note:: This is not guaranteed to provide exactly the fraction specified of the total

count of the given :class:`DataFrame`.

>>> rdd = sc.parallelize(range(100), 4)

>>> 6 <= rdd.sample(False, 0.1, 81).count() <= 14

True

"""

assert fraction >= 0.0, "Negative fraction value: %s" % fraction

return self.mapPartitionsWithIndex(RDDSampler(withReplacement, fraction, seed).func, True)

def randomSplit(self, weights, seed=None):

"""

Randomly splits this RDD with the provided weights.

:param weights: weights for splits, will be normalized if they don't sum to 1

:param seed: random seed

:return: split RDDs in a list

>>> rdd = sc.parallelize(range(500), 1)

>>> rdd1, rdd2 = rdd.randomSplit([2, 3], 17)

>>> len(rdd1.collect() + rdd2.collect())

500

>>> 150 < rdd1.count() < 250

True

>>> 250 < rdd2.count() < 350

True

"""

s = float(sum(weights))

cweights = [0.0]

for w in weights:

cweights.append(cweights[-1] + w / s)

if seed is None:

seed = random.randint(0, 2 ** 32 - 1)

return [self.mapPartitionsWithIndex(RDDRangeSampler(lb, ub, seed).func, True)

for lb, ub in zip(cweights, cweights[1:])]

# this is ported from scala/spark/RDD.scala

def takeSample(self, withReplacement, num, seed=None):

"""

Return a fixed-size sampled subset of this RDD.

.. note:: This method should only be used if the resulting array is expected

to be small, as all the data is loaded into the driver's memory.

>>> rdd = sc.parallelize(range(0, 10))

>>> len(rdd.takeSample(True, 20, 1))

20

>>> len(rdd.takeSample(False, 5, 2))

5

>>> len(rdd.takeSample(False, 15, 3))

10

"""

numStDev = 10.0

if num < 0:

raise ValueError("Sample size cannot be negative.")

elif num == 0:

return []

initialCount = self.count()

if initialCount == 0:

return []

rand = random.Random(seed)

if (not withReplacement) and num >= initialCount:

# shuffle current RDD and return

samples = self.collect()

rand.shuffle(samples)

return samples

maxSampleSize = sys.maxsize - int(numStDev * sqrt(sys.maxsize))

if num > maxSampleSize:

raise ValueError(

"Sample size cannot be greater than %d." % maxSampleSize)

fraction = RDD._computeFractionForSampleSize(

num, initialCount, withReplacement)

samples = self.sample(withReplacement, fraction, seed).collect()

# If the first sample didn't turn out large enough, keep trying to take samples;

# this shouldn't happen often because we use a big multiplier for their initial size.

# See: scala/spark/RDD.scala

while len(samples) < num:

# TODO: add log warning for when more than one iteration was run

seed = rand.randint(0, sys.maxsize)

samples = self.sample(withReplacement, fraction, seed).collect()

rand.shuffle(samples)

return samples[0:num]

@staticmethod

def _computeFractionForSampleSize(sampleSizeLowerBound, total, withReplacement):

"""

Returns a sampling rate that guarantees a sample of

size >= sampleSizeLowerBound 99.99% of the time.

How the sampling rate is determined:

Let p = num / total, where num is the sample size and total is the

total number of data points in the RDD. We're trying to compute

q > p such that

- when sampling with replacement, we're drawing each data point

with prob_i ~ Pois(q), where we want to guarantee

Pr[s < num] < 0.0001 for s = sum(prob_i for i from 0 to

total), i.e. the failure rate of not having a sufficiently large

sample < 0.0001. Setting q = p + 5 * sqrt(p/total) is sufficient

to guarantee 0.9999 success rate for num > 12, but we need a

slightly larger q (9 empirically determined).

- when sampling without replacement, we're drawing each data point

with prob_i ~ Binomial(total, fraction) and our choice of q

guarantees 1-delta, or 0.9999 success rate, where success rate is

defined the same as in sampling with replacement.

"""

fraction = float(sampleSizeLowerBound) / total

if withReplacement:

numStDev = 5

if (sampleSizeLowerBound < 12):

numStDev = 9

return fraction + numStDev * sqrt(fraction / total)

else:

delta = 0.00005

gamma = - log(delta) / total

return min(1, fraction + gamma + sqrt(gamma * gamma + 2 * gamma * fraction))

def union(self, other):

"""

Return the union of this RDD and another one.

>>> rdd = sc.parallelize([1, 1, 2, 3])

>>> rdd.union(rdd).collect()

[1, 1, 2, 3, 1, 1, 2, 3]

"""

if self._jrdd_deserializer == other._jrdd_deserializer:

rdd = RDD(self._jrdd.union(other._jrdd), self.ctx,

self._jrdd_deserializer)

else:

# These RDDs contain data in different serialized formats, so we

# must normalize them to the default serializer.

self_copy = self._reserialize()

other_copy = other._reserialize()

rdd = RDD(self_copy._jrdd.union(other_copy._jrdd), self.ctx,

self.ctx.serializer)

if (self.partitioner == other.partitioner and

self.getNumPartitions() == rdd.getNumPartitions()):

rdd.partitioner = self.partitioner

return rdd

def intersection(self, other):

"""

Return the intersection of this RDD and another one. The output will

not contain any duplicate elements, even if the input RDDs did.

.. note:: This method performs a shuffle internally.

>>> rdd1 = sc.parallelize([1, 10, 2, 3, 4, 5])

>>> rdd2 = sc.parallelize([1, 6, 2, 3, 7, 8])

>>> rdd1.intersection(rdd2).collect()

[1, 2, 3]

"""

return self.map(lambda v: (v, None)) \

.cogroup(other.map(lambda v: (v, None))) \

.filter(lambda k_vs: all(k_vs[1])) \

.keys()

def _reserialize(self, serializer=None):

serializer = serializer or self.ctx.serializer

if self._jrdd_deserializer != serializer:

self = self.map(lambda x: x, preservesPartitioning=True)

self._jrdd_deserializer = serializer

return self

def __add__(self, other):

"""

Return the union of this RDD and another one.

>>> rdd = sc.parallelize([1, 1, 2, 3])

>>> (rdd + rdd).collect()

[1, 1, 2, 3, 1, 1, 2, 3]

"""

if not isinstance(other, RDD):

raise TypeError

return self.union(other)

def repartitionAndSortWithinPartitions(self, numPartitions=None, partitionFunc=portable_hash,

ascending=True, keyfunc=lambda x: x):

"""

Repartition the RDD according to the given partitioner and, within each resulting partition,

sort records by their keys.

>>> rdd = sc.parallelize([(0, 5), (3, 8), (2, 6), (0, 8), (3, 8), (1, 3)])

>>> rdd2 = rdd.repartitionAndSortWithinPartitions(2, lambda x: x % 2, 2)

>>> rdd2.glom().collect()

[[(0, 5), (0, 8), (2, 6)], [(1, 3), (3, 8), (3, 8)]]

"""

if numPartitions is None:

numPartitions = self._defaultReducePartitions()

memory = _parse_memory(self.ctx._conf.get("spark.python.worker.memory", "512m"))

serializer = self._jrdd_deserializer

def sortPartition(iterator):

sort = ExternalSorter(memory * 0.9, serializer).sorted

return iter(sort(iterator, key=lambda k_v: keyfunc(k_v[0]), reverse=(not ascending)))

return self.partitionBy(numPartitions, partitionFunc).mapPartitions(sortPartition, True)

def sortByKey(self, ascending=True, numPartitions=None, keyfunc=lambda x: x):

"""

Sorts this RDD, which is assumed to consist of (key, value) pairs.

# noqa

>>> tmp = [('a', 1), ('b', 2), ('1', 3), ('d', 4), ('2', 5)]

>>> sc.parallelize(tmp).sortByKey().first()

('1', 3)

>>> sc.parallelize(tmp).sortByKey(True, 1).collect()

[('1', 3), ('2', 5), ('a', 1), ('b', 2), ('d', 4)]

>>> sc.parallelize(tmp).sortByKey(True, 2).collect()

[('1', 3), ('2', 5), ('a', 1), ('b', 2), ('d', 4)]

>>> tmp2 = [('Mary', 1), ('had', 2), ('a', 3), ('little', 4), ('lamb', 5)]

>>> tmp2.extend([('whose', 6), ('fleece', 7), ('was', 8), ('white', 9)])

>>> sc.parallelize(tmp2).sortByKey(True, 3, keyfunc=lambda k: k.lower()).collect()

[('a', 3), ('fleece', 7), ('had', 2), ('lamb', 5),...('white', 9), ('whose', 6)]

"""

if numPartitions is None:

numPartitions = self._defaultReducePartitions()

memory = self._memory_limit()

serializer = self._jrdd_deserializer

def sortPartition(iterator):

sort = ExternalSorter(memory * 0.9, serializer).sorted

return iter(sort(iterator, key=lambda kv: keyfunc(kv[0]), reverse=(not ascending)))

if numPartitions == 1:

if self.getNumPartitions() > 1:

self = self.coalesce(1)

return self.mapPartitions(sortPartition, True)

# first compute the boundary of each part via sampling: we want to partition

# the key-space into bins such that the bins have roughly the same

# number of (key, value) pairs falling into them

rddSize = self.count()

if not rddSize:

return self # empty RDD

maxSampleSize = numPartitions * 20.0 # constant from Spark's RangePartitioner

fraction = min(maxSampleSize / max(rddSize, 1), 1.0)

samples = self.sample(False, fraction, 1).map(lambda kv: kv[0]).collect()

samples = sorted(samples, key=keyfunc)

# we have numPartitions many parts but one of the them has

# an implicit boundary

bounds = [samples[int(len(samples) * (i + 1) / numPartitions)]

for i in range(0, numPartitions - 1)]

def rangePartitioner(k):

p = bisect.bisect_left(bounds, keyfunc(k))

if ascending:

return p

else:

return numPartitions - 1 - p

return self.partitionBy(numPartitions, rangePartitioner).mapPartitions(sortPartition, True)

def sortBy(self, keyfunc, ascending=True, numPartitions=None):

"""

Sorts this RDD by the given keyfunc

>>> tmp = [('a', 1), ('b', 2), ('1', 3), ('d', 4), ('2', 5)]

>>> sc.parallelize(tmp).sortBy(lambda x: x[0]).collect()

[('1', 3), ('2', 5), ('a', 1), ('b', 2), ('d', 4)]

>>> sc.parallelize(tmp).sortBy(lambda x: x[1]).collect()

[('a', 1), ('b', 2), ('1', 3), ('d', 4), ('2', 5)]

"""

return self.keyBy(keyfunc).sortByKey(ascending, numPartitions).values()

def glom(self):

"""

Return an RDD created by coalescing all elements within each partition

into a list.

>>> rdd = sc.parallelize([1, 2, 3, 4], 2)

>>> sorted(rdd.glom().collect())

[[1, 2], [3, 4]]

"""

def func(iterator):

yield list(iterator)

return self.mapPartitions(func)

def cartesian(self, other):

"""

Return the Cartesian product of this RDD and another one, that is, the

RDD of all pairs of elements C{(a, b)} where C{a} is in C{self} and

C{b} is in C{other}.

>>> rdd = sc.parallelize([1, 2])

>>> sorted(rdd.cartesian(rdd).collect())

[(1, 1), (1, 2), (2, 1), (2, 2)]

"""

# Due to batching, we can't use the Java cartesian method.

deserializer = CartesianDeserializer(self._jrdd_deserializer,

other._jrdd_deserializer)

return RDD(self._jrdd.cartesian(other._jrdd), self.ctx, deserializer)

def groupBy(self, f, numPartitions=None, partitionFunc=portable_hash):

"""

Return an RDD of grouped items.

>>> rdd = sc.parallelize([1, 1, 2, 3, 5, 8])

>>> result = rdd.groupBy(lambda x: x % 2).collect()

>>> sorted([(x, sorted(y)) for (x, y) in result])

[(0, [2, 8]), (1, [1, 1, 3, 5])]

"""

return self.map(lambda x: (f(x), x)).groupByKey(numPartitions, partitionFunc)

@ignore_unicode_prefix

def pipe(self, command, env=None, checkCode=False):

"""

Return an RDD created by piping elements to a forked external process.

>>> sc.parallelize(['1', '2', '', '3']).pipe('cat').collect()

[u'1', u'2', u'', u'3']

:param checkCode: whether or not to check the return value of the shell command.

"""

if env is None:

env = dict()

def func(iterator):

pipe = Popen(

shlex.split(command), env=env, stdin=PIPE, stdout=PIPE)

def pipe_objs(out):

for obj in iterator:

s = str(obj).rstrip('\n') + '\n'

out.write(s.encode('utf-8'))

out.close()

Thread(target=pipe_objs, args=[pipe.stdin]).start()

def check_return_code():

pipe.wait()

if checkCode and pipe.returncode:

raise Exception("Pipe function `%s' exited "

"with error code %d" % (command, pipe.returncode))

else:

for i in range(0):

yield i

return (x.rstrip(b'\n').decode('utf-8') for x in

chain(iter(pipe.stdout.readline, b''), check_return_code()))

return self.mapPartitions(func)

def foreach(self, f):

"""

Applies a function to all elements of this RDD.

>>> def f(x): print(x)

>>> sc.parallelize([1, 2, 3, 4, 5]).foreach(f)

"""

def processPartition(iterator):

for x in iterator:

f(x)

return iter([])

self.mapPartitions(processPartition).count() # Force evaluation

def foreachPartition(self, f):

"""

Applies a function to each partition of this RDD.

>>> def f(iterator):

... for x in iterator:

... print(x)

>>> sc.parallelize([1, 2, 3, 4, 5]).foreachPartition(f)

"""

def func(it):

r = f(it)

try:

return iter(r)

except TypeError:

return iter([])

self.mapPartitions(func).count() # Force evaluation

def collect(self):

"""

Return a list that contains all of the elements in this RDD.

.. note:: This method should only be used if the resulting array is expected

to be small, as all the data is loaded into the driver's memory.

"""

with SCCallSiteSync(self.context) as css:

port = self.ctx._jvm.PythonRDD.collectAndServe(self._jrdd.rdd())

return list(_load_from_socket(port, self._jrdd_deserializer))

def reduce(self, f):

"""

Reduces the elements of this RDD using the specified commutative and

associative binary operator. Currently reduces partitions locally.

>>> from operator import add

>>> sc.parallelize([1, 2, 3, 4, 5]).reduce(add)

15

>>> sc.parallelize((2 for _ in range(10))).map(lambda x: 1).cache().reduce(add)

10

>>> sc.parallelize([]).reduce(add)

Traceback (most recent call last):

...

ValueError: Can not reduce() empty RDD

"""

def func(iterator):

iterator = iter(iterator)

try:

initial = next(iterator)

except StopIteration:

return

yield reduce(f, iterator, initial)

vals = self.mapPartitions(func).collect()

if vals:

return reduce(f, vals)

raise ValueError("Can not reduce() empty RDD")

def treeReduce(self, f, depth=2):

"""

Reduces the elements of this RDD in a multi-level tree pattern.

:param depth: suggested depth of the tree (default: 2)

>>> add = lambda x, y: x + y

>>> rdd = sc.parallelize([-5, -4, -3, -2, -1, 1, 2, 3, 4], 10)

>>> rdd.treeReduce(add)

-5

>>> rdd.treeReduce(add, 1)

-5

>>> rdd.treeReduce(add, 2)

-5

>>> rdd.treeReduce(add, 5)

-5

>>> rdd.treeReduce(add, 10)

-5

"""

if depth < 1:

raise ValueError("Depth cannot be smaller than 1 but got %d." % depth)

zeroValue = None, True # Use the second entry to indicate whether this is a dummy value.

def op(x, y):

if x[1]:

return y

elif y[1]:

return x

else:

return f(x[0], y[0]), False

reduced = self.map(lambda x: (x, False)).treeAggregate(zeroValue, op, op, depth)

if reduced[1]:

raise ValueError("Cannot reduce empty RDD.")

return reduced[0]

def fold(self, zeroValue, op):

"""

Aggregate the elements of each partition, and then the results for all

the partitions, using a given associative function and a neutral "zero value."

The function C{op(t1, t2)} is allowed to modify C{t1} and return it

as its result value to avoid object allocation; however, it should not

modify C{t2}.

This behaves somewhat differently from fold operations implemented

for non-distributed collections in functional languages like Scala.

This fold operation may be applied to partitions individually, and then

fold those results into the final result, rather than apply the fold

to each element sequentially in some defined ordering. For functions

that are not commutative, the result may differ from that of a fold

applied to a non-distributed collection.

>>> from operator import add

>>> sc.parallelize([1, 2, 3, 4, 5]).fold(0, add)

15

"""

def func(iterator):

acc = zeroValue

for obj in iterator:

acc = op(acc, obj)

yield acc

# collecting result of mapPartitions here ensures that the copy of

# zeroValue provided to each partition is unique from the one provided

# to the final reduce call

vals = self.mapPartitions(func).collect()

return reduce(op, vals, zeroValue)

def aggregate(self, zeroValue, seqOp, combOp):

"""

Aggregate the elements of each partition, and then the results for all

the partitions, using a given combine functions and a neutral "zero

value."

The functions C{op(t1, t2)} is allowed to modify C{t1} and return it

as its result value to avoid object allocation; however, it should not

modify C{t2}.

The first function (seqOp) can return a different result type, U, than

the type of this RDD. Thus, we need one operation for merging a T into

an U and one operation for merging two U

>>> seqOp = (lambda x, y: (x[0] + y, x[1] + 1))

>>> combOp = (lambda x, y: (x[0] + y[0], x[1] + y[1]))

>>> sc.parallelize([1, 2, 3, 4]).aggregate((0, 0), seqOp, combOp)

(10, 4)

>>> sc.parallelize([]).aggregate((0, 0), seqOp, combOp)

(0, 0)

"""

def func(iterator):

acc = zeroValue

for obj in iterator:

acc = seqOp(acc, obj)

yield acc

# collecting result of mapPartitions here ensures that the copy of

# zeroValue provided to each partition is unique from the one provided

# to the final reduce call

vals = self.mapPartitions(func).collect()

return reduce(combOp, vals, zeroValue)

def treeAggregate(self, zeroValue, seqOp, combOp, depth=2):

"""

Aggregates the elements of this RDD in a multi-level tree

pattern.

:param depth: suggested depth of the tree (default: 2)

>>> add = lambda x, y: x + y

>>> rdd = sc.parallelize([-5, -4, -3, -2, -1, 1, 2, 3, 4], 10)

>>> rdd.treeAggregate(0, add, add)

-5

>>> rdd.treeAggregate(0, add, add, 1)

-5

>>> rdd.treeAggregate(0, add, add, 2)

-5

>>> rdd.treeAggregate(0, add, add, 5)

-5

>>> rdd.treeAggregate(0, add, add, 10)

-5

"""

if depth < 1:

raise ValueError("Depth cannot be smaller than 1 but got %d." % depth)

if self.getNumPartitions() == 0:

return zeroValue

def aggregatePartition(iterator):

acc = zeroValue

for obj in iterator:

acc = seqOp(acc, obj)

yield acc

partiallyAggregated = self.mapPartitions(aggregatePartition)

numPartitions = partiallyAggregated.getNumPartitions()

scale = max(int(ceil(pow(numPartitions, 1.0 / depth))), 2)

# If creating an extra level doesn't help reduce the wall-clock time, we stop the tree

# aggregation.

while numPartitions > scale + numPartitions / scale:

numPartitions /= scale

curNumPartitions = int(numPartitions)

def mapPartition(i, iterator):

for obj in iterator:

yield (i % curNumPartitions, obj)

partiallyAggregated = partiallyAggregated \

.mapPartitionsWithIndex(mapPartition) \

.reduceByKey(combOp, curNumPartitions) \

.values()

return partiallyAggregated.reduce(combOp)

def max(self, key=None):

"""

Find the maximum item in this RDD.

:param key: A function used to generate key for comparing

>>> rdd = sc.parallelize([1.0, 5.0, 43.0, 10.0])

>>> rdd.max()