Merge pull request #278 from google/map

Map operator
google · Oct 10, 2023 · 95ce3ff · 95ce3ff
2 parents 267cccb + 1673611
commit 95ce3ff
Show file tree

Hide file tree

Showing 44 changed files with 904 additions and 2,305 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,13 +6,15 @@
 
 - Added `EventSet.tick_calendar()` operator.
 - Added `EventSet.where()` operator.
-- Add `filter_moving_count` operator.
+- Added `EventSet.filter_moving_count()` operator.
+- Added `EventSet.map()` operator.
 
 ### Improvements
 
 - Print `EventSet` timestamps as datetimes instead of float.
 - Support `sampling` argument in `EventSet.cumsum()` operator.
 - Using utf-8 codec to support non-ascii in string values.
+- New `tp.types` module to facilitate access to types used throughout the API.
 
 ### Fixes
 

diff --git a/docs/public_api_test.py b/docs/public_api_test.py
@@ -22,7 +22,6 @@
     "EventSet",
     "IndexData",
     "EventSetNode",
-    "EventSetOrNode",
     "Schema",
     "duration",
     "run",
@@ -32,6 +31,7 @@
     "plot",
     "compile",
     "config",
+    "types",
     # IO
     "to_csv",
     "from_csv",

diff --git a/docs/src/reference/index.md b/docs/src/reference/index.md
@@ -54,6 +54,7 @@ Check the index on the left for a more detailed description of any symbol.
 | [`EventSet.join()`][temporian.EventSet.join]                                                               | Join [`EventSets`][temporian.EventSet] with different samplings but the same index together.                 |
 | [`EventSet.lag()`][temporian.EventSet.lag]                                                                 | Adds a delay to an [`EventSet`][temporian.EventSet]'s timestamps.                                            |
 | [`EventSet.leak()`][temporian.EventSet.leak]                                                               | Subtracts a duration from an [`EventSet`][temporian.EventSet]'s timestamps.                                  |
+| [`EventSet.map()`][temporian.EventSet.map]                                                                 | Applies a function on each of an [`EventSet`][temporian.EventSet]'s values.                                  |
 | [`EventSet.prefix()`][temporian.EventSet.prefix]                                                           | Adds a prefix to the names of the features in an [`EventSet`][temporian.EventSet].                           |
 | [`EventSet.propagate()`][temporian.EventSet.propagate]                                                     | Propagates feature values over a sub index.                                                                  |
 | [`EventSet.rename()`][temporian.EventSet.rename]                                                           | Renames an [`EventSet`][temporian.EventSet]'s features and index.                                            |

diff --git a/docs/src/reference/temporian/internals/typing/WindowLength.md b/docs/src/reference/temporian/internals/typing/WindowLength.md
diff --git a/docs/src/reference/temporian/operators/map.md b/docs/src/reference/temporian/operators/map.md
@@ -0,0 +1 @@
+::: temporian.EventSet.map
diff --git a/docs/src/reference/temporian/other/typing/EventSetCollection.md b/docs/src/reference/temporian/other/typing/EventSetCollection.md
diff --git a/docs/src/reference/temporian/other/typing/EventSetNodeCollection.md b/docs/src/reference/temporian/other/typing/EventSetNodeCollection.md
diff --git a/docs/src/reference/temporian/other/typing/IndexKey.md b/docs/src/reference/temporian/other/typing/IndexKey.md
diff --git a/docs/src/reference/temporian/other/typing/IndexKeyItem.md b/docs/src/reference/temporian/other/typing/IndexKeyItem.md
diff --git a/docs/src/reference/temporian/other/typing/IndexKeyList.md b/docs/src/reference/temporian/other/typing/IndexKeyList.md
diff --git a/docs/src/reference/temporian/other/typing/NodeToEventSetMapping.md b/docs/src/reference/temporian/other/typing/NodeToEventSetMapping.md
diff --git a/.../temporian/other/typing/EventSetOrNode.md → docs/src/reference/temporian/types.md b/.../temporian/other/typing/EventSetOrNode.md → docs/src/reference/temporian/types.md
diff --git a/docs/src/tutorials/bank_fraud_detection_with_tfdf.ipynb b/docs/src/tutorials/bank_fraud_detection_with_tfdf.ipynb
@@ -581,7 +581,7 @@
    ],
    "source": [
     "@tp.compile\n",
-    "def augment_transactions(transactions: tp.EventSetOrNode) -> Dict[str, tp.EventSetOrNode]:\n",
+    "def augment_transactions(transactions: tp.types.EventSetOrNode) -> Dict[str, tp.types.EventSetOrNode]:\n",
     "    \"\"\"Temporian function to augment the transactions with temporal features.\"\"\"\n",
     "    print(\"TRANSACTIONS:\\n\", transactions.schema, sep = '')\n",
     "\n",

diff --git a/docs/src/tutorials/heart_rate_analysis.ipynb b/docs/src/tutorials/heart_rate_analysis.ipynb
diff --git a/docs/src/user_guide.ipynb b/docs/src/user_guide.ipynb
@@ -382,7 +382,7 @@
    "outputs": [],
    "source": [
     "@tp.compile\n",
-    "def my_function(x : tp.EventSetOrNode) -> tp.EventSetOrNode:\n",
+    "def my_function(x : tp.types.EventSetOrNode) -> tp.types.EventSetOrNode:\n",
     "    return x.simple_moving_average(window_length=0.5)\n",
     "\n",
     "# Feeding an EventSet\n",
@@ -450,7 +450,7 @@
    "outputs": [],
    "source": [
     "@tp.compile\n",
-    "def my_function(x : tp.EventSetOrNode, a:bool) -> tp.EventSetOrNode:\n",
+    "def my_function(x : tp.types.EventSetOrNode, a:bool) -> tp.types.EventSetOrNode:\n",
     "    if a:\n",
     "        return x.rename(\"a_branch\")\n",
     "    else:\n",
@@ -485,7 +485,7 @@
    "outputs": [],
    "source": [
     "@tp.compile\n",
-    "def my_function(x : tp.EventSetOrNode) -> tp.EventSetOrNode:\n",
+    "def my_function(x : tp.types.EventSetOrNode) -> tp.types.EventSetOrNode:\n",
     "    return x[\"value\"].filter(x[\"condition\"])\n",
     "\n",
     "my_function(tp.event_set(\n",

diff --git a/temporian/__init__.py b/temporian/__init__.py
@@ -85,7 +85,7 @@
 from temporian.core.compilation import compile
 
 # Types
-from temporian.core.typing import EventSetOrNode
+from temporian.core import types
 
 # Runtime check
 from temporian.utils.typecheck import runtime_check_raise_exception

diff --git a/temporian/beam/evaluation.py b/temporian/beam/evaluation.py
@@ -139,7 +139,7 @@ def run_multi_io(
 
     num_steps = len(schedule.steps)
     for step_idx, step in enumerate(schedule.steps):
-        operator_def = step.op.definition()
+        operator_def = step.op.definition
 
         if verbose > 0:
             print("=============================", file=sys.stderr)

diff --git a/temporian/core/BUILD b/temporian/core/BUILD
@@ -100,3 +100,19 @@ py_library(
         "//temporian/implementation/numpy/data:event_set",
     ],
 )
+
+py_library(
+    name = "dataclasses",
+    srcs = ["dataclasses.py"],
+    srcs_version = "PY3",
+)
+
+py_library(
+    name = "types",
+    srcs = ["types.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":dataclasses",
+        ":typing",
+    ],
+)
diff --git a/temporian/core/compilation.py b/temporian/core/compilation.py
@@ -33,8 +33,8 @@ def compile(fn: Optional[F] = None, *, verbose: int = 0) -> F:
     """Compiles a Temporian function.
 
     A Temporian function is a function that takes
-    [`EventSetOrNodes`][temporian.EventSetOrNode] as arguments and returns
-    [`EventSetOrNodes`][temporian.EventSetOrNode] as outputs.
+    [`EventSetOrNodes`][temporian.types.EventSetOrNode] as arguments and returns
+    [`EventSetOrNodes`][temporian.types.EventSetOrNode] as outputs.
 
     Compiling a function allows Temporian to optimize the underlying graph
     defined by the operators used inside the function, making it run on

diff --git a/temporian/core/data/dtype.py b/temporian/core/data/dtype.py
@@ -14,11 +14,9 @@
 
 """Data types declaration."""
 
-from typing import Any
-
 import math
 from enum import Enum
-from typing import Union
+from typing import Any, Union
 
 
 class DType(Enum):
@@ -80,23 +78,23 @@ def from_python_value(cls, value: Any) -> "DType":
         """
 
         try:
-            return _PY_TYPE_TO_DTYPE[type(value)]
+            return PY_TYPE_TO_DTYPE[type(value)]
         except KeyError as e:
             raise ValueError(
                 f"Couldn't find a dtype to store a value of type {type(value)}."
                 f" Value is: {value}"
             ) from e
 
 
-_PY_TYPE_TO_DTYPE = {
+PY_TYPE_TO_DTYPE = {
     float: DType.FLOAT64,
     int: DType.INT64,
     str: DType.STRING,
     bytes: DType.STRING,
     bool: DType.BOOLEAN,
 }
 
-_DTYPE_TO_PY_TYPES = {
+DTYPE_TO_PY_TYPES = {
     DType.FLOAT64: float,
     DType.FLOAT32: float,
     DType.INT64: int,
@@ -107,7 +105,7 @@ def from_python_value(cls, value: Any) -> "DType":
 
 
 def tp_dtype_to_py_type(dtype: DType) -> Any:
-    return _DTYPE_TO_PY_TYPES[dtype]
+    return DTYPE_TO_PY_TYPES[dtype]
 
 
 # The dtype of indexes.

diff --git a/temporian/core/data/schema.py b/temporian/core/data/schema.py
@@ -15,10 +15,9 @@
 """Schema and related classes."""
 
 from __future__ import annotations
-
+from dataclasses import dataclass
 from typing import List, Tuple, Dict, Union
 
-from dataclasses import dataclass
 from temporian.core.data.dtype import DType, IndexDType
 
 

diff --git a/temporian/core/dataclasses.py b/temporian/core/dataclasses.py
@@ -0,0 +1,35 @@
+# Copyright 2021 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
+#
+#     https://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.
+
+"""Data class definitions."""
+
+from dataclasses import dataclass
+
+from temporian.core.typing import IndexKey
+
+
+@dataclass
+class MapExtras:
+    """Object containing information about a value's position in an
+    [`EventSet`][temporian.EventSet].
+
+    Attributes:
+        index_key: The index the value belongs to.
+        timestamp: The timestamp of the value's event.
+        feature_name: The name of the feature the value belongs to.
+    """
+
+    index_key: IndexKey
+    timestamp: float
+    feature_name: str