[Do not submit] Timeseries backfill avro

pipeline/pom.xml

@@ -22,6 +22,7 @@
     <module>differ</module>
     <module>ingestion</module>
     <module>spanner</module>
+    <module>timeseries-backfill</module>
     <module>util</module>
   </modules>
 

pipeline/timeseries-backfill/pom.xml

@@ -0,0 +1,114 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <parent>
+        <groupId>org.datacommons</groupId>
+        <artifactId>pipeline</artifactId>
+        <version>${revision}</version>
+    </parent>
+
+    <groupId>org.datacommons</groupId>
+    <artifactId>timeseries-backfill</artifactId>
+    <version>${revision}</version>
+    <name>Data Commons - Timeseries Backfill</name>
+
+    <dependencies>
+        <dependency>
+            <artifactId>spanner</artifactId>
+            <groupId>org.datacommons</groupId>
+            <version>${revision}</version>
+        </dependency>
+        <dependency>
+            <artifactId>data</artifactId>
+            <groupId>org.datacommons</groupId>
+            <version>${revision}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-sdks-java-core</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-sdks-java-io-google-cloud-platform</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-runners-google-cloud-dataflow-java</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-runners-direct-java</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.beam</groupId>
+            <artifactId>beam-sdks-java-extensions-avro</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.cloud</groupId>
+            <artifactId>google-cloud-spanner</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.code.gson</groupId>
+            <artifactId>gson</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-jdk14</artifactId>
+            <version>${slf4j.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <version>${junit.version}</version>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.codehaus.mojo</groupId>
+                <artifactId>exec-maven-plugin</artifactId>
+                <configuration>
+                    <skip>false</skip>
+                    <cleanupDaemonThreads>false</cleanupDaemonThreads>
+                </configuration>
+            </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-shade-plugin</artifactId>
+                <version>3.5.3</version>
+                <executions>
+                    <execution>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>shade</goal>
+                        </goals>
+                        <configuration>
+                            <finalName>${project.artifactId}-bundled-${project.version}</finalName>
+                            <transformers>
+                                <transformer implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
+                                    <mainClass>org.datacommons.ingestion.timeseries.TimeseriesBackfillPipeline</mainClass>
+                                </transformer>
+                                <transformer implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
+                            </transformers>
+                            <filters>
+                                <filter>
+                                    <artifact>*:*</artifact>
+                                    <excludes>
+                                        <exclude>META-INF/*.SF</exclude>
+                                        <exclude>META-INF/*.DSA</exclude>
+                                        <exclude>META-INF/*.RSA</exclude>
+                                    </excludes>
+                                </filter>
+                            </filters>
+                        </configuration>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+    </build>
+</project>

pipeline/timeseries-backfill/recreate_timeseries_tables.sh

@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ $# -ne 3 ]]; then
+  echo "Usage: $0 <project_id> <instance_id> <database_id>" >&2
+  exit 1
+fi
+
+project_id="$1"
+instance_id="$2"
+database_id="$3"
+
+script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+schema_file="${script_dir}/../../../rk-experiments/mixer/spanner/bq_spanner_ingestion/timeseries_schema.sql"
+schema_with_suffix_file="$(mktemp)"
+
+if [[ ! -f "${schema_file}" ]]; then
+  echo "Schema file not found: ${schema_file}" >&2
+  exit 1
+fi
+
+current_schema_file="$(mktemp)"
+drop_ddl_file="$(mktemp)"
+trap 'rm -f "${current_schema_file}" "${drop_ddl_file}" "${schema_with_suffix_file}"' EXIT
+
+perl -0pe '
+  s/\bTimeSeriesAttributePropertyValue\b/TimeSeriesAttributePropertyValue_rk/g;
+  s/\bTimeSeriesAttributeValue\b/TimeSeriesAttributeValue_rk/g;
+  s/\bTimeSeriesByProvenance\b/TimeSeriesByProvenance_rk/g;
+  s/\bTimeSeriesByVariableMeasured\b/TimeSeriesByVariableMeasured_rk/g;
+  s/\bObservationAttribute\b/ObservationAttribute_rk/g;
+  s/\bStatVarObservation\b/StatVarObservation_rk/g;
+  s/\bTimeSeriesAttribute\b/TimeSeriesAttribute_rk/g;
+  s/\bTimeSeries\b/TimeSeries_rk/g;
+' "${schema_file}" > "${schema_with_suffix_file}"
+
+gcloud spanner databases ddl describe "${database_id}" \
+  --project="${project_id}" \
+  --instance="${instance_id}" \
+  > "${current_schema_file}"
+
+append_if_present() {
+  local pattern="$1"
+  local ddl="$2"
+  if rg -q "${pattern}" "${current_schema_file}"; then
+    printf '%s\n' "${ddl}" >> "${drop_ddl_file}"
+  fi
+}
+
+append_if_present "CREATE INDEX TimeSeriesByProvenance_rk " "DROP INDEX TimeSeriesByProvenance_rk"
+append_if_present "CREATE INDEX TimeSeriesByVariableMeasured_rk " "DROP INDEX TimeSeriesByVariableMeasured_rk"
+append_if_present "CREATE INDEX TimeSeriesAttributePropertyValue_rk " "DROP INDEX TimeSeriesAttributePropertyValue_rk"
+append_if_present "CREATE INDEX TimeSeriesAttributeValue_rk " "DROP INDEX TimeSeriesAttributeValue_rk"
+append_if_present "CREATE TABLE ObservationAttribute_rk " "DROP TABLE ObservationAttribute_rk"
+append_if_present "CREATE TABLE StatVarObservation_rk " "DROP TABLE StatVarObservation_rk"
+append_if_present "CREATE TABLE TimeSeriesAttribute_rk " "DROP TABLE TimeSeriesAttribute_rk"
+append_if_present "CREATE TABLE TimeSeries_rk " "DROP TABLE TimeSeries_rk"
+
+if [[ -s "${drop_ddl_file}" ]]; then
+  gcloud spanner databases ddl update "${database_id}" \
+    --project="${project_id}" \
+    --instance="${instance_id}" \
+    --ddl-file="${drop_ddl_file}"
+fi
+
+gcloud spanner databases ddl update "${database_id}" \
+  --project="${project_id}" \
+  --instance="${instance_id}" \
+  --ddl-file="${schema_with_suffix_file}"

pipeline/timeseries-backfill/src/main/java/org/datacommons/ingestion/timeseries/LocalProgressTracker.java

@@ -0,0 +1,146 @@
+package org.datacommons.ingestion.timeseries;
+
+import com.google.cloud.Timestamp;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.LongSupplier;
+import org.slf4j.Logger;
+
+/** Emits local progress logs for validator and DirectRunner runs. */
+final class LocalProgressTracker implements AutoCloseable {
+  interface LogSink {
+    void info(String message);
+  }
+
+  private final String runnerName;
+  private final int progressEverySourceRows;
+  private final int heartbeatSeconds;
+  private final LogSink logSink;
+  private final LongSupplier currentTimeMillis;
+  private final ScheduledExecutorService heartbeatExecutor;
+  private final AtomicLong sourceRows;
+  private final AtomicLong timeSeriesRows;
+  private final AtomicLong timeSeriesAttributeRows;
+  private final AtomicLong statVarObservationRows;
+  private final AtomicLong lastRowProgressMillis;
+  private final AtomicBoolean closed;
+
+  LocalProgressTracker(
+      String runnerName, int progressEverySourceRows, int heartbeatSeconds, Logger logger) {
+    this(
+        runnerName,
+        progressEverySourceRows,
+        heartbeatSeconds,
+        logger::info,
+        System::currentTimeMillis,
+        createHeartbeatExecutor(runnerName, heartbeatSeconds));
+  }
+
+  LocalProgressTracker(
+      String runnerName,
+      int progressEverySourceRows,
+      int heartbeatSeconds,
+      LogSink logSink,
+      LongSupplier currentTimeMillis,
+      ScheduledExecutorService heartbeatExecutor) {
+    this.runnerName = runnerName;
+    this.progressEverySourceRows = progressEverySourceRows;
+    this.heartbeatSeconds = heartbeatSeconds;
+    this.logSink = logSink;
+    this.currentTimeMillis = currentTimeMillis;
+    this.heartbeatExecutor = heartbeatExecutor;
+    this.sourceRows = new AtomicLong();
+    this.timeSeriesRows = new AtomicLong();
+    this.timeSeriesAttributeRows = new AtomicLong();
+    this.statVarObservationRows = new AtomicLong();
+    this.lastRowProgressMillis = new AtomicLong(currentTimeMillis.getAsLong());
+    this.closed = new AtomicBoolean();
+
+    if (heartbeatExecutor != null) {
+      heartbeatExecutor.scheduleAtFixedRate(
+          this::logHeartbeatNow, heartbeatSeconds, heartbeatSeconds, TimeUnit.SECONDS);
+    }
+  }
+
+  boolean isEnabled() {
+    return progressEverySourceRows > 0 || heartbeatSeconds > 0;
+  }
+
+  void recordRow(int timeSeriesAttributeRowsDelta, int statVarObservationRowsDelta) {
+    long sourceRowCount = sourceRows.incrementAndGet();
+    timeSeriesRows.incrementAndGet();
+    timeSeriesAttributeRows.addAndGet(timeSeriesAttributeRowsDelta);
+    statVarObservationRows.addAndGet(statVarObservationRowsDelta);
+    lastRowProgressMillis.set(currentTimeMillis.getAsLong());
+
+    if (progressEverySourceRows > 0 && sourceRowCount % progressEverySourceRows == 0) {
+      logSink.info(buildProgressMessage("progress"));
+    }
+  }
+
+  void recordValidatorFlush(int mutationCount, Timestamp writeTimestamp) {
+    logSink.info(
+        runnerName + " flush: mutations=" + mutationCount + ", write_timestamp=" + writeTimestamp);
+  }
+
+  void logHeartbeatNow() {
+    if (heartbeatSeconds <= 0 || closed.get()) {
+      return;
+    }
+    logSink.info(buildProgressMessage("heartbeat"));
+  }
+
+  @Override
+  public void close() {
+    if (!closed.compareAndSet(false, true)) {
+      return;
+    }
+    if (heartbeatExecutor != null) {
+      heartbeatExecutor.shutdownNow();
+    }
+  }
+
+  private String buildProgressMessage(String kind) {
+    long sourceRowCount = sourceRows.get();
+    long idleSeconds =
+        TimeUnit.MILLISECONDS.toSeconds(
+            currentTimeMillis.getAsLong() - lastRowProgressMillis.get());
+    StringBuilder message =
+        new StringBuilder()
+            .append(runnerName)
+            .append(' ')
+            .append(kind)
+            .append(": source_rows=")
+            .append(sourceRowCount)
+            .append(", timeseries_rows=")
+            .append(timeSeriesRows.get())
+            .append(", timeseries_attribute_rows=")
+            .append(timeSeriesAttributeRows.get())
+            .append(", stat_var_observation_rows=")
+            .append(statVarObservationRows.get())
+            .append(", seconds_since_last_row=")
+            .append(idleSeconds);
+    if (sourceRowCount == 0) {
+      message.append(", no_source_rows_yet=true");
+    }
+    return message.toString();
+  }
+
+  private static ScheduledExecutorService createHeartbeatExecutor(
+      String runnerName, int heartbeatSeconds) {
+    if (heartbeatSeconds <= 0) {
+      return null;
+    }
+    ThreadFactory threadFactory =
+        runnable -> {
+          Thread thread = new Thread(runnable, runnerName + "-progress-heartbeat");
+          thread.setDaemon(true);
+          return thread;
+        };
+    return Executors.newSingleThreadScheduledExecutor(threadFactory);
+  }
+}