Composite primary key for linked tables

[khazana]

I have two tables user and events. user is my subject table and the primary key is id. events has a composite primary key id and user_id. I've set the foreign key config for the events table as:

  "foreign_keys": [
            {"column": "user_id", "referenced_table": "user",
             "is_context": is_context}],

However, in the generated events data there are multiple rows with the same composite primary key (id, user_id). I'm passing in the data as a pandas dataframe to the generator. Is there any way to specify composite primary key for linked tables? My training data had 500K data points.

[shuangwu5]

Hi @khazana,
Currently we don't support composite primary key directly.
Could you please try to preprocess your dataset to concatenate those two keys into another column, and set it as the primary key?
And later, postprocess the generated synthetic dataset to transform columns back to your original format.
Hope it helps and please let us know if it works. Thanks!

[khazana]

Thanks for the reply @shuangwu5!

Just so that I understand clearly:

1. Add a new column where value is id - user_id and set that as primary key.
2. Remove columns id and user_id from existing data frame?
3. How can I reference the foreign key user_id then? I need the correlations between the user and events table to be maintained.

[shuangwu5]

Hi @khazana,
If I understand your use case correctly, you would only need to add a new column where value is id - user_id and remove the original id column.

For user table:

• Set id as primary key as before

For event table:

• set primary key to the new column with values id - user_id
• set user_id as foreign key referencing to id of user table as before

[khazana]

@shuangwu5 I'm still having trouble with the generated data. Here's what my event data looks like:

id	user_id	type
new_user_event	1	internal_event
sign_up	1	internal_event
random_event_1234	1	user_event
new_user_event	2	internal_event
sign_up	2	internal_event
random_event_4566	2	user_event

For every user, the table will contain new_user_event and sign_up rows. A user can also have a bunch of other rows with random event uuids. The table has 60+ columns whose data depends on the id, specifically for the non-uuid ones. The generated data however looks like the following:

id	user_id	type
new_user_event	1	internal_event
new_user_event	1	user_event
random_event_1234	1	user_event
sign_up	2	user_event
sign_up 5	2	internal_event
random_event_4566	2	internal_event

I expected the model to be able to understand this tabular relation...but maybe I'm not setting some field correctly?

The encoding type is set to ModelEncodingType.tabular_categorical for the id field. I did try adding a composite key as you suggested but the resulting data still produces duplicate events for a given user - determined by other fields such as the type.

[mplatzer]

Hi. If I understand correct, than events contains

1. internal events, with the id column contains the specific event type
2. user events, with the id containing a random UUID (with no statistical value)

So, it seems like a mixture of ID and Type, but we would need to separate those 2 concerns.

I thus suggest the following pre-processing:

events['type'] = events['id'].where(events['id'].isin(['new_user_event', 'sign_up']), 'user_event')
events = events.drop(columns=['id'])

and configuring the generator via

g = mostly.train(config={
    'tables': [{
        'name': 'users',
        'data': users,
        'primary_key': 'user_id',
    }, {
        'name': 'events',
        'data': events,
        'foreign_keys': [
            {'column': 'user_id', 'referenced_table': 'users', 'is_context': True}
        ]
    }]
}

If you have sufficient number of users, then the model should be able to properly learn that the first 2 events are new_user_event and sign_up, which is then being followed by a number of user events.

Once you have generated the synthetic data, you can transform back via:

import uuid
events['id'] = events['type'].where(events['type'].isin(['new_user_event', 'sign_up']), [uuid.uuid4() for i in range(len(events))])
events['type'] = events['type'].where(events['type'] == 'user_event', 'internal_event')
events

See also attached. Please let us know whether this works for you.

[khazana]

Thanks for the message @mplatzer.

Just for my understanding purpose, I split my events table into two dataframes. One with just internal_event events and the other with user_event. Here's the size of my training data:

1. user - 10K
2. events_internal - 23348
3. events_user - 499439

In the events_internal data frame, I confirmed that the new_user_event and sign_up events are present for every single user. We have a few more internal events type that don't necessarily have to be present for every user, which contribute to the extra 3.3K events in the data.

Now, I trained my generator with just user and events_internal dataframes with user_id foreign key set up as before. I then generated data using sd = mostly.generate(g.id, size=1000) and expected to see at least 2k rows in the synthetic sd.data()['events_internal'] table but only 183 rows were generated. Why is the model not capturing that 1:1 relationship correctly i.e a new_user_event has to be present for every single user?

An interesting thing I noticed from the logs is that for the user table, I see:

[2025-05-16 18:58:52] INFO   : tgt_table=ParquetDataTable(name=user, is_output=None, path=/home/ubuntu/mostlyai/in_progress/c159af3d-0b4c-4f26-90e7-77a562f526c0/user:tabular/__PULL_FETCH/user) ctx_table=None
[2025-05-16 18:58:52] INFO   : tgt_n_rows=10000 ctx_n_rows=None

but I don't see the correct number of rows imported for the events_internal table:

[2025-05-16 19:00:47] INFO   : tgt_table=ParquetDataTable(name=event_internal, is_output=None, path=/home/ubuntu/mostlyai/in_progress/c159af3d-0b4c-4f26-90e7-77a562f526c0/event_internal:tabular/__PULL_FETCH/device_internal) ctx_table=ParquetDataTable(name=user, is_output=None, path=/home/ubuntu/mostlyai/in_progress/c159af3d-0b4c-4f26-90e7-77a562f526c0/event_internal:tabular/__PULL_FETCH/user)
[2025-05-16 19:00:47] INFO   : tgt_n_rows=1838 ctx_n_rows=10000

[2025-05-16 19:04:18] INFO   : tgt_table=ParquetDataTable(name=event_internal, is_output=False, path=/home/ubuntu/mostlyai/in_progress/f8db5ef2-542a-4089-83a7-69f941908138/event_internal:tabular/SyntheticData) ctx_table=ParquetDataTable(name=user, is_output=False, path=/home/ubuntu/mostlyai/in_progress/f8db5ef2-542a-4089-83a7-69f941908138/user:tabular/SyntheticData)
[2025-05-16 19:04:18] INFO   : tgt_n_rows=0 ctx_n_rows=20

Only 1838 rows? 🤔 And in the generated report, I see 755 original samples, 630 synthetic samples for the event_internal table whereas 10,000 original samples, 1,000 synthetic samples as expected for the subject table. Is the model not pulling all the training data for some reason?

[mplatzer]

Indeed, it seems that not all records are pulled for events_internal, and then the model is never aware of these, thus also generates too few.

Might it be that the keys are not matching, or of different type? Could you please share the output of these lines:

print(f"users: {len(users)} rows, {len(set(users.user_id))} unique user ids")
print(f"events_internal: {len(events_internal)} rows, {len(set(events_internal.user_id))} unique user ids")
print(f"{len(set(events_internal.user_id) - set(users.user_id))} user ids in events_internal but not in users")
print(f"{len(set(users.user_id) - set(events_internal.user_id))} user ids in users but not in events_internal")
print(f"{users.user_id.dtype=}, {events_internal.user_id.dtype=}")

Plus, also share the config that you are passing to mostly.train()? Do you have a (non-unique) primary key defined for events_internal?

And any other log lines during PULL that could provide us a clue?

[khazana]

@mplatzer thank you so much for the suggestion, it helped me find my mistake 🤦‍♀️ . This was my code:

import petl as etl

table = etl.fromdb(prod_rds, 'select * from account.user ORDER BY RAND() limit 10000')
user_df = table.todf()   ---> used for subject table training
user_ids = etl.values(table, 'id')
events_df = get_data_in_batches("account.event", prod_rds, user_ids) ---> used for linked table training

It turned out that etl.values(table, 'id') was making another underlying query to fetch user ids so the events_df ended up having events for a completely different set of users. The generated data looks much better now - thank you so much for your time!

[mplatzer]

Great to hear!