Merge pull request #635 from qinzzz/issue-606

implement add_annotation_raw

Author: qinzzz

forte/data/data_store.py

@@ -11,11 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-from typing import List, Iterator, Tuple, Optional, Any
+from typing import Dict, List, Iterator, Tuple, Optional, Any
 import uuid
 from bisect import bisect_left
 from heapq import heappush, heappop
+from sortedcontainers import SortedList
 
 from forte.utils import get_class
 from forte.data.base_store import BaseStore
@@ -29,7 +29,9 @@ class DataStore(BaseStore):
     # TODO: temporarily disable this for development purposes.
     # pylint: disable=pointless-string-statement
 
-    def __init__(self, onto_file_path: Optional[str] = None):
+    def __init__(
+        self, onto_file_path: Optional[str] = None, dynamically_add_type=True
+    ):
         r"""An implementation of the data store object that mainly uses
         primitive types. This class will be used as the internal data
         representation behind data pack. The usage of primitive types provides
@@ -121,15 +123,23 @@ def __init__(self, onto_file_path: Optional[str] = None):
             onto_file_path (str, optional): the path to the ontology file.
         """
         super().__init__()
-        self.onto_file_path = onto_file_path
+
+        if onto_file_path is None and not dynamically_add_type:
+            raise RuntimeError(
+                "DataStore is initialized with no existing types. Setting"
+                "dynamically_add_type to False without providing onto_file_path"
+                "will lead to no usable type in DataStore."
+            )
+        self._onto_file_path = onto_file_path
+        self._dynamically_add_type = dynamically_add_type
 
         """
         The ``_type_attributes`` is a private dictionary that provides
         ``type_name``, their parent entry, and the order of corresponding attributes.
         The keys are fully qualified names of every type; The value is a dictionary with
         two keys. Key ``attribute`` provides an inner dictionary with all valid attributes
-        for this type and the indices of attributes among these lists. Key ``parent_entry``
-        is a string representing the direct parent of this type.
+        for this type and the indices of attributes among these lists. Key ``parent_class``
+        is a string representing the ancesters of this type.
 
         This structure is supposed to be built dynamically. When a user adds new entries,
         data_store will check unknown types and add them to ``_type_attributes``.
@@ -144,20 +154,20 @@ def __init__(self, onto_file_path: Optional[str] = None):
             #       "attributes": {"pos": 4, "ud_xpos": 5,
             #               "lemma": 6, "chunk": 7, "ner": 8, "sense": 9,
             #               "is_root": 10, "ud_features": 11, "ud_misc": 12},
-            #       "parent_entry": "forte.data.ontology.top.Annotation", },
+            #       "parent_class": "forte.data.ontology.top.Annotation", },
             #     "ft.onto.base_ontology.Document": {
             #       "attributes": {"document_class": 4,
             #               "sentiment": 5, "classifications": 6},
-            #       "parent_entry": "forte.data.ontology.top.Annotation", },
+            #       "parent_class": "forte.data.ontology.top.Annotation", },
             #     "ft.onto.base_ontology.Sentence": {
             #       "attributes": {"speaker": 4,
             #               "part_id": 5, "sentiment": 6,
             #               "classification": 7, "classifications": 8},
-            #       "parent_entry": "forte.data.ontology.top.Annotation", }
+            #       "parent_class": "forte.data.ontology.top.Annotation", }
             # }
         """
         self._type_attributes: dict = {}
-        if self.onto_file_path:
+        if self._onto_file_path:
             self._parse_onto_file()
 
         """
@@ -190,6 +200,104 @@ def _new_tid(self) -> int:
         r"""This function generates a new ``tid`` for an entry."""
         return uuid.uuid4().int
 
+    def _get_type_info(self, type_name: str) -> Dict[str, Any]:
+        """
+        Get the dictionary containing type information from ``self._type_attributes``.
+        If the ``type_name`` does not currecntly exists and dynamic import is enabled,
+        this function will add a new key-value pair into ``self._type_attributes``. The
+        value consists of a full attribute-to-index dictionary and an empty parent set.
+
+        This function returns a dictionary containing an attribute dict and a set of parent
+        entries of the given type. For example:
+
+        .. code-block:: python
+
+            "ft.onto.base_ontology.Sentence": {
+                    "attributes": {
+                        "speaker": 4,
+                        "part_id": 5,
+                        "sentiment": 6,
+                        "classification": 7,
+                        "classifications": 8,
+                    },
+                    "parent_class": set(),
+                }
+
+        Args:
+            type_name (str): The fully qualified type name of a type.
+        Returns:
+            attr_dict (dict): The dictionary containing an attribute dict and a set of parent
+            entries of the given type.
+        Raises:
+            RuntimeError: When the type is not provided by ontology file and
+            dynamic import is disabled.
+        """
+        # check if type is in dictionary
+        if type_name in self._type_attributes:
+            return self._type_attributes[type_name]
+        if not self._dynamically_add_type:
+            raise ValueError(
+                f"{type_name} is not an existing type in current data store."
+                f"Dynamically add type is disabled."
+                f"Set dynamically_add_type=True if you need to use types other than"
+                f"types specified in the ontology file."
+            )
+        # get attribute dictionary
+        attributes = self._get_entry_attributes_by_class(type_name)
+
+        attr_dict = {}
+        attr_idx = constants.ENTRY_TYPE_INDEX + 1
+        for attr_name in attributes:
+            attr_dict[attr_name] = attr_idx
+            attr_idx += 1
+
+        new_entry_info = {
+            "attributes": attr_dict,
+            "parent_class": set(),
+        }
+        self._type_attributes[type_name] = new_entry_info
+
+        return new_entry_info
+
+    def _get_type_attribute_dict(self, type_name: str) -> Dict[str, int]:
+        """Get the attribute dict of an entry type. The attribute dict maps
+        attribute names to a list of consecutive integers as indicies. For example:
+        .. code-block:: python
+
+            "attributes": {
+                        "speaker": 4,
+                        "part_id": 5,
+                        "sentiment": 6,
+                        "classification": 7,
+                        "classifications": 8,
+            },
+
+        Args:
+            type_name (str): The fully qualified type name of a type.
+        Returns:
+            attr_dict (dict): The attribute-to-index dictionary of an entry.
+        """
+        return self._get_type_info(type_name)["attributes"]
+
+    def _get_type_parent(self, type_name: str) -> str:
+        """Get a set of parent names of an entry type. The set is a subset of all
+        ancestors of the given type.
+        Args:
+            type_name (str): The fully qualified type name of a type.
+        Returns:
+            parent_class (str): The parent entry name of an entry.
+        """
+        return self._get_type_info(type_name)["parent_class"]
+
+    def _num_attributes_for_type(self, type_name: str) -> int:
+        """Get the length of the attribute dict of an entry type.
+        Args:
+            type_name (str): The fully qualified type name of the new entry.
+        Returns:
+            attr_dict (dict): The attributes-to-index dict of an entry.
+        """
+        return len(self._get_type_attribute_dict(type_name))
+
     def _new_annotation(self, type_name: str, begin: int, end: int) -> List:
         r"""This function generates a new annotation with default fields.
         All default fields are filled with None.
@@ -207,8 +315,10 @@ def _new_annotation(self, type_name: str, begin: int, end: int) -> List:
 
         tid: int = self._new_tid()
         entry: List[Any]
+
         entry = [begin, end, tid, type_name]
-        entry += len(self._type_attributes[type_name]) * [None]
+        entry += self._num_attributes_for_type(type_name) * [None]
+
         return entry
 
     def _new_link(
@@ -230,8 +340,10 @@ def _new_link(
 
         tid: int = self._new_tid()
         entry: List[Any]
+
         entry = [parent_tid, child_tid, tid, type_name]
-        entry += len(self._type_attributes[type_name]) * [None]
+        entry += self._num_attributes_for_type(type_name) * [None]
+
         return entry
 
     def _new_group(self, type_name: str, member_type: str) -> List:
@@ -249,21 +361,22 @@ def _new_group(self, type_name: str, member_type: str) -> List:
         """
 
         tid: int = self._new_tid()
+
         entry = [member_type, [], tid, type_name]
-        entry += len(self._type_attributes[type_name]) * [None]
+        entry += self._num_attributes_for_type(type_name) * [None]
+
         return entry
 
     def _is_annotation(self, type_name: str) -> bool:
         r"""This function takes a type_id and returns whether a type
         is an annotation type or not.
-
         Args:
             type_name (str): The name of type in `self.__elements`.
-
         Returns:
             A boolean value whether this type_id belongs to an annotation
             type or not.
         """
+        # TODO: use is_subclass() in DataStore to replace this
         entry_class = get_class(type_name)
         return issubclass(entry_class, (Annotation, AudioAnnotation))
 
@@ -286,7 +399,15 @@ def add_annotation_raw(self, type_name: str, begin: int, end: int) -> int:
         # annotation type entry data with default fields.
         # A reference to the entry should be store in both self.__elements and
         # self.__entry_dict.
-        raise NotImplementedError
+        entry = self._new_annotation(type_name, begin, end)
+        try:
+            self.__elements[type_name].add(entry)
+        except KeyError:
+            self.__elements[type_name] = SortedList(key=lambda s: (s[0], s[1]))
+            self.__elements[type_name].add(entry)
+        tid = entry[constants.TID_INDEX]
+        self.__entry_dict[tid] = entry
+        return tid
 
     def add_link_raw(
         self, type_name: str, parent_tid: int, child_tid: int
@@ -340,16 +461,17 @@ def set_attribute(self, tid: int, attr_name: str, attr_value: Any):
             KeyError: when ``tid`` or ``attr_name`` is not found.
         """
         try:
-            entry_type = self.__entry_dict[tid][constants.ENTRY_TYPE_INDEX]
+            entry = self.__entry_dict[tid]
+            entry_type = entry[constants.ENTRY_TYPE_INDEX]
         except KeyError as e:
             raise KeyError(f"Entry with tid {tid} not found.") from e
 
         try:
-            attr_id = self._type_attributes[entry_type][attr_name]
+            attr_id = self._get_type_attribute_dict(entry_type)[attr_name]
         except KeyError as e:
             raise KeyError(f"{entry_type} has no {attr_name} attribute.") from e
 
-        self._set_attr(tid, attr_id, attr_value)
+        entry[attr_id] = attr_value
 
     def _set_attr(self, tid: int, attr_id: int, attr_value: Any):
         r"""This function locates the entry data with ``tid`` and sets its
@@ -381,16 +503,17 @@ def get_attribute(self, tid: int, attr_name: str) -> Any:
             KeyError: when ``tid`` or ``attr_name`` is not found.
         """
         try:
-            entry_type = self.__entry_dict[tid][constants.ENTRY_TYPE_INDEX]
+            entry = self.__entry_dict[tid]
+            entry_type = entry[constants.ENTRY_TYPE_INDEX]
         except KeyError as e:
             raise KeyError(f"Entry with tid {tid} not found.") from e
 
         try:
-            attr_id = self._type_attributes[entry_type][attr_name]
+            attr_id = self._get_type_attribute_dict(entry_type)[attr_name]
         except KeyError as e:
             raise KeyError(f"{entry_type} has no {attr_name} attribute.") from e
 
-        return self._get_attr(tid, attr_id)
+        return entry[attr_id]
 
     def _get_attr(self, tid: int, attr_id: int) -> Any:
         r"""This function locates the entry data with ``tid`` and gets the value
@@ -439,7 +562,7 @@ def delete_entry(self, tid: int):
         if self._is_annotation(type_name):
             entry_index = bisect_left(target_list, entry_data)
         else:  # if it's group or link, use the index in entry_list
-            entry_index = entry_data[-1]
+            entry_index = entry_data[constants.ENTRY_INDEX_INDEX]
 
         if (
             entry_index >= len(target_list)
@@ -455,8 +578,6 @@ def delete_entry(self, tid: int):
     def _delete_entry_by_loc(self, type_name: str, index_id: int):
         r"""It removes an entry of `index_id` by taking both the `type_id`
         and `index_id`. Called by `delete_entry()`.
-        This function will raise an IndexError if the `type_id` or `index_id`
-        is invalid.
 
         Args:
             type_id (int): The index of the list in ``self.__elements``.
@@ -769,7 +890,7 @@ def _parse_onto_file(self):
         A user can use classes both in the ontology specification file and their parent
         entries's paths.
         """
-        if self.onto_file_path is None:
+        if self._onto_file_path is None:
             return
         raise NotImplementedError