add embedding geneator to ingest embedding with BQ

[shixiao-coder]

No description provided.

[codacy-production]

Not up to standards ⛔

🔴 Issues 6 high · 8 medium · 6 minor

Alerts:
⚠ 20 issues (≤ 0 issues of at least minor severity)

Results:
20 new issues

Category	Results
UnusedCode	3 medium 3 minor
Documentation	1 minor
ErrorProne	6 high
Security	3 medium
CodeStyle	2 minor
Complexity	2 medium

View in Codacy

🟢 Metrics 40 complexity

Metric	Results
Complexity	40

View in Codacy

_{NEW Get contextual insights on your PRs based on Codacy's metrics, along with PR and Jira context, without leaving GitHub. Enable AI reviewer
TIP This summary will be updated as you push new changes.}

[gemini-code-assist]

Code Review

This pull request introduces a new EmbeddingGenerator to generate node embeddings using Vertex AI remote models and save them to Spanner, integrating it into the orchestrator pipeline with corresponding configuration schemas, tests, and dependencies. The review feedback highlights critical improvements: first, to comply with Spanner Data Boost requirements, the EXISTS subquery in the EXTERNAL_QUERY should be replaced with a list of OR conditions; second, to prevent exceeding Spanner's query size limits, the massive statistical variable filter should be offloaded from Spanner to BigQuery; third, GCS file reading can be simplified by using storage.Client directly; and finally, redundant executions of the generator across active imports in the orchestrator should be addressed.

I am having trouble creating individual review comments. Click here to see my feedback.

pipeline/workflow/aggregation-helper/aggregation/embedding_generator.py (131-132)

To comply with Spanner Data Boost requirements (which prohibit subqueries/joins in EXTERNAL_QUERY), we should avoid using the EXISTS subquery in the Spanner query. We can prepare a list of OR conditions using IN UNNEST(types) instead. Additionally, we should use the imported _escape_sql_literal helper to safely escape the node types.

        node_types_conds = [f"'{_escape_sql_literal(nt)}' IN UNNEST(types)" for nt in node_types]
        node_types_sql = f"({ ' OR '.join(node_types_conds) })"

References

1. When executing Spanner queries via BigQuery's EXTERNAL_QUERY with Data Boost enabled, do not use joins or subqueries as they are not supported by Spanner Data Boost.

pipeline/workflow/aggregation-helper/aggregation/embedding_generator.py (155-157)

Replace the EXISTS subquery with the pre-constructed node_types_sql condition to avoid subqueries and ensure compatibility with Spanner Data Boost.

              AND {node_types_sql}

References

1. When executing Spanner queries via BigQuery's EXTERNAL_QUERY with Data Boost enabled, do not use joins or subqueries as they are not supported by Spanner Data Boost.

pipeline/workflow/aggregation-helper/aggregation/embedding_generator.py (137-140)

When node_filter_type is "NLStatisticalVariable", _extract_nl_stat_var() returns thousands of statistical variables. Embedding this massive array literal directly inside the Spanner query via EXTERNAL_QUERY can exceed Spanner's query size/complexity limits and severely degrade performance. Instead, we should use _escape_sql_literal to safely escape the variables, and we can offload this filter to BigQuery by applying it on the temporary table created from the EXTERNAL_QUERY results.

        elif node_filter_type == "NLStatisticalVariable":
            nl_stat_vars = _extract_nl_stat_var()
            safe_vars = [f"'{_escape_sql_literal(var)}'" for var in nl_stat_vars]
            filter_condition_sql = f"subject_id IN UNNEST([{', '.join(safe_vars)}])"

pipeline/workflow/aggregation-helper/aggregation/embedding_generator.py (153-154)

Remove {filter_condition_sql} from the Spanner query since it is now applied in BigQuery.

              AND %s

pipeline/workflow/aggregation-helper/aggregation/embedding_generator.py (182-188)

Apply the filter_condition_sql filter in BigQuery on the temporary table raw_nodes instead of inside the Spanner query. This keeps the Spanner query lightweight and compliant with Spanner's limits.

        EXECUTE IMMEDIATE FORMAT('''
          -- 1. Get raw text from Spanner
          CREATE TEMP TABLE raw_nodes AS
          SELECT * FROM EXTERNAL_QUERY("{conn_id}", {spanner_query_str})
          WHERE {filter_condition_sql};
        ''', 
        timelock_condition
        );

pipeline/workflow/aggregation-helper/aggregation/embedding_generator.py (55-65)

Since google-cloud-storage is a direct dependency of this package and storage.Client is already imported, we can simplify _extract_nl_stat_var by using storage.Client directly to read the GCS file. This is more robust, secure, and idiomatic than performing raw HTTP requests via urllib.request.

    if path.startswith("gs://"):
        parts = path[5:].split("/", 1)
        client = storage.Client()
        content = client.bucket(parts[0]).blob(parts[1]).download_as_text()

pipeline/workflow/aggregation-helper/aggregation/orchestrator.py (400-406)

The EmbeddingGenerator is triggered for each active import in the orchestrator's loop. However, since EmbeddingGenerator.run_all ignores the import_names parameter and generates embeddings for all nodes based on the global latest_lock_timestamp, running it for every active import in the same orchestrator run is highly redundant and expensive. Consider either updating EmbeddingGenerator to filter by the current import (if possible), or keeping track of whether embeddings have already been generated in the current orchestrator session to avoid duplicate executions.