pyantz.jobs.analysis package

pyantz.jobs.analysis package#

Submodules#

pyantz.jobs.analysis.external module#

Jobs that call user-defined Python functions for data extraction and analysis.

save_from_external_extraction() calls a user function and saves the returned Polars DataFrame to a Parquet file.

external_simple() calls a function and returns its boolean result.

external_analysis() loads one or more Parquet files into Polars frames, then passes them as keyword arguments to a user function.

class ExternalAnalysisParams(*, variable_names_to_parquets: ~collections.abc.Mapping[str, ~pathlib.Path], lazy_or_eager: ~typing.Literal['lazy', 'eager'] | ~collections.abc.Mapping[str, ~typing.Literal['lazy', 'eager']], function_args: list[~typing.Any] = <factory>, function_kwargs: dict[str, ~typing.Any] = <factory>, external_function: ~collections.abc.Callable[..., ~typing.Any])[source]#

Bases: BaseModel

Params to pass to an external function which will perform analysis on data.

Parameters:

variable_names_to_parquets (Mapping[str, Path])
lazy_or_eager (Literal['lazy', 'eager'] | ~collections.abc.Mapping[str, ~typing.Literal['lazy', 'eager']])
function_args (list[Any])
function_kwargs (dict[str, Any])
external_function (Annotated[Callable[[...], Any], BeforeValidator(func=~pyantz.infrastructure.config.fn_utils.import_function_by_name, json_schema_input_type=PydanticUndefined), PlainSerializer(func=~pyantz.infrastructure.config.fn_utils.serialize_function, return_type=str, when_used=always), WithJsonSchema(json_schema={'type': 'string', 'format': 'base64'}, mode=None)])

_abc_impl = <_abc._abc_data object>#

external_function: Callable[..., Any]#: Fully resolved name of the function, e.g. package.module.function.

function_args: list[Any]#: Positional arguments to pass directly to the function.

function_kwargs: dict[str, Any]#: Keyword arguments to pass directly to the function.

lazy_or_eager: Literal['lazy', 'eager'] | Mapping[str, Literal['lazy', 'eager']]#: If lazy, all dataframes are scanned lazily (polars LazyFrames); if eager, all dataframes are read into memory (polars DataFrames); if a mapping, maps each name of variable_names_to_parquets to one of those.

variable_names_to_parquets: Mapping[str, Path]#: Maps keyword argument names to parquet files; the parquet files are read into polars LazyFrames or DataFrames and passed as keyword arguments to the analysis function.

class ExternalExtractParams(*, function_args: list[~typing.Any] = <factory>, function_kwargs: dict[str, ~typing.Any] = <factory>, external_function: ~collections.abc.Callable[..., ~polars.dataframe.frame.DataFrame | ~polars.lazyframe.frame.LazyFrame], result_parquet_location: str)[source]#

Bases: BaseModel

Parameters to call an external function returning a polars dataframe.

Parameters:

function_args (list[Any])
function_kwargs (dict[str, Any])
external_function (Annotated[Callable[[...], DataFrame | LazyFrame], BeforeValidator(func=~pyantz.infrastructure.config.fn_utils.import_function_by_name, json_schema_input_type=PydanticUndefined), PlainSerializer(func=~pyantz.infrastructure.config.fn_utils.serialize_function, return_type=str, when_used=always), WithJsonSchema(json_schema={'type': 'string', 'format': 'base64'}, mode=None)])
result_parquet_location (str)

_abc_impl = <_abc._abc_data object>#

external_function: Callable[..., DataFrame | LazyFrame]#: Fully resolved name of the function, e.g. package.module.function; it must return a polars DataFrame or LazyFrame.

function_args: list[Any]#: Positional arguments to pass directly to the function.

function_kwargs: dict[str, Any]#: Keyword arguments to pass directly to the function.

result_parquet_location: str#: Location to save the result from the function.

class ExternalSimpleParams(*, function_args: list[~typing.Any] = <factory>, function_kwargs: dict[str, ~typing.Any] = <factory>, external_function: ~collections.abc.Callable[..., ~typing.Any])[source]#

Bases: BaseModel

Call an external function.

Parameters:

function_args (list[Any])
function_kwargs (dict[str, Any])
external_function (Annotated[Callable[[...], Any], BeforeValidator(func=~pyantz.infrastructure.config.fn_utils.import_function_by_name, json_schema_input_type=PydanticUndefined), PlainSerializer(func=~pyantz.infrastructure.config.fn_utils.serialize_function, return_type=str, when_used=always), WithJsonSchema(json_schema={'type': 'string', 'format': 'base64'}, mode=None)])

_abc_impl = <_abc._abc_data object>#

external_function: Callable[..., Any]#: Fully resolved name of the function, e.g. package.module.function.

function_args: list[Any]#: Positional arguments to pass directly to the function.

function_kwargs: dict[str, Any]#: Keyword arguments to pass directly to the function.

external_analysis(params: ExternalAnalysisParams)[source]#

Call some external function with the data loaded into polars.

Each parquet of variable_names_to_parquets is loaded and passed to the function as the keyword argument of its name, after function_args and along with function_kwargs. The function’s result is discarded.

Parameters:

params (ExternalAnalysisParams) – the job’s parameters, validated as ExternalAnalysisParams.

Job parameters:

variable_names_to_parquets (Mapping[str, Path]) – keyword argument names mapped to the parquet files to load for them.
lazy_or_eager (str | Mapping[str, str]) – lazy (LazyFrames), eager (DataFrames), or one of those for each keyword argument name.
function_args (list[Any]) – positional arguments for the function. Defaults to none.
function_kwargs (dict[str, Any]) – more keyword arguments for the function. Defaults to none.
external_function (str) – import path of the function to call.

Returns:

True if the function ran, False if the data can’t be loaded (or a name has no lazy_or_eager entry) or the function raises.

Return type:

bool

Raises:

None – failures are logged and reported by returning False.

Example

import polars as pl

from pyantz.jobs.analysis.external import external_analysis

pl.DataFrame({"a": [1, 2, 3]}).write_parquet("numbers.parquet")

# in a configuration, the function is given by its import path
def report_total(numbers, label):
    print(label, numbers["a"].sum())

external_analysis(
    {
        "variable_names_to_parquets": {"numbers": "numbers.parquet"},
        "lazy_or_eager": "eager",
        "function_kwargs": {"label": "total:"},
        "external_function": report_total,
    },
    lambda job: None,
)

total: 6

external_simple(params: ExternalSimpleParams)[source]#

Call the external function.

Parameters:

params (ExternalSimpleParams) – the job’s parameters, validated as ExternalSimpleParams.

Job parameters:

function_args (list[Any]) – positional arguments for the function. Defaults to none.
function_kwargs (dict[str, Any]) – keyword arguments for the function. Defaults to none.
external_function (str) – import path of the function to call.

Returns:

the function’s result if it returns a bool, False if it raises, and True otherwise.

Return type:

bool

Raises:

None – failures are logged and reported by returning False.

Example

from pyantz.jobs.analysis.external import external_simple

def submit(job):
    pass

for args in ([1, 1], [1, 2]):
    params = {"external_function": "operator.eq", "function_args": args}
    print(external_simple(params, submit))

True
False

save_from_external_extraction(params: ExternalExtractParams)[source]#

Call a user defined function and save the results off.

While the user could directly call their function, the utility of this function is saving it in the parquet file, allowing the user to ignore the file system fun. Missing parent directories of the parquet are created.

Parameters:

params (ExternalExtractParams) – the job’s parameters, validated as ExternalExtractParams.

Job parameters:

function_args (list[Any]) – positional arguments for the function. Defaults to none.
function_kwargs (dict[str, Any]) – keyword arguments for the function. Defaults to none.
external_function (str) – import path of the function, which must return a polars DataFrame or LazyFrame.
result_parquet_location (str) – parquet file to save the result to.

Returns:

True if the result was saved, False if the function raises, returns something other than a dataframe, or the parquet can’t be written.

Return type:

bool

Raises:

None – failures are logged and reported by returning False.

Example

import polars as pl

from pyantz.jobs.analysis.external import save_from_external_extraction

save_from_external_extraction(
    {
        "external_function": "polars.DataFrame",
        "function_kwargs": {"data": {"a": [1, 2, 3]}},
        "result_parquet_location": "out/extracted.parquet",
    },
    lambda job: None,
)
print(pl.read_parquet("out/extracted.parquet")["a"].to_list())

[1, 2, 3]

pyantz.jobs.analysis.simple module#

Simple columnar operations on Parquet data files.

columnar_operation() applies an arithmetic operation between two columns and stores the result in a new column.

join_parquets() performs a Polars join between two Parquet files and writes the output to a third file.

class ColumnarOperationParameters(*, source_parquet_file: Path, result_parquet_file: str, left_col: str, right_col: str, result_col: str, operation: Literal['+', '-', '*', '/'])[source]#

Bases: BaseModel

Perform an operation and save the resulting dataframe in a parquet.

Parameters:

source_parquet_file (Path)
result_parquet_file (str)
left_col (str)
right_col (str)
result_col (str)
operation (Literal['+', '-', '*', '/'])

_abc_impl = <_abc._abc_data object>#

left_col: str#: Column on the left side of the expression.

operation: Literal['+', '-', '*', '/']#: Operation to be performed between the two columns.

result_col: str#: Column to save the results.

result_parquet_file: str#: Location to save the results after performing the operation.

right_col: str#: Column on the right side of the expression.

source_parquet_file: Path#: Parquet file to read dataframe from for these operations.

class JoinParquetParams(*, left_parquet: Path, right_parquet: Path, how: Literal['inner', 'left', 'right', 'full', 'semi', 'anti', 'cross'] = 'inner', output_parquet: Path, on_: str)[source]#

Bases: BaseModel

Parameters to join the dataframes from parquet files.

Parameters:

left_parquet (Path)
right_parquet (Path)
how (Literal['inner', 'left', 'right', 'full', 'semi', 'anti', 'cross'])
output_parquet (Path)
on_ (str)

_abc_impl = <_abc._abc_data object>#

check_file_paths()[source]#

Check that output is not the same as the input file.

Parameters:: None.
Returns:: the validated parameters.
Return type:: Self
Raises:: ValueError – if the output parquet is the left or right parquet.

how: Literal['inner', 'left', 'right', 'full', 'semi', 'anti', 'cross']#: Supported styles of joins (see polars docs for details).

left_parquet: Path#: Parquet with tabular data to join, on the left of the join operation.

on_: str#: Column name to join on.

output_parquet: Path#: Parquet file to save the results, cannot be the same as left/right.

right_parquet: Path#: Parquet with tabular data to join, on the right of the join operation.

OP_MAPPING = {'*': <function Expr.mul>, '+': <function Expr.add>, '-': <function Expr.sub>, '/': <function Expr.truediv>}#: Polars expression method for each operation of columnar_operation().

columnar_operation(params: ColumnarOperationParameters)[source]#

Perform a simple operation between two columns.

The source parquet is read lazily, the result of left_col <operation> right_col is added as result_col, and the dataframe with the new column is written to result_parquet_file.

Parameters:

params (ColumnarOperationParameters) – the job’s parameters, validated as ColumnarOperationParameters.

Job parameters:

source_parquet_file (Path) – parquet file to read the dataframe from.
result_parquet_file (str) – parquet file to write the result to.
left_col (str) – column on the left side of the expression.
right_col (str) – column on the right side of the expression.
result_col (str) – column to save the results in.
operation (str) – +, -, * or /.

Returns:

True if the result was written, False if the source doesn’t exist, lacks a column, or can’t be read or written.

Return type:

bool

Raises:

None – failures are logged and reported by returning False.

Example

import os

import polars as pl

from pyantz.jobs.analysis.simple import columnar_operation


assert os.path.exists("test_input.parquet")
columnar_operation({
    "source_parquet_file": "test_input.parquet",
    "result_parquet_file": "test_output.parquet",
    "left_col": "a",
    "right_col": "b",
    "result_col": "c",
    "operation": "+",
}, lambda *_: None)

result = pl.read_parquet("test_output.parquet")
print(result)

shape: (2, 3)
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 3   ┆ 4   │
│ 2   ┆ 4   ┆ 6   │
└─────┴─────┴─────┘

join_parquets(params: JoinParquetParams)[source]#

Join two dataframes read from a parquet.

The parquets are read lazily and the joined dataframe is streamed into the output parquet. A join doesn’t keep the order of the rows.

Parameters:

params (JoinParquetParams) – the job’s parameters, validated as JoinParquetParams.

Job parameters:

left_parquet (Path) – parquet on the left of the join.
right_parquet (Path) – parquet on the right of the join.
how (str) – style of join, one of inner, left, right, full, semi, anti or cross (see polars’ docs). Defaults to inner.
output_parquet (Path) – parquet to save the result to; it can’t be the left or right parquet.
on_ (str) – name of the column to join on.

Returns:

True if the result was written, False if the join fails (for example, a parquet or column doesn’t exist).

Return type:

bool

Raises:

None – failures are logged and reported by returning False.

Example

from pyantz import start
config = {
    "jobs": [
        {
            "function": "pyantz.jobs.analysis.simple.join_parquets",
            "parameters": {
                "left_parquet": "left.parquet",
                "right_parquet": "right.parquet",
                "how": "inner",
                "output_parquet": "output.parquet",
                "on_": "idx",
            },
        }
    ],
    "submitter": {
        "type_": "local_proc",
        "working_directory": ".",
    },
}
start(config)

# a join doesn't keep the order of the rows
result = pl.read_parquet("output.parquet").sort("idx")
print(result)

shape: (2, 3)
┌─────┬─────┬───────┐
│ idx ┆ a   ┆ b     │
│ --- ┆ --- ┆ ---   │
│ i64 ┆ i64 ┆ str   │
╞═════╪═════╪═══════╡
│ 1   ┆ 10  ┆ hello │
│ 4   ┆ 13  ┆ world │
└─────┴─────┴───────┘

Module contents#

Jobs for performing simple data analysis.

pyantz.jobs.analysis package

Contents

pyantz.jobs.analysis package#

Submodules#

pyantz.jobs.analysis.external module#

pyantz.jobs.analysis.simple module#

Module contents#