qtype 0.0.16__py3-none-any.whl → 0.1.1__py3-none-any.whl

qtype/interpreter/batch/flow.py

@@ -1,95 +0,0 @@
-from __future__ import annotations
-
-import logging
-from typing import Any, Tuple
-
-import pandas as pd
-
-from qtype.interpreter.batch.step import batch_execute_step
-from qtype.interpreter.batch.types import BatchConfig
-from qtype.interpreter.batch.utils import reconcile_results_and_errors
-from qtype.semantic.model import Flow, Sink
-
-logger = logging.getLogger(__name__)
-
-
-def batch_execute_flow(
-    flow: Flow,
-    inputs: pd.DataFrame,
-    batch_config: BatchConfig,
-    **kwargs: dict[Any, Any],
-) -> Tuple[pd.DataFrame, pd.DataFrame]:
-    """Executes a flow in a batch context.
-
-    Args:
-        flow: The flow to execute.
-        batch_config: The batch configuration to use.
-        **kwargs: Additional keyword arguments to pass to the flow.
-
-    Returns:
-        A list of output variables produced by the flow.
-    """
-
-    previous_outputs = inputs
-
-    all_errors = []
-
-    # Iterate over each step in the flow
-    for step in flow.steps:
-        results: list[pd.DataFrame] = []
-        errors: list[pd.DataFrame] = []
-
-        if isinstance(step, Sink):
-            # Send the entire batch to the sink
-            batch_results, batch_errors = batch_execute_step(
-                step, previous_outputs, batch_config
-            )
-            results.append(batch_results)
-            if len(batch_errors) > 1:
-                errors.append(batch_errors)
-        else:
-            # batch the current data into dataframes of max size batch_size
-            batch_size = batch_config.batch_size
-            for start in range(0, len(previous_outputs), batch_size):
-                end = start + batch_size
-                batch = previous_outputs.iloc[start:end].copy()
-                # Execute the step with the current batch
-                batch_results, batch_errors = batch_execute_step(
-                    step, batch, batch_config
-                )
-
-                results.append(batch_results)
-                if len(batch_errors) > 1:
-                    errors.append(batch_errors)
-
-        previous_outputs, errors_df = reconcile_results_and_errors(
-            results, errors
-        )
-
-        if len(errors_df):
-            all_errors.append(errors_df)
-            if batch_config.write_errors_to:
-                output_file = (
-                    f"{batch_config.write_errors_to}/{step.id}.errors.parquet"
-                )
-                try:
-                    errors_df.to_parquet(
-                        output_file, engine="pyarrow", compression="snappy"
-                    )
-                    logging.info(
-                        f"Saved errors for step {step.id} to {output_file}"
-                    )
-                except Exception as e:
-                    logging.warning(
-                        f"Could not save errors step {step.id} to {output_file}",
-                        exc_info=e,
-                        stack_info=True,
-                    )
-
-    # Return the last steps results and errors
-    rv_errors = (
-        pd.concat(all_errors, ignore_index=True)
-        if len(all_errors)
-        else pd.DataFrame({})
-    )
-    return previous_outputs, rv_errors

qtype/interpreter/batch/sql_source.py

@@ -1,92 +0,0 @@
-from typing import Any, Tuple
-
-import boto3  # type: ignore[import-untyped]
-import pandas as pd
-import sqlalchemy
-from sqlalchemy import create_engine
-from sqlalchemy.exc import SQLAlchemyError
-
-from qtype.base.exceptions import InterpreterError
-from qtype.interpreter.auth.generic import auth
-from qtype.interpreter.batch.types import BatchConfig, ErrorMode
-from qtype.interpreter.batch.utils import reconcile_results_and_errors
-from qtype.semantic.model import SQLSource
-
-
-def to_output_columns(
-    df: pd.DataFrame, output_columns: set[str]
-) -> pd.DataFrame:
-    """Filters the DataFrame to only include specified output columns.
-
-    Args:
-        df: The input DataFrame.
-        output_columns: A set of column names to retain in the DataFrame.
-
-    Returns:
-        A DataFrame containing only the specified output columns.
-    """
-    if len(df) == 0:
-        return df
-    missing = output_columns - set(df.columns)
-    if missing:
-        raise InterpreterError(
-            f"SQL Result was missing expected columns: {','.join(missing)}, it has columns: {','.join(df.columns)}"
-        )
-
-    return df[[col for col in df.columns if col in output_columns]]
-
-
-def execute_sql_source(
-    step: SQLSource,
-    inputs: pd.DataFrame,
-    batch_config: BatchConfig,
-    **kwargs: dict[Any, Any],
-) -> Tuple[pd.DataFrame, pd.DataFrame]:
-    """Executes a SQLSource step to retrieve data from a SQL database.
-
-    Args:
-        step: The SQLSource step to execute.
-
-    Returns:
-        A tuple containing two DataFrames:
-            - The first DataFrame contains the successfully retrieved data.
-            - The second DataFrame contains rows that encountered errors with an 'error' column.
-    """
-    # Create a database engine
-    connect_args = {}
-    if step.auth:
-        with auth(step.auth) as creds:
-            if isinstance(creds, boto3.Session):
-                connect_args["session"] = creds
-    engine = create_engine(step.connection, connect_args=connect_args)
-
-    output_columns = {output.id for output in step.outputs}
-
-    results = []
-    errors = []
-    step_inputs = {i.id for i in step.inputs}
-    for _, row in inputs.iterrows():
-        try:
-            # Make a dictionary of column_name: value from row
-            params = {col: row[col] for col in row.index if col in step_inputs}
-            # Execute the query and fetch the results into a DataFrame
-            with engine.connect() as connection:
-                result = connection.execute(
-                    sqlalchemy.text(step.query),
-                    parameters=params if len(params) else None,
-                )
-                df = pd.DataFrame(
-                    result.fetchall(), columns=list(result.keys())
-                )
-            df = to_output_columns(df, output_columns)
-            # Augment with all input row columns (fan-out-right)
-            df = df.assign(**row.to_dict())
-            results.append(df)
-        except SQLAlchemyError as e:
-            if batch_config.error_mode == ErrorMode.FAIL:
-                raise e
-            # If there's an error, return an empty DataFrame and the error message
-            error_df = pd.DataFrame([{"error": str(e)}])
-            errors.append(error_df)
-
-    return reconcile_results_and_errors(results, errors)

qtype/interpreter/batch/step.py

@@ -1,74 +0,0 @@
-from functools import partial
-from typing import Any, Tuple
-
-import pandas as pd
-
-from qtype.interpreter.batch.file_sink_source import (
-    execute_file_sink,
-    execute_file_source,
-)
-from qtype.interpreter.batch.sql_source import execute_sql_source
-from qtype.interpreter.batch.types import BatchConfig
-from qtype.interpreter.batch.utils import (
-    batch_iterator,
-    single_step_adapter,
-    validate_inputs,
-)
-from qtype.interpreter.exceptions import InterpreterError
-from qtype.semantic.model import (
-    Condition,
-    Decoder,
-    FileSink,
-    FileSource,
-    Flow,
-    PromptTemplate,
-    Search,
-    SQLSource,
-    Step,
-    Tool,
-)
-
-SINGLE_WRAP_STEPS = {Decoder, Condition, PromptTemplate, Search, Tool}
-
-
-def batch_execute_step(
-    step: Step,
-    inputs: pd.DataFrame,
-    batch_config: BatchConfig,
-    **kwargs: dict[str, Any],
-) -> Tuple[pd.DataFrame, pd.DataFrame]:
-    """
-    Executes a given step in a batch processing pipeline.
-
-    Args:
-        step (Step): The step to be executed.
-        inputs (pd.DataFrame): The input data for the step.
-        batch_config (BatchConfig): Configuration for batch processing.
-        **kwargs: Additional keyword arguments.
-
-    Returns:
-        Tuple[pd.DataFrame, pd.DataFrame]: A tuple containing the output results and any rows that returned errors.
-    """
-
-    # ensure the inputs to step are included in the current data frame
-    validate_inputs(inputs, step)
-
-    if isinstance(step, Flow):
-        from qtype.interpreter.batch.flow import batch_execute_flow
-
-        return batch_execute_flow(step, inputs, batch_config, **kwargs)
-    elif isinstance(step, SQLSource):
-        return execute_sql_source(step, inputs, batch_config, **kwargs)
-    elif isinstance(step, FileSource):
-        return execute_file_source(step, inputs, batch_config, **kwargs)
-    elif isinstance(step, FileSink):
-        return execute_file_sink(step, inputs, batch_config, **kwargs)
-    elif type(step) in SINGLE_WRAP_STEPS:
-        return batch_iterator(
-            f=partial(single_step_adapter, step=step),
-            batch=inputs,
-            batch_config=batch_config,
-        )
-    # TODO: implement batching for multi-row steps. For example, llm inference can be sped up in batch...
-    else:
-        raise InterpreterError(f"Unsupported step type: {type(step).__name__}")

qtype/interpreter/batch/types.py

@@ -1,41 +0,0 @@
-from __future__ import annotations
-
-from enum import Enum
-
-from pydantic import BaseModel, Field
-
-
-class ErrorMode(str, Enum):
-    """Error handling mode for batch processing."""
-
-    FAIL = "fail"
-    DROP = "drop"
-
-
-class BatchConfig(BaseModel):
-    """Configuration for batch execution.
-
-    Attributes:
-        num_workers: Number of async workers for batch operations.
-        batch_size: Maximum number of rows to send to a step at a time.
-        error_mode: Error handling mode for batch processing.
-    """
-
-    num_workers: int = Field(
-        default=4,
-        description="Number of async workers for batch operations",
-        gt=0,
-    )
-    batch_size: int = Field(
-        default=512,
-        description="Max number of rows to send to a step at a time",
-        gt=0,
-    )
-    error_mode: ErrorMode = Field(
-        default=ErrorMode.FAIL,
-        description="Error handling mode for batch processing",
-    )
-    write_errors_to: str | None = Field(
-        default=None,
-        description="If error mode is DROP, the errors for any step are saved to this directory",
-    )

qtype/interpreter/batch/utils.py

@@ -1,178 +0,0 @@
-from __future__ import annotations
-
-import copy
-from typing import Any, Callable, Tuple
-
-import pandas as pd
-
-from qtype.interpreter.batch.types import BatchConfig, ErrorMode
-from qtype.semantic.model import Step
-
-
-class InputMissingError(Exception):
-    """Raised when a required input variable is missing from the DataFrame."""
-
-    pass
-
-
-def validate_inputs(batch: pd.DataFrame, step: Step) -> None:
-    """Ensures all input variables for the step are columns in the DataFrame.
-    If not, an Exception is raised.
-
-    Args:
-        batch: The input DataFrame to decode.
-        step: The step to validate.
-    Raises:
-        InputMissingError: If any input variable is missing from the DataFrame.
-    """
-    input_ids = [input_var.id for input_var in step.inputs]
-    for input_var in input_ids:
-        if input_var not in batch.columns:
-            raise InputMissingError(
-                f"Input DataFrame must contain column '{input_var}' for step {step.id}."
-            )
-
-
-def fail_mode_wrapper(
-    f: Callable[..., dict[str, Any]],
-    row: pd.Series,
-    batch_config: BatchConfig,
-    **kwargs: dict[str, Any],
-) -> dict | Exception:
-    """Executes a function with error handling based on the batch configuration.
-
-    Args:
-        f: The function to execute that can take any arguments and returns a dict of results.
-        row: The input row as a dictionary.
-        batch_config: Configuration for error handling.
-        **kwargs: Additional keyword arguments.
-
-    Returns:
-        The result of the function or an Exception if an error occurs and the error mode is set to CONTINUE.
-    """
-    try:
-        # turn row into a dict and merge with kwargs
-        merged_kwargs = {**row.to_dict(), **kwargs}
-        return {**f(**merged_kwargs), **row.to_dict()}
-    except Exception as e:
-        if batch_config.error_mode == ErrorMode.FAIL:
-            raise e
-        else:
-            return e
-
-
-def single_step_adapter(
-    step: Step, **inputs: dict[str, Any]
-) -> dict[str, Any]:
-    """A batch adapter for steps that have no side effects or access shared resources."""
-    from qtype.interpreter.step import execute_step
-
-    step_clone = copy.deepcopy(step)
-    for input_var in step_clone.inputs:
-        if input_var.id in inputs:
-            input_var.value = inputs[input_var.id]
-        else:
-            raise ValueError(
-                f"Input variable '{input_var.id}' not found in inputs."
-            )
-    execute_step(step_clone)
-    return {
-        output_var.id: output_var.value for output_var in step_clone.outputs
-    }
-
-
-def to_series(
-    rv: dict | Exception, error_col_name: str = "error"
-) -> pd.Series:
-    # If rv is an exception, return a series with index "error"
-    if isinstance(rv, Exception):
-        return pd.Series({error_col_name: str(rv)})
-    return pd.Series(rv)  # type: ignore[no-any-return]
-
-
-def batch_iterator(
-    f: Callable[..., dict[str, Any]],
-    batch: pd.DataFrame,
-    batch_config: BatchConfig,
-    **kwargs: Any,
-) -> Tuple[pd.DataFrame, pd.DataFrame]:
-    """Executes a step over a batch of inputs with error handling.
-
-    Args:
-        step: The step to execute.
-        batch: The input DataFrame to process.
-        batch_config: Configuration for error handling.
-        **kwargs: Additional keyword arguments to pass to the step.
-
-    Returns:
-        A tuple containing two DataFrames:
-            - The first DataFrame contains successful results with output columns.
-            - The second DataFrame contains rows that encountered errors with an 'error' column.
-    """
-
-    # Use a unique column name for errors
-    error_col = "error_" + str(id(f))
-
-    # If error_col is already in the dataframe, throw an exception
-    if error_col in batch.columns:
-        raise ValueError(
-            f"Error column name '{error_col}' already exists in the batch DataFrame."
-        )
-
-    def the_pipe(row: pd.Series) -> pd.Series:
-        return to_series(
-            fail_mode_wrapper(f, row, batch_config=batch_config, **kwargs),
-            error_col_name=error_col,
-        )
-
-    results = batch.apply(the_pipe, axis=1)
-
-    # If error column doesn't exist, add it with NaN values
-    if error_col not in results.columns:
-        results[error_col] = pd.NA
-
-    # Split the results into two dataframes, one where error_col is not defined, and one where it is.
-    success_mask = ~results[error_col].notna()
-    failed_mask = results[error_col].notna()
-
-    # Create success DataFrame (drop the error column)
-    success_df = results[success_mask].drop(columns=[error_col])
-
-    # Create failed DataFrame (keep only original columns plus error)
-    original_columns = batch.columns.tolist()
-
-    if failed_mask.any():
-        failed_df = results[failed_mask]
-        # Drop all the output columns from failed_df, keep only original input columns + error
-        failed_df = failed_df[original_columns + [error_col]]
-    else:
-        # No failed rows, create empty DataFrame with expected structure
-        failed_df = pd.DataFrame(columns=original_columns + [error_col])
-
-    return success_df, failed_df
-
-
-def reconcile_results_and_errors(
-    results: list[pd.DataFrame], errors: list[pd.DataFrame]
-) -> Tuple[pd.DataFrame, pd.DataFrame]:
-    """
-    Concatenates lists of pandas DataFrames containing results and errors into single DataFrames.
-
-    If the input lists are empty, creates empty DataFrames as placeholders.
-
-    Args:
-        results (list[pd.DataFrame]): List of DataFrames containing results.
-        errors (list[pd.DataFrame]): List of DataFrames containing errors.
-
-    Returns:
-        Tuple[pd.DataFrame, pd.DataFrame]: A tuple containing:
-            - A single DataFrame with all results concatenated.
-            - A single DataFrame with all errors concatenated.
-    """
-    if not results:
-        results = [pd.DataFrame({})]
-    if not errors:
-        errors = [pd.DataFrame({})]
-    return pd.concat(results, ignore_index=True), pd.concat(
-        errors, ignore_index=True
-    )