cursus.steps.scripts.temporal_split_preprocessing

Temporal Split Preprocessing Script

Comprehensive preprocessing with temporal splitting capabilities: 1. Temporal cutoff (date-based split for OOT test) 2. Customer-level random split (train/validation) 3. Ensures no customer leakage between train and OOT 4. Parallel processing for large datasets 5. Signature file support 6. Memory-efficient batch concatenation 7. Multiple output formats (CSV, TSV, Parquet)

load_signature_columns(signature_path)

Load column names from signature file.

Parameters:

signature_path (str) – Path to the signature file directory

Returns:

List of column names if signature file exists, None otherwise

Return type:

list | None

peek_json_format(file_path, open_func=<built-in function open>)

Check if the JSON file is in JSON Lines or regular format.

combine_shards(input_dir, signature_columns=None, max_workers=None, batch_size=10)

Detect and combine all supported data shards in a directory using parallel processing.

Uses parallel shard reading and batch concatenation for improved performance. Memory-efficient approach avoids PyArrow’s 2GB column limit error.

Parameters:

• input_dir (str) – Directory containing data shards

• signature_columns (list | None) – Optional column names for CSV/TSV files

• max_workers (int | None) – Maximum number of parallel workers (default: cpu_count)

• batch_size (int) – Number of DataFrames to concatenate at once (default: 10)

Returns:

Combined DataFrame from all shards

Return type:

DataFrame

generate_main_task_label(df, targets, main_task_index=0, logger=None)

Generate main task label based on subtasks by taking the maximum value across subtasks.

The main task label is set as the maximum value of all subtasks for each sample. For example, if targets=[‘is_abuse’,’is_abusive_dnr’,’is_abusive_pda’,’is_abusive_rr’] and main_task_index=0, then ‘is_abuse’ will be set as the max of the other subtasks.

Parameters:

• df (DataFrame) – Input DataFrame

• targets (list) – List of target column names (main task + subtasks)

• main_task_index (int) – Index of the main task in the targets list (default: 0)

• logger (Callable[[str], None] | None) – Optional logger function

Returns:

DataFrame with updated main task labels

Return type:

DataFrame

Example

>>> targets = ['is_abuse','is_abusive_dnr','is_abusive_pda','is_abusive_rr']
>>> df = generate_main_task_label(df, targets, main_task_index=0)
# 'is_abuse' will be set as max('is_abusive_dnr', 'is_abusive_pda', 'is_abusive_rr')

temporal_customer_split(df, date_column, group_id_column, split_date, train_ratio=0.9, random_seed=42, logger=None)

Split data temporally and by group ID.

Parameters:

• df (DataFrame) – Input DataFrame

• date_column (str) – Name of the date column

• group_id_column (str) – Name of the group ID column

• split_date (str) – Date string for temporal split (format: YYYY-MM-DD)

• train_ratio (float) – Ratio of groups for training (rest go to validation)

• random_seed (int) – Random seed for reproducibility

• logger (Callable[[str], None] | None) – Optional logger function

Returns:

Dictionary with ‘train’, ‘val’, and ‘oot’ DataFrames

Return type:

Dict[str, DataFrame]

process_batch_mode_temporal_split(input_data_dir, signature_columns, date_column, group_id_column, split_date, train_ratio, random_seed, targets_str, main_task_index, label_field, output_format, max_workers, batch_size, output_paths, log_func)

Process temporal split in batch mode.

Loads all data into memory, applies temporal split logic, and saves outputs.

Parameters:

• input_data_dir (str) – Input data directory

• signature_columns (list | None) – Optional signature columns

• date_column (str) – Name of date column

• group_id_column (str) – Name of group ID column

• split_date (str) – Temporal split date

• train_ratio (float) – Train/val ratio

• random_seed (int) – Random seed

• targets_str (str | None) – Optional targets string for multi-task

• main_task_index (int | None) – Optional main task index

• label_field (str | None) – Optional label field

• output_format (str) – Output format

• max_workers (int) – Max parallel workers

• batch_size (int) – Batch concatenation size

• output_paths (Dict[str, str]) – Output path dictionary

• log_func (Callable) – Logging function

Returns:

Dictionary with training_data and oot_data DataFrames

Return type:

Dict[str, DataFrame]

find_input_shards(input_dir, log_func)

Find all input shards in directory.

extract_shard_number(shard_path)

Extract shard number from filename like part-00042.csv.

Handles various formats: - part-00042.csv → 42 - part-00042.csv.gz → 42 - part-00042.parquet → 42 - part-00042.snappy.parquet → 42

Parameters:

shard_path (Path) – Path to shard file

Returns:

Integer shard number

Raises:

ValueError – If shard number cannot be extracted

Return type:

int

Example

>>> extract_shard_number(Path("part-00042.csv"))
42
>>> extract_shard_number(Path("part-00001.csv.gz"))
1

collect_customer_allocation(all_shards, signature_columns, date_column, group_id_column, split_date, train_ratio, random_seed, log_func)

Pass 1: Scan pre-split data and build complete customer allocation map.

Memory: O(unique_customers) - typically ~50MB for 1M customers Scans all shards to collect unique customer IDs from pre-split period, then randomly allocates them to train vs validation sets.

Parameters:

• all_shards (List[Path]) – List of all input shard paths

• signature_columns (list | None) – Optional column names

• date_column (str) – Name of date column for temporal filtering

• group_id_column (str) – Name of customer/group ID column

• split_date (str) – Date string for temporal cutoff (YYYY-MM-DD)

• train_ratio (float) – Proportion of customers for training

• random_seed (int) – Random seed for reproducibility

• log_func (Callable) – Logging function

Returns:

Dictionary mapping customer_id to split assignment (“train” or “val”) Example: {“customer_1”: “train”, “customer_2”: “val”, …}

Return type:

Dict[str, str]

write_single_shard(df, output_dir, shard_number, output_format='csv')

Write a single data shard in the specified format.

Note: Kept for potential future use, but not currently used by streaming mode.

process_shard_temporal_split(args)

Process single shard with global customer allocation map. Uses VECTORIZED operations for performance (not df.apply).

This function runs in parallel across multiple workers.

Parameters:

args (tuple) – Tuple of (shard_path, shard_num, customer_split_map, config)

Returns:

Statistics dict with row counts per split

Return type:

Dict[str, int]

process_streaming_temporal_split_parallel(all_shards, training_output_path, customer_split_map, config, max_workers, log_func)

Pass 2: Process all shards in parallel using customer map.

Parameters:

• all_shards (List[Path]) – List of all input shard paths

• training_output_path (Path) – Base output directory

• customer_split_map (Dict[str, str]) – Customer→split mapping from Pass 1

• config (Dict) – Configuration dictionary

• max_workers (int) – Number of parallel workers

• log_func (Callable) – Logging function

consolidate_shards_to_single_files(training_output_path, output_format, log_func)

Consolidate temporary shards into single files per split.

main(input_paths, output_paths, environ_vars, job_args, logger=None)

Main logic for temporal split preprocessing with comprehensive features, refactored for testability.

Parameters:

• input_paths (Dict[str, str]) – Dictionary of input paths with logical names

• output_paths (Dict[str, str]) – Dictionary of output paths with logical names

• environ_vars (Dict[str, str]) – Dictionary of environment variables

• job_args (Namespace) – Command line arguments

• logger (Callable[[str], None] | None) – Optional logger object (defaults to print if None)

Returns:

Dictionary of DataFrames by split name (e.g., ‘train’, ‘val’, ‘oot’)

Return type:

Dict[str, DataFrame]