From 8a928d26740c60e60efc3a83caa70eb2b2a68bc3 Mon Sep 17 00:00:00 2001 From: rohit kumar Date: Mon, 20 Jul 2026 14:25:14 +0000 Subject: [PATCH 1/3] optimizise prov summary for larger imports --- .../aggregation/aggregation_test.py | 6 + .../provenance_summary_generator.py | 88 ++--- .../provenance_summary_one_import.sql | 312 ++++++++++++++++++ .../aggregation/aggregation_test.py | 6 + .../provenance_summary_generator.py | 86 ++--- 5 files changed, 414 insertions(+), 84 deletions(-) create mode 100644 pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql diff --git a/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py b/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py index 3d5d61332..527e8537e 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py @@ -259,6 +259,12 @@ 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.assertNotIn("temp_obs_flat", query) class TestPlaceAggregationGenerator(unittest.TestCase): diff --git a/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py b/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py index 272513472..412c27725 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py +++ b/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py @@ -62,42 +62,40 @@ def run_provenance_summary_aggregation( DECLARE place_count INT64; DECLARE sample_dcids_str STRING; - -- 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 '''); -- 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`); @@ -127,10 +125,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 @@ -177,8 +179,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 @@ -225,13 +227,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 ( diff --git a/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql b/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql new file mode 100644 index 000000000..29106f3b3 --- /dev/null +++ b/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql @@ -0,0 +1,312 @@ +-- Copyright 2026 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +-- BigQuery UI query for previewing provenance summaries for one import. +-- +-- Before running, replace both placeholders below throughout this file: +-- datcom-store.us-central1.rk-prod-spanner +-- dc/base/CensusACS5YearSurvey +-- +-- The query returns one row per (variable_measured, provenance), not one row +-- for the entire import. It does not write the results to Spanner. + +DECLARE place_dcids_str STRING; +DECLARE place_count INT64; +DECLARE sample_dcids_str STRING; + +-- Aggregate in Spanner before transferring data to BigQuery. This changes the +-- federated result from one row per observation to one row per time series. +CREATE OR REPLACE TEMPORARY TABLE `temp_series_summary` AS +SELECT * +FROM EXTERNAL_QUERY( + "datcom-store.us-central1.rk-prod-spanner", + '''SELECT + 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 = 'dc/base/CensusACS5YearSurvey' + GROUP BY + ts.variable_measured, + ts.entity1, + ts.extra_entities_id, + ts.facet_id, + ts.provenance''' +); + +CREATE OR REPLACE TEMPORARY TABLE `temp_dataset_places` AS +SELECT DISTINCT observation_about +FROM `temp_series_summary`; + +SET place_count = (SELECT COUNT(*) FROM `temp_dataset_places`); + +-- Spanner limits IN to 10,000 values. For larger datasets, retain the existing +-- fallback that reads all typeOf edges and filters them in BigQuery. +IF place_count <= 10000 THEN + SET place_dcids_str = ( + SELECT IFNULL( + STRING_AGG(FORMAT("'%s'", REPLACE(observation_about, "'", "\\'")), ','), + "''" + ) + FROM `temp_dataset_places` + ); + + EXECUTE IMMEDIATE FORMAT(''' + CREATE OR REPLACE TEMPORARY TABLE `temp_type_edges_filtered` AS + SELECT subject_id, object_id AS place_type + FROM EXTERNAL_QUERY( + "datcom-store.us-central1.rk-prod-spanner", + "SELECT subject_id, object_id FROM Edge WHERE predicate = 'typeOf' AND subject_id IN (%s)" + ); + ''', place_dcids_str); +ELSE + CREATE OR REPLACE TEMPORARY TABLE `temp_type_edges_filtered` AS + SELECT subject_id, object_id AS place_type + FROM EXTERNAL_QUERY( + "datcom-store.us-central1.rk-prod-spanner", + "SELECT subject_id, object_id FROM Edge WHERE predicate = 'typeOf'" + ); +END IF; + +CREATE OR REPLACE TEMPORARY TABLE `temp_prepared` AS +SELECT + series.variable_measured, + series.observation_about, + series.facet_id, + series.provenance, + series.min_value, + series.max_value, + edges.place_type +FROM `temp_series_summary` AS series +LEFT JOIN `temp_type_edges_filtered` AS edges + ON series.observation_about = edges.subject_id; + +CREATE OR REPLACE TEMPORARY TABLE `temp_top_place_dcids` AS +WITH distinct_places AS ( + SELECT DISTINCT + variable_measured, + provenance, + facet_id, + place_type, + observation_about AS dcid + FROM `temp_prepared` + WHERE place_type IS NOT NULL +) +SELECT + variable_measured, + provenance, + facet_id, + place_type, + ARRAY_AGG(dcid ORDER BY dcid LIMIT 3) AS top_dcids +FROM distinct_places +GROUP BY variable_measured, provenance, facet_id, place_type; + +SET sample_dcids_str = ( + SELECT IFNULL( + STRING_AGG(DISTINCT FORMAT("'%s'", REPLACE(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( + "datcom-store.us-central1.rk-prod-spanner", + "SELECT subject_id, name FROM Node WHERE subject_id IN (%s)" + ); +''', sample_dcids_str); + +CREATE OR REPLACE TEMPORARY TABLE `temp_place_type_summary` AS +WITH place_stats AS ( + SELECT + variable_measured, + provenance, + facet_id, + place_type, + 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 + GROUP BY variable_measured, provenance, facet_id, place_type +), +aggregated_places AS ( + SELECT + samples.variable_measured, + samples.provenance, + samples.facet_id, + samples.place_type, + ARRAY_AGG( + STRUCT(dcid, names.name) + ORDER BY dcid + ) AS top_places + FROM `temp_top_place_dcids` AS samples + CROSS JOIN UNNEST(samples.top_dcids) AS dcid + LEFT JOIN `temp_node_names_filtered` AS names + ON dcid = names.subject_id + GROUP BY + samples.variable_measured, + samples.provenance, + samples.facet_id, + samples.place_type +) +SELECT + stats.variable_measured, + stats.provenance AS provenance_dcid, + stats.facet_id, + stats.place_type, + stats.place_count, + stats.min_val, + stats.max_val, + places.top_places +FROM place_stats AS stats +JOIN aggregated_places AS places + USING (variable_measured, provenance, facet_id, place_type); + +-- Preview the rows that the production generator exports to KeyValueStore. +WITH facet_base AS ( + SELECT + variable_measured, + provenance AS provenance_dcid, + facet_id, + ANY_VALUE( + IF(provenance LIKE 'dc/base/%', SUBSTR(provenance, 9), provenance) + ) AS import_name, + ANY_VALUE(measurement_method) AS measurement_method, + ANY_VALUE(observation_period) AS observation_period, + ANY_VALUE(unit) AS unit, + ANY_VALUE(scaling_factor) AS scaling_factor, + ANY_VALUE(is_dc_aggregate) AS is_dc_aggregate, + 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 ( + SELECT + facets.variable_measured, + facets.provenance_dcid, + facets.facet_id, + facets.import_name, + facets.measurement_method, + facets.observation_period, + facets.unit, + facets.scaling_factor, + facets.is_dc_aggregate, + facets.min_date, + facets.max_date, + facets.facet_min, + facets.facet_max, + facets.facet_obs_count, + facets.facet_ts_count, + ARRAY_AGG( + STRUCT( + places.place_type, + places.place_count, + places.min_val, + places.max_val, + places.top_places + ) + ) AS pt_summaries + FROM facet_base AS facets + LEFT JOIN `temp_place_type_summary` AS places + USING (variable_measured, provenance_dcid, facet_id) + GROUP BY + facets.variable_measured, + facets.provenance_dcid, + facets.facet_id, + facets.import_name, + facets.measurement_method, + facets.observation_period, + facets.unit, + facets.scaling_factor, + facets.is_dc_aggregate, + facets.min_date, + facets.max_date, + facets.facet_min, + facets.facet_max, + facets.facet_obs_count, + facets.facet_ts_count +) +SELECT + 'ProvenanceSummary' AS type, + variable_measured AS key, + provenance_dcid AS provenance, + JSON_OBJECT( + 'import_name', ANY_VALUE(import_name), + 'observation_count', CAST(SUM(facet_obs_count) AS FLOAT64), + 'time_series_count', CAST(SUM(facet_ts_count) AS FLOAT64), + 'series_summary', ARRAY_AGG( + JSON_OBJECT( + 'series_key', JSON_OBJECT( + 'measurement_method', measurement_method, + 'observation_period', observation_period, + 'unit', unit, + 'scaling_factor', scaling_factor, + 'is_dc_aggregate', COALESCE(is_dc_aggregate, FALSE) + ), + 'earliest_date', min_date, + 'latest_date', max_date, + 'min_value', facet_min, + 'max_value', facet_max, + 'observation_count', CAST(facet_obs_count AS FLOAT64), + 'time_series_count', CAST(facet_ts_count AS FLOAT64), + 'place_type_summary', ( + SELECT IF(ARRAY_LENGTH(keys) > 0, JSON_OBJECT(keys, vals), NULL) + FROM ( + SELECT + ARRAY_AGG(place_type) AS keys, + ARRAY_AGG( + JSON_OBJECT( + 'place_count', place_count, + 'min_value', min_val, + 'max_value', max_val, + 'top_places', ( + SELECT ARRAY_AGG( + JSON_OBJECT('dcid', top_place.dcid, 'name', top_place.name) + ) + FROM UNNEST(top_places) AS top_place + ) + ) + ) AS vals + FROM UNNEST(pt_summaries) + WHERE place_type IS NOT NULL + ) + ) + ) + ) + ) AS value +FROM facet_summaries +GROUP BY variable_measured, provenance_dcid +ORDER BY key; diff --git a/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py b/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py index f9d93a681..d78caab66 100644 --- a/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py +++ b/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py @@ -240,6 +240,12 @@ 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.assertNotIn("temp_obs_flat", query) class TestPlaceAggregationGenerator(unittest.TestCase): diff --git a/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py b/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py index 70f2165db..c324a6292 100644 --- a/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py +++ b/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py @@ -63,40 +63,40 @@ def run_provenance_summary_aggregation( DECLARE place_count INT64; DECLARE sample_dcids_str STRING; - -- 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, - 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, - 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 '''); -- 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`); @@ -126,10 +126,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 @@ -176,8 +180,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 @@ -224,13 +228,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 observation_about) 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 ( From f9fda26ccefa28228c99b3b45c671037a8714b0d Mon Sep 17 00:00:00 2001 From: rohit kumar Date: Mon, 20 Jul 2026 15:24:06 +0000 Subject: [PATCH 2/3] Batch provenance summary name lookups. Avoid oversized Spanner federation responses --- .../aggregation/aggregation_test.py | 3 ++ .../provenance_summary_generator.py | 46 +++++++++++++---- .../provenance_summary_one_import.sql | 50 ++++++++++++++----- .../aggregation/aggregation_test.py | 3 ++ .../provenance_summary_generator.py | 46 +++++++++++++---- 5 files changed, 118 insertions(+), 30 deletions(-) diff --git a/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py b/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py index 527e8537e..1a2fd97cd 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py @@ -264,6 +264,9 @@ def test_run_all(self): 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 1000", 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) diff --git a/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py b/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py index 412c27725..e909835b5 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py +++ b/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py @@ -61,6 +61,9 @@ 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 1000; -- Step 1: Aggregate observations per time series in Spanner so that only -- one row per series is transferred to BigQuery. @@ -156,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 diff --git a/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql b/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql index 29106f3b3..1feae1073 100644 --- a/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql +++ b/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql @@ -24,6 +24,9 @@ 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 1000; -- Aggregate in Spanner before transferring data to BigQuery. This changes the -- federated result from one row per observation to one row per time series. @@ -125,23 +128,46 @@ SELECT FROM distinct_places GROUP BY variable_measured, provenance, facet_id, place_type; -SET sample_dcids_str = ( - SELECT IFNULL( - STRING_AGG(DISTINCT FORMAT("'%s'", REPLACE(dcid, "'", "\\'")), ','), - "''" - ) +-- Fetch optional sample labels in bounded batches. This avoids exceeding the +-- 200 MiB 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( - "datcom-store.us-central1.rk-prod-spanner", - "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( + "datcom-store.us-central1.rk-prod-spanner", + "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; CREATE OR REPLACE TEMPORARY TABLE `temp_place_type_summary` AS WITH place_stats AS ( diff --git a/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py b/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py index d78caab66..b64cc53a1 100644 --- a/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py +++ b/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py @@ -245,6 +245,9 @@ def test_run_all(self): 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 1000", 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) diff --git a/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py b/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py index c324a6292..991b028c5 100644 --- a/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py +++ b/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py @@ -62,6 +62,9 @@ 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 1000; -- Step 1: Aggregate observations per time series in Spanner so that only -- one row per series is transferred to BigQuery. @@ -157,20 +160,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 From ef19af9d90b896ec357c8f1fdb4872ccb92ece18 Mon Sep 17 00:00:00 2001 From: rohit kumar Date: Mon, 20 Jul 2026 16:49:04 +0000 Subject: [PATCH 3/3] reduced batch size to 100 --- .../aggregation/aggregation_test.py | 2 +- .../aggregation/provenance_summary_generator.py | 2 +- .../provenance_summary_one_import.sql | 12 ++++++------ .../ingestion-helper/aggregation/aggregation_test.py | 2 +- .../aggregation/provenance_summary_generator.py | 2 +- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py b/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py index 1a2fd97cd..186a184d1 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py @@ -264,7 +264,7 @@ def test_run_all(self): 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 1000", 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) diff --git a/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py b/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py index e909835b5..f279bb3ac 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py +++ b/pipeline/workflow/aggregation-helper/aggregation/provenance_summary_generator.py @@ -63,7 +63,7 @@ def run_provenance_summary_aggregation( DECLARE sample_dcids_str STRING; DECLARE sample_batch_index INT64 DEFAULT 0; DECLARE sample_batch_count INT64; - DECLARE sample_dcid_batch_size INT64 DEFAULT 1000; + DECLARE sample_dcid_batch_size INT64 DEFAULT 100; -- Step 1: Aggregate observations per time series in Spanner so that only -- one row per series is transferred to BigQuery. diff --git a/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql b/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql index 1feae1073..17c339a99 100644 --- a/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql +++ b/pipeline/workflow/aggregation-helper/provenance_summary_one_import.sql @@ -15,7 +15,7 @@ -- BigQuery UI query for previewing provenance summaries for one import. -- -- Before running, replace both placeholders below throughout this file: --- datcom-store.us-central1.rk-prod-spanner +-- projects/datcom-store/locations/us-central1/connections/rk-prod-spanner -- dc/base/CensusACS5YearSurvey -- -- The query returns one row per (variable_measured, provenance), not one row @@ -26,14 +26,14 @@ 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 1000; +DECLARE sample_dcid_batch_size INT64 DEFAULT 100; -- Aggregate in Spanner before transferring data to BigQuery. This changes the -- federated result from one row per observation to one row per time series. CREATE OR REPLACE TEMPORARY TABLE `temp_series_summary` AS SELECT * FROM EXTERNAL_QUERY( - "datcom-store.us-central1.rk-prod-spanner", + "projects/datcom-store/locations/us-central1/connections/rk-prod-spanner", '''SELECT ts.variable_measured, ts.entity1 AS observation_about, @@ -82,7 +82,7 @@ IF place_count <= 10000 THEN CREATE OR REPLACE TEMPORARY TABLE `temp_type_edges_filtered` AS SELECT subject_id, object_id AS place_type FROM EXTERNAL_QUERY( - "datcom-store.us-central1.rk-prod-spanner", + "projects/datcom-store/locations/us-central1/connections/rk-prod-spanner", "SELECT subject_id, object_id FROM Edge WHERE predicate = 'typeOf' AND subject_id IN (%s)" ); ''', place_dcids_str); @@ -90,7 +90,7 @@ ELSE CREATE OR REPLACE TEMPORARY TABLE `temp_type_edges_filtered` AS SELECT subject_id, object_id AS place_type FROM EXTERNAL_QUERY( - "datcom-store.us-central1.rk-prod-spanner", + "projects/datcom-store/locations/us-central1/connections/rk-prod-spanner", "SELECT subject_id, object_id FROM Edge WHERE predicate = 'typeOf'" ); END IF; @@ -161,7 +161,7 @@ WHILE sample_batch_index < sample_batch_count DO INSERT INTO `temp_node_names_filtered` (subject_id, name) SELECT subject_id, name FROM EXTERNAL_QUERY( - "datcom-store.us-central1.rk-prod-spanner", + "projects/datcom-store/locations/us-central1/connections/rk-prod-spanner", "SELECT subject_id, SUBSTR(name, 1, 1024) AS name FROM Node WHERE subject_id IN (%s)" ); ''', sample_dcids_str); diff --git a/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py b/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py index b64cc53a1..75b3dcb8f 100644 --- a/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py +++ b/pipeline/workflow/ingestion-helper/aggregation/aggregation_test.py @@ -245,7 +245,7 @@ def test_run_all(self): 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 1000", 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) diff --git a/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py b/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py index 991b028c5..833559843 100644 --- a/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py +++ b/pipeline/workflow/ingestion-helper/aggregation/provenance_summary_generator.py @@ -64,7 +64,7 @@ def run_provenance_summary_aggregation( DECLARE sample_dcids_str STRING; DECLARE sample_batch_index INT64 DEFAULT 0; DECLARE sample_batch_count INT64; - DECLARE sample_dcid_batch_size INT64 DEFAULT 1000; + DECLARE sample_dcid_batch_size INT64 DEFAULT 100; -- Step 1: Aggregate observations per time series in Spanner so that only -- one row per series is transferred to BigQuery.