Enhance partition_by to support strings (#1191)

Author: dmpetrov

src/datachain/lib/dc/datachain.py

@@ -21,6 +21,7 @@
 import orjson
 import sqlalchemy
 from pydantic import BaseModel
+from sqlalchemy.sql.elements import ColumnElement
 from tqdm import tqdm
 
 from datachain import semver
@@ -806,11 +807,35 @@ def agg_sum(
             chain.save("new_dataset")
             ```
         """
+        # Convert string partition_by parameters to Column objects
+        processed_partition_by = partition_by
+        if partition_by is not None:
+            if isinstance(partition_by, (str, Function, ColumnElement)):
+                list_partition_by = [partition_by]
+            else:
+                list_partition_by = list(partition_by)
+
+            processed_partition_columns: list[ColumnElement] = []
+            for col in list_partition_by:
+                if isinstance(col, str):
+                    col_db_name = ColumnMeta.to_db_name(col)
+                    col_type = self.signals_schema.get_column_type(col_db_name)
+                    column = Column(col_db_name, python_to_sql(col_type))
+                    processed_partition_columns.append(column)
+                elif isinstance(col, Function):
+                    column = col.get_column(self.signals_schema)
+                    processed_partition_columns.append(column)
+                else:
+                    # Assume it's already a ColumnElement
+                    processed_partition_columns.append(col)
+
+            processed_partition_by = processed_partition_columns
+
         udf_obj = self._udf_to_obj(Aggregator, func, params, output, signal_map)
         return self._evolve(
             query=self._query.generate(
                 udf_obj.to_udf_wrapper(),
-                partition_by=partition_by,
+                partition_by=processed_partition_by,
                 **self._settings.to_dict(),
             ),
             signal_schema=udf_obj.output,

src/datachain/query/dataset.py

@@ -82,7 +82,10 @@
 INSERT_BATCH_SIZE = 10000
 
 PartitionByType = Union[
-    Function, ColumnElement, Sequence[Union[Function, ColumnElement]]
+    str,
+    Function,
+    ColumnElement,
+    Sequence[Union[str, Function, ColumnElement]],
 ]
 JoinPredicateType = Union[str, ColumnClause, ColumnElement]
 DatasetDependencyType = tuple["DatasetRecord", str]

tests/unit/lib/test_datachain.py

@@ -3595,3 +3595,60 @@ def test_save_create_project_not_allowed(test_session, allow_create_project):
         dc.read_values(fib=[1, 1, 2, 3, 5, 8], session=test_session).save(
             "dev.numbers.fibonacci"
         )
+
+
+def test_agg_partition_by_string_notation(test_session):
+    """Test that agg method supports string notation for partition_by."""
+
+    class _ImageGroup(BaseModel):
+        name: str
+        size: int
+
+    def func(key, val) -> Iterator[tuple[File, _ImageGroup]]:
+        n = "-".join(key)
+        v = sum(val)
+        yield File(path=n), _ImageGroup(name=n, size=v)
+
+    keys = ["n1", "n2", "n1"]
+    values = [1, 5, 9]
+
+    # Test using string notation (NEW functionality)
+    ds = dc.read_values(key=keys, val=values, session=test_session).agg(
+        x=func,
+        partition_by="key",  # String notation instead of C("key")
+    )
+
+    assert ds.order_by("x_1.name").to_values("x_1.name") == ["n1-n1", "n2"]
+    assert ds.order_by("x_1.size").to_values("x_1.size") == [5, 10]
+
+
+def test_agg_partition_by_string_sequence(test_session):
+    """Test that agg method supports sequence of strings for partition_by."""
+
+    class _ImageGroup(BaseModel):
+        name: str
+        size: int
+
+    def func(key1, key2, val) -> Iterator[tuple[File, _ImageGroup]]:
+        n = f"{key1[0]}-{key2[0]}"
+        v = sum(val)
+        yield File(path=n), _ImageGroup(name=n, size=v)
+
+    key1_values = ["a", "a", "b"]
+    key2_values = ["x", "y", "x"]
+    values = [1, 5, 9]
+
+    # Test using sequence of strings (NEW functionality)
+    ds = dc.read_values(
+        key1=key1_values, key2=key2_values, val=values, session=test_session
+    ).agg(
+        x=func,
+        partition_by=["key1", "key2"],  # Sequence of strings
+    )
+
+    result_names = ds.order_by("x_1.name").to_values("x_1.name")
+    result_sizes = ds.order_by("x_1.size").to_values("x_1.size")
+
+    # Should have 3 partitions: (a,x), (a,y), (b,x)
+    assert len(result_names) == 3
+    assert len(result_sizes) == 3