feat: object tracking

eva/binder/binder_utils.py

@@ -18,7 +18,11 @@
 from typing import TYPE_CHECKING, List
 
 from eva.catalog.catalog_type import TableType
-from eva.catalog.catalog_utils import is_string_col, is_video_table
+from eva.catalog.catalog_utils import (
+    get_video_table_column_definitions,
+    is_string_col,
+    is_video_table,
+)
 from eva.expression.function_expression import FunctionExpression
 from eva.parser.alias import Alias
 
@@ -132,3 +136,74 @@ def resolve_alias_table_value_expression(node: FunctionExpression):
     assert len(node.alias.col_names) == len(
         node.output_objs
     ), f"""Expected {len(node.output_objs)} output columns for {node.alias.alias_name}, got {len(node.alias.col_names)}."""
+
+
+def handle_bind_extract_object_function(
+    node: FunctionExpression, binder_context: StatementBinderContext
+):
+    """Handles the binding of extract_object function.
+        1. Bind the source video data
+        2. Create and bind the detector function expression using the provided name.
+        3. Create and bind the tracker function expression.
+            Its inputs are id, data, output of detector.
+        4. Bind the EXTRACT_OBJECT function expression and append the new children.
+        5. Handle the alias and populate the outputs of the EXTRACT_OBJECT function
+
+    Args:
+        node (FunctionExpression): The function expression representing the extract object operation.
+        binder_context (StatementBinderContext): The context object used to bind expressions in the statement.
+
+    Raises:
+        AssertionError: If the number of children in the `node` is not equal to 3.
+    """
+    assert (
+        len(node.children) == 3
+    ), f"Invalid arguments provided to {node}. Example correct usage, (data, Detector, Tracker)"
+
+    # 1. Bind the source video
+    video_data = node.children[0]
+    binder_context.bind(video_data)
+
+    # 2. Construct the detector
+    # convert detector to FunctionExpression before binding
+    # eg. YoloV5 -> YoloV5(data)
+    detector = FunctionExpression(None, node.children[1].col_name)
+    detector.append_child(video_data.copy())
+    binder_context.bind(detector)
+
+    # 3. Construct the tracker
+    # convert tracker to FunctionExpression before binding
+    # eg. ByteTracker -> ByteTracker(id, data, labels, bboxes, scores)
+    tracker = FunctionExpression(None, node.children[2].col_name)
+    # create the video id expression
+    columns = get_video_table_column_definitions()
+    tracker.append_child(
+        TupleValueExpression(
+            col_name=columns[1].name, table_alias=video_data.table_alias
+        )
+    )
+    tracker.append_child(video_data.copy())
+    binder_context.bind(tracker)
+    # append the bound output of detector
+    for obj in detector.output_objs:
+        col_alias = "{}.{}".format(obj.udf_name.lower(), obj.name.lower())
+        child = TupleValueExpression(
+            obj.name,
+            table_alias=obj.udf_name.lower(),
+            col_object=obj,
+            col_alias=col_alias,
+        )
+        tracker.append_child(child)
+
+    # 4. Bind the EXTRACT_OBJECT expression and append the new children.
+    node.children = []
+    node.children = [video_data, detector, tracker]
+
+    # 5. assign the outputs of tracker to the output of extract_object
+    node.output_objs = tracker.output_objs
+    node.projection_columns = [obj.name.lower() for obj in node.output_objs]
+
+    # 5. resolve alias based on the what user provided
+    # we assign the alias to tracker as it governs the output of the extract object
+    resolve_alias_table_value_expression(node)
+    tracker.alias = node.alias

eva/binder/statement_binder.py

@@ -22,6 +22,7 @@
     check_groupby_pattern,
     check_table_object_is_video,
     extend_star,
+    handle_bind_extract_object_function,
     resolve_alias_table_value_expression,
 )
 from eva.binder.statement_binder_context import StatementBinderContext
@@ -40,6 +41,7 @@
 from eva.parser.select_statement import SelectStatement
 from eva.parser.statement import AbstractStatement
 from eva.parser.table_ref import TableRef
+from eva.parser.types import UDFType
 from eva.third_party.huggingface.binder import assign_hf_udf
 from eva.utils.generic_utils import get_file_checksum, load_udf_class_from_file
 from eva.utils.logging_manager import logger
@@ -234,6 +236,10 @@ def _bind_tuple_expr(self, node: TupleValueExpression):
 
     @bind.register(FunctionExpression)
     def _bind_func_expr(self, node: FunctionExpression):
+        # handle the special case of "extract_object"
+        if node.name.upper() == str(UDFType.EXTRACT_OBJECT):
+            handle_bind_extract_object_function(node, self)
+            return
         # bind all the children
         for child in node.children:
             self.bind(child)

eva/executor/apply_and_merge_executor.py

@@ -39,16 +39,11 @@ def __init__(self, node: ApplyAndMergePlan):
     def exec(self, *args, **kwargs) -> Iterator[Batch]:
         child_executor = self.children[0]
         for batch in child_executor.exec(**kwargs):
-            res = self.func_expr.evaluate(batch)
-            if not res.empty():
-                if self.do_unnest:
-                    res.unnest()
+            func_result = self.func_expr.evaluate(batch)
+            output = Batch.merge_column_wise([batch, func_result])
+            if self.do_unnest:
+                output.unnest(func_result.columns)
+                # we reset the index as after unnest there can be duplicate index
+                output.reset_index()
 
-                # Merge the results to the input.
-                # This assumes that the batch index is preserved by the function
-                # call. Since both the batch and the results are sorted, we could
-                # perform a sorted merge, though the typical small size of the
-                # batch and results should not significantly impact performance.
-                merged_batch = Batch.join(batch, res)
-                merged_batch.reset_index()
-                yield merged_batch
+            yield output

eva/expression/abstract_expression.py

@@ -91,6 +91,10 @@ def get_child(self, index: int):
     def children(self):
         return self._children
 
+    @children.setter
+    def children(self, children):
+        self._children = children
+
     def append_child(self, child):
         self._children.append(child)
 

eva/optimizer/operators.py

@@ -62,6 +62,7 @@ class OperatorType(IntEnum):
     LOGICALEXPLAIN = auto()
     LOGICALCREATEINDEX = auto()
     LOGICAL_APPLY_AND_MERGE = auto()
+    LOGICAL_EXTRACT_OBJECT = auto()
     LOGICALFAISSINDEXSCAN = auto()
     LOGICALDELIMITER = auto()
 
@@ -918,6 +919,45 @@ def __hash__(self) -> int:
         return hash((super().__hash__(), self.func_expr, self.do_unnest, self.alias))
 
 
+class LogicalExtractObject(Operator):
+    def __init__(
+        self,
+        detector: FunctionExpression,
+        tracker: FunctionExpression,
+        alias: Alias,
+        do_unnest: bool = False,
+        children: List = None,
+    ):
+        super().__init__(OperatorType.LOGICAL_EXTRACT_OBJECT, children)
+        self.detector = detector
+        self.tracker = tracker
+        self.do_unnest = do_unnest
+        self.alias = alias
+
+    def __eq__(self, other):
+        is_subtree_equal = super().__eq__(other)
+        if not isinstance(other, LogicalExtractObject):
+            return False
+        return (
+            is_subtree_equal
+            and self.detector == other.detector
+            and self.tracker == other.tracker
+            and self.do_unnest == other.do_unnest
+            and self.alias == other.alias
+        )
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                super().__hash__(),
+                self.detector,
+                self.tracker,
+                self.do_unnest,
+                self.alias,
+            )
+        )
+
+
 class LogicalJoin(Operator):
     """
     Logical node for join operators

eva/optimizer/rules/rules.py

@@ -61,6 +61,7 @@
     LogicalDrop,
     LogicalDropUDF,
     LogicalExplain,
+    LogicalExtractObject,
     LogicalFaissIndexScan,
     LogicalFilter,
     LogicalFunctionScan,
@@ -346,7 +347,7 @@ def apply(self, before: LogicalFilter, context: OptimizerContext):
 class XformLateralJoinToLinearFlow(Rule):
     """If the inner node of a lateral join is a function-valued expression, we
     eliminate the join node and make the inner node the parent of the outer node. This
-    produces a linear #data flow path. Because this scenario is common in our system,
+    produces a linear data flow path. Because this scenario is common in our system,
     we chose to explicitly convert it to a linear flow, which simplifies the
     implementation of other optimizations such as UDF reuse and parallelized plans by
     removing the join."""
@@ -442,6 +443,50 @@ def apply(self, before: LogicalFilter, context: OptimizerContext):
         yield root_node
 
 
+class XformExtractObjectToLinearFlow(Rule):
+    """If the inner node of a lateral join is a Extract_Object function-valued
+    expression, we eliminate the join node and make the inner node the parent of the
+    outer node. This produces a linear data flow path.
+    TODO: We need to add a sorting operation after detector to ensure we always provide tracker data in order.
+    """
+
+    #                                          LogicalApplyAndMerge(tracker)
+    #     LogicalJoin(Lateral)                         |
+    #     /           \                 ->    LogicalApplyAndMerge(detector)
+    #    A        LogicalExtractObject                 |
+    #                                                  A
+
+    def __init__(self):
+        pattern = Pattern(OperatorType.LOGICALJOIN)
+        pattern.append_child(Pattern(OperatorType.DUMMY))
+        pattern.append_child(Pattern(OperatorType.LOGICAL_EXTRACT_OBJECT))
+        super().__init__(RuleType.XFORM_EXTRACT_OBJECT_TO_LINEAR_FLOW, pattern)
+
+    def promise(self):
+        return Promise.XFORM_EXTRACT_OBJECT_TO_LINEAR_FLOW
+
+    def check(self, before: LogicalJoin, context: OptimizerContext):
+        if before.join_type == JoinType.LATERAL_JOIN:
+            return True
+        return False
+
+    def apply(self, before: LogicalJoin, context: OptimizerContext):
+        A: Dummy = before.children[0]
+        logical_extract_obj: LogicalExtractObject = before.children[1]
+
+        detector = LogicalApplyAndMerge(
+            logical_extract_obj.detector, alias=logical_extract_obj.detector.alias
+        )
+        tracker = LogicalApplyAndMerge(
+            logical_extract_obj.tracker,
+            alias=logical_extract_obj.alias,
+            do_unnest=logical_extract_obj.do_unnest,
+        )
+        detector.append_child(A)
+        tracker.append_child(detector)
+        yield tracker
+
+
 class CombineSimilarityOrderByAndLimitToFaissIndexScan(Rule):
     """
     This rule currently rewrites Order By + Limit to a Faiss index scan.

eva/optimizer/rules/rules_base.py

@@ -34,6 +34,7 @@ class RuleType(Flag):
 
     # REWRITE RULES TOP DOWN APPLY FIRST (LOGICAL -> LOGICAL)
     XFORM_LATERAL_JOIN_TO_LINEAR_FLOW = auto()
+    XFORM_EXTRACT_OBJECT_TO_LINEAR_FLOW = auto()
     TOP_DOWN_DELIMETER = auto()
 
     # REWRITE RULES BOTTOM UP APPLY SECOND (LOGICAL -> LOGICAL)
@@ -82,7 +83,6 @@ class RuleType(Flag):
     LOGICAL_CREATE_INDEX_TO_FAISS = auto()
     LOGICAL_APPLY_AND_MERGE_TO_PHYSICAL = auto()
     LOGICAL_FAISS_INDEX_SCAN_TO_PHYSICAL = auto()
-
     IMPLEMENTATION_DELIMETER = auto()
 
     NUM_RULES = auto()
@@ -139,6 +139,7 @@ class Promise(IntEnum):
     # REWRITE RULES
     EMBED_FILTER_INTO_GET = auto()
     EMBED_SAMPLE_INTO_GET = auto()
+    XFORM_EXTRACT_OBJECT_TO_LINEAR_FLOW = auto()
     XFORM_LATERAL_JOIN_TO_LINEAR_FLOW = auto()
     PUSHDOWN_FILTER_THROUGH_JOIN = auto()
     PUSHDOWN_FILTER_THROUGH_APPLY_AND_MERGE = auto()

eva/optimizer/rules/rules_manager.py

@@ -70,6 +70,7 @@
     PushDownFilterThroughApplyAndMerge,
     PushDownFilterThroughJoin,
     ReorderPredicates,
+    XformExtractObjectToLinearFlow,
     XformLateralJoinToLinearFlow,
 )
 from eva.optimizer.rules.rules_base import Rule
@@ -86,6 +87,7 @@ def __init__(self):
 
         self._stage_one_rewrite_rules = [
             XformLateralJoinToLinearFlow(),
+            XformExtractObjectToLinearFlow(),
         ]
 
         self._stage_two_rewrite_rules = [

eva/optimizer/statement_to_opr_convertor.py

@@ -22,6 +22,7 @@
     LogicalDrop,
     LogicalDropUDF,
     LogicalExplain,
+    LogicalExtractObject,
     LogicalFilter,
     LogicalFunctionScan,
     LogicalGet,
@@ -57,6 +58,7 @@
 from eva.parser.show_statement import ShowStatement
 from eva.parser.statement import AbstractStatement
 from eva.parser.table_ref import TableRef
+from eva.parser.types import UDFType
 from eva.utils.logging_manager import logger
 
 
@@ -78,11 +80,19 @@ def visit_table_ref(self, table_ref: TableRef):
 
         elif table_ref.is_table_valued_expr():
             tve = table_ref.table_valued_expr
-            self._plan = LogicalFunctionScan(
-                func_expr=tve.func_expr,
-                alias=table_ref.alias,
-                do_unnest=tve.do_unnest,
-            )
+            if tve.func_expr.name.lower() == str(UDFType.EXTRACT_OBJECT).lower():
+                self._plan = LogicalExtractObject(
+                    detector=tve.func_expr.children[1],
+                    tracker=tve.func_expr.children[2],
+                    alias=table_ref.alias,
+                    do_unnest=tve.do_unnest,
+                )
+            else:
+                self._plan = LogicalFunctionScan(
+                    func_expr=tve.func_expr,
+                    alias=table_ref.alias,
+                    do_unnest=tve.do_unnest,
+                )
 
         elif table_ref.is_select():
             # NestedQuery

eva/parser/types.py

@@ -67,3 +67,7 @@ class FileFormatType(EVAEnum):
 class ShowType(EVAEnum):
     UDFS  # noqa: F821
     TABLES  # noqa: F821
+
+
+class UDFType(EVAEnum):
+    EXTRACT_OBJECT  # noqa: F821

eva/udfs/decorators/utils.py

@@ -31,7 +31,22 @@ def load_io_from_udf_decorators(
         Type[UdfIOCatalogEntry]: UdfIOCatalogEntry object created from the input decorator in setup
     """
     tag_key = "input" if is_input else "output"
-    io_signature = udf.forward.tags[tag_key]
+    io_signature = None
+    if hasattr(udf.forward, "tags") and tag_key in udf.forward.tags:
+        io_signature = udf.forward.tags[tag_key]
+    else:
+        # Attempt to populate from the parent class and stop at the first parent class
+        # where the required tags are found.
+        for base_class in udf.__bases__:
+            if hasattr(base_class, "forward") and hasattr(base_class.forward, "tags"):
+                if tag_key in base_class.forward.tags:
+                    io_signature = base_class.forward.tags[tag_key]
+                    break
+
+    assert (
+        io_signature is not None
+    ), f"Cannot infer io signature from the decorator for {udf}."
+
     result_list = []
     for io in io_signature:
         result_list.extend(io.generate_catalog_entries(is_input))