Spark, Delta Lake, and Databricks utility library for Python.
The patterns you keep re-implementing across PySpark jobs — Delta table management, streaming foreachBatch wiring, schema migrations, data quality checks — packaged as a library. Runs on Databricks and locally with open-source Spark + Delta Lake.
pip install databricks4pyRequirements: Python >= 3.10, Java 17+ (for PySpark). pyspark.dbutils only available on Databricks Runtime.
Wraps the delta-spark API into a single object that handles table creation, reads, writes, metadata, and schema validation.
from databricks4py.io import DeltaTable, GeneratedColumn
# Create a table with generated columns and partitioning
table = DeltaTable(
table_name="catalog.schema.events",
schema={"id": "int", "name": "string", "event_ts": "timestamp", "event_date": "date"},
location="/data/events",
partition_by="event_date",
generated_columns=[GeneratedColumn("event_date", "DATE", "CAST(event_ts AS DATE)")],
)
df = table.dataframe() # Read
table.write(df, mode="append") # Write (schema_check=True by default)
table.detail() # Metadata DataFrame
table.partition_columns() # ["event_date"]
table.size_in_bytes() # Physical size in bytesConvenience wrappers:
from databricks4py.io import DeltaTableAppender, DeltaTableOverwriter
appender = DeltaTableAppender("target", schema=my_schema)
appender.append(df)
overwriter = DeltaTableOverwriter("target", schema=my_schema)
overwriter.overwrite(df)Table maintenance:
from databricks4py.io.delta import optimize_table, vacuum_table
optimize_table("catalog.schema.events", z_order_by=["id"])
vacuum_table("catalog.schema.events", retention_hours=168)Fluent builder for Delta MERGE INTO. Chain conditions, execute, get back row counts.
from databricks4py.io import MergeBuilder
result = (
MergeBuilder("catalog.schema.target", source_df, spark)
.on("id", "date") # Join keys
.when_matched_update(["name", "value"]) # Update specific columns
.when_not_matched_insert() # Insert new rows
.when_not_matched_by_source_delete() # Remove stale rows
.execute()
)
print(f"Inserted: {result.rows_inserted}, Updated: {result.rows_updated}")Custom join conditions:
result = (
MergeBuilder("target", source_df, spark)
.on_condition("target.id = source.id AND target.region = source.region")
.when_matched_update()
.when_not_matched_insert()
.execute()
)SCD Type 2 is built into DeltaTable:
table.scd_type2(source_df, keys=["id"])Subclass StreamingTableReader, implement process_batch, and the base class handles the foreachBatch wiring, empty-batch skipping, row filtering, metrics, and DLQ routing.
from databricks4py.io import StreamingTableReader, StreamingTriggerOptions
class EventProcessor(StreamingTableReader):
def process_batch(self, df, batch_id):
clean = df.where("status IS NOT NULL")
clean.write.format("delta").mode("append").saveAsTable("output")
reader = EventProcessor(
source_table="catalog.schema.raw_events",
trigger=StreamingTriggerOptions.PROCESSING_TIME_1M,
checkpoint_location="/checkpoints/events",
dead_letter_table="catalog.schema.dlq", # Failed batches go here
)
query = reader.start()
reader.stop(timeout_seconds=30) # Graceful shutdownIf dead_letter_table is set and process_batch raises, the batch DataFrame is written to the DLQ table with _dlq_error_message, _dlq_error_timestamp, and _dlq_batch_id columns appended. The stream keeps running.
Checkpoint management:
from databricks4py.io import CheckpointManager
mgr = CheckpointManager(base_path="/checkpoints")
reader = EventProcessor(
source_table="input",
checkpoint_manager=mgr, # Auto-generates checkpoint path
)Ordered migration runner, same idea as Flyway or Alembic but for Delta tables. Steps are Python callables, history is tracked in a Delta table, and each version runs exactly once.
from databricks4py.migrations import MigrationRunner, MigrationStep
def add_audit_columns(spark):
spark.sql("ALTER TABLE catalog.schema.events ADD COLUMNS (created_at TIMESTAMP)")
def backfill_defaults(spark):
spark.sql("UPDATE catalog.schema.events SET created_at = current_timestamp()")
runner = MigrationRunner(history_table="catalog.schema._migrations")
runner.register(
MigrationStep(version="V001", description="Add audit columns", up=add_audit_columns),
MigrationStep(version="V002", description="Backfill defaults", up=backfill_defaults),
)
# Check what needs to run
pending = runner.pending() # [MigrationStep(V002, ...)]
# Execute (idempotent — already-applied steps are skipped)
result = runner.run()
print(result.applied) # ["V002"]
print(result.skipped) # ["V001"]
# Dry run (logs what would run, no side effects)
result = runner.run(dry_run=True)Pre/post validation per step:
MigrationStep(
version="V003",
description="Rename column",
up=lambda spark: spark.sql("ALTER TABLE t RENAME COLUMN old TO new"),
pre_validate=lambda spark: spark.catalog.tableExists("t"),
post_validate=lambda spark: "new" in spark.read.table("t").columns,
)Queue up ALTER TABLE operations and apply them in one go.
from databricks4py.migrations import TableAlter
TableAlter("catalog.schema.events") \
.add_column("region", "STRING", comment="ISO-3166 region code") \
.set_property("delta.enableChangeDataFeed", "true") \
.apply()Rename and drop require Delta column mapping:
TableAlter("catalog.schema.events") \
.set_property("delta.columnMapping.mode", "name") \
.set_property("delta.minReaderVersion", "2") \
.set_property("delta.minWriterVersion", "5") \
.apply()
TableAlter("catalog.schema.events") \
.rename_column("old_name", "new_name") \
.drop_column("deprecated_col") \
.apply()Compare two schemas (or a live table vs. an incoming DataFrame) and get back a list of column-level changes with severity.
from databricks4py.migrations import SchemaDiff
diff = SchemaDiff.from_tables("catalog.schema.events", new_df)
for change in diff.changes():
print(f"{change.column}: {change.change_type} [{change.severity}]")
if diff.has_breaking_changes():
raise RuntimeError(diff.summary())Check that a Delta table matches expected columns, partitions, and location before or after a migration.
from databricks4py.migrations import TableValidator
validator = TableValidator(
table_name="catalog.schema.events",
expected_columns=["id", "name", "event_date"],
expected_partition_columns=["event_date"],
expected_location_contains="/data/events",
)
result = validator.validate()
if not result.is_valid:
print(result.errors) # ["Missing required columns: ['event_date']"]
print(result.warnings) # ["Unexpected extra columns: ['debug_flag']"]
result.raise_if_invalid("catalog.schema.events") # Raises MigrationErrorRow-level expectations you can run individually or bundle into a gate that raises, warns, or quarantines bad rows.
from databricks4py.quality.expectations import NotNull, InRange, Unique, RowCount, MatchesRegex
from databricks4py.quality.gate import QualityGate
# Individual expectations
result = NotNull("id", "name").validate(df)
result = InRange("score", min_val=0, max_val=100).validate(df)
result = Unique("id").validate(df)
result = RowCount(min_count=1, max_count=1_000_000).validate(df)
result = MatchesRegex("email", r"^.+@.+\..+$").validate(df)
# result.passed, result.failing_rows, result.total_rows
# Quality gate — enforce multiple expectations
gate = QualityGate(
NotNull("id"),
InRange("score", min_val=0, max_val=100),
on_fail="raise", # or "warn" or "quarantine"
)
clean_df = gate.enforce(df) # Raises QualityError if checks failQuarantine mode splits bad rows and routes them to a handler:
gate = QualityGate(
NotNull("id"),
on_fail="quarantine",
quarantine_handler=lambda bad_df: bad_df.write.saveAsTable("quarantine_table"),
)
clean_df = gate.enforce(df) # Returns only clean rowsChain DataFrame transformations. Each filter is a callable that takes and returns a DataFrame.
from databricks4py.filters import FilterPipeline, DropDuplicates, WhereFilter, ColumnFilter
pipeline = FilterPipeline([
DropDuplicates(subset=["id"]),
WhereFilter("status = 'active'"),
ColumnFilter(columns=["id", "name", "status"]),
])
clean_df = pipeline(raw_df)Base class for Databricks jobs. Handles SparkSession init, config application, lifecycle metrics (job_start/job_complete/job_failed), quality gates, and optional retry.
from databricks4py import Workflow
from databricks4py.quality.expectations import NotNull
from databricks4py.quality.gate import QualityGate
class MyETL(Workflow):
def run(self):
source = self.spark.read.table("raw_events")
# Quality check with metrics emission
gate = QualityGate(NotNull("id", "event_ts"), on_fail="raise")
clean = self.quality_check(source, gate, table_name="raw_events")
clean.write.format("delta").mode("append").saveAsTable("clean_events")
self.emit_metric("write_complete", row_count=clean.count())
# With config and metrics
from databricks4py.config import JobConfig
from databricks4py.metrics import DeltaMetricsSink
config = JobConfig(tables={"source": "raw_events"}, spark_configs={"spark.sql.shuffle.partitions": "8"})
sink = DeltaMetricsSink("catalog.schema.job_metrics")
MyETL(config=config, metrics=sink).execute()Environment auto-detected from Databricks widgets or ENV/ENVIRONMENT env vars, defaults to DEV.
from databricks4py.config import JobConfig, Environment
config = JobConfig(
tables={"events": "catalog.bronze.events", "users": "catalog.silver.users"},
secret_scope="my-scope",
spark_configs={"spark.sql.shuffle.partitions": "8"},
)
config.env # Environment.DEV (auto-detected)
config.table("events") # "catalog.bronze.events"Unity Catalog — environment-aware catalog resolution:
from databricks4py.config import UnityConfig
config = UnityConfig(catalog_prefix="myapp", schemas=["bronze", "silver"])
config.catalog # "myapp_dev" (or myapp_prod in production)
config.table("bronze.events") # "myapp_dev.bronze.events"Buffer events and write them to a Delta table, log them as JSON, or both.
from databricks4py.metrics import DeltaMetricsSink, LoggingMetricsSink, CompositeMetricsSink, MetricEvent
# Write metrics to a Delta table
delta_sink = DeltaMetricsSink("catalog.schema.metrics", buffer_size=50)
# Log metrics as JSON
log_sink = LoggingMetricsSink()
# Fan out to multiple sinks
sink = CompositeMetricsSink(delta_sink, log_sink)
sink.emit(MetricEvent(job_name="etl", event_type="batch_complete", timestamp=now, row_count=1000))
sink.flush()Structured batch logging — JSON log records per batch with correlation IDs, queryable in any log aggregation system.
from databricks4py.observability import BatchContext, BatchLogger
logger = BatchLogger(extra_fields={"pipeline": "events", "env": "prod"})
# Inside your StreamingTableReader.process_batch:
ctx = BatchContext.create(batch_id=batch_id, source_table="catalog.schema.events")
logger.batch_start(ctx)
# ... process ...
logger.batch_complete(ctx, row_count=df.count(), duration_ms=ctx.elapsed_ms())
# On error:
logger.batch_error(ctx, error=str(exc))Each log line is single-line JSON: {"event": "batch_complete", "batch_id": 42, "correlation_id": "a1b2c3d4e5f6", "row_count": 1000, ...}
Query progress listener — wraps PySpark 3.4+ StreamingQueryListener to collect progress snapshots and route them to a MetricsSink.
from databricks4py.observability import QueryProgressObserver
observer = QueryProgressObserver(metrics_sink=my_sink, query_name_filter="events_processor")
observer.attach()
# After the stream runs:
latest = observer.latest_progress()
print(f"Batch {latest.batch_id}: {latest.processed_rows_per_second} rows/sec")
history = observer.history(limit=10)
observer.detach()Health checks — poll a streaming query for stuck detection, slow batches, and low throughput.
from databricks4py.observability import StreamingHealthCheck, HealthStatus
check = StreamingHealthCheck(
query,
max_batch_duration_ms=60_000, # DEGRADED if batch > 60s
min_processing_rate=100.0, # DEGRADED if < 100 rows/sec
stale_timeout_seconds=300, # UNHEALTHY if no progress for 5min
)
result = check.evaluate()
if result.status != HealthStatus.HEALTHY:
print(result.summary())from databricks4py.retry import retry, RetryConfig
@retry(RetryConfig(max_attempts=5, base_delay_seconds=2.0, backoff_factor=3.0))
def fetch_from_api():
return requests.get(url).json()Session-scoped SparkSession (one JVM per test run), function-scoped cleanup, and helpers for building test data.
# conftest.py — register fixtures
from databricks4py.testing.fixtures import * # noqa: F401,F403DataFrameBuilder — fluent test data construction:
def test_my_transform(spark_session_function):
df = (
DataFrameBuilder(spark_session_function)
.with_columns({"id": "int", "name": "string", "score": "int"})
.with_rows((1, "Alice", 95), (2, "Bob", 80))
.build()
)
assert df.count() == 2TempDeltaTable — ephemeral Delta tables for test isolation:
def test_merge(spark_session_function, tmp_path):
with TempDeltaTable(spark_session_function, schema={"id": "int"}, data=[(1,), (2,)]) as table:
assert table.dataframe().count() == 2
# Table is auto-dropped after the context exitsAssertions:
from databricks4py.testing.assertions import assert_frame_equal, assert_schema_equal
assert_frame_equal(actual_df, expected_df, check_order=False)
assert_schema_equal(actual_df.schema, expected_schema, check_nullable=False)Mock Databricks utilities:
from databricks4py.testing.mocks import MockDBUtils, MockDBUtilsModule
mock = MockDBUtils()
mock.secrets.put("scope", "key", "secret-value")
assert mock.secrets.get("scope", "key") == "secret-value"src/databricks4py/
├── __init__.py # Top-level exports
├── spark_session.py # get_active(), active_fallback(), get_or_create_local_session()
├── catalog.py # CatalogSchema for schema-qualified table access
├── logging.py # configure_logging(), get_logger()
├── secrets.py # SecretFetcher with injectable dbutils
├── retry.py # retry() decorator with exponential backoff
├── workflow.py # Workflow ABC for Databricks job entry points
├── config/
│ ├── base.py # JobConfig, Environment
│ └── unity.py # UnityConfig (Unity Catalog-aware)
├── io/
│ ├── delta.py # DeltaTable, Appender, Overwriter, optimize, vacuum
│ ├── merge.py # MergeBuilder, MergeResult
│ ├── streaming.py # StreamingTableReader, StreamingTriggerOptions
│ ├── checkpoint.py # CheckpointManager, CheckpointInfo
│ └── dbfs.py # DBFS file operations (Databricks only)
├── filters/
│ └── base.py # Filter, FilterPipeline, DropDuplicates, WhereFilter, ColumnFilter
├── migrations/
│ ├── runner.py # MigrationRunner, MigrationStep, MigrationRunResult
│ ├── alter.py # TableAlter (fluent DDL builder)
│ ├── validators.py # TableValidator, ValidationResult, MigrationError
│ └── schema_diff.py # SchemaDiff, ColumnChange, SchemaEvolutionError
├── quality/
│ ├── base.py # Expectation, ExpectationResult, QualityReport
│ ├── expectations.py # NotNull, InRange, Unique, RowCount, MatchesRegex, ColumnExists
│ └── gate.py # QualityGate, QualityError
├── metrics/
│ ├── base.py # MetricEvent, MetricsSink, CompositeMetricsSink
│ ├── delta_sink.py # DeltaMetricsSink (buffered Delta table writer)
│ └── logging_sink.py # LoggingMetricsSink (JSON to logger)
├── observability/
│ ├── batch_context.py # BatchContext, BatchLogger (structured per-batch JSON logging)
│ ├── query_listener.py # QueryProgressObserver (StreamingQueryListener wrapper)
│ └── health.py # StreamingHealthCheck, HealthStatus, HealthResult
└── testing/
├── fixtures.py # spark_session, spark_session_function, df_builder, temp_delta
├── builders.py # DataFrameBuilder (fluent test data)
├── temp_table.py # TempDeltaTable (auto-cleanup context manager)
├── assertions.py # assert_frame_equal, assert_schema_equal
└── mocks.py # MockDBUtils, MockDBUtilsModule
git clone https://github.com/kirankbs/databricks4py.git
cd databricks4py
pip install -e ".[dev]"
# Lint
ruff check src/ tests/ docs/
ruff format --check src/ tests/ docs/
# Tests
pytest -m no_pyspark --timeout=30 # Fast, no Spark/Java
pytest -m "integration or unit" --timeout=120 # Integration (requires Java 17+)
pytest -v --timeout=120 # Everything| PySpark | delta-spark | Python |
|---|---|---|
| 3.5.x | 3.2.x | >= 3.10 |
| 3.4.x | 2.4.x | >= 3.10 |
| 4.x | 4.x | >= 3.10 |
MIT