litestar-org / litestar-org/polyfactory

Enhancement: Add a rejection sampler

Open
#174 5 comments 0 reactions 0 assignees View on GitHub
enhancement good first issue help wanted
Dominant language
Python
Stars
1.5k
Forks
120
PR merge metrics
No merged PRs in 30d

Description

### Summary

Currently the `batch` method fails with a validation error if any of the generated rows fail the schema validators. To allow use of the package in a testing environment, it would be useful to be able to generate a dataframe of any size using a rejection sampler method. This method should store the random seeds of successful builds in order to reproduce the same dataframe each time.

I have created a class that performs these actions included below. Given this is something I have needed for my project, it could be a useful feature for others wanting to use Polyfactory for testing. I built it based off the original pydantic factories package, but I imagine it would be pretty similar for the additional Factory options in Polyfactory.

### Basic Example

```
import time
import json
import pandas as pd
from polyfactory.factories.pydantic_factory import ModelFactory

class RejectionSampler:
"""Function to create a synthetic dataset based off the pydantic schema,
dropping rows that do not meet the validation set up in the schema.

Parameters
----------

factory (ModelFactory): pydantic factories ModelFactory created from pydantic schema
size (int): Length of dataset to create
"""

def __init__(self, factory: ModelFactory, size: int) -> None:

self.factory = factory
self.size = size
self.used_seeds = []

def setup_seeds(self):

start = time.time()

synthetic_data = pd.DataFrame()

# start seed at 1, increase seed by 1 each pass/fail of factory.build() to ensure reproducibility
seed_no = 1

for _ in range(self.size):
result = None
while not result:
try:
self.factory.seed_random(seed_no)
result = self.factory.build()
result_dict = json.loads(result.json())
synthetic_data = synthetic_data.append(
pd.DataFrame(result_dict, index=[0])
)
self.used_seeds += [seed_no]
seed_no += 1
result = True
except ValidationError:
seed_no += 1

end = time.time()

print(f"finished, took {seed_no-1} attempts to generate {self.size} rows")
print(f"took {end-start} seconds to setup seeds")

def generate(self):

start = time.time()

synthetic_data = pd.DataFrame()

for seed in self.used_seeds:
self.factory.seed_random(seed)
result = self.factory.build()
result_dict = json.loads(result.json())
synthetic_data = synthetic_data.append(pd.DataFrame(result_dict, index=[0]))

end = time.time()

print(f"took {end-start} seconds to generate new data")

return synthetic_data
```

### Drawbacks and Impact

_No response_

### Unresolved questions

_No response_

Contributor guide

Open the contributing guide

Research direction

Start with the existing `batch` method and `polyfactory.factories.pydantic_factory.ModelFactory`, then compare them with the proposed `RejectionSampler` example. Clarify how rejected rows, successful seeds, and replay should work across the package’s other Factory options. Done means a documented sampler API can produce the requested size and reproduce the same dataframe from stored seeds.

Written by the indexing model from the issue text.

Assessment

Tech stack
pandas, python
Domain
data, testing-qa
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
28/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.