Skip to content

Databricks

For Spark and Databricks users. It teaches SparkGenerator, which builds a users table natively in Spark.

Prerequisite and run

Use a Databricks runtime, or a machine with PySpark and Java available. rand-engine itself does not install PySpark. Scaffold the file, then import it into Databricks or run it where that Spark environment is already available:

Get the file:

pip install rand-engine
rand-engine get-started databricks
python databricks-example.py

It writes databricks-example.py to the current directory. --out DIR writes it to DIR instead, created when missing. An existing file is never overwritten. The native path builds a 1,000-row Spark DataFrame, then shows its first rows and schema; the final cell also shows the 100-row pandas-to-Spark route.

# Databricks notebook source
# Spark: a users table generated natively with SparkGenerator. Runs in Databricks, or under plain python with PySpark installed.

# COMMAND ----------
# Get the session. In Databricks it returns the notebook's session; anywhere else it returns the one already running.

from pyspark.sql import SparkSession, functions as F

from rand_engine import SparkGenerator

spark = SparkSession.builder.getOrCreate()

# COMMAND ----------
# Describe the users with supported methods only.

spec = {
  "user_id": {"method": "int_zfilled", "kwargs": {"length": 8}},
  "age": {"method": "integers", "kwargs": {"min": 18, "max": 80}},
  "score": {"method": "floats", "kwargs": {"min": 0, "max": 100, "decimals": 2}},
  "is_active": {"method": "booleans", "kwargs": {"true_prob": 0.7}},
  "plan": {"method": "distincts", "kwargs": {"distincts": ["free", "pro", "team"]}},
  "signup": {"method": "dates", "kwargs": {"start": "2020-01-01", "end": "2023-12-31", "date_format": "%Y-%m-%d"}},
  "last_seen": {"method": "unix_timestamps", "kwargs": {"start": "2024-01-01", "end": "2024-12-31"}},
}

# COMMAND ----------
# Generate 1 000 rows natively into df, then look at them.

df = SparkGenerator(spark, F, spec).size(1_000).get_df()
df.show(5)
df.printSchema()

# COMMAND ----------
# Note: pk, fk, the distributions (exponential, lognormal, poisson, zipf), constant, null_rate and
# anomaly_rate are DataGenerator-only; SparkGenerator refuses them. For a small volume, generate with
# DataGenerator(...).get_df() and wrap the pandas frame in spark.createDataFrame(pdf).
# The cell below does that for 100 rows; df above stays the native one.
# A text column's nulls are pd.NA, so Spark's Arrow option is enabled before createDataFrame.
# On a self-managed Spark on Java 17, start the JVM with -Dio.netty.tryReflectionSetAccessible=true.

from rand_engine import DataGenerator

pdf = DataGenerator(spec, seed=7).size(100).get_df()
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
small = spark.createDataFrame(pdf)
small.show(5)

The file is in Databricks notebook source format: its first line is # Databricks notebook source, and its cells are separated by # COMMAND ----------. Import it into a Databricks workspace as a notebook, or run it with python where PySpark and Java are available.

DataGenerator-only features

SparkGenerator accepts stateless pk and fk through the same key definitions as DataGenerator; an explicit seed applies to those key columns only. It still refuses the distributions (exponential, lognormal, poisson, zipf), constant, null_rate and anomaly_rate. For those, or for a small volume, generate with DataGenerator and wrap the pandas frame with spark.createDataFrame. The file's last cell does that for 100 rows. A text column's nulls are pd.NA, so enable Spark's Arrow option before createDataFrame: spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true"). On a self-managed Spark on Java 17, also start the JVM with -Dio.netty.tryReflectionSetAccessible=true.

What it shows

  • Session: the Spark session, the one the notebook already runs or the one created where there is none.
  • Describe: the users spec, with supported methods only.
  • Generate: 1 000 rows built natively into df, then shown with show and printSchema.
  • Small volume: the DataGenerator path and spark.createDataFrame (Arrow option on, see above), beside the native df.

Next steps