[SPARK-8777] [SQL] Add random data generator test utilities to Spark SQL #7176
Prev
Previous commit
Roll back ScalaCheck usage
- Loading branch information
commit f71634d73470189cfe45a89d2a69ea9c5ffa9e29
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,17 +17,54 @@ | |
|
|
||
| package org.apache.spark.sql | ||
|
|
||
| import java.lang.Double.longBitsToDouble | ||
| import java.lang.Float.intBitsToFloat | ||
| import java.math.MathContext | ||
|
|
||
| import scala.util.Random | ||
|
|
||
| import org.apache.spark.sql.types._ | ||
|
|
||
| /** | ||
| * Random data generators for Spark SQL DataTypes. These generators do not generate uniformly random | ||
| * values; instead, they're biased to return "interesting" values (such as maximum / minimum values) | ||
| * with higher probability. | ||
| */ | ||
| object RandomDataGenerator { | ||
|
|
||
| /** | ||
| * The conditional probability of a non-null value being drawn from a set of "interesting" values | ||
| * instead of being chosen uniformly at random. | ||
| */ | ||
| private val PROBABILITY_OF_INTERESTING_VALUE: Float = 0.5f | ||
|
|
||
| /** | ||
| * The probability of the generated value being null | ||
| */ | ||
| private val PROBABILITY_OF_NULL: Float = 0.1f | ||
|
|
||
| private val MAX_STR_LEN: Int = 1024 | ||
| private val MAX_ARR_SIZE: Int = 128 | ||
| private val MAX_MAP_SIZE: Int = 128 | ||
|
|
||
| /** | ||
| * Helper function for constructing a biased random number generator which returns "interesting" | ||
| * values with a higher probability. | ||
| */ | ||
| private def randomNumeric[T]( | ||
| rand: Random, | ||
| uniformRand: Random => T, | ||
| interestingValues: Seq[T]): Some[() => T] = { | ||
| val f = () => { | ||
| if (rand.nextFloat() <= PROBABILITY_OF_INTERESTING_VALUE) { | ||
| interestingValues(rand.nextInt(interestingValues.length)) | ||
| } else { | ||
| uniformRand(rand) | ||
| } | ||
| } | ||
| Some(f) | ||
| } | ||
|
|
||
| /** | ||
| * Returns a function which generates random values for the given [[DataType]], or `None` if no | ||
| * random data generator is defined for that data type. The generated values will use an external | ||
|
|
@@ -37,58 +74,85 @@ object RandomDataGenerator { | |
| * | ||
| * @param dataType the type to generate values for | ||
| * @param nullable whether null values should be generated | ||
| * @param seed an optional seed for the random number generator | ||
| * @return a function which can be called to generate random values. | ||
| */ | ||
| def forType( | ||
| dataType: DataType, | ||
| nullable: Boolean = true, | ||
| seed: Option[Long] = None): Option[() => Any] = { | ||
| val rand = new Random() | ||
| seed.foreach(rand.setSeed) | ||
|
|
||
| val valueGenerator: Option[() => Any] = dataType match { | ||
| case StringType => Some(() => rand.nextString(rand.nextInt(MAX_STR_LEN))) | ||
| case BinaryType => Some(() => { | ||
| val arr = new Array[Byte](rand.nextInt(MAX_STR_LEN)) | ||
| rand.nextBytes(arr) | ||
| arr | ||
| }) | ||
| case BooleanType => Some(() => rand.nextBoolean()) | ||
| case DateType => Some(() => new java.sql.Date(rand.nextInt())) | ||
| case TimestampType => Some(() => new java.sql.Timestamp(rand.nextLong())) | ||
| case DecimalType.Unlimited => Some( | ||
| () => BigDecimal.apply(rand.nextLong, rand.nextInt, MathContext.UNLIMITED)) | ||
| case DoubleType => randomNumeric[Double]( | ||
| rand, r => longBitsToDouble(r.nextLong()), Seq(Double.MinValue, Double.MinPositiveValue, | ||
|
Contributor
There was a problem hiding this comment. Are we using This implies more chances than expected to generate NaNs. But it's probably OK?
Contributor
Author
There was a problem hiding this comment. The goal here was to produce doubles that were uniformly distributed over the range of possible double values (`rand.nextDouble() just returns doubles in the range 0.0 to 1.0). Empirically, the number of NaNs produced by this seems to be quite small and a back-of-the-envelope calculation seems to back this up; I think that ((0x7fffffffffffffff - 0x7ff0000000000001) + (0xffffffffffffffff - 0xfff0000000000001)) / 2^64 works out to be a roughly 0.05% chance of producing a NaN through this method (see https://www.wolframalpha.com/input/?i=%28%280x7fffffffffffffff+-+0x7ff0000000000001%29+%2B+%280xffffffffffffffff+-+0xfff0000000000001%29%29+%2F+2%5E64). I think this is small enough to ignore for our purposes, but we can revisit later if it's a problem. |
||
| Double.MaxValue, Double.PositiveInfinity, Double.NegativeInfinity, Double.NaN, 0.0)) | ||
| case FloatType => randomNumeric[Float]( | ||
| rand, r => intBitsToFloat(r.nextInt()), Seq(Float.MinValue, Float.MinPositiveValue, | ||
|
|
||
| Float.MaxValue, Float.PositiveInfinity, Float.NegativeInfinity, Float.NaN, 0.0f)) | ||
| case ByteType => randomNumeric[Byte]( | ||
| rand, _.nextInt().toByte, Seq(Byte.MinValue, Byte.MaxValue, 0.toByte)) | ||
| case IntegerType => randomNumeric[Int]( | ||
| rand, _.nextInt(), Seq(Int.MinValue, Int.MaxValue, 0)) | ||
| case LongType => randomNumeric[Long]( | ||
| rand, _.nextLong(), Seq(Long.MinValue, Long.MaxValue, 0L)) | ||
| case ShortType => randomNumeric[Short]( | ||
| rand, _.nextInt().toShort, Seq(Short.MinValue, Short.MaxValue, 0.toShort)) | ||
| case NullType => Some(() => null) | ||
| case ArrayType(elementType, containsNull) => { | ||
| forType(elementType, nullable = containsNull, seed = Some(rand.nextLong())).map { | ||
| elementGenerator => () => Array.fill(rand.nextInt(MAX_ARR_SIZE))(elementGenerator()) | ||
| } | ||
| } | ||
| case MapType(keyType, valueType, valueContainsNull) => { | ||
| for ( | ||
| keyGenerator <- forType(keyType, nullable = false, seed = Some(rand.nextLong())); | ||
| valueGenerator <- | ||
| forType(valueType, nullable = valueContainsNull, seed = Some(rand.nextLong())) | ||
| ) yield { | ||
| () => { | ||
| Seq.fill(rand.nextInt(MAX_MAP_SIZE))((keyGenerator(), valueGenerator())).toMap | ||
| } | ||
| } | ||
| } | ||
| case StructType(fields) => { | ||
| val maybeFieldGenerators: Seq[Option[() => Any]] = fields.map { field => | ||
| forType(field.dataType, nullable = field.nullable, seed = Some(rand.nextLong())) | ||
| } | ||
| if (maybeFieldGenerators.forall(_.isDefined)) { | ||
| val fieldGenerators: Seq[() => Any] = maybeFieldGenerators.map(_.get) | ||
| Some(() => Row.fromSeq(fieldGenerators.map(_.apply()))) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| case unsupportedType => None | ||
| } | ||
| // Handle nullability by wrapping the non-null value generator: | ||
| valueGenerator.map { valueGenerator => | ||
| if (nullable) { | ||
| () => { | ||
| if (rand.nextFloat() <= PROBABILITY_OF_NULL) { | ||
| null | ||
| } else { | ||
| valueGenerator() | ||
| } | ||
| } | ||
| } else { | ||
| valueGenerator | ||
| } | ||
| } | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does
nextStringcover unicode characters?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes.