fastflowtransform.executors.bigquery.bigframes¶

BigQueryBFExecutor ¶

Bases: BigQueryBaseExecutor[BFDataFrame]

Source code in src/fastflowtransform/executors/bigquery/bigframes.py

class BigQueryBFExecutor(BigQueryBaseExecutor[BFDataFrame]):
    ENGINE_NAME = "bigquery"

    def __init__(
        self,
        project: str,
        dataset: str,
        location: str | None = None,
        allow_create_dataset: bool = False,
    ):
        if not project:
            raise RuntimeError("BigFrames executor requires FF_BQ_PROJECT to be set.")
        if not location:
            raise RuntimeError(
                "BigFrames executor requires FF_BQ_LOCATION to be set. "
                "Use the dataset's region (e.g., EU or US)."
            )
        super().__init__(
            project=project,
            dataset=dataset,
            location=location,
            allow_create_dataset=allow_create_dataset,
        )

        try:
            ctx = BigQueryOptions(
                project=project,
                location=location,
            )
            self.session = bigframes.Session(context=ctx)
        except Exception as exc:
            raise RuntimeError(
                "Failed to initialize BigFrames session. Verify FF_BQ_PROJECT, "
                "FF_BQ_DATASET, and FF_BQ_LOCATION are set for the active profile. "
                f"{exc}"
            ) from exc

    def run_python(self, node: Node) -> None:
        """
        Execute Python models with a session scoped to this executor.

        We avoid mutating the process-wide default session; instead we
        temporarily set the executor session as the active global session so
        model code using bpd.DataFrame(...) picks up the configured location,
        then restore afterward.
        """
        ctx = bf_global_session._GlobalSessionContext(self.session)
        with ctx:
            super().run_python(node)

    # ---------- Python (Frames) ----------
    def _read_relation(self, relation: str, node: Node, deps: Iterable[str]) -> BFDataFrame:
        table_id = f"{self.project}.{self.dataset}.{relation}"
        try:
            return self.session.read_gbq(table_id)
        except NotFound as e:
            existing = [
                t.table_id for t in self.client.list_tables(f"{self.project}.{self.dataset}")
            ]
            raise RuntimeError(
                f"Dependency table not found: {table_id}\n"
                f"Deps: {list(deps)}\nExisting in dataset: {existing}\n"
                "Hinweis: Seeds/Upstream-Modelle erzeugt? DATASET korrekt?"
            ) from e

    def _materialize_relation(self, relation: str, df: BFDataFrame, node: Node) -> None:
        self._ensure_dataset()
        table_id = f"{self.project}.{self.dataset}.{relation}"

        to_gbq = getattr(df, "to_gbq", None)
        if callable(to_gbq):
            to_gbq(table_id, if_exists="replace")
            return

        # Fallback only when it is truly a method (not a column name!)
        mat = getattr(df, "materialize", None)
        if callable(mat):
            mat(table=table_id, mode="overwrite")
            return

        raise RuntimeError(
            "BigQuery DataFrames: Ergebnis nicht materialisierbar. "
            "Erwarte df.to_gbq(...) oder df.materialize(...)."
        )

    # ---- Required-columns validation tuned for BigFrames ----
    def _validate_required(
        self,
        node_name: str,
        inputs: Any,
        requires: dict[str, set[str]],
    ) -> None:
        if not requires:
            return

        def cols(bf_df: BFDataFrame) -> set[str]:
            if hasattr(bf_df, "columns"):
                return set(map(str, list(bf_df.columns)))
            if hasattr(bf_df, "schema") and hasattr(bf_df.schema, "names"):
                return set(bf_df.schema.names)
            return set()

        errs: list[str] = []
        if self._is_frame(inputs):
            # Single input frame case
            need = next(iter(requires.values()), set())
            miss = need - cols(inputs)
            if miss:
                errs.append(f"- missing columns: {sorted(miss)}")
        else:
            # Mapping {rel -> frame}
            for rel, need in requires.items():
                if rel not in inputs:
                    errs.append(f"- missing dependency key '{rel}'")
                    continue
                miss = need - cols(inputs[rel])
                if miss:
                    errs.append(f"- [{rel}] missing: {sorted(miss)}")

        if errs:
            raise ValueError(
                f"Required columns check failed for BigQuery DataFrames model '{node_name}'.\n"
                + "\n".join(errs)
            )

    def _columns_of(self, frame: BFDataFrame) -> list[str]:
        if hasattr(frame, "columns"):
            return [str(c) for c in list(frame.columns)]
        if hasattr(frame, "schema") and hasattr(frame.schema, "names"):
            return list(frame.schema.names)
        return []

    def _is_frame(self, obj: Any) -> bool:
        if obj is None:
            return False
        return (
            callable(getattr(obj, "to_gbq", None))
            or callable(getattr(obj, "materialize", None))
            or hasattr(obj, "columns")
        )

    def _frame_name(self) -> str:
        return "BigQuery DataFrame (BigFrames)"

        # ---- Unit-test helpers (pandas-facing) --------------------------------

    def utest_read_relation(self, relation: str) -> pd.DataFrame:
        """
        Read a relation into a pandas DataFrame for unit-test assertions.

        Even though this executor uses BigFrames for normal execution,
        utests compare pandas DataFrames, so we convert.
        """
        q = f"SELECT * FROM {self._qualified_identifier(relation)}"
        job = self.client.query(q, location=self.location)
        return job.result().to_dataframe(create_bqstorage_client=True)

    def utest_load_relation_from_rows(self, relation: str, rows: list[dict]) -> None:
        """
        Load rows into a BigQuery table for unit tests (replace if exists).

        Implementation uses the raw BigQuery client with pandas, which is
        perfectly fine for test input setup.
        """
        self._ensure_dataset()
        table_id = f"{self.project}.{self.dataset}.{relation}"
        df = pd.DataFrame(rows)

        job_config = LoadJobConfig(write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE)

        try:
            job = self.client.load_table_from_dataframe(
                df,
                table_id,
                job_config=job_config,
                location=self.location,
            )
            job.result()
        except BadRequest as e:
            raise RuntimeError(f"BigQuery utest write failed: {table_id}\n{e}") from e

    def utest_clean_target(self, relation: str) -> None:
        """
        For unit tests: drop any table/view with this name in the configured dataset.
        """
        table_id = f"{self.project}.{self.dataset}.{relation}"

        try:
            self.client.delete_table(table_id, not_found_ok=True)
        except NotFound:
            pass
        except TypeError:
            with suppress(NotFound):
                self.client.delete_table(table_id)
        except Exception:
            # Best-effort; don't make the whole test run fail because cleanup hiccupped.
            pass

run_python ¶

run_python(node)

Execute Python models with a session scoped to this executor.

We avoid mutating the process-wide default session; instead we temporarily set the executor session as the active global session so model code using bpd.DataFrame(...) picks up the configured location, then restore afterward.

Source code in src/fastflowtransform/executors/bigquery/bigframes.py

def run_python(self, node: Node) -> None:
    """
    Execute Python models with a session scoped to this executor.

    We avoid mutating the process-wide default session; instead we
    temporarily set the executor session as the active global session so
    model code using bpd.DataFrame(...) picks up the configured location,
    then restore afterward.
    """
    ctx = bf_global_session._GlobalSessionContext(self.session)
    with ctx:
        super().run_python(node)

utest_read_relation ¶

utest_read_relation(relation)

Read a relation into a pandas DataFrame for unit-test assertions.

Even though this executor uses BigFrames for normal execution, utests compare pandas DataFrames, so we convert.

Source code in src/fastflowtransform/executors/bigquery/bigframes.py

def utest_read_relation(self, relation: str) -> pd.DataFrame:
    """
    Read a relation into a pandas DataFrame for unit-test assertions.

    Even though this executor uses BigFrames for normal execution,
    utests compare pandas DataFrames, so we convert.
    """
    q = f"SELECT * FROM {self._qualified_identifier(relation)}"
    job = self.client.query(q, location=self.location)
    return job.result().to_dataframe(create_bqstorage_client=True)

utest_load_relation_from_rows ¶

utest_load_relation_from_rows(relation, rows)

Load rows into a BigQuery table for unit tests (replace if exists).

Implementation uses the raw BigQuery client with pandas, which is perfectly fine for test input setup.

Source code in src/fastflowtransform/executors/bigquery/bigframes.py

def utest_load_relation_from_rows(self, relation: str, rows: list[dict]) -> None:
    """
    Load rows into a BigQuery table for unit tests (replace if exists).

    Implementation uses the raw BigQuery client with pandas, which is
    perfectly fine for test input setup.
    """
    self._ensure_dataset()
    table_id = f"{self.project}.{self.dataset}.{relation}"
    df = pd.DataFrame(rows)

    job_config = LoadJobConfig(write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE)

    try:
        job = self.client.load_table_from_dataframe(
            df,
            table_id,
            job_config=job_config,
            location=self.location,
        )
        job.result()
    except BadRequest as e:
        raise RuntimeError(f"BigQuery utest write failed: {table_id}\n{e}") from e

utest_clean_target ¶

utest_clean_target(relation)

For unit tests: drop any table/view with this name in the configured dataset.

Source code in src/fastflowtransform/executors/bigquery/bigframes.py

def utest_clean_target(self, relation: str) -> None:
    """
    For unit tests: drop any table/view with this name in the configured dataset.
    """
    table_id = f"{self.project}.{self.dataset}.{relation}"

    try:
        self.client.delete_table(table_id, not_found_ok=True)
    except NotFound:
        pass
    except TypeError:
        with suppress(NotFound):
            self.client.delete_table(table_id)
    except Exception:
        # Best-effort; don't make the whole test run fail because cleanup hiccupped.
        pass

run_sql ¶

run_sql(node, env)

Orchestrate SQL models

1) Render Jinja (ref/source/this) and strip leading {{ config(...) }}. 2) If the SQL is full DDL (CREATE …), execute it verbatim (passthrough). 3) Otherwise, normalize to CREATE OR REPLACE {TABLE|VIEW} AS . The body is CTE-aware (keeps WITH … SELECT … intact).

On failure, raise ModelExecutionError with a helpful snippet.

Source code in src/fastflowtransform/executors/base.py

def run_sql(self, node: Node, env: Environment) -> None:
    """
    Orchestrate SQL models:
      1) Render Jinja (ref/source/this) and strip leading {{ config(...) }}.
      2) If the SQL is full DDL (CREATE …), execute it verbatim (passthrough).
      3) Otherwise, normalize to CREATE OR REPLACE {TABLE|VIEW} AS <body>.
         The body is CTE-aware (keeps WITH … SELECT … intact).
    On failure, raise ModelExecutionError with a helpful snippet.
    """
    meta = getattr(node, "meta", {}) or {}
    if self._meta_is_incremental(meta):
        # Delegates to incremental engine: render, schema sync, merge/insert, etc.
        return _ff_incremental.run_or_dispatch(self, node, env)

    if self._meta_is_snapshot(meta):
        # Snapshots are executed via the dedicated CLI: `fft snapshot run`.
        raise ModelExecutionError(
            node_name=node.name,
            relation=relation_for(node.name),
            message=(
                "Snapshot models cannot be executed via 'fft run'. "
                "Use 'fft snapshot run' instead."
            ),
            sql_snippet="",
        )

    sql_rendered = self.render_sql(
        node,
        env,
        ref_resolver=lambda name: self._resolve_ref(name, env),
        source_resolver=self._resolve_source,
    )
    sql = self._strip_leading_config(sql_rendered).strip()

    materialization = (node.meta or {}).get("materialized", "table")
    if materialization == "ephemeral":
        return

    # 1) Direct DDL passthrough (CREATE [OR REPLACE] {TABLE|VIEW} …)
    if self._looks_like_direct_ddl(sql):
        try:
            self._execute_sql_direct(sql, node)
            return
        except NotImplementedError:
            # Engine doesn't implement direct DDL → fall back to normalized materialization.
            pass
        except Exception as e:
            raise ModelExecutionError(
                node_name=node.name,
                relation=relation_for(node.name),
                message=str(e),
                sql_snippet=sql,
            ) from e

    # 2) Normalized materialization path (CTE-safe body)
    body = self._selectable_body(sql).rstrip(" ;\n\t")
    target_sql = self._format_relation_for_ref(node.name)

    # Centralized SQL preview logging (applies to ALL engines)
    preview = (
        f"=== MATERIALIZE ===\n"
        f"-- model: {node.name}\n"
        f"-- materialized: {materialization}\n"
        f"-- target: {target_sql}\n"
        f"{body}\n"
    )
    echo_debug(preview)

    try:
        self._apply_sql_materialization(node, target_sql, body, materialization)
    except Exception as e:
        preview = f"-- materialized={materialization}\n-- target={target_sql}\n{body}"
        raise ModelExecutionError(
            node_name=node.name,
            relation=relation_for(node.name),
            message=str(e),
            sql_snippet=preview,
        ) from e

configure_query_budget_limit ¶

configure_query_budget_limit(limit)

Inject a configured per-query byte limit (e.g. from budgets.yml).

Source code in src/fastflowtransform/executors/base.py

def configure_query_budget_limit(self, limit: int | None) -> None:
    """
    Inject a configured per-query byte limit (e.g. from budgets.yml).
    """
    if limit is None:
        self._ff_configured_query_limit = None
        return
    try:
        iv = int(limit)
    except Exception:
        self._ff_configured_query_limit = None
        return
    self._ff_configured_query_limit = iv if iv > 0 else None

reset_node_stats ¶

reset_node_stats()

Reset per-node statistics buffer.

The run engine calls this before executing a model so that all stats recorded via _record_query_stats(...) belong to that node.

Source code in src/fastflowtransform/executors/base.py

def reset_node_stats(self) -> None:
    """
    Reset per-node statistics buffer.

    The run engine calls this before executing a model so that all
    stats recorded via `_record_query_stats(...)` belong to that node.
    """
    # just clear the buffer; next recording will re-create it
    self._ff_query_stats_buffer = []

get_node_stats ¶

get_node_stats()

Aggregate buffered QueryStats into a simple dict:

{
  "bytes_scanned": <sum>,
  "rows": <sum>,
  "query_duration_ms": <sum>,
}

Called by the run engine after a node finishes.

Source code in src/fastflowtransform/executors/base.py

def get_node_stats(self) -> dict[str, int]:
    """
    Aggregate buffered QueryStats into a simple dict:

        {
          "bytes_scanned": <sum>,
          "rows": <sum>,
          "query_duration_ms": <sum>,
        }

    Called by the run engine after a node finishes.
    """
    stats_list = self._drain_query_stats()
    if not stats_list:
        return {}

    total_bytes = 0
    total_rows = 0
    total_duration = 0

    for s in stats_list:
        if s.bytes_processed is not None:
            total_bytes += int(s.bytes_processed)
        if s.rows is not None:
            total_rows += int(s.rows)
        if s.duration_ms is not None:
            total_duration += int(s.duration_ms)

    return {
        "bytes_scanned": total_bytes,
        "rows": total_rows,
        "query_duration_ms": total_duration,
    }