data

DataSource declarations, tokenizer fitting, and entity-sharded Parquet generation for fastpragma.

path = untar_data(URLs.ML_100k)

DataSource

DataSource declares profile or event columns for a Polars LazyFrame.

• cats, conts, signed_conts, texts, and lifelong declare field roles.
• entity_col identifies the entity shared across sources.
• Event sources require time_col; profile sources use is_profile=True and must be unique by entity.
• DataSource.from_df adapts a pandas DataFrame; DataSource.from_file selects a Polars scanner from the file suffix.
• show_summary later reports rows, entities, field counts, and event-time ranges.

---

source

DataSource

def DataSource(
    df:LazyFrame, cats:list=None, conts:list=None, signed_conts:list=None, texts:list=None, time_col:str=None,
    lifelong:list=None, num_buckets:int=100, cardinality_threshold:int=100, entity_col:str='entity_id',
    is_profile:bool=False, name:str=None, uri:str=None
):

Declare a lazy local or cloud-backed data source.

# Test 1: Explicit column types
df = pl.LazyFrame({'cat_col': ['a', 'b', 'a'], 'num_col': [1.0, 2.0, 3.0], 'id_col': [10, 20, 30]})
s = DataSource(df, cats=['cat_col'], conts=['num_col'], entity_col='id_col')
test_eq(sorted(s.cats), ['cat_col'])
test_eq(s.conts, ['num_col'])

# Test 2: time_col is excluded from cats/conts
df3 = pl.LazyFrame({'id_col': [1,2], 'evt': ['x', 'y'], 'val': [1, 2], 'ts': [100, 200]})
s3 = DataSource(df3, cats=['evt'], conts=['val'], time_col='ts', entity_col='id_col')
test_eq('ts' not in s3.cats + s3.conts, True)

# Test 3: Float32
df4 = pl.LazyFrame({'f32': [1.0, 2.0], 'f33': [1.0, 2.0]}, schema={'f32': pl.Float32, 'f33': pl.Float32})
s4 = DataSource(df4, conts=['f32'], entity_col='f33')
test_eq(s4.conts, ['f32'])

profile = DataSource(
    pl.scan_csv(path/'u.user', separator='|', has_header=False, new_columns=['user_id','age','gender','occupation','zip_code']),
    cats=['gender','occupation'], conts=['age'], entity_col='user_id', name='users', uri=str(path/'u.user'), is_profile=True
)
profile

DataSource(name='users', uri='/app/data/.fastai/data/ml-100k/u.user', columns=['user_id', 'age', 'gender', 'occupation', 'zip_code'], cats=['gender', 'occupation'], conts=['age'], texts=[], time_col=None, is_profile=True)

pdf = pd.DataFrame({'cat': ['x','y','x'], 'val': [1,2,3], 'id': [10,20,30]})
s5 = DataSource.from_df(pdf, cats=['cat'], conts=['val'], entity_col='id')
test_eq(sorted(s5.cats), ['cat'])
test_eq(s5.conts, ['val'])
test_eq(s5.entity_col, 'id')
test_eq(s5.time_col, None)

Key-Value-Time Conversion

The tokenizer converts wide sources into a common long format with one row per entity field value:

entity_col, key, value, vtype, time

Profile fields use time=0; event fields use the declared timestamp; lifelong fields use their own milestone time. Signed continuous fields are split into magnitude and a categorical sign field.

---

source

Tokenizer

def Tokenizer(
    num_buckets:int=100, cardinality_threshold:int=100, bpe_vocab_size:int=100
):

Initialize self. See help(type(self)) for accurate signature.

Tokenizer

Tokenizer.fit builds shared key/value vocabularies across all sources.

• Categorical values receive direct vocabulary IDs.
• Continuous values are mapped to percentile buckets, with explicit zero and unknown handling.
• Low-cardinality text is treated as categorical.
• Higher-cardinality text uses the fitted BPE tokenizer.
• encode_source maps a source to model-ready IDs and positions.
• save and load persist vocabularies, numeric bounds, low-cardinality text fields, and BPE state.

Keys use a single shared vocabulary across all sources. replace_strict maps known keys to their index; any unknown key gets the [UNK] fallback.

Scalable key fitting

_fit_keys builds the key vocabulary from DataSource declarations alone, so it does not scan any source rows. It adds declared categorical, continuous, and text field names, generated <key>_sign fields for signed continuous values, and the shared lifelong key when any source declares milestones. Sorting before insertion makes the resulting IDs deterministic.

Bounded categorical fitting

_cat_counts scans only declared categorical columns, converts them to the common string representation, reshapes them into (key, value) rows, and counts each pair lazily.

_cat_vocab_counts combines those counts across sources, drops rare values, applies a per-key maximum, and uses deterministic sorting. Only this bounded table needs to be collected. _fit_cat_vals then adds the retained values plus declared lifelong values and the fixed sign labels used by signed continuous fields.

Numeric sampling

Fitting exact percentiles over multi-terabyte sources would require processing every continuous value. _num_sample instead builds a bounded, deterministic sample for one source:

1. Read the source row count and choose a stride targeting at most sample_size rows.
2. Project only continuous columns and cast them to Float64.
3. Convert signed fields to absolute magnitude because sign is tokenized separately.
4. Sample rows before unpivoting, avoiding a potentially much larger long table.
5. Remove null, NaN, and zero values; zero already has its own <bucket_0> token.

The source-specific offset changes the sampling phase. _num_samples applies this to every source with continuous fields and concatenates the resulting lazy samples.

Numeric percentile boundaries

_fit_num_bounds operates on the bounded sample after collection. For each key, it calculates num_buckets-1 percentile cut points and stores them in num_bounds. During encoding, np.searchsorted uses these boundaries to map each non-zero value into <bucket_1> through <bucket_{num_buckets}>; zero remains <bucket_0> and missing or unparsable values become [UNK].

---

source

Tokenizer.fit

def fit(
    sources:L, eval_points:LazyFrame=None, min_freq:int=2, max_values_per_key:int=100000, sample_size:int=1000000
):

Call self as a function.

Time Features

Tokenizer._tokenize adds the temporal features consumed by the model.

• Event logsec is the log-scaled distance from the entity’s global latest event across all event sources.
• Lifelong profile fields use distance from eval_time.
• Regular profile fields remain at zero.
• Events also receive hour, dow, and dom calendar features.

Truncation is deliberately deferred to the dataloader so experiments can change context limits without re-tokenizing.

Truncation (max events per user, max tokens per event) belongs in the dataloader, not in tokenization. Tokenization produces the full flat parquet per source; the dataloader applies truncation at load time so it can be tuned per experiment without re-tokenizing.

df = pl.LazyFrame(dict(
    uid=[1, 2, 3, 4],
    plan_note=["student discount", "student discount", "family plan", "family plan",],
    bio=[
        "loves sci-fi movies and long documentaries",
        "mostly watches comedy shows after work",
        "interested in independent films and animation",
        "rewatches classic action movies on weekends",
    ],
))

src = DataSource(
    df,
    texts=["plan_note", "bio"],
    entity_col="uid",
    is_profile=True,
)

tok = Tokenizer(num_buckets=10, cardinality_threshold=2, bpe_vocab_size=50)
tok.fit(L([src]))

print("Low-cardinality text keys:", tok.lowcard_text_keys)
print("Has BPE tokenizer:", tok.bpe_tokenizer is not None)

encoded = tok._tokenize(src, eval_time="2024-01-01").head(10).collect()

Low-cardinality text keys: {'plan_note'}
Has BPE tokenizer: True

/tmp/ipymini_487/276104125.py:16: DeprecationWarning: In Polars 2.0, the default behavior for `empty_as_null` will change to `False`. To keep the current behavior, explicitly set `empty_as_null=True`.
  high = high.explode('_ids', 'val_pos').rename({'_ids':'val_id'})

test_eq("plan_note" in tok.lowcard_text_keys, True)
test_eq("bio" in tok.lowcard_text_keys, False)
test_eq(tok.bpe_tokenizer is not None, True)

tok.n_keys, tok.n_keys

(7, 7)

Tokenizer Persistence

The tokenizer is part of the dataset contract. Tokenizer.save(path) writes JSON state and, when present, a companion BPE model. Tokenizer.load(path) reconstructs the same key/value vocabularies, numeric boundaries, low-cardinality text configuration, BPE offset, and tokenizer model for reproducible encoding.

---

source

Tokenizer.to_dict

def to_dict():

Call self as a function.

---

source

Tokenizer.from_dict

def from_dict(
    state
):

Call self as a function.

---

source

Tokenizer.save

def save(
    dest:str='tokenizer.json', storage_options:NoneType=None
):

Call self as a function.

---

source

Tokenizer.load

def load(
    src:str='tokenizer.json', storage_options:NoneType=None
):

Call self as a function.

tmp = Path(tempfile.mkdtemp())/'tokenizer.json'
tok.save(tmp)
tok2 = Tokenizer.load(tmp)
test_eq(tok2.key_vocab, tok.key_vocab)
test_eq(tok2.val_vocab, tok.val_vocab)
test_eq(set(tok2.num_bounds), set(tok.num_bounds))
test_eq(tok2.bpe_offset, len(tok2.val_vocab))

PRAGMADataset

PRAGMADataset combines one optional profile source with one or more event sources sharing entity_col.

fit_tokenizer fits and saves the tokenizer. write_kv(eval_time, n_shards) then tokenizes all sources, applies a shared latest-event reference, adds source/event/calendar metadata, assigns entities deterministically to shards, and writes shard_*.parquet. It returns the output path, key-vocabulary size, and value-vocabulary size.

---

source

PRAGMADataset

def PRAGMADataset(
    profile:NoneType=None, events:NoneType=None, entity_col:str='entity_id', out_path:NoneType=None,
    storage_options:NoneType=None, credential_provider:str='auto', fs_options:NoneType=None, engine:str='streaming'
):

Initialize self. See help(type(self)) for accurate signature.

---

source

PRAGMADataset.write_kv

def write_kv(
    eval_time, n_shards:int=100
):

Tokenize sources and write to entity-sharded parquet files.

tmp = Path(tempfile.mkdtemp())
ds = PRAGMADataset(profile=psrc, events=[esrc], entity_col='uid', out_path=tmp)
tok = ds.fit_tokenizer(num_buckets=4)

test_eq(tok is ds.tokenizer, True)
test_eq((tmp/'tokenizer.json').exists(), True)
test_eq('plan' in tok.key_vocab, True)
test_eq('evt' in tok.key_vocab, True)
test_eq('amt' in tok.key_vocab, True)

Keys: 8, Vals: 9, BPE: none

---

source

PRAGMADataset.show_summary

def show_summary():

Call self as a function.

out = Path(tempfile.mkdtemp())/'ml100k_tok'
ds = PRAGMADataset(profile=profile, events=[events], entity_col='user_id', out_path=out)
ds.show_summary()

shape: (2, 14)

i	name	kind	rows	entities	entity_col	time_col	cats	conts	signed_conts	texts	lifelong	min_time	max_time
i64	str	str	i64	i64	str	str	i64	i64	i64	i64	i64	i64	i64
0	"users"	"profile"	943	943	"user_id"	null	2	1	0	0	0	null	null
1	"ratings"	"event"	100000	943	"user_id"	"timestamp"	1	1	0	0	0	874724710	893286638

import nbdev; nbdev.nbdev_export()

End-to-End Tokenization Check

The MovieLens 100K check verifies the current contract: profile and event sources are tokenized with one shared tokenizer, output is split into entity shards, profile rows use source_idx=-1 and event_idx=-1, event rows carry event and calendar features, and token IDs contain no unexpected unknowns.

out = Path(tempfile.mkdtemp())/'ml100k_tok'
ds = PRAGMADataset(profile=profile, events=[events], entity_col='user_id', out_path=out)
print(ds.show_summary())
# tok = ds.fit_tokenizer(num_buckets=10, cardinality_threshold=100)
ds.write_kv(eval_time='1998-04-01T00:00:00', n_shards=4)
sorted(o.name for o in out.iterdir())

shape: (2, 14)
┌─────┬─────────┬─────────┬────────┬───┬───────┬──────────┬───────────┬───────────┐
│ i   ┆ name    ┆ kind    ┆ rows   ┆ … ┆ texts ┆ lifelong ┆ min_time  ┆ max_time  │
│ --- ┆ ---     ┆ ---     ┆ ---    ┆   ┆ ---   ┆ ---      ┆ ---       ┆ ---       │
│ i64 ┆ str     ┆ str     ┆ i64    ┆   ┆ i64   ┆ i64      ┆ i64       ┆ i64       │
╞═════╪═════════╪═════════╪════════╪═══╪═══════╪══════════╪═══════════╪═══════════╡
│ 0   ┆ users   ┆ profile ┆ 943    ┆ … ┆ 0     ┆ 0        ┆ null      ┆ null      │
│ 1   ┆ ratings ┆ event   ┆ 100000 ┆ … ┆ 0     ┆ 0        ┆ 874724710 ┆ 893286638 │
└─────┴─────────┴─────────┴────────┴───┴───────┴──────────┴───────────┴───────────┘
Keys: 10, Vals: 1668, BPE: none
tokenizing profile
█

 |----------------------------------------| 0.00% [0/1 00:00<?]tokenizing event source 0: ratings

 |████████████████████████████████████████| 100.00% [1/1 00:00<00:00]combining sources

['shard_id=0', 'shard_id=1', 'shard_id=2', 'shard_id=3', 'tokenizer.json']

tok = ds.fit_tokenizer(num_buckets=10, cardinality_threshold=100)

tokenizer found

tok.n_keys, tok.n_vals, tok.num_bounds.keys(), tok.bpe_tokenizer

(10, 1668, dict_keys(['age', 'rating']), None)

files = sorted(out.rglob('*.parquet'))
toks = pl.scan_parquet([str(o) for o in files], hive_partitioning=False).collect()
prof,evts = toks.filter(pl.col('source_idx') == -1),toks.filter(pl.col('source_idx') == 0)

test_eq(len(files), 4)
test_eq(toks.height, 943*3 + 100_000*2)
test_eq((prof.height, evts.height), (943*3, 100_000*2))
test_eq(prof['event_idx'].unique().to_list(), [-1])
test_eq(evts['event_idx'].n_unique(), 100_000)
test_eq(evts.group_by('event_idx').len()['len'].unique().to_list(), [2])
test_eq(toks.group_by('user_id').agg(pl.col('shard_id').n_unique())['shard_id'].unique().to_list(), [1])
test_eq(toks.select(pl.col('key_id', 'val_id').null_count()).row(0), (0,0))
test_eq(toks['key_id'].max() < tok.n_keys, True)
test_eq(toks['val_id'].max() < tok.n_vals, True)
test_eq(prof['logsec'].unique().to_list(), [0.])
test_eq(evts.group_by('user_id').agg(pl.col('logsec').min())['logsec'].unique().to_list(), [0.])
test_eq(evts['hour'].is_between(0, 23).all(), True)
test_eq(evts['dow'].is_between(1, 7).all(), True)
test_eq(evts['dom'].is_between(1, 31).all(), True)

toks.shape, prof.shape, evts.shape

((202829, 15), (2829, 15), (200000, 15))

u = 1
prof.filter(pl.col('user_id') == u).sort('key')

shape: (3, 15)

user_id	key	key_id	value	val_id	val_pos	vtype	time	logsec	source_idx	event_idx	hour	dow	dom	shard_id
i64	str	i64	str	i64	i64	str	i64	f64	i32	i64	f32	f32	f32	u64
1	"age"	5	"24"	1589	0	"num"	0	0.0	-1	-1	0.0	0.0	0.0	0
1	"gender"	6	"M"	1545	0	"cat"	0	0.0	-1	-1	0.0	0.0	0.0	0
1	"occupation"	8	"technician"	1565	0	"cat"	0	0.0	-1	-1	0.0	0.0	0.0	0

evts.filter(pl.col('user_id') == u).sort('event_idx','key','val_pos').head(12)

shape: (12, 15)

user_id	key	key_id	value	val_id	val_pos	vtype	time	logsec	source_idx	event_idx	hour	dow	dom	shard_id
i64	str	i64	str	i64	i64	str	i64	f64	i32	i64	f32	f32	f32	u64
1	"movie_id"	7	"61"	1118	0	"cat"	878542420	113.22252	0	202	7.0	1.0	3.0	0
1	"rating"	9	"4"	1612	0	"num"	878542420	113.22252	0	202	7.0	1.0	3.0	0
1	"movie_id"	7	"189"	651	0	"cat"	888732928	94.037681	0	305	6.0	7.0	1.0	0
1	"rating"	9	"3"	1585	0	"num"	888732928	94.037681	0	305	6.0	7.0	1.0	0
1	"movie_id"	7	"33"	808	0	"cat"	878542699	113.22232	0	333	7.0	1.0	3.0	0
…	…	…	…	…	…	…	…	…	…	…	…	…	…	…
1	"rating"	9	"4"	1612	0	"num"	875072547	115.380003	0	334	3.0	3.0	24.0	0
1	"movie_id"	7	"20"	664	0	"cat"	887431883	100.62061	0	478	4.0	6.0	14.0	0
1	"rating"	9	"4"	1612	0	"num"	887431883	100.62061	0	478	4.0	6.0	14.0	0
1	"movie_id"	7	"202"	667	0	"cat"	875072442	115.38006	0	639	3.0	3.0	24.0	0
1	"rating"	9	"5"	1646	0	"num"	875072442	115.38006	0	639	3.0	3.0	24.0	0

chk = evts.filter(pl.col('user_id') == 1, pl.col('key') == 'movie_id').sort('time')
chk.select('user_id','event_idx','time','logsec','hour','dow','dom').head(), chk.select('user_id','event_idx','time','logsec','hour','dow','dom').tail()

(shape: (5, 7)
 ┌─────────┬───────────┬───────────┬────────────┬──────┬─────┬──────┐
 │ user_id ┆ event_idx ┆ time      ┆ logsec     ┆ hour ┆ dow ┆ dom  │
 │ ---     ┆ ---       ┆ ---       ┆ ---        ┆ ---  ┆ --- ┆ ---  │
 │ i64     ┆ i64       ┆ i64       ┆ f64        ┆ f32  ┆ f32 ┆ f32  │
 ╞═════════╪═══════════╪═══════════╪════════════╪══════╪═════╪══════╡
 │ 1       ┆ 59972     ┆ 874965478 ┆ 115.438142 ┆ 21.0 ┆ 1.0 ┆ 22.0 │
 │ 1       ┆ 92487     ┆ 874965478 ┆ 115.438142 ┆ 21.0 ┆ 1.0 ┆ 22.0 │
 │ 1       ┆ 74577     ┆ 874965518 ┆ 115.438121 ┆ 21.0 ┆ 1.0 ┆ 22.0 │
 │ 1       ┆ 48214     ┆ 874965556 ┆ 115.4381   ┆ 21.0 ┆ 1.0 ┆ 22.0 │
 │ 1       ┆ 15764     ┆ 874965677 ┆ 115.438035 ┆ 22.0 ┆ 1.0 ┆ 22.0 │
 └─────────┴───────────┴───────────┴────────────┴──────┴─────┴──────┘,
 shape: (5, 7)
 ┌─────────┬───────────┬───────────┬───────────┬──────┬─────┬──────┐
 │ user_id ┆ event_idx ┆ time      ┆ logsec    ┆ hour ┆ dow ┆ dom  │
 │ ---     ┆ ---       ┆ ---       ┆ ---       ┆ ---  ┆ --- ┆ ---  │
 │ i64     ┆ i64       ┆ i64       ┆ f64       ┆ f32  ┆ f32 ┆ f32  │
 ╞═════════╪═══════════╪═══════════╪═══════════╪══════╪═════╪══════╡
 │ 1       ┆ 88259     ┆ 889751711 ┆ 11.336528 ┆ 1.0  ┆ 5.0 ┆ 13.0 │
 │ 1       ┆ 30479     ┆ 889751712 ┆ 11.090355 ┆ 1.0  ┆ 5.0 ┆ 13.0 │
 │ 1       ┆ 47638     ┆ 889751712 ┆ 11.090355 ┆ 1.0  ┆ 5.0 ┆ 13.0 │
 │ 1       ┆ 3248      ┆ 889751736 ┆ 0.0       ┆ 1.0  ┆ 5.0 ┆ 13.0 │
 │ 1       ┆ 19699     ┆ 889751736 ┆ 0.0       ┆ 1.0  ┆ 5.0 ┆ 13.0 │
 └─────────┴───────────┴───────────┴───────────┴──────┴─────┴──────┘)

chk = evts.filter(pl.col('user_id') == 1, pl.col('key') == 'movie_id')
test_eq(chk.filter(pl.col('time') == pl.col('time').max())['logsec'].unique().item(), 0.0)
test_eq((chk['logsec'] >= 0).all(), True)
test_eq(chk.sort('time')['logsec'].to_list()[0] >= chk.sort('time')['logsec'].to_list()[-1], True)

for GPU

ds = PRAGMADataset(..., out_path='gs://bucket/fastpragma/data', engine='streaming')
ds = PRAGMADataset(..., out_path='gs://bucket/fastpragma/data', engine=pl.GPUEngine(raise_on_fail=True))

gpu = pl.GPUEngine(device=0, raise_on_fail=True)

ds = PRAGMADataset(
    profile=profile,
    events=[events],
    entity_col='user_id',
    out_path='gs://my-bucket/fastpragma/ml100k',
    storage_options=None,
    fs_options=dict(token='google_default'),
    credential_provider='auto',
    engine=gpu
)

out,n_keys,n_vals = ds.write_kv(eval_time='1998-04-01T00:00:00', n_shards=100)
out,n_keys,n_vals