Add op tests

temporian/core/operators/add_index.py

@@ -37,10 +37,10 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self._output_feature_schemas = self._get_output_feature_schemas(
+        output_feature_schemas = self._get_output_feature_schemas(
             input, indexes
         )
-        self._output_indexes = self._get_output_indexes(input, indexes)
+        output_indexes = self._get_output_indexes(input, indexes)
         self._indexes = indexes
 
         self.add_input("input", input)
@@ -50,8 +50,8 @@ def __init__(
         self.add_output(
             "output",
             create_node_new_features_new_sampling(
-                features=self._output_feature_schemas,
-                indexes=self._output_indexes,
+                features=output_feature_schemas,
+                indexes=output_indexes,
                 is_unix_timestamp=input.schema.is_unix_timestamp,
                 creator=self,
             ),
@@ -75,12 +75,12 @@ def _get_output_indexes(
 
         new_indexes: List[IndexSchema] = []
         for index in indexes:
-            if index not in feature_dict:
-                raise ValueError(f"{index} is not a feature in input.")
-
             if index in index_dict:
                 raise ValueError(f"{index} is already an index in input.")
 
+            if index not in feature_dict:
+                raise ValueError(f"{index} is not a feature in input.")
+
             new_indexes.append(
                 IndexSchema(name=index, dtype=feature_dict[index])
             )
@@ -101,14 +101,6 @@ def build_op_definition(cls) -> pb.OperatorDef:
             outputs=[pb.OperatorDef.Output(key="output")],
         )
 
-    @property
-    def output_feature_schemas(self) -> List[FeatureSchema]:
-        return self._output_feature_schemas
-
-    @property
-    def output_indexes(self) -> List[IndexSchema]:
-        return self._output_indexes
-
     @property
     def indexes(self) -> List[str]:
         return self._indexes

temporian/core/operators/drop_index.py

@@ -50,7 +50,7 @@ def __init__(
             input, indexes, keep
         )
 
-        self._output_indexes = [
+        output_indexes = [
             index_level
             for index_level in input.schema.indexes
             if index_level.name not in indexes
@@ -60,7 +60,7 @@ def __init__(
             "output",
             create_node_new_features_new_sampling(
                 features=self._output_feature_schemas,
-                indexes=self._output_indexes,
+                indexes=output_indexes,
                 is_unix_timestamp=input.schema.is_unix_timestamp,
                 creator=self,
             ),
@@ -74,19 +74,9 @@ def _get_output_feature_schemas(
             return input.schema.features
 
         index_dict = input.schema.index_name_to_dtype()
-        feature_dict = input.schema.feature_name_to_dtype()
 
         new_features: List[FeatureSchema] = []
         for index in indexes:
-            if index not in index_dict:
-                raise ValueError(f"{index} is not an index in input.")
-
-            if index in feature_dict:
-                raise ValueError(
-                    f"{index} already exists in input's features. If you"
-                    " want to drop the index, specify `keep=False`."
-                )
-
             new_features.append(
                 FeatureSchema(name=index, dtype=index_dict[index])
             )
@@ -98,10 +88,6 @@ def _get_output_feature_schemas(
     def output_feature_schemas(self) -> List[FeatureSchema]:
         return self._output_feature_schemas
 
-    @property
-    def output_indexes(self) -> List[IndexSchema]:
-        return self._output_indexes
-
     @property
     def indexes(self) -> List[str]:
         return self._indexes
@@ -150,7 +136,7 @@ def _normalize_indexes(
     index_dict = input.schema.index_name_to_dtype()
     for index in indexes:
         if index not in index_dict:
-            raise KeyError(
+            raise ValueError(
                 f"{index} is not an index in {input.schema.indexes}."
             )
 

temporian/core/operators/glue.py

@@ -41,11 +41,11 @@ def __init__(
         # serialization.
 
         if len(inputs) < 2:
-            raise ValueError("At least two arguments should be provided")
+            raise ValueError("At least two arguments should be provided.")
 
         if len(inputs) >= MAX_NUM_ARGUMENTS:
             raise ValueError(
-                f"Too many (>{MAX_NUM_ARGUMENTS}) arguments provided"
+                f"Too many (>{MAX_NUM_ARGUMENTS}) arguments provided."
             )
 
         # inputs

temporian/core/operators/test/test_add_index.py

@@ -23,65 +23,65 @@ class AddIndexTest(TestCase):
     def setUp(self) -> None:
         self.timestamps = [0.4, 0.5, 0.6, 0.1, 0.2, 0.3, 0.4, 0.3]
         self.features = {
-            "state_id": ["A", "A", "A", "B", "B", "B", "B", "B"],
-            "store_id": [0, 0, 0, 0, 0, 1, 1, 1],
-            "item_id": [1, 1, 1, 2, 2, 2, 2, 3],
-            "sales": [10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0],
+            "a": ["A", "A", "A", "B", "B", "B", "B", "B"],
+            "b": [0, 0, 0, 0, 0, 1, 1, 1],
+            "c": [1, 1, 1, 2, 2, 2, 2, 3],
+            "d": [10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0],
         }
         self.evset = event_set(
             timestamps=self.timestamps,
             features=self.features,
-            indexes=["state_id"],
+            indexes=["a"],
         )
 
     def test_add_index_single(self):
-        result = self.evset.add_index("store_id")
+        result = self.evset.add_index("b")
 
         expected = event_set(
             timestamps=self.timestamps,
             features=self.features,
-            indexes=["state_id", "store_id"],
+            indexes=["a", "b"],
         )
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
     def test_add_index_multiple(self):
-        result = self.evset.add_index(["store_id", "item_id"])
+        result = self.evset.add_index(["b", "c"])
 
         expected = event_set(
             timestamps=self.timestamps,
             features=self.features,
-            indexes=["state_id", "store_id", "item_id"],
+            indexes=["a", "b", "c"],
         )
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
     def test_set_index_single(self):
-        result = self.evset.set_index("store_id")
+        result = self.evset.set_index("b")
 
         expected = event_set(
             timestamps=self.timestamps,
             # Old index is now the last feature
             features={
-                **{k: v for k, v in self.features.items() if k != "state_id"},
-                "state_id": self.features["state_id"],
+                **{k: v for k, v in self.features.items() if k != "a"},
+                "a": self.features["a"],
             },
-            indexes=["store_id"],
+            indexes=["b"],
         )
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
     def test_set_index_multiple(self):
-        result = self.evset.set_index(["store_id", "item_id"])
+        result = self.evset.set_index(["b", "c"])
 
         expected = event_set(
             timestamps=self.timestamps,
             # Old index is now the last feature
             features={
-                **{k: v for k, v in self.features.items() if k != "state_id"},
-                "state_id": self.features["state_id"],
+                **{k: v for k, v in self.features.items() if k != "a"},
+                "a": self.features["a"],
             },
-            indexes=["store_id", "item_id"],
+            indexes=["b", "c"],
         )
 
         assertOperatorResult(self, result, expected, check_sampling=False)
@@ -178,6 +178,26 @@ def test_mix(self):
         self.assertEqual(result.data[(2, b"X")].timestamps.tolist(), [4, 5])
         self.assertEqual(result.data[(3, b"Z")].timestamps.tolist(), [6])
 
+    def test_target_doesnt_exist(self):
+        with self.assertRaisesRegex(ValueError, "is not a feature in input"):
+            self.evset.add_index("e")
+
+    def test_target_is_index(self):
+        with self.assertRaisesRegex(ValueError, "is already an index in input"):
+            self.evset.add_index("a")
+
+    def test_empty_list(self):
+        with self.assertRaisesRegex(
+            ValueError, "Cannot specify empty list as `indexes` argument"
+        ):
+            self.evset.add_index([])
+
+    def test_empty_list_set(self):
+        with self.assertRaisesRegex(
+            ValueError, "Cannot specify empty list as `indexes` argument"
+        ):
+            self.evset.set_index([])
+
 
 if __name__ == "__main__":
     absltest.main()

temporian/core/operators/test/test_begin.py

@@ -35,6 +35,15 @@ def test_basic(self):
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
+    def test_empty(self):
+        evset = event_set(timestamps=[], features={"a": []})
+
+        result = evset.begin()
+
+        expected = event_set(timestamps=[])
+
+        assertOperatorResult(self, result, expected, check_sampling=False)
+
 
 if __name__ == "__main__":
     absltest.main()

temporian/core/operators/test/test_drop_index.py

@@ -23,26 +23,26 @@ class DropIndexTest(TestCase):
     def setUp(self) -> None:
         self.timestamps = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
         self.features = {
-            "state_id": ["A", "A", "A", "B", "B", "B", "B", "B"],
-            "store_id": [0, 0, 0, 0, 0, 1, 1, 1],
-            "item_id": [1, 1, 1, 2, 2, 2, 2, 3],
-            "sales": [10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0],
+            "a": ["A", "A", "A", "B", "B", "B", "B", "B"],
+            "b": [0, 0, 0, 0, 0, 1, 1, 1],
+            "c": [1, 1, 1, 2, 2, 2, 2, 3],
+            "d": [10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0],
         }
         self.evset = event_set(
             timestamps=self.timestamps,
             features=self.features,
-            indexes=["store_id", "item_id"],
+            indexes=["b", "c"],
         )
 
     def test_drop_all(self) -> None:
         expected = event_set(
             timestamps=self.timestamps,
             # Old indexes are now the last features
             features={
-                "state_id": self.features["state_id"],
-                "sales": self.features["sales"],
-                "store_id": self.features["store_id"],
-                "item_id": self.features["item_id"],
+                "a": self.features["a"],
+                "d": self.features["d"],
+                "b": self.features["b"],
+                "c": self.features["c"],
             },
         )
 
@@ -55,15 +55,15 @@ def test_drop_single_first(self) -> None:
             timestamps=self.timestamps,
             # Old indexes are now the last features
             features={
-                "item_id": self.features["item_id"],
-                "state_id": self.features["state_id"],
-                "sales": self.features["sales"],
-                "store_id": self.features["store_id"],
+                "c": self.features["c"],
+                "a": self.features["a"],
+                "d": self.features["d"],
+                "b": self.features["b"],
             },
-            indexes=["item_id"],
+            indexes=["c"],
         )
 
-        result = self.evset.drop_index("store_id")
+        result = self.evset.drop_index("b")
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
@@ -72,15 +72,15 @@ def test_drop_single_second(self) -> None:
             timestamps=self.timestamps,
             # Old indexes are now the last features
             features={
-                "store_id": self.features["store_id"],
-                "state_id": self.features["state_id"],
-                "sales": self.features["sales"],
-                "item_id": self.features["item_id"],
+                "b": self.features["b"],
+                "a": self.features["a"],
+                "d": self.features["d"],
+                "c": self.features["c"],
             },
-            indexes=["store_id"],
+            indexes=["b"],
         )
 
-        result = self.evset.drop_index("item_id")
+        result = self.evset.drop_index("c")
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
@@ -89,14 +89,14 @@ def test_drop_single_keep_false(self) -> None:
             timestamps=self.timestamps,
             # Old indexes are now the last features
             features={
-                "store_id": self.features["store_id"],
-                "state_id": self.features["state_id"],
-                "sales": self.features["sales"],
+                "b": self.features["b"],
+                "a": self.features["a"],
+                "d": self.features["d"],
             },
-            indexes=["store_id"],
+            indexes=["b"],
         )
 
-        result = self.evset.drop_index("item_id", keep=False)
+        result = self.evset.drop_index("c", keep=False)
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
@@ -122,6 +122,16 @@ def test_str_index(self) -> None:
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
+    def test_wrong_index(self):
+        with self.assertRaisesRegex(ValueError, "x is not an index in"):
+            self.evset.drop_index("x")
+
+    def test_empty_list(self):
+        with self.assertRaisesRegex(
+            ValueError, "Cannot specify empty list as `indexes` argument"
+        ):
+            self.evset.drop_index([])
+
 
 if __name__ == "__main__":
     absltest.main()

temporian/core/operators/test/test_end.py

@@ -35,6 +35,15 @@ def test_basic(self):
 
         assertOperatorResult(self, result, expected, check_sampling=False)
 
+    def test_empty(self):
+        evset = event_set(timestamps=[], features={"a": []})
+
+        result = evset.end()
+
+        expected = event_set(timestamps=[])
+
+        assertOperatorResult(self, result, expected, check_sampling=False)
+
 
 if __name__ == "__main__":
     absltest.main()