Move TimeseriesPadAndCap into padding stage.

dgf/src/api/data.py

@@ -49,6 +49,7 @@
 
 from dgf.src.data.padding import Padding
 from dgf.src.data.padding import EdgeSetPadding
+from dgf.src.data.padding import FeaturePadding
 from dgf.src.data.padding import NodeSetPadding
 
 from dgf.src.data.graph_snapshots_metadata import GraphSnapshotsFormat

dgf/src/data/padding.py

@@ -15,17 +15,24 @@
 """Padding of graphs."""
 
 import dataclasses
-from typing import Dict
+from typing import Dict, Optional
+
+
+@dataclasses.dataclass
+class FeaturePadding:
+  max_timeseries_len: Optional[int] = None
 
 
 @dataclasses.dataclass
 class NodeSetPadding:
-  num_nodes: int
+  num_nodes: Optional[int] = None
+  features: Dict[str, FeaturePadding] = dataclasses.field(default_factory=dict)
 
 
 @dataclasses.dataclass
 class EdgeSetPadding:
-  num_edges: int
+  num_edges: Optional[int] = None
+  features: Dict[str, FeaturePadding] = dataclasses.field(default_factory=dict)
 
 
 @dataclasses.dataclass

dgf/src/learning/ten_lines/dataset.py

@@ -24,7 +24,6 @@
 from dgf.src.sampling import config as sampling_config_lib
 from dgf.src.sampling import in_memory_sampler as in_memory_sampler_lib
 from dgf.src.transform import merge as merge_lib
-from dgf.src.transform import timeseries as timeseries_transform
 from dgf.src.util import temporal as temporal_util
 from dgf.src.util import util
 import numpy as np

dgf/src/transform/BUILD

@@ -28,10 +28,12 @@ py_library(
     name = "merge",
     srcs = ["merge.py"],
     deps = [
+        ":timeseries_padding",
         "//dgf/src/data:in_memory_graph",
         "//dgf/src/data:padding",
         "//dgf/src/data:schema",
         "//dgf/src/data:tf_in_memory_graph",
+        "//dgf/src/util:temporal",
         "//dgf/src/util/weak_dep:weak_dep_tensorflow",
         # numpy dep,
     ],
@@ -117,6 +119,19 @@ py_library(
     ],
 )
 
+py_library(
+    name = "timeseries_padding",
+    srcs = ["timeseries_padding.py"],
+    deps = [
+        "//dgf/src/data:in_memory_graph",
+        "//dgf/src/data:padding",
+        "//dgf/src/data:schema",
+        "//dgf/src/io:feature_format",
+        "//dgf/src/util:temporal",
+        # numpy dep,
+    ],
+)
+
 py_library(
     name = "timeseries",
     srcs = ["timeseries.py"],
@@ -154,8 +169,10 @@ py_test(
         # apache_beam dep,
         "//dgf/src/data:in_memory_graph",
         "//dgf/src/data:padding",
+        "//dgf/src/data:schema",
         "//dgf/src/data:tf_in_memory_graph",
         "//dgf/src/util:gen_test_graph",
+        "//dgf/src/util:temporal",
         "//dgf/src/util:test_util",
         # numpy dep,
         # tensorflow:tensorflow_no_contrib dep,
@@ -261,3 +278,18 @@ py_test(
         # numpy dep,
     ],
 )
+
+py_test(
+    name = "timeseries_padding_test",
+    srcs = ["timeseries_padding_test.py"],
+    deps = [
+        ":timeseries_padding",
+        # absl/testing:absltest dep,
+        # absl/testing:parameterized dep,
+        "//dgf/src/data:in_memory_graph",
+        "//dgf/src/data:padding",
+        "//dgf/src/data:schema",
+        "//dgf/src/util:test_util",
+        # numpy dep,
+    ],
+)