Skip to content

optimizise prov summary for larger imports #668

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
rohitkumarbhagat wants to merge 3 commits into
base: master
Choose a base branch
from
Draft

optimizise prov summary for larger imports #668

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8a928d2
optimizise prov summary for larger imports
rohitkumarbhagat Jul 20, 2026
f9fda26
Batch provenance summary name lookups. Avoid oversized Spanner federa…
rohitkumarbhagat Jul 20, 2026
ef19af9
reduced batch size to 100
rohitkumarbhagat Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,15 @@ def test_run_all(self):
self.assertIn("spanner-uri", query)
self.assertIn("'dc/base/import1'",
query) # Since is_base_dc=True
self.assertIn("CREATE OR REPLACE TEMPORARY TABLE `temp_series_summary`",
query)
self.assertIn("COUNT(*) AS observation_count", query)
self.assertIn("SUM(observation_count) as facet_obs_count", query)
self.assertIn("COUNT(*) as facet_ts_count", query)
self.assertIn("sample_dcid_batch_size INT64 DEFAULT 100", query)
self.assertIn("WHILE sample_batch_index < sample_batch_count", query)
self.assertIn("SUBSTR(name, 1, 1024) AS name", query)
self.assertNotIn("temp_obs_flat", query)


class TestPlaceAggregationGenerator(unittest.TestCase):
Expand Down
134 changes: 82 additions & 52 deletions pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,43 +61,44 @@ def run_provenance_summary_aggregation(
DECLARE place_dcids_str STRING;
DECLARE place_count INT64;
DECLARE sample_dcids_str STRING;
DECLARE sample_batch_index INT64 DEFAULT 0;
DECLARE sample_batch_count INT64;
DECLARE sample_dcid_batch_size INT64 DEFAULT 100;

-- Step 1: Fetch joined TimeSeries and Observation data from Spanner
-- We filter by provenance (which corresponds to import_name with prefix)
CREATE OR REPLACE TEMPORARY TABLE `temp_obs_flat` AS
SELECT
variable_measured,
entity1 as observation_about,
extra_entities_id,
facet_id,
provenance,
observation_period,
measurement_method,
unit,
JSON_VALUE(facet, '$.scalingFactor') as scaling_factor,
CAST(JSON_VALUE(facet, '$.isDcAggregate') AS BOOL) as is_dc_aggregate,
date as date_val,
SAFE_CAST(value AS FLOAT64) as value_num
-- Step 1: Aggregate observations per time series in Spanner so that only
-- one row per series is transferred to BigQuery.
CREATE OR REPLACE TEMPORARY TABLE `temp_series_summary` AS
SELECT *
FROM EXTERNAL_QUERY("{connection_id}",
'''SELECT
variable_measured,
entity1,
extra_entities_id,
facet_id,
provenance,
observation_period,
measurement_method,
unit,
facet,
date,
value
FROM TimeSeries
JOIN Observation USING (variable_measured, entity1, extra_entities_id, facet_id)
WHERE provenance IN ({provenances_str}) ''');
ts.variable_measured,
ts.entity1 AS observation_about,
ts.facet_id,
ts.provenance,
ANY_VALUE(ts.observation_period) AS observation_period,
ANY_VALUE(ts.measurement_method) AS measurement_method,
ANY_VALUE(ts.unit) AS unit,
ANY_VALUE(JSON_VALUE(ts.facet, '$.scalingFactor')) AS scaling_factor,
ANY_VALUE(SAFE_CAST(JSON_VALUE(ts.facet, '$.isDcAggregate') AS BOOL)) AS is_dc_aggregate,
MIN(obs.date) AS min_date,
MAX(obs.date) AS max_date,
MIN(SAFE_CAST(obs.value AS FLOAT64)) AS min_value,
MAX(SAFE_CAST(obs.value AS FLOAT64)) AS max_value,
COUNT(*) AS observation_count
FROM TimeSeries AS ts
JOIN Observation AS obs
USING (variable_measured, entity1, extra_entities_id, facet_id)
WHERE ts.provenance IN ({provenances_str})
GROUP BY
ts.variable_measured,
ts.entity1,
ts.extra_entities_id,
ts.facet_id,
ts.provenance ''');
Comment on lines 73 to +97

@gemini-code-assist gemini-code-assist Bot Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

In GoogleSQL (used by both BigQuery and Spanner), columns specified in the USING clause of a JOIN (such as variable_measured, entity1, extra_entities_id, and facet_id) cannot be qualified with a table alias (like ts.) in the SELECT list or GROUP BY clause. Doing so will result in a query compilation error. Please remove the ts. prefix from these columns.

          '''SELECT 
               variable_measured,
               entity1 AS observation_about,
               facet_id,
               ts.provenance,
               ANY_VALUE(ts.observation_period) AS observation_period,
               ANY_VALUE(ts.measurement_method) AS measurement_method,
               ANY_VALUE(ts.unit) AS unit,
               ANY_VALUE(JSON_VALUE(ts.facet, '$.scalingFactor')) AS scaling_factor,
               ANY_VALUE(SAFE_CAST(JSON_VALUE(ts.facet, '$.isDcAggregate') AS BOOL)) AS is_dc_aggregate,
               MIN(obs.date) AS min_date,
               MAX(obs.date) AS max_date,
               MIN(SAFE_CAST(obs.value AS FLOAT64)) AS min_value,
               MAX(SAFE_CAST(obs.value AS FLOAT64)) AS max_value,
               COUNT(*) AS observation_count
             FROM TimeSeries AS ts
             JOIN Observation AS obs
               USING (variable_measured, entity1, extra_entities_id, facet_id)
             WHERE ts.provenance IN ({provenances_str})
             GROUP BY
               variable_measured,
               entity1,
               extra_entities_id,
               facet_id,
               ts.provenance ''';


-- Step 2: Extract distinct place IDs in this dataset
CREATE OR REPLACE TEMPORARY TABLE `temp_dataset_places` AS
SELECT DISTINCT observation_about FROM `temp_obs_flat`;
SELECT DISTINCT observation_about FROM `temp_series_summary`;

SET place_count = (SELECT COUNT(*) FROM `temp_dataset_places`);

Expand Down Expand Up @@ -127,10 +128,14 @@ def run_provenance_summary_aggregation(
-- Step 4: Join observations with filtered place_type only
CREATE OR REPLACE TEMPORARY TABLE `temp_prepared` AS
SELECT
raw.*,
IF(raw.provenance LIKE 'dc/base/%', SUBSTR(raw.provenance, 9), raw.provenance) as import_name,
raw.variable_measured,
raw.observation_about,
raw.facet_id,
raw.provenance,
raw.min_value,
raw.max_value,
edges.place_type
FROM `temp_obs_flat` raw
FROM `temp_series_summary` raw
LEFT JOIN `temp_type_edges_filtered` edges ON raw.observation_about = edges.subject_id;

-- Step 5: Extract top 3 sample place DCIDs per summary group
Expand All @@ -154,20 +159,45 @@ def run_provenance_summary_aggregation(
FROM distinct_places
GROUP BY variable_measured, provenance, facet_id, place_type;

-- Step 6: Fetch ONLY place names for the selected top sample places from Spanner
SET sample_dcids_str = (
SELECT IFNULL(STRING_AGG(DISTINCT FORMAT("'%s'", REPLACE(dcid, "'", "\\'")), ','), "''")
-- Step 6: Fetch place names in bounded batches. A single large result
-- can exceed the Spanner federation message limit.
CREATE OR REPLACE TEMPORARY TABLE `temp_sample_dcids` AS
SELECT
dcid,
DIV(ROW_NUMBER() OVER (ORDER BY dcid) - 1, sample_dcid_batch_size) AS batch_index
FROM (
SELECT DISTINCT dcid
FROM `temp_top_place_dcids`
CROSS JOIN UNNEST(top_dcids) as dcid
);

EXECUTE IMMEDIATE FORMAT('''
CREATE OR REPLACE TEMPORARY TABLE `temp_node_names_filtered` AS
SELECT subject_id, name
FROM EXTERNAL_QUERY("{connection_id}",
"SELECT subject_id, name FROM Node WHERE subject_id IN (%s)"
SET sample_batch_count = (
SELECT IFNULL(MAX(batch_index) + 1, 0)
FROM `temp_sample_dcids`
);

CREATE OR REPLACE TEMPORARY TABLE `temp_node_names_filtered` (
subject_id STRING,
name STRING
);

WHILE sample_batch_index < sample_batch_count DO
SET sample_dcids_str = (
SELECT STRING_AGG(FORMAT("'%s'", REPLACE(dcid, "'", "\\'")), ',')
FROM `temp_sample_dcids`
WHERE batch_index = sample_batch_index
);
''', sample_dcids_str);

EXECUTE IMMEDIATE FORMAT('''
INSERT INTO `temp_node_names_filtered` (subject_id, name)
SELECT subject_id, name
FROM EXTERNAL_QUERY("{connection_id}",
"SELECT subject_id, SUBSTR(name, 1, 1024) AS name FROM Node WHERE subject_id IN (%s)"
);
''', sample_dcids_str);

SET sample_batch_index = sample_batch_index + 1;
END WHILE;

-- Step 7: Aggregate Place Type Summaries and attach names to top 3 sample places
CREATE OR REPLACE TEMPORARY TABLE `temp_place_type_summary` AS
Expand All @@ -177,8 +207,8 @@ def run_provenance_summary_aggregation(
provenance,
facet_id,
place_type,
MIN(value_num) as min_val,
MAX(value_num) as max_val,
MIN(min_value) as min_val,
MAX(max_value) as max_val,
COUNT(DISTINCT observation_about) as place_count
FROM `temp_prepared`
WHERE place_type IS NOT NULL
Expand Down Expand Up @@ -225,13 +255,13 @@ def run_provenance_summary_aggregation(
ANY_VALUE(unit) as unit,
ANY_VALUE(scaling_factor) as scaling_factor,
ANY_VALUE(is_dc_aggregate) as is_dc_aggregate,
MIN(date_val) as min_date,
MAX(date_val) as max_date,
MIN(value_num) as facet_min,
MAX(value_num) as facet_max,
COUNT(*) as facet_obs_count,
COUNT(DISTINCT CONCAT(observation_about, '|', extra_entities_id)) as facet_ts_count
FROM `temp_obs_flat`
MIN(min_date) as min_date,
MAX(max_date) as max_date,
MIN(min_value) as facet_min,
MAX(max_value) as facet_max,
SUM(observation_count) as facet_obs_count,
COUNT(*) as facet_ts_count
FROM `temp_series_summary`
GROUP BY variable_measured, provenance, facet_id
),
facet_summaries AS (
Expand Down
Loading