cursus.steps.scripts.tabular_preprocessing

optimize_dtypes(df, log_func=None)

Optimize DataFrame dtypes to reduce memory usage.

Applies the following optimizations: - Downcast numeric types (int64->int32, float64->float32) - Convert object columns with low cardinality to category

Parameters:

• df (DataFrame) – Input DataFrame

• log_func (Callable | None) – Optional logging function

Returns:

DataFrame with optimized dtypes

Return type:

DataFrame

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

process_label_column(df, label_field, log_func)

Process label column: convert to numeric and handle missing values.

Used by both batch and streaming modes.

Parameters:

• df (DataFrame) – DataFrame with label column

• label_field (str) – Name of label column

• log_func (Callable) – Logging function

Returns:

DataFrame with processed labels

Return type:

DataFrame

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, streaming_batch_size=None)

Detect and combine all supported data shards from one or more directories.

Used by BATCH MODE only.

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

Streaming Mode: When streaming_batch_size is set, processes shards in batches to avoid loading all DataFrames into memory simultaneously. This is the most memory-efficient mode.

Parameters:

• input_dir – Directory or list of directories 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)

• streaming_batch_size (int | None) – Number of shards to process per batch (enables streaming mode) - If None: Loads all shards into memory (original behavior) - If set: Processes shards in batches, concatenating incrementally - Recommended: 10-20 shards per batch for memory-constrained environments

Returns:

Combined DataFrame from all shards

Return type:

DataFrame

process_batch_mode_preprocessing(input_data_dir, input_signature_dir, output_dir, signature_columns, job_type, label_field, train_ratio, test_val_ratio, output_format, max_workers, batch_size, streaming_batch_size, optimize_memory, logger=None)

Batch mode for tabular preprocessing.

Loads full DataFrame into memory, applies transformations, and splits data. Uses stratified splits for training jobs when labels are available.

Parameters:

• input_data_dir – Directory or list of directories containing input shards

• input_signature_dir (str) – Directory containing signature file

• output_dir (str) – Base output directory

• signature_columns (list | None) – Optional column names from signature file

• job_type (str) – “training”, “validation”, “testing”, or “calibration”

• label_field (str | None) – Name of label column (optional)

• train_ratio (float) – Training set ratio (for training jobs)

• test_val_ratio (float) – Test/val split ratio (for training jobs)

• output_format (str) – “csv”, “tsv”, or “parquet”

• max_workers (int | None) – Max parallel workers

• batch_size (int) – Batch size for concatenation

• streaming_batch_size (int | None) – Optional incremental loading batch size

• optimize_memory (bool) – Whether to optimize dtypes

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

Returns:

Dictionary of DataFrames by split name

Return type:

Dict[str, DataFrame]

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

Write a single data shard in the specified format.

Used by STREAMING MODE to write temporary shards.

Parameters:

• df (DataFrame) – DataFrame to write

• output_dir (Path) – Output directory

• shard_number (int) – Shard index number

• output_format (str) – “csv”, “tsv”, or “parquet”

Returns:

Path to the written shard file

Return type:

Path

assign_random_splits(df, train_ratio, test_val_ratio)

Randomly assign rows to train/test/val splits.

Vectorized implementation for performance and reliability. For large datasets, random assignment approximates stratified splits well.

Parameters:

• df (DataFrame) – Input DataFrame

• train_ratio (float) – Proportion for training (e.g., 0.7)

• test_val_ratio (float) – Proportion of non-train for test vs val (e.g., 0.5)

Returns:

DataFrame with added ‘_split’ column

Return type:

DataFrame

write_splits_to_shards(df, output_base, split_counters, shard_size, output_format, log_func)

Write DataFrame to separate split directories based on ‘_split’ column.

Parameters:

• df (DataFrame) – DataFrame with ‘_split’ column

• output_base (Path) – Base output directory

• split_counters (Dict[str, int]) – Dictionary tracking shard numbers per split (modified in place)

• shard_size (int) – Rows per shard

• output_format (str) – “csv”, “tsv”, or “parquet”

• log_func (Callable) – Logging function

find_input_shards(input_dir, log_func)

Find all input shards in one or more directories.

process_single_batch(shard_files, signature_columns, batch_size, optimize_memory, label_field, log_func)

Process a single batch of shards.

process_training_splits_streaming(all_shards, output_path, signature_columns, label_field, train_ratio, test_val_ratio, output_format, streaming_batch_size, shard_size, batch_size, optimize_memory, consolidate_shards, log_func)

Process training data with random splits in streaming mode.

Supports two output modes via consolidate_shards parameter: - consolidate_shards=True: Direct write to single files (train/val/test_processed_data.*) - consolidate_shards=False: Write to shards (train/part-., val/part-., test/part-.)

process_single_split_streaming(all_shards, output_path, job_type, signature_columns, label_field, output_format, streaming_batch_size, shard_size, batch_size, optimize_memory, consolidate_shards, log_func)

Process non-training data as single split in streaming mode.

Supports two output modes via consolidate_shards parameter: - consolidate_shards=True: Direct write to single file (job_type_processed_data.*) - consolidate_shards=False: Write to shards (job_type/part-.)

consolidate_shards_to_single_files(output_path, job_type, output_format, log_func)

Consolidate temporary shards into single files per split.

process_streaming_mode_preprocessing(input_dir, output_dir, signature_columns, job_type, label_field, train_ratio, test_val_ratio, output_format, streaming_batch_size, shard_size, max_workers, batch_size, optimize_memory, consolidate_shards=True, logger=None)

True streaming mode for tabular preprocessing.

Processes data in batches, never loading the full DataFrame into memory. For training jobs, uses random split instead of stratified split. For non-training jobs, processes as a single split.

Memory usage: Fixed at ~2GB per batch regardless of total data size.

Parameters:

• input_dir (str) – Directory containing input shards

• output_dir (str) – Base output directory

• signature_columns (list | None) – Optional column names from signature file

• job_type (str) – “training”, “validation”, “testing”, or “calibration”

• label_field (str | None) – Name of label column (optional)

• train_ratio (float) – Training set ratio (for training jobs)

• test_val_ratio (float) – Test/val split ratio (for training jobs)

• output_format (str) – “csv”, “tsv”, or “parquet”

• streaming_batch_size (int) – Number of shards per batch

• shard_size (int) – Rows per output shard

• max_workers (int | None) – Max parallel workers

• batch_size (int) – Batch size for concatenation

• optimize_memory (bool) – Whether to optimize dtypes

• consolidate_shards (bool) – Whether to consolidate output into single files (default: True)

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

Returns:

Empty dictionary (data was written incrementally)

Return type:

Dict[str, DataFrame]

extract_shard_number(shard_path)

Extract numeric shard number from filename.

Parameters:

shard_path (Path) – Path to shard file

Returns:

Shard number as integer

Return type:

int

Examples

part-00042.csv → 42 part-00042.csv.gz → 42 part-00000-68e1f319-…-c000.snappy.parquet → 0

write_shard_file(df, output_path, output_format)

Write DataFrame to file in specified format.

Parameters:

• df (DataFrame) – DataFrame to write

• output_path (Path) – Full path to output file

• output_format (str) – “csv”, “tsv”, or “parquet”

assign_stratified_splits_approximate(df, label_field, train_ratio, test_val_ratio)

Approximate stratified split using per-label deterministic random assignment.

Vectorized implementation for performance and reliability. Each class gets the same split ratios independently, providing approximate stratification without requiring global coordination.

Parameters:

• df (DataFrame) – Input DataFrame with label column

• label_field (str) – Name of label column

• train_ratio (float) – Proportion for training

• test_val_ratio (float) – Proportion of non-train for test vs val

Returns:

DataFrame with ‘_split’ column added

Return type:

DataFrame

process_shard_end_to_end_generic(args)

Process a single shard completely: read → preprocess → split → write.

Generic version with approximate stratification or random splits. No domain-specific preprocessing or global coordination required.

Parameters:

args (tuple) – Tuple of (shard_path, shard_index, config_dict, output_base, signature_columns, label_field, optimize_memory, output_format, train_ratio, test_val_ratio)

Returns:

Statistics dict with row counts per split

Return type:

Dict[str, int]

process_training_streaming_fully_parallel_generic(all_shards, output_path, config_dict, signature_columns, label_field, optimize_memory, output_format, train_ratio, test_val_ratio, max_workers, log_func)

Fully parallel streaming preprocessing for training jobs.

Generic version with approximate stratification or random splits. No global Pass 1 needed - each shard processes independently.

Parameters:

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

• output_path (Path) – Base output directory

• config_dict (Dict) – Configuration dictionary

• signature_columns (list | None) – Optional column names

• label_field (str | None) – Name of label column

• optimize_memory (bool) – Whether to optimize dtypes

• output_format (str) – Output format

• train_ratio (float) – Training set ratio

• test_val_ratio (float) – Test/val split ratio

• max_workers (int) – Number of parallel workers

• log_func (Callable) – Logging function

process_single_split_streaming_fully_parallel_generic(all_shards, output_path, job_type, signature_columns, label_field, optimize_memory, output_format, max_workers, log_func)

Fully parallel preprocessing for single split (validation/testing/calibration).

Parameters:

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

• output_path (Path) – Base output directory

• job_type (str) – Job type (validation/testing/calibration)

• signature_columns (list | None) – Optional column names

• label_field (str | None) – Name of label column

• optimize_memory (bool) – Whether to optimize dtypes

• output_format (str) – Output format

• max_workers (int) – Number of parallel workers

• log_func (Callable) – Logging function

process_fully_parallel_mode_preprocessing_generic(input_dir, output_dir, signature_columns, job_type, label_field, train_ratio, test_val_ratio, output_format, max_workers, optimize_memory, logger=None)

Fully parallel streaming mode with 1:1 shard mapping (generic version).

Uses approximate stratification when labels available, random splits otherwise. No global coordination needed - each shard processes independently.

Parameters:

• input_dir – Directory or list of directories containing input shards

• output_dir (str) – Base output directory

• signature_columns (list | None) – Optional column names

• job_type (str) – Job type

• label_field (str | None) – Name of label column

• train_ratio (float) – Training set ratio

• test_val_ratio (float) – Test/val split ratio

• output_format (str) – Output format

• max_workers (int | None) – Max parallel workers

• optimize_memory (bool) – Whether to optimize dtypes

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

Returns:

Empty dictionary (data written to disk)

Return type:

Dict[str, DataFrame]

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

Main logic for preprocessing data, 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’, ‘test’, ‘val’)

Return type:

Dict[str, DataFrame]