Bases: PipelineAction
Add or update multiple columns in the DataFrame using SQL expressions.
This action uses PySpark's expr() function to evaluate SQL expressions and
create or update one or more columns in the DataFrame.
Examples:
Source code in src/cloe_nessy/pipeline/actions/transform_with_column.py
| class TransformWithColumnAction(PipelineAction):
"""Add or update multiple columns in the DataFrame using SQL expressions.
This action uses PySpark's expr() function to evaluate SQL expressions and
create or update one or more columns in the DataFrame.
Examples:
=== "Single column transformation"
```yaml
Create Full Name:
action: TRANSFORM_WITH_COLUMN
options:
columns:
full_name: concat(first_name, ' ', last_name)
```
=== "Multiple column transformations"
```yaml
Transform Multiple Columns:
action: TRANSFORM_WITH_COLUMN
options:
columns:
full_name: concat(first_name, ' ', last_name)
email_lower: lower(email)
total_price: price * quantity
year: year(order_date)
```
=== "Complex expressions"
```yaml
Calculate Fields:
action: TRANSFORM_WITH_COLUMN
options:
columns:
total_amount: price * quantity
discount_amount: total_amount * discount_rate
final_amount: total_amount - discount_amount
status_flag: case when final_amount > 1000 then 'high' else 'normal' end
```
=== "Date transformations"
```yaml
Extract Date Parts:
action: TRANSFORM_WITH_COLUMN
options:
columns:
order_year: year(order_date)
order_month: month(order_date)
order_day: day(order_date)
order_quarter: quarter(order_date)
```
"""
name: str = "TRANSFORM_WITH_COLUMN"
def run(
self,
context: PipelineContext,
*,
columns: dict[str, str] | None = None,
**_: Any,
) -> PipelineContext:
"""Add or update one or more columns using SQL expressions.
Args:
context: The pipeline context containing the DataFrame
columns: Dictionary mapping column names to SQL expressions.
Format: {column_name: sql_expression}
**_: Additional unused keyword arguments
Returns:
PipelineContext: Updated context with the modified DataFrame
Raises:
ValueError: If columns is not provided or is empty
ValueError: If context.data is None
Exception: If any SQL expression is invalid
Example:
```python
action.run(
context=pipeline_context,
columns={
"full_name": "concat(first_name, ' ', last_name)",
"year": "year(order_date)",
"total": "price * quantity"
}
)
```
"""
if not columns:
raise ValueError("No columns provided. Expected a dictionary of {column_name: expression}.")
if not isinstance(columns, dict):
raise ValueError("columns must be a dictionary mapping column names to SQL expressions.")
if context.data is None:
raise ValueError("Data from context is required for transform_with_column")
df = context.data
column_count = len(columns)
self._console_logger.info(f"Adding/updating {column_count} column(s)")
for column_name, expression in columns.items():
if not column_name or not isinstance(column_name, str):
raise ValueError(f"Invalid column name: {column_name}. Column names must be non-empty strings.")
if not expression or not isinstance(expression, str):
raise ValueError(
f"Invalid expression for column '{column_name}': {expression}. "
"Expressions must be non-empty strings."
)
try:
self._console_logger.info(f" - Adding/updating column '{column_name}' with expression: {expression}")
df = df.withColumn(column_name, F.expr(expression))
except Exception as e:
self._console_logger.error(
f"Failed to evaluate expression '{expression}' for column '{column_name}': {e}"
)
raise
self._console_logger.info(f"Successfully added/updated {column_count} column(s)")
return context.from_existing(data=df)
|
Add or update one or more columns using SQL expressions.
Parameters:
| Name |
Type |
Description |
Default |
context
|
PipelineContext
|
The pipeline context containing the DataFrame
|
required
|
columns
|
dict[str, str] | None
|
Dictionary mapping column names to SQL expressions.
Format: {column_name: sql_expression}
|
None
|
**_
|
Any
|
Additional unused keyword arguments
|
{}
|
Returns:
| Name | Type |
Description |
PipelineContext |
PipelineContext
|
Updated context with the modified DataFrame
|
Raises:
| Type |
Description |
ValueError
|
If columns is not provided or is empty
|
ValueError
|
|
Exception
|
If any SQL expression is invalid
|
Example
action.run(
context=pipeline_context,
columns={
"full_name": "concat(first_name, ' ', last_name)",
"year": "year(order_date)",
"total": "price * quantity"
}
)
Source code in src/cloe_nessy/pipeline/actions/transform_with_column.py
| def run(
self,
context: PipelineContext,
*,
columns: dict[str, str] | None = None,
**_: Any,
) -> PipelineContext:
"""Add or update one or more columns using SQL expressions.
Args:
context: The pipeline context containing the DataFrame
columns: Dictionary mapping column names to SQL expressions.
Format: {column_name: sql_expression}
**_: Additional unused keyword arguments
Returns:
PipelineContext: Updated context with the modified DataFrame
Raises:
ValueError: If columns is not provided or is empty
ValueError: If context.data is None
Exception: If any SQL expression is invalid
Example:
```python
action.run(
context=pipeline_context,
columns={
"full_name": "concat(first_name, ' ', last_name)",
"year": "year(order_date)",
"total": "price * quantity"
}
)
```
"""
if not columns:
raise ValueError("No columns provided. Expected a dictionary of {column_name: expression}.")
if not isinstance(columns, dict):
raise ValueError("columns must be a dictionary mapping column names to SQL expressions.")
if context.data is None:
raise ValueError("Data from context is required for transform_with_column")
df = context.data
column_count = len(columns)
self._console_logger.info(f"Adding/updating {column_count} column(s)")
for column_name, expression in columns.items():
if not column_name or not isinstance(column_name, str):
raise ValueError(f"Invalid column name: {column_name}. Column names must be non-empty strings.")
if not expression or not isinstance(expression, str):
raise ValueError(
f"Invalid expression for column '{column_name}': {expression}. "
"Expressions must be non-empty strings."
)
try:
self._console_logger.info(f" - Adding/updating column '{column_name}' with expression: {expression}")
df = df.withColumn(column_name, F.expr(expression))
except Exception as e:
self._console_logger.error(
f"Failed to evaluate expression '{expression}' for column '{column_name}': {e}"
)
raise
self._console_logger.info(f"Successfully added/updated {column_count} column(s)")
return context.from_existing(data=df)
|