Add the optional dependencies parameters in the transform/partitioned…

…_transform methods (#4891) * Add the optional dependencies parameters * Apply suggestions from code review Co-authored-by: Ryan Caudy <[email protected]> * Respond to review comments --------- Co-authored-by: Ryan Caudy <[email protected]>
deephaven · Nov 29, 2023 · 1745a66 · 1745a66
1 parent fb87423
commit 1745a66
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 6 deletions.
diff --git a/py/server/deephaven/table.py b/py/server/deephaven/table.py
@@ -2663,13 +2663,18 @@ def constituent_tables(self) -> List[Table]:
         """Returns all the current constituent tables."""
         return list(map(Table, self.j_partitioned_table.constituents()))
 
-    def transform(self, func: Callable[[Table], Table]) -> PartitionedTable:
+    def transform(self, func: Callable[[Table], Table],
+                  dependencies: Optional[Sequence[Union[Table, PartitionedTable]]] = None) -> PartitionedTable:
         """Apply the provided function to all constituent Tables and produce a new PartitionedTable with the results
         as its constituents, with the same data for all other columns in the underlying partitioned Table. Note that
         if the Table underlying this PartitionedTable changes, a corresponding change will propagate to the result.
 
         Args:
             func (Callable[[Table], Table]): a function which takes a Table as input and returns a new Table
+            dependencies (Optional[Sequence[Union[Table, PartitionedTable]]]): additional dependencies that must be
+                satisfied before applying the provided transform function to added or modified constituents during
+                update processing. If the transform function uses any other refreshing Table or refreshing Partitioned
+                Table, they must be included in this argument. Defaults to None.
 
         Returns:
             a PartitionedTable
@@ -2679,13 +2684,18 @@ def transform(self, func: Callable[[Table], Table]) -> PartitionedTable:
         """
         try:
             j_operator = j_unary_operator(func, dtypes.from_jtype(Table.j_object_type.jclass))
-            with auto_locking_ctx(self):
-                j_pt = self.j_partitioned_table.transform(j_operator)
+            dependencies = to_sequence(dependencies, wrapped=True)
+            j_dependencies = [d.j_table for d in dependencies if isinstance(d, Table) and d.is_refreshing]
+            j_dependencies.extend([d.table.j_table for d in dependencies if isinstance(d, PartitionedTable) and d.is_refreshing])
+            with auto_locking_ctx(self, *dependencies):
+                j_pt = self.j_partitioned_table.transform(j_operator, j_dependencies)
                 return PartitionedTable(j_partitioned_table=j_pt)
         except Exception as e:
             raise DHError(e, "failed to transform the PartitionedTable.") from e
 
-    def partitioned_transform(self, other: PartitionedTable, func: Callable[[Table, Table], Table]) -> PartitionedTable:
+    def partitioned_transform(self, other: PartitionedTable, func: Callable[[Table, Table], Table],
+                              dependencies: Optional[Sequence[Union[Table, PartitionedTable]]] = None) -> \
+            PartitionedTable:
         """Join the underlying partitioned Tables from this PartitionedTable and other on the key columns, then apply
         the provided function to all pairs of constituent Tables with the same keys in order to produce a new
         PartitionedTable with the results as its constituents, with the same data for all other columns in the
@@ -2698,6 +2708,10 @@ def partitioned_transform(self, other: PartitionedTable, func: Callable[[Table,
             other (PartitionedTable): the other Partitioned table whose constituent tables will be passed in as the 2nd
                 argument to the provided function
             func (Callable[[Table, Table], Table]): a function which takes two Tables as input and returns a new Table
+            dependencies (Optional[Sequence[Union[Table, PartitionedTable]]]): additional dependencies that must be
+                satisfied before applying the provided transform function to added, modified, or newly-matched
+                constituents during update processing. If the transform function uses any other refreshing Table or
+                refreshing Partitioned Table, they must be included in this argument. Defaults to None.
 
         Returns:
             a PartitionedTable
@@ -2707,8 +2721,12 @@ def partitioned_transform(self, other: PartitionedTable, func: Callable[[Table,
         """
         try:
             j_operator = j_binary_operator(func, dtypes.from_jtype(Table.j_object_type.jclass))
-            with auto_locking_ctx(self, other):
-                j_pt = self.j_partitioned_table.partitionedTransform(other.j_partitioned_table, j_operator)
+            dependencies = to_sequence(dependencies, wrapped=True)
+            j_dependencies = [d.j_table for d in dependencies if isinstance(d, Table) and d.is_refreshing]
+            j_dependencies.extend([d.table.j_table for d in dependencies if isinstance(d, PartitionedTable) and d.is_refreshing])
+            with auto_locking_ctx(self, other, *dependencies):
+                j_pt = self.j_partitioned_table.partitionedTransform(other.j_partitioned_table, j_operator,
+                                                                     j_dependencies)
                 return PartitionedTable(j_partitioned_table=j_pt)
         except Exception as e:
             raise DHError(e, "failed to transform the PartitionedTable with another PartitionedTable.") from e

diff --git a/py/server/tests/test_partitioned_table.py b/py/server/tests/test_partitioned_table.py
@@ -11,6 +11,7 @@
 from deephaven.filters import Filter
 
 from deephaven import read_csv, DHError, new_table, update_graph, time_table, empty_table
+from deephaven.update_graph import shared_lock
 from tests.testbase import BaseTestCase
 from deephaven.execution_context import get_exec_ctx
 
@@ -128,6 +129,13 @@ def test_transform(self):
             pt = self.partitioned_table.transform(Transformer)
             self.assertIn("f", [col.name for col in pt.constituent_table_columns])
 
+            ticking_t = time_table("PT00:00:01")
+            pt = self.partitioned_table.transform(Transformer, dependencies=[ticking_t])
+            self.assertIn("f", [col.name for col in pt.constituent_table_columns])
+
+            pt = self.partitioned_table.transform(Transformer, dependencies=[self.test_table])
+            self.assertIn("f", [col.name for col in pt.constituent_table_columns])
+
             with self.assertRaises(DHError) as cm:
                 pt = self.partitioned_table.transform(lambda t, t1: t.join(t1))
             self.assertRegex(str(cm.exception), r"missing .* argument")
@@ -141,6 +149,14 @@ def test_partitioned_transform(self):
             pt = self.partitioned_table.partitioned_transform(other_pt, PartitionedTransformer())
             self.assertIn("f", [col.name for col in pt.constituent_table_columns])
 
+            ticking_pt = time_table("PT00:00:01").update(["X= i % 10", "Y = String.valueOf(i)"]).partition_by("X")
+            pt = self.partitioned_table.partitioned_transform(other_pt, PartitionedTransformer(),
+                                                              dependencies=[ticking_pt])
+            self.assertIn("f", [col.name for col in pt.constituent_table_columns])
+
+            pt = self.partitioned_table.partitioned_transform(other_pt, PartitionedTransformer(), dependencies=[other_pt])
+            self.assertIn("f", [col.name for col in pt.constituent_table_columns])
+
     def test_partition_agg(self):
         with update_graph.shared_lock(self.test_update_graph):
             test_table = time_table("PT00:00:00.001").update(["X=i", "Y=i%13", "Z=X*Y"])