| # |
| # 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. |
| # |
| |
| """ |
| Python package for random data generation. |
| """ |
| |
| import sys |
| from functools import wraps |
| from typing import Any, Callable, Optional |
| |
| import numpy as np |
| |
| from pyspark.mllib.common import callMLlibFunc |
| from pyspark.core.context import SparkContext |
| from pyspark.core.rdd import RDD |
| from pyspark.mllib.linalg import Vector |
| |
| |
| __all__ = [ |
| "RandomRDDs", |
| ] |
| |
| |
| def toArray(f: Callable[..., RDD[Vector]]) -> Callable[..., RDD[np.ndarray]]: |
| @wraps(f) |
| def func(sc: SparkContext, *a: Any, **kw: Any) -> RDD[np.ndarray]: |
| rdd = f(sc, *a, **kw) |
| return rdd.map(lambda vec: vec.toArray()) |
| |
| return func |
| |
| |
| class RandomRDDs: |
| """ |
| Generator methods for creating RDDs comprised of i.i.d samples from |
| some distribution. |
| |
| .. versionadded:: 1.1.0 |
| """ |
| |
| @staticmethod |
| def uniformRDD( |
| sc: SparkContext, size: int, numPartitions: Optional[int] = None, seed: Optional[int] = None |
| ) -> RDD[float]: |
| """ |
| Generates an RDD comprised of i.i.d. samples from the |
| uniform distribution U(0.0, 1.0). |
| |
| To transform the distribution in the generated RDD from U(0.0, 1.0) |
| to U(a, b), use |
| ``RandomRDDs.uniformRDD(sc, n, p, seed).map(lambda v: a + (b - a) * v)`` |
| |
| .. versionadded:: 1.1.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| used to create the RDD. |
| size : int |
| Size of the RDD. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of float comprised of i.i.d. samples ~ `U(0.0, 1.0)`. |
| |
| Examples |
| -------- |
| >>> x = RandomRDDs.uniformRDD(sc, 100).collect() |
| >>> len(x) |
| 100 |
| >>> max(x) <= 1.0 and min(x) >= 0.0 |
| True |
| >>> RandomRDDs.uniformRDD(sc, 100, 4).getNumPartitions() |
| 4 |
| >>> parts = RandomRDDs.uniformRDD(sc, 100, seed=4).getNumPartitions() |
| >>> parts == sc.defaultParallelism |
| True |
| """ |
| return callMLlibFunc("uniformRDD", sc._jsc, size, numPartitions, seed) |
| |
| @staticmethod |
| def normalRDD( |
| sc: SparkContext, size: int, numPartitions: Optional[int] = None, seed: Optional[int] = None |
| ) -> RDD[float]: |
| """ |
| Generates an RDD comprised of i.i.d. samples from the standard normal |
| distribution. |
| |
| To transform the distribution in the generated RDD from standard normal |
| to some other normal N(mean, sigma^2), use |
| ``RandomRDDs.normal(sc, n, p, seed).map(lambda v: mean + sigma * v)`` |
| |
| .. versionadded:: 1.1.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| used to create the RDD. |
| size : int |
| Size of the RDD. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of float comprised of i.i.d. samples ~ N(0.0, 1.0). |
| |
| Examples |
| -------- |
| >>> x = RandomRDDs.normalRDD(sc, 1000, seed=1) |
| >>> stats = x.stats() |
| >>> stats.count() |
| 1000 |
| >>> bool(abs(stats.mean() - 0.0) < 0.1) |
| True |
| >>> bool(abs(stats.stdev() - 1.0) < 0.1) |
| True |
| """ |
| return callMLlibFunc("normalRDD", sc._jsc, size, numPartitions, seed) |
| |
| @staticmethod |
| def logNormalRDD( |
| sc: SparkContext, |
| mean: float, |
| std: float, |
| size: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[float]: |
| """ |
| Generates an RDD comprised of i.i.d. samples from the log normal |
| distribution with the input mean and standard distribution. |
| |
| .. versionadded:: 1.3.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| used to create the RDD. |
| mean : float |
| mean for the log Normal distribution |
| std : float |
| std for the log Normal distribution |
| size : int |
| Size of the RDD. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| RDD of float comprised of i.i.d. samples ~ log N(mean, std). |
| |
| Examples |
| -------- |
| >>> from math import sqrt, exp |
| >>> mean = 0.0 |
| >>> std = 1.0 |
| >>> expMean = exp(mean + 0.5 * std * std) |
| >>> expStd = sqrt((exp(std * std) - 1.0) * exp(2.0 * mean + std * std)) |
| >>> x = RandomRDDs.logNormalRDD(sc, mean, std, 1000, seed=2) |
| >>> stats = x.stats() |
| >>> stats.count() |
| 1000 |
| >>> bool(abs(stats.mean() - expMean) < 0.5) |
| True |
| >>> from math import sqrt |
| >>> bool(abs(stats.stdev() - expStd) < 0.5) |
| True |
| """ |
| return callMLlibFunc( |
| "logNormalRDD", sc._jsc, float(mean), float(std), size, numPartitions, seed |
| ) |
| |
| @staticmethod |
| def poissonRDD( |
| sc: SparkContext, |
| mean: float, |
| size: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[float]: |
| """ |
| Generates an RDD comprised of i.i.d. samples from the Poisson |
| distribution with the input mean. |
| |
| .. versionadded:: 1.1.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| mean : float |
| Mean, or lambda, for the Poisson distribution. |
| size : int |
| Size of the RDD. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of float comprised of i.i.d. samples ~ Pois(mean). |
| |
| Examples |
| -------- |
| >>> mean = 100.0 |
| >>> x = RandomRDDs.poissonRDD(sc, mean, 1000, seed=2) |
| >>> stats = x.stats() |
| >>> stats.count() |
| 1000 |
| >>> abs(stats.mean() - mean) < 0.5 |
| True |
| >>> from math import sqrt |
| >>> bool(abs(stats.stdev() - sqrt(mean)) < 0.5) |
| True |
| """ |
| return callMLlibFunc("poissonRDD", sc._jsc, float(mean), size, numPartitions, seed) |
| |
| @staticmethod |
| def exponentialRDD( |
| sc: SparkContext, |
| mean: float, |
| size: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[float]: |
| """ |
| Generates an RDD comprised of i.i.d. samples from the Exponential |
| distribution with the input mean. |
| |
| .. versionadded:: 1.3.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| mean : float |
| Mean, or 1 / lambda, for the Exponential distribution. |
| size : int |
| Size of the RDD. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of float comprised of i.i.d. samples ~ Exp(mean). |
| |
| Examples |
| -------- |
| >>> mean = 2.0 |
| >>> x = RandomRDDs.exponentialRDD(sc, mean, 1000, seed=2) |
| >>> stats = x.stats() |
| >>> stats.count() |
| 1000 |
| >>> abs(stats.mean() - mean) < 0.5 |
| True |
| >>> from math import sqrt |
| >>> bool(abs(stats.stdev() - sqrt(mean)) < 0.5) |
| True |
| """ |
| return callMLlibFunc("exponentialRDD", sc._jsc, float(mean), size, numPartitions, seed) |
| |
| @staticmethod |
| def gammaRDD( |
| sc: SparkContext, |
| shape: float, |
| scale: float, |
| size: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[float]: |
| """ |
| Generates an RDD comprised of i.i.d. samples from the Gamma |
| distribution with the input shape and scale. |
| |
| .. versionadded:: 1.3.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| shape : float |
| shape (> 0) parameter for the Gamma distribution |
| scale : float |
| scale (> 0) parameter for the Gamma distribution |
| size : int |
| Size of the RDD. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of float comprised of i.i.d. samples ~ Gamma(shape, scale). |
| |
| Examples |
| -------- |
| >>> from math import sqrt |
| >>> shape = 1.0 |
| >>> scale = 2.0 |
| >>> expMean = shape * scale |
| >>> expStd = sqrt(shape * scale * scale) |
| >>> x = RandomRDDs.gammaRDD(sc, shape, scale, 1000, seed=2) |
| >>> stats = x.stats() |
| >>> stats.count() |
| 1000 |
| >>> bool(abs(stats.mean() - expMean) < 0.5) |
| True |
| >>> bool(abs(stats.stdev() - expStd) < 0.5) |
| True |
| """ |
| return callMLlibFunc( |
| "gammaRDD", sc._jsc, float(shape), float(scale), size, numPartitions, seed |
| ) |
| |
| @staticmethod |
| @toArray |
| def uniformVectorRDD( |
| sc: SparkContext, |
| numRows: int, |
| numCols: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[Vector]: |
| """ |
| Generates an RDD comprised of vectors containing i.i.d. samples drawn |
| from the uniform distribution U(0.0, 1.0). |
| |
| .. versionadded:: 1.1.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| numRows : int |
| Number of Vectors in the RDD. |
| numCols : int |
| Number of elements in each Vector. |
| numPartitions : int, optional |
| Number of partitions in the RDD. |
| seed : int, optional |
| Seed for the RNG that generates the seed for the generator in each partition. |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of Vector with vectors containing i.i.d samples ~ `U(0.0, 1.0)`. |
| |
| Examples |
| -------- |
| >>> import numpy as np |
| >>> mat = np.matrix(RandomRDDs.uniformVectorRDD(sc, 10, 10).collect()) |
| >>> mat.shape |
| (10, 10) |
| >>> bool(mat.max() <= 1.0 and mat.min() >= 0.0) |
| True |
| >>> RandomRDDs.uniformVectorRDD(sc, 10, 10, 4).getNumPartitions() |
| 4 |
| """ |
| return callMLlibFunc("uniformVectorRDD", sc._jsc, numRows, numCols, numPartitions, seed) |
| |
| @staticmethod |
| @toArray |
| def normalVectorRDD( |
| sc: SparkContext, |
| numRows: int, |
| numCols: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[Vector]: |
| """ |
| Generates an RDD comprised of vectors containing i.i.d. samples drawn |
| from the standard normal distribution. |
| |
| .. versionadded:: 1.1.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| numRows : int |
| Number of Vectors in the RDD. |
| numCols : int |
| Number of elements in each Vector. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of Vector with vectors containing i.i.d. samples ~ `N(0.0, 1.0)`. |
| |
| Examples |
| -------- |
| >>> import numpy as np |
| >>> mat = np.matrix(RandomRDDs.normalVectorRDD(sc, 100, 100, seed=1).collect()) |
| >>> mat.shape |
| (100, 100) |
| >>> bool(abs(mat.mean() - 0.0) < 0.1) |
| True |
| >>> bool(abs(mat.std() - 1.0) < 0.1) |
| True |
| """ |
| return callMLlibFunc("normalVectorRDD", sc._jsc, numRows, numCols, numPartitions, seed) |
| |
| @staticmethod |
| @toArray |
| def logNormalVectorRDD( |
| sc: SparkContext, |
| mean: float, |
| std: float, |
| numRows: int, |
| numCols: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[Vector]: |
| """ |
| Generates an RDD comprised of vectors containing i.i.d. samples drawn |
| from the log normal distribution. |
| |
| .. versionadded:: 1.3.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| mean : float |
| Mean of the log normal distribution |
| std : float |
| Standard Deviation of the log normal distribution |
| numRows : int |
| Number of Vectors in the RDD. |
| numCols : int |
| Number of elements in each Vector. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of Vector with vectors containing i.i.d. samples ~ log `N(mean, std)`. |
| |
| Examples |
| -------- |
| >>> import numpy as np |
| >>> from math import sqrt, exp |
| >>> mean = 0.0 |
| >>> std = 1.0 |
| >>> expMean = exp(mean + 0.5 * std * std) |
| >>> expStd = sqrt((exp(std * std) - 1.0) * exp(2.0 * mean + std * std)) |
| >>> m = RandomRDDs.logNormalVectorRDD(sc, mean, std, 100, 100, seed=1).collect() |
| >>> mat = np.matrix(m) |
| >>> mat.shape |
| (100, 100) |
| >>> bool(abs(mat.mean() - expMean) < 0.1) |
| True |
| >>> bool(abs(mat.std() - expStd) < 0.1) |
| True |
| """ |
| return callMLlibFunc( |
| "logNormalVectorRDD", |
| sc._jsc, |
| float(mean), |
| float(std), |
| numRows, |
| numCols, |
| numPartitions, |
| seed, |
| ) |
| |
| @staticmethod |
| @toArray |
| def poissonVectorRDD( |
| sc: SparkContext, |
| mean: float, |
| numRows: int, |
| numCols: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[Vector]: |
| """ |
| Generates an RDD comprised of vectors containing i.i.d. samples drawn |
| from the Poisson distribution with the input mean. |
| |
| .. versionadded:: 1.1.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| mean : float |
| Mean, or lambda, for the Poisson distribution. |
| numRows : float |
| Number of Vectors in the RDD. |
| numCols : int |
| Number of elements in each Vector. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`) |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of Vector with vectors containing i.i.d. samples ~ Pois(mean). |
| |
| Examples |
| -------- |
| >>> import numpy as np |
| >>> mean = 100.0 |
| >>> rdd = RandomRDDs.poissonVectorRDD(sc, mean, 100, 100, seed=1) |
| >>> mat = np.asmatrix(rdd.collect()) |
| >>> mat.shape |
| (100, 100) |
| >>> bool(abs(mat.mean() - mean) < 0.5) |
| True |
| >>> from math import sqrt |
| >>> bool(abs(mat.std() - sqrt(mean)) < 0.5) |
| True |
| """ |
| return callMLlibFunc( |
| "poissonVectorRDD", sc._jsc, float(mean), numRows, numCols, numPartitions, seed |
| ) |
| |
| @staticmethod |
| @toArray |
| def exponentialVectorRDD( |
| sc: SparkContext, |
| mean: float, |
| numRows: int, |
| numCols: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[Vector]: |
| """ |
| Generates an RDD comprised of vectors containing i.i.d. samples drawn |
| from the Exponential distribution with the input mean. |
| |
| .. versionadded:: 1.3.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| mean : float |
| Mean, or 1 / lambda, for the Exponential distribution. |
| numRows : int |
| Number of Vectors in the RDD. |
| numCols : int |
| Number of elements in each Vector. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`) |
| seed : int, optional |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of Vector with vectors containing i.i.d. samples ~ Exp(mean). |
| |
| Examples |
| -------- |
| >>> import numpy as np |
| >>> mean = 0.5 |
| >>> rdd = RandomRDDs.exponentialVectorRDD(sc, mean, 100, 100, seed=1) |
| >>> mat = np.asmatrix(rdd.collect()) |
| >>> mat.shape |
| (100, 100) |
| >>> bool(abs(mat.mean() - mean) < 0.5) |
| True |
| >>> from math import sqrt |
| >>> bool(abs(mat.std() - sqrt(mean)) < 0.5) |
| True |
| """ |
| return callMLlibFunc( |
| "exponentialVectorRDD", sc._jsc, float(mean), numRows, numCols, numPartitions, seed |
| ) |
| |
| @staticmethod |
| @toArray |
| def gammaVectorRDD( |
| sc: SparkContext, |
| shape: float, |
| scale: float, |
| numRows: int, |
| numCols: int, |
| numPartitions: Optional[int] = None, |
| seed: Optional[int] = None, |
| ) -> RDD[Vector]: |
| """ |
| Generates an RDD comprised of vectors containing i.i.d. samples drawn |
| from the Gamma distribution. |
| |
| .. versionadded:: 1.3.0 |
| |
| Parameters |
| ---------- |
| sc : :py:class:`pyspark.SparkContext` |
| SparkContext used to create the RDD. |
| shape : float |
| Shape (> 0) of the Gamma distribution |
| scale : float |
| Scale (> 0) of the Gamma distribution |
| numRows : int |
| Number of Vectors in the RDD. |
| numCols : int |
| Number of elements in each Vector. |
| numPartitions : int, optional |
| Number of partitions in the RDD (default: `sc.defaultParallelism`). |
| seed : int, optional, |
| Random seed (default: a random long integer). |
| |
| Returns |
| ------- |
| :py:class:`pyspark.RDD` |
| RDD of Vector with vectors containing i.i.d. samples ~ Gamma(shape, scale). |
| |
| Examples |
| -------- |
| >>> import numpy as np |
| >>> from math import sqrt |
| >>> shape = 1.0 |
| >>> scale = 2.0 |
| >>> expMean = shape * scale |
| >>> expStd = sqrt(shape * scale * scale) |
| >>> mat = np.matrix(RandomRDDs.gammaVectorRDD(sc, shape, scale, 100, 100, seed=1).collect()) |
| >>> mat.shape |
| (100, 100) |
| >>> bool(abs(mat.mean() - expMean) < 0.1) |
| True |
| >>> bool(abs(mat.std() - expStd) < 0.1) |
| True |
| """ |
| return callMLlibFunc( |
| "gammaVectorRDD", |
| sc._jsc, |
| float(shape), |
| float(scale), |
| numRows, |
| numCols, |
| numPartitions, |
| seed, |
| ) |
| |
| |
| def _test() -> None: |
| import doctest |
| from pyspark.sql import SparkSession |
| |
| globs = globals().copy() |
| # The small batch size here ensures that we see multiple batches, |
| # even in these small test examples: |
| spark = SparkSession.builder.master("local[2]").appName("mllib.random tests").getOrCreate() |
| globs["sc"] = spark.sparkContext |
| (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) |
| spark.stop() |
| if failure_count: |
| sys.exit(-1) |
| |
| |
| if __name__ == "__main__": |
| _test() |
