nxcals.api.extraction.data.builders.DataFrame.sampleBy

DataFrame.sampleBy(col: ColumnOrName, fractions: Dict[Any, float], seed: Optional[int] = None) → DataFrame

Returns a stratified sample without replacement based on the fraction given on each stratum.

New in version 1.5.0.

Changed in version 3.4.0: Supports Spark Connect.

Parameters:

  • col (Column or str) –

    column that defines strata

    Changed in version 3.0.0: Added sampling by a column of Column

  • fractions (dict) – sampling fraction for each stratum. If a stratum is not specified, we treat its fraction as zero.

  • seed (int, optional) – random seed

Return type:

a new DataFrame that represents the stratified sample

Examples

>>> from pyspark.sql.functions import col
>>> dataset = spark.range(0, 100).select((col("id") % 3).alias("key"))
>>> sampled = dataset.sampleBy("key", fractions={0: 0.1, 1: 0.2}, seed=0)
>>> sampled.groupBy("key").count().orderBy("key").show()
+---+-----+
|key|count|
+---+-----+
|  0|    3|
|  1|    6|
+---+-----+
>>> dataset.sampleBy(col("key"), fractions={2: 1.0}, seed=0).count()
33