Schema Annotations in Bauplan 0.3.x
This week, we released Bauplan 0.3.0, which offers a new syntax for defining and validating pipelines. This week we’re deprecating the old SDK, so we thought it’d be a good time to step back and offer an overview of the new features.
Our goal with the new SDK is to offer the same sort of tradeoff a compiler or linter does: invest more time up front, and the system will catch bugs early for you. As we wrote previously, we think that this tradeoff is even more valuable if you’re coding primarily with agents. Either way, you’ll find lots to like about the new annotations. Even if the new syntax can take a bit of getting used to!
How to write pipelines with the new SDK
At the heart of the new syntax is the TableSchema type:
class RawTrips(TableSchema): # (1)
"""Columns read from taxi_fhvhv."""
PULocationID: Int64 | None
request_datetime: TimestampMicroUTC | None
on_scene_datetime: TimestampMicroUTC | None
class WaitTimes(TableSchema): # (2)
"""How long riders waited for their taxi."""
PULocationID: Annotated[Int64, TableField(lineage=RawTrips["PULocationID"])] # (3)
wait_seconds: Annotated[
Int64,
TableField(doc="Seconds between the ride request and the driver's arrival."),
]
@bauplan.model()
@bauplan.python("3.12", pip={"polars": "2.0.0"})
def wait_times(
trips: Annotated[pa.Table, Model("taxi_fhvhv", projection_schema=RawTrips)], # (4)
) -> Annotated[pa.Table, WaitTimes]: # (5)
import polars as pl
trips = pl.from_arrow(trips).drop_nulls()
return compute_wait_seconds(trips).to_arrow() # elided for clarityHere we’ve defined a model in a pipeline, an input schema RawTrips, and an output schema WaitTimes. The types are attached to the model using the Annotated syntax:
RawTrips[1] is used as theprojection_schema[4] for reading fromtaxi_fhvhv. Note that any valid subset of columns with compatible types (any projection) can be used here. The system automatically pushes the projection down, and transparently casts input types to the types inRawTripsif that can be done in a lossless way.WaitTimes[2] is defined as the output of the function [5].- We use
lineage[3] to track thatPULocationIDis the same field as in the raw table.
Why do we require this?
Each Bauplan model function represents the transformation between two data types, for example, a transformation from raw data in the lake to a cleaned-up representation.
If you take the transformations as the edges in a DAG, then these schemas are tagging the nodes. Because we don’t look inside your transformations and can’t reason about the Python code inside, this extra metadata gives us the ability to offer new error checks and planning-time optimizations that weren’t available before.
Using schemas also benefits the code itself; both the types and the docstrings serve as documentation for anyone working on the code. And this documentation outlives the code; it’s saved as part of the metadata for the table and versioned along with it, in each commit.
Anyone with access to the table can then fetch it (and previous versions!) via the API; as standardized Iceberg metadata; and of course via the CLI, Python SDK, and MCP:
$ bauplan table get wait_times
Table Documentation:
Dispatched trips and how long the rider waited for the driver.
COLUMN TYPE NULLABLE DOC
PULocationID long false -
wait_seconds long false Seconds between the ride requ...
Note: some documentation was truncated. Use `-O json` to see it in full.What Bauplan will validate for you
At planning time
If any of the schemas don’t match the actual input, we’ll reject that immediately, before anything runs.
For example, if you try to read from taxi_fhvhv with the following schema instead:
class RawTrips(TableSchema):
"""Columns read from taxi_fhvhv."""
PULocationID: Int64 # Missing `| None` to indicate the column is nullable
request_datetime: TimestampMicroUTC | None
on_scene_datetime: TimestampMicroUTC | NoneYou will get the following error:
- "PULocationID" has incompatible projection:
- constraint mismatch: cannot cast from optional to requiredThis may seem like a small thing, but if you’ve ever spent hours hunting down the source of nulls in a complicated pipeline, then you know the value of this type of static check.
At runtime
We can’t verify ahead of time that the code inside wait_times produces a table without nulls. But because we set WaitTimes as the output schema, the system will validate the output efficiently as it runs. If you remove the call to drop_nulls(), you’ll get the following error at runtime:
- "wait_seconds" has incompatible projection:
- constraint mismatch: cannot cast from optional to requiredThis serves two purposes:
- It prevents bugs in code from corrupting the data. This serves as a test that your code does what you say it does, efficiently and without having to run separate SQL queries or expensive data checks.
- It lets downstream code rely on these constraints. If you say a column doesn’t have nulls in it, the system enforces that, so other pipelines using this table don’t have to be written defensively.
Upcoming work
Not everything we wanted to build made it in time for the cutoff, but we have a bunch of relevant features in the pipeline.
More types
Currently, most primitive Arrow types are supported, but we’ll be adding more soon. In particular:
Decimal128, with configurablescale/precision- Arbitrarily nested struct, list, and map types
In the meantime, you can use bauplan.Any for types that aren’t currently expressible, and we’ll skip validation for those.
More constraints
Today, the only constraint we can validate at runtime is nullability. We’re planning to extend that to arbitrary row-level checks, for example that an integer value is nonzero.
Better ergonomics
Although schema annotations are required in this release, we’re working to loosen that constraint, so that it’s easier to do quick, ad-hoc workflows or add types gradually to complicated pipelines.
We also recognize that the Annotated[...] syntax is quite verbose. We are aiming for correctness with this release, but we’re looking for ways to make the syntax more beautiful and easier for humans to read. To that end, we’re keeping a close eye on PEP 835.
Freely mixing SQL and Python models
Because SQL is a typed language, we can feed the type annotations from upstream Python models into the SQL planner. This will allow something we don’t currently: using SQL models in the middle of a DAG, downstream of Python models (without jinja templating!).
Better diagnostics, faster
In parallel with the new SDK, we’ve been working on a rewrite of our planner from the ground up, with the extra benefits of these type annotations built in from the get-go.
It’s still a work in progress; stay tuned for news in the coming months, but we can share a few details:
- It’s much faster than the current planner. Excluding reads from the catalog, we can plan a DAG with hundreds of nodes in less than a millisecond. With all the fetching we need to do, we’re hopeful we can get to a worst-case time of around 50 ms for the vast majority of pipelines, so that you can iterate faster.
- We’ve put a lot of effort into making errors more legible, and adding pointers to the source. For example, here’s the error above presented a different way:
Error: invalid `projection_schema`
╭─[ models.py:42:50 ]
│
42 │ bauplan.Model('trips', projection_schema=MismatchedSchema),
│ ────────┬───────
│ ╰───────── here
│
├─[ models.py:42:50 ]
│
7 │ class PickupsSchema(bauplan.TableSchema):
│ ──────┬──────
│ ╰──────── the output of `trips` is defined here
│
18 │ PULocationID: bauplan.Int64
│ ──────┬──────
│ ╰──────── this must be `Int64 | None`,
│ because `PULocationID` may contain nullsFurther resources
We have detailed documentation available for the new SDK on our documentation page. If you're migrating from the old syntax, we've also published some skills that make the process quick and painless.
If you'd like help, or if you have any questions or feedback, we'd be delighted to hear from you!
