Metadata-Version: 2.4
Name: vlab_prepro
Version: 0.6.0
Summary: Preprocessing for Virtual Lab survey data, on pandas or Polars DataFrames
License-File: LICENSE
Author: Nandan Rao
Author-email: nandanmarkrao@gmail.com
Requires-Python: >=3.10,<4
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Python :: 3.15
Requires-Dist: pandas (>=2.2,<3)
Requires-Dist: polars (>=1.30,<2.0)
Requires-Dist: pyarrow (>=16)
Requires-Dist: pyfarmhash (>=0.3.2,<0.4.0)
Requires-Dist: toolz (>=0.11.1)
Description-Content-Type: text/markdown

# vlab_prepro

Preprocessing library for survey data from the Virtual Lab chatbot platform. It transforms long-format response data (one row per answer) into wide-format datasets (one row per respondent) with enrichment for timing, metadata, data quality indicators, and randomization group assignment.

Every function works on a pandas or a Polars DataFrame and returns the type it was given.

## Installation

```shell
pip install vlab-prepro
```

## Dependencies

| Package | Purpose |
|---|---|
| `polars` | Computation engine for every step, and an accepted input/output type |
| `pandas` (2.2 or later) | Accepted input/output type |
| `pyarrow` | Converts between pandas and Polars (loaded only for pandas input, as is pandas itself) |
| `toolz` | `@curry` decorator enabling partial application / pipe composition |
| `pyfarmhash` | FarmHash fingerprint for `compute_seed` (matches JS implementation) |

## Input Data Format

The library expects two DataFrames:

**Response data** — one row per answer:

| Column | Description |
|---|---|
| `surveyid` | Identifies which survey form was taken |
| `userid` | Unique respondent identifier |
| `flowid` | Identifies a specific survey session/attempt |
| `question_ref` | The question identifier (becomes a column after pivoting) |
| `question_idx` | Numeric position of the question |
| `response` | Raw response value |
| `timestamp` | ISO 8601 string when the answer was submitted |
| `metadata` | JSON string of per-response metadata (e.g. `{"clusterid": "A"}`) |

**Form metadata** — one row per survey form:

| Column | Description |
|---|---|
| `surveyid` | Matches `surveyid` in response data |
| `shortcode` | Human-readable form identifier |
| `version` | Form version number |
| `survey_created` | Timestamp of form creation |
| `metadata` | JSON string of form-level metadata (e.g. `{"wave": "1"}`) |

## Quick Start

```python
from vlab_prepro import Preprocessor
from toolz import pipe
import pandas as pd

responses = pd.read_csv('responses.csv')
forms = pd.read_csv('forms.csv')

p = Preprocessor()

result = pipe(
    responses,
    p.add_form_data(forms),
    p.add_metadata(['clusterid']),
    p.add_duration,
    p.add_time_indicators(['week', 'month']),
    p.count_invalid,
    p.keep_final_answer,
    p.drop_users_without('clusterid'),
    p.drop_duplicated_users(['wave']),
    p.pivot('response'),
)
```

The output is a wide-format DataFrame with one row per respondent per survey, where each `question_ref` value becomes a column.

## pandas or Polars: type in, type out

Every function and method accepts either a `pd.DataFrame` or a `pl.DataFrame` and returns the same type.

- **Polars in, Polars out.** The work is done in Polars and nothing is converted to pandas. This is the fast, low-memory path: use it for large exports.

  ```python
  import polars as pl

  responses = pl.read_csv('responses.csv')
  forms = pl.read_csv('forms.csv')
  result = pipe(responses, p.add_form_data(forms), p.keep_final_answer, p.pivot('response'))
  # result is a pl.DataFrame
  ```

- **pandas in, pandas out.** The columns a step reads are converted to Polars, the step runs there, and the result is assembled from your own DataFrame's rows. Values, dtypes, missing-value markers (`None` vs `NaN`) and index labels are the same as in vlab_prepro 0.5.x, so existing pandas code keeps working.

`add_form_data` accepts a form table of either type; it is converted to the type of the responses.

Differences from 0.5.x, for both types:

- `pivot` orders columns deterministically: the key columns sorted by name, then the question columns sorted.
- `add_duration` and `pivot` order rows deterministically (by the key columns, `userid` first for `pivot`). 0.5.x's order there depended on Python's set iteration order and could change from one run to the next. The rows themselves are the same.
- Functions no longer modify the DataFrame you pass in (0.5.x's `add_metadata` added its columns to its input as well as returning them).
- Timestamps carrying a UTC offset other than `+00:00` are converted to UTC; 0.5.x kept the offset as a fixed-offset time zone.

## API Reference

### `Preprocessor`

The main class. All methods are curried, so they can be called with partial arguments to produce a single-argument function suitable for use in `toolz.pipe`.

```python
p = Preprocessor()

# Fully applied — transforms df immediately
p.add_metadata(['clusterid'], df)

# Partially applied — returns a function that takes df
step = p.add_metadata(['clusterid'])
result = step(df)

# In a pipe
pipe(df, p.add_metadata(['clusterid']), ...)
```

---

#### `add_form_data(form_df, df, prefix=None)`

Merges form-level metadata into the response DataFrame. The `metadata` JSON column in `form_df` is flattened — each key becomes its own column. Merged on `surveyid`.

```python
# Without prefix: metadata key "wave" becomes column "wave"
p.add_form_data(forms, df)

# With prefix: metadata key "wave" becomes column "form_wave"
p.add_form_data(forms, df, prefix='form')
```

Adds all merged columns (including `surveyid`, `shortcode`, `version`, and any metadata keys) to the internal key set used when pivoting.

---

#### `remove_form_data(df)`

Reverses `add_form_data`, dropping all columns that were added by it. Raises `PreprocessingError` if no form data has been added. Useful when you want to pivot across all surveys combined rather than per-survey.

---

#### `add_metadata(keys, df)`

Extracts specific keys from the per-row `metadata` JSON column into individual columns.

```python
p.add_metadata(['clusterid', 'region'], df)
```

Each key is added as a column and registered in the internal key set for pivoting. A key that is also a `question_ref` is added as `<key>_metadata`, so it cannot collide with the question's column after pivoting. A row without the key gets a missing value.

---

#### `parse_timestamp(df)`

Converts the `timestamp` column from ISO 8601 strings (such as CockroachDB's `2026-10-14 13:31:45.123456+00:00`) to datetimes. Strings with a UTC offset become time-zone-aware (UTC) datetimes; strings without one stay naive. Called automatically by `add_duration` and `add_time_indicators` if needed.

---

#### `add_duration(df)`

Calculates timing metrics for each respondent's survey session and adds them as columns. Groups by `userid` (and `surveyid` if form data has been added).

New columns added:

| Column | Description |
|---|---|
| `survey_start_time` | Timestamp of the first answer |
| `survey_end_time` | Timestamp of the last answer |
| `survey_duration` | Total seconds between first and last answer |
| `answer_time_min` | Minimum seconds between consecutive answers |
| `answer_time_median` | Median seconds between consecutive answers |
| `answer_time_75` | 75th percentile seconds between answers |
| `answer_time_90` | 90th percentile seconds between answers |

All timing columns are added to the internal key set for pivoting.

Rows come back sorted by the key columns (by name), then by time.

---

#### `add_time_indicators(indicators, df)`

Adds calendar-based grouping columns derived from `survey_start_time`. Requires `add_duration` to have been called first.

```python
p.add_time_indicators(['week', 'month'], df)
```

Supported indicators:

| Value | Column added | Description |
|---|---|---|
| `'week'` | `week` | ISO week number of the year |
| `'month'` | `month` | Month number (1–12) |

---

#### `add_final_answer(df)`

Adds a boolean `final_answer` column. In chatbot surveys, respondents may re-answer a question multiple times (e.g. after an invalid entry). This marks the chronologically last answer for each `(userid, surveyid, question_ref)` combination as `True` and earlier answers as `False`.

---

#### `keep_final_answer(df)`

Filters the DataFrame to only final answers, removing all earlier attempts. Calls `add_final_answer` automatically if the column does not exist. Use this before pivoting to ensure each question has at most one answer per respondent per survey.

---

#### `count_invalid(df)`

Calculates answer correction rates per user across their entire history in the dataset.

New columns added:

| Column | Description |
|---|---|
| `invalid_answer_count` | Number of non-final answers submitted by the user |
| `invalid_answer_percentage` | Fraction of all answers that were not final (0.0–1.0) |

Both columns are added to the internal key set for pivoting.

---

#### `drop_users_without(metadata_key, df)`

Removes all rows for any user who ever submitted an answer without a value for `metadata_key`. This catches test users or misconfigured sessions — if a user answered even once without the key, all of their data is removed.

```python
p.drop_users_without('clusterid', df)
```

Raises `PreprocessingError` if the column does not exist (call `add_metadata` first). Logs a warning with the count of removed users.

---

#### `drop_duplicated_users(form_keys, df)`

Removes all rows for any user who has more than one distinct `flowid` for the same form. Identifies "the same form" using the columns in `form_keys`.

```python
# Remove users who took the same wave twice
p.drop_duplicated_users(['wave'], df)

# Remove users who took the same shortcode twice
p.drop_duplicated_users(['shortcode'], df)
```

Logs a warning with the count of removed users.

---

#### `pivot(answer_column, df)`

Pivots the long-format DataFrame to wide format. Each unique `question_ref` value becomes a column, populated with values from `answer_column`. The index is all columns in the internal key set.

```python
p.pivot('response', df)
```

The result has one row per unique combination of key columns (e.g. one row per `userid`/`surveyid` pair if form data was added). Rows are sorted by `userid` (then by the other key columns); columns are the key columns sorted by name, then the question columns sorted.

Raises `PreprocessingError` if the pivot fails, which typically means there are duplicate `(userid, surveyid, question_ref)` entries — use `keep_final_answer` first.

Logs a warning if `surveyid` is not in the key set, since `question_ref` values must then be unique across all surveys.

---

#### `map_columns(cols, fn, df)`

Applies a function to one or more columns.

```python
p.map_columns(['age', 'income'], int, df)
```

---

#### `hash_userid(df)`

Replaces the `userid` column with its SHA-256 hash (hex string). Use for anonymization before sharing data.

---

### `compute_seed(seed, n=None, m=0, *, key=None)`

Calculates a respondent's randomization group assignment from the base seed included in the data export. Replicates the JavaScript implementation used by the survey platform.

```python
from vlab_prepro import compute_seed

# 2-arm trial (treatment/control)
df['arm'] = df['seed'].apply(lambda s: compute_seed(s, n=2))

# 3-arm trial using key format
df['arm'] = df['seed'].apply(lambda s: compute_seed(s, key='seed_3'))
```

**Parameters:**

| Parameter | Type | Description |
|---|---|---|
| `seed` | `int` | Base seed from the data export (32-bit integer) |
| `n` | `int` | Number of groups; result is in range `[1, n]` |
| `m` | `int` | Rehash count for independent randomizations (default `0`) |
| `key` | `str` | Key string in format `"seed_N"` or `"seed_N_M"` |

Either `n` or `key` must be provided. `key` is parsed to extract `n` and `m` — `"seed_3_1"` means `n=3, m=1`.

**Multiple independent randomizations:**

When a survey uses multiple randomized branches, each uses a different `m` value to generate independent assignments from the same base seed:

```python
# Independent coin flip
df['treatment'] = df['seed'].apply(lambda s: compute_seed(s, key='seed_2'))

# Independent 3-way split
df['variant'] = df['seed'].apply(lambda s: compute_seed(s, key='seed_3_1'))

# Another independent coin flip
df['image'] = df['seed'].apply(lambda s: compute_seed(s, key='seed_2_2'))
```

---

### `parse_number(s)`

Parses a number string using loose validation that matches the chatbot's own number validator. Strips commas and periods before converting to integer.

```python
from vlab_prepro import parse_number

parse_number('500')       # 500
parse_number('500,000')   # 500000  (commas stripped)
parse_number('500.00')    # 50000   (periods stripped)
parse_number(500)         # 500     (already an int)
parse_number('abc')       # None
parse_number(None)        # None
```

> Note: this intentionally does not distinguish thousands separators from decimal points — `"500.00"` and `"500,00"` both become `50000`.

---

### `PreprocessingError`

Custom exception raised for unrecoverable preprocessing problems:

- Attempting to remove form data that was never added
- Calling `drop_users_without` with a column that does not exist
- Pivot failures due to duplicate `(userid, question_ref)` entries

```python
from vlab_prepro import PreprocessingError

try:
    result = p.pivot('response', df)
except PreprocessingError as e:
    print(f"Preprocessing failed: {e}")
```

