Add temporal support to graph schema.

dgf/src/data/BUILD

@@ -294,6 +294,7 @@ py_test(
     name = "schema_cc_test",
     srcs = ["schema_cc_test.py"],
     deps = [
+        ":schema",
         ":schema_ext",
         # absl/testing:absltest dep,
         "//dgf/src/util:gen_test_graph",

dgf/src/data/schema.cc

@@ -41,9 +41,13 @@ std::string GraphSchema::Feature::to_string(int indent) const {
   auto shape_formatter = [](std::string* out, int dim) {
     absl::StrAppend(out, dim == -1 ? "None" : std::to_string(dim));
   };
-  return absl::StrCat(prefix, "Feature(name='", name, "', shape=[",
-                      absl::StrJoin(shape, ", ", shape_formatter),
-                      "], format=", FormatToString(format), ")");
+  return absl::StrCat(
+      prefix, "Feature(name='", name, "', shape=[",
+      absl::StrJoin(shape, ", ", shape_formatter),
+      "], format=", FormatToString(format),
+      is_timeseries ? ", is_timeseries=true" : "",
+      timestamps.empty() ? "" : absl::StrCat(", timestamps='", timestamps, "'"),
+      ")");
 }
 
 std::string GraphSchema::Nodeset::to_string(int indent) const {

dgf/src/data/schema.h

@@ -28,6 +28,8 @@ struct GraphSchema {
     // Shape of the feature. -1 (in cc) is equivalent to None (in python).
     std::vector<int> shape;
     eFormat format;
+    bool is_timeseries = false;
+    std::string timestamps;
 
     std::string to_string(int indent) const;
 

dgf/src/data/schema.py

@@ -90,13 +90,21 @@ class FeatureSchema:
     is_utf8_string: Whether the feature is a UTF-8 string. This is only relevant
       when feature_format is BYTES, to distinguish between Spanner STRING (True)
       and Spanner BYTES (False).
+    is_timeseries: Whether the feature represents a temporal series / sequence.
+    timestamps: For temporal sequence features, the name of the feature
+      containing the corresponding timestamp sequence (e.g., "time"). The
+      length of the corresponding timestamps feature must equal the length of
+      the timeseries feature along the 0th dimension. Cannot be set for non
+      timeseries features.
   """
 
   format: FeatureFormat
   semantic: FeatureSemantic = FeatureSemantic.UNKNOWN
   shape: Shape = None
   num_categorical_values: Optional[int] = None
   is_utf8_string: Optional[bool] = False
+  is_timeseries: Optional[bool] = False
+  timestamps: Optional[str] = None
 
   def is_static_shape(self) -> bool:
     """Returns true if the feature has a fully static shape."""

dgf/src/data/schema_cc_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from absl.testing import absltest
+from dgf.src.data import schema as schema_lib
 from dgf.src.data import schema_ext as lib
 from dgf.src.util import gen_test_graph
 
@@ -42,6 +43,101 @@ def test_parse_schema(self):
 ])""",
     )
 
+  def test_parse_temporal_schema(self):
+    schema = schema_lib.GraphSchema(
+        node_sets={
+            "n1": schema_lib.NodeSchema(
+                features={
+                    "#id": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.BYTES,
+                        semantic=schema_lib.FeatureSemantic.PRIMARY_ID,
+                    ),
+                    "#creation_time": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.INTEGER_64,
+                        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+                    ),
+                    "time": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.INTEGER_64,
+                        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+                        shape=(None,),
+                        is_timeseries=True,
+                    ),
+                    "f1_seq": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.FLOAT_32,
+                        semantic=schema_lib.FeatureSemantic.NUMERICAL,
+                        shape=(None,),
+                        is_timeseries=True,
+                        timestamps="time",
+                    ),
+                    "f2_seq": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.FLOAT_32,
+                        semantic=schema_lib.FeatureSemantic.EMBEDDING,
+                        shape=(None, 4),
+                        is_timeseries=True,
+                        timestamps="time",
+                    ),
+                }
+            ),
+            "n2": schema_lib.NodeSchema(
+                features={
+                    "#id": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.INTEGER_64,
+                        semantic=schema_lib.FeatureSemantic.PRIMARY_ID,
+                    ),
+                    "sensor_ts": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.FLOAT_32,
+                        semantic=schema_lib.FeatureSemantic.TIMESERIES,
+                        shape=(20, 8),
+                        is_timeseries=True,
+                    ),
+                }
+            ),
+        },
+        edge_sets={
+            "e1": schema_lib.EdgeSchema(
+                source="n1",
+                target="n2",
+                features={
+                    "edge_time": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.INTEGER_64,
+                        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+                        shape=(None,),
+                        is_timeseries=True,
+                    ),
+                    "edge_val": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.INTEGER_32,
+                        semantic=schema_lib.FeatureSemantic.NUMERICAL,
+                        shape=(None,),
+                        is_timeseries=True,
+                        timestamps="edge_time",
+                    ),
+                },
+            )
+        },
+    )
+    self.assertEqual(
+        lib.ParseAndDebugPrintSchema(schema),
+        """\
+GraphSchema(nodesets=[
+  Nodeset(name='n1', features=[
+    Feature(name='#creation_time', shape=[], format=INTEGER_64),
+    Feature(name='#id', shape=[], format=BYTES),
+    Feature(name='f1_seq', shape=[None], format=FLOAT_32, is_timeseries=true, timestamps='time'),
+    Feature(name='f2_seq', shape=[None, 4], format=FLOAT_32, is_timeseries=true, timestamps='time'),
+    Feature(name='time', shape=[None], format=INTEGER_64, is_timeseries=true)
+  ]),
+  Nodeset(name='n2', features=[
+    Feature(name='#id', shape=[], format=INTEGER_64),
+    Feature(name='sensor_ts', shape=[20, 8], format=FLOAT_32, is_timeseries=true)
+  ])
+], edgesets=[
+  Edgeset(name='e1', source_nodeset=0, target_nodeset=1, features=[
+    Feature(name='edge_time', shape=[None], format=INTEGER_64, is_timeseries=true),
+    Feature(name='edge_val', shape=[None], format=INTEGER_32, is_timeseries=true, timestamps='edge_time')
+  ])
+])""",
+    )
+
 
 if __name__ == "__main__":
   absltest.main()

dgf/src/data/schema_nb.cc

@@ -62,6 +62,22 @@ absl::StatusOr<GraphSchema::Feature> ParseFeatureSchema(
     return absl::InvalidArgumentError(
         absl::StrCat("Invalid shape type for feature '", feature_name, "'"));
   }
+
+  if (nb::hasattr(py_feature_schema, "is_timeseries")) {
+    nb::object py_is_timeseries = py_feature_schema.attr("is_timeseries");
+    if (!py_is_timeseries.is_none() &&
+        nb::isinstance<nb::bool_>(py_is_timeseries)) {
+      feature.is_timeseries = nb::cast<bool>(py_is_timeseries);
+    }
+  }
+
+  if (nb::hasattr(py_feature_schema, "timestamps")) {
+    nb::object py_timestamps = py_feature_schema.attr("timestamps");
+    if (!py_timestamps.is_none() && nb::isinstance<nb::str>(py_timestamps)) {
+      feature.timestamps = nb::cast<std::string>(py_timestamps);
+    }
+  }
+
   return feature;
 }
 

dgf/src/io/hgraph_in_beam.py

@@ -83,7 +83,6 @@ def read_graphai_hgraph(
       node_id_column=node_id_column,
       edge_id_column=edge_id_column,
       override_schema=override_schema,
-      research_node_format=research_node_format,
       remove_dangling_edges=remove_dangling_edges,
   )
 

dgf/src/validate/BUILD

@@ -54,5 +54,6 @@ py_test(
         "//dgf/src/data:schema",
         "//dgf/src/util:gen_test_graph",
         "//dgf/src/util:test_util",
+        # numpy dep,
     ],
 )

dgf/src/validate/in_memory_graph.py

@@ -122,6 +122,101 @@ def feature_set_issues(
             )
         )
 
+    if feature_schema.timestamps is not None:
+      if not feature_schema.is_timeseries:
+        items.append(
+            Issue.error(
+                f"The feature {feature_name!r} in {source} has timestamps set"
+                f" to {feature_schema.timestamps!r}, but is_timeseries is"
+                " False."
+            )
+        )
+      ts_name = feature_schema.timestamps
+      if ts_name not in featureset_schema:
+        items.append(
+            Issue.error(
+                f"The feature {feature_name!r} in {source} references"
+                f" timestamps feature {ts_name!r} which is not defined in the"
+                " schema."
+            )
+        )
+      else:
+        ts_schema = featureset_schema[ts_name]
+        if (
+            not ts_schema.is_timeseries
+            or ts_schema.semantic != schema_lib.FeatureSemantic.TIMESTAMP
+        ):
+          if not ts_schema.is_timeseries:
+            items.append(
+                Issue.error(
+                    f"The feature {feature_name!r} in {source} references"
+                    f" timestamps feature {ts_name!r}, but {ts_name!r} does not"
+                    " have is_timeseries=True."
+                )
+            )
+          if ts_schema.semantic != schema_lib.FeatureSemantic.TIMESTAMP:
+            items.append(
+                Issue.error(
+                    f"The feature {feature_name!r} in {source} references"
+                    f" timestamps feature {ts_name!r}, but {ts_name!r} does not"
+                    " have semantic=TIMESTAMP."
+                )
+            )
+          continue
+        ts_shape = ts_schema.shape or ()
+        feat_shape = feature_schema.shape or ()
+        if len(ts_shape) != 1:
+          items.append(
+              Issue.error(
+                  f"The feature {feature_name!r} in {source} references"
+                  f" timestamps feature {ts_name!r}, but {ts_name!r} must have"
+                  " exactly 1 sequence dimension in schema shape."
+              )
+          )
+        if len(feat_shape) < 1:
+          items.append(
+              Issue.error(
+                  f"The feature {feature_name!r} in {source} references"
+                  f" timestamps feature {ts_name!r}, but {feature_name!r} must"
+                  " have at least 1 sequence dimension in schema shape."
+              )
+          )
+        if (
+            len(ts_shape) == 1
+            and len(feat_shape) >= 1
+            and ts_shape[0] != feat_shape[0]
+        ):
+          items.append(
+              Issue.error(
+                  f"The feature {feature_name!r} in {source} has schema shape"
+                  f" {feat_shape} whose 0th dimension ({feat_shape[0]}) does not"
+                  f" match timestamps feature {ts_name!r} schema shape 0th"
+                  f" dimension ({ts_shape[0]})."
+              )
+          )
+        if feature_name in featureset_data and ts_name in featureset_data:
+          if feat_shape and feat_shape[0] is None:
+            feature_data = featureset_data[feature_name]
+            ts_data = featureset_data[ts_name]
+            if len(feature_data) == len(ts_data):
+              for i in range(len(feature_data)):
+                f_val = feature_data[i]
+                t_val = ts_data[i]
+                if f_val is not None and t_val is not None:
+                  f_len = len(f_val) if hasattr(f_val, "__len__") else 1
+                  t_len = len(t_val) if hasattr(t_val, "__len__") else 1
+                  if f_len != t_len:
+                    items.append(
+                        Issue.error(
+                            f"The feature {feature_name!r} in {source} has a"
+                            f" variable-length timeseries at index {i} of"
+                            f" length {f_len}, which does not match the"
+                            f" timestamps sequence {ts_name!r} of length"
+                            f" {t_len}."
+                        )
+                    )
+                    break
+
   return items