Tables#

Traveler uses JaxTable for rectangular (tabular) model data that will be consumed by JAX. A JaxTable is similar in shape to a Polars DataFrame: both have named columns and a common number of rows. The difference is what the columns contain and what kind of work each object is designed to do.

A Polars DataFrame is well suited to reading files, selecting and deriving columns, joining tables, grouping records, and other data preparation. A JaxTable stores its columns as JAX arrays so numerical model code can use JAX transformations and device execution.

Important

A JaxTable is not a Polars DataFrame subclass, and the two objects do not share an API. They are representations used on opposite sides of a conversion boundary.

From Polars to JAX and back#

A JaxTable column is normally a one-dimensional JAX array. All such columns must have the same length. Scalar, zero-dimensional columns are also allowed, and categorical columns use Traveler’s CategoricalArray.

Conversion is explicit:

import polars as pl

from traveler import JaxTable


frame = pl.DataFrame(
    {
        "household_id": pl.Series([1, 2], dtype=pl.Int32),
        "income": pl.Series([45_000, 90_000], dtype=pl.Int32),
    }
)

table = JaxTable.from_polars(frame)
table.income  # Array([45000, 90000], dtype=int32)

frame_again = table.to_polars()

Compatible numeric columns can often pass through Arrow and DLPack without a copy. Some layouts or data types require a copy instead. Polars enum columns become CategoricalArray columns, while types that JAX cannot represent—such as arbitrary strings or nullable values—must be encoded or otherwise handled before conversion.

A typical workflow therefore looks like this:

  1. Read and prepare source data with Polars.

  2. Convert the prepared frame to a JaxTable.

  3. Put the JaxTable in a Store and run model calculations.

  4. Convert results back to Polars for inspection, further data processing, or writing to disk.

JaxTable deliberately offers a smaller set of table operations than Polars. It supports column access, row selection, metadata, and assignment of derived columns, but Polars remains the better tool for general-purpose table manipulation.

Data and views#

Traveler’s other table classes do not make additional copies of the table’s columns. They describe how model code should view a JaxTable held in a store.

Type

Owns table columns?

Attached to a store?

Purpose

Polars DataFrame

Yes, as Polars/Arrow data

No

Input, output, and data preparation

JaxTable

Yes, as JAX arrays

No

Numerical model data

WorkingTable

No

Depends on its concrete form

Common interface for attaching and detaching a table view

StoredTable

No

Yes

Typed view of a named JaxTable, optionally restricted to selected rows

ReferencedTable

No

No

Lightweight, detachable reference to a table type and row selection

For a typed store, the same name can expose the underlying data or its richer application-specific view:

store["households"]  # the JaxTable that owns the columns
store.households     # a Households StoredTable view

The view knows that its data is stored under the name "households"; it does not contain a second household table.

WorkingTable#

WorkingTable is the abstract base class for the two forms of table view. It defines the small protocol that Traveler steps can rely on:

  • with_store(store) produces a store-attached StoredTable;

  • without_store() produces a detached ReferencedTable.

Application code normally works with one of those concrete forms rather than constructing a WorkingTable directly.

StoredTable#

StoredTable connects an application-specific table definition to a store. A subclass declares the store key that holds the underlying JaxTable, column fields, and optionally relationships or derived attributes:

import numpy as np

from traveler.core.stored import StoredTable


class Households(StoredTable):
    table_name = "households"
    table_id_col = "household_id"

    household_id: np.int32
    income: np.int32

The annotations declare the columns available through the typed view and the dtypes used for loading and validation. No class-level value is assigned: the arrays themselves remain in the store’s underlying JaxTable. Declared columns are read-only on the view; use store operations such as assign_on() to create a table with updated data. Columns that are present in the underlying table but are not part of the declared interface remain available with item access, such as households["temporary_column"].

The application’s store declares that typed view and holds the JaxTable created earlier:

from traveler import AbstractStore


class Store(AbstractStore):
    households: Households


store = Store(households=table)

When connected to a store, Households retrieves those columns from store["households"]. The annotation alone declares this connection; no class-level default or descriptor is needed. An optional index array makes it a view of particular rows:

import jax.numpy as jnp


all_households = Households(store=store)
second_household = Households(store=store, index_array=jnp.array([1]))

all_households.income     # every income in the underlying JaxTable
second_household.income  # Array([90000], dtype=int32)

A StoredTable can also define validation rules, lazy derived columns, and relationships to other stored tables. For example, a person view can return a household view indexed to each person’s household.

Because it retains a reference to the entire store, a StoredTable cannot be flattened as a JAX pytree. Detach it before passing the selection through code that expects a pytree.

ReferencedTable#

ReferencedTable is the detached form of a StoredTable. It keeps only the information needed to reconstruct the view:

  • the name of the table in a store;

  • the concrete StoredTable subclass;

  • the row index array.

It does not contain the underlying JaxTable and does not retain the store. This makes the reference a small JAX pytree whose dynamic content is the row selection.

reference = second_household.without_store()

# The reference can travel separately from the model's data.
second_household_again = reference.with_store(store)
second_household_again.income  # Array([90000], dtype=int32)

Traveler steps commonly accept a ReferencedTable alongside a store. The step reattaches the reference with with_store(store), uses the resulting typed StoredTable to read columns and relationships, and leaves the actual arrays in the store’s JaxTable.

In short, the conversion and view relationships are:

        block
    columns 3
    polars["Polars DataFrame"] conversion["from_polars() →<br/>← to_polars()"] jax["JaxTable<br/>stored in Store"]
    space:2 reads["reads selected rows ↑"]
    working["WorkingTable<br/>(shared interface)"] space stored["StoredTable"]
    space:2 attachment["without_store() ↓<br/>↑ with_store(store)"]
    space:2 referenced["ReferencedTable"]

    polars --- conversion
    conversion --- jax
    stored --> reads
    reads --> jax
    stored --- attachment
    attachment --- referenced
    working -.-> stored
    working -.-> referenced

    style conversion fill:transparent,stroke:transparent
    style reads fill:transparent,stroke:transparent
    style attachment fill:transparent,stroke:transparent