diff --git a/py/server/deephaven/table.py b/py/server/deephaven/table.py index e3a2da77423..eab32badecd 100644 --- a/py/server/deephaven/table.py +++ b/py/server/deephaven/table.py @@ -1430,16 +1430,17 @@ def sort(self, order_by: Union[str, Sequence[str]], except Exception as e: raise DHError(e, "failed to sort the partitioned table.") from e - def get_constituent(self, key_values: Sequence[Any]) -> Optional[Table]: + def get_constituent(self, key_values: [Any, Sequence[Any]]) -> Optional[Table]: """Gets a single constituent table by its corresponding key column values. Args: - key_values (Sequence[Any]): the values of the key columns + key_values (Any, Sequence[Any]): the value(s) of the key column(s) Returns: a Table or None """ - j_table = self.j_partitioned_table.constituentFor(*key_values) + key_values = to_sequence(key_values) + j_table = self.j_partitioned_table.constituentFor(key_values) if j_table: return Table(j_table=j_table) else: diff --git a/py/server/tests/test_partitioned_table.py b/py/server/tests/test_partitioned_table.py index 2c65afad56d..45cc8f20ebf 100644 --- a/py/server/tests/test_partitioned_table.py +++ b/py/server/tests/test_partitioned_table.py @@ -5,7 +5,7 @@ from deephaven.filters import Filter -from deephaven import read_csv, DHError +from deephaven import read_csv, DHError, new_table from tests.testbase import BaseTestCase @@ -72,6 +72,22 @@ def test_get_constituent(self): keys = [917, 167] self.assertIsNotNone(self.partitioned_table.get_constituent(keys)) + from deephaven.column import string_col, int_col, double_col + + houses = new_table([ + string_col("HomeType", ["Colonial", "Contemporary", "Contemporary", "Condo", "Colonial", "Apartment"]), + int_col("HouseNumber", [1, 3, 4, 15, 4, 9]), + string_col("StreetName", ["Test Drive", "Test Drive", "Test Drive", "Deephaven Road", "Community Circle", + "Community Circle"]), + int_col("SquareFeet", [2251, 1914, 4266, 1280, 3433, 981]), + int_col("Price", [450000, 400000, 1250000, 300000, 600000, 275000]), + double_col("LotSizeAcres", [0.41, 0.26, 1.88, 0.11, 0.95, 0.10]) + ]) + + houses_by_type = houses.partition_by("HomeType") + colonial_homes = houses_by_type.get_constituent("Colonial") + self.assertIsNotNone(colonial_homes) + def test_constituents(self): constituent_tables = self.partitioned_table.constituent_tables self.assertGreater(len(constituent_tables), 0)