Flowfile 0.2.2__py3-none-any.whl

flowfile_frame/group_frame.py

@@ -0,0 +1,199 @@
+
+from flowfile_frame.expr import col, Expr
+from flowfile_frame.selectors import Selector
+from flowfile_frame.utils import _parse_inputs_as_iterable
+from flowfile_core.schemas import transform_schema, input_schema
+from typing import TYPE_CHECKING
+
+# Corrected TYPE_CHECKING block as provided by user
+if TYPE_CHECKING:
+    from flowfile_frame.flow_frame import FlowFrame
+else:
+    FlowFrame = None
+
+
+class GroupByFrame:
+    """Represents a grouped DataFrame for aggregation operations."""
+
+    def __init__(self, node_id: int, parent_frame, by_cols, maintain_order=False, description: str = None):
+        self.parent = parent_frame
+        self.by_cols = _parse_inputs_as_iterable(by_cols)
+        self.maintain_order = maintain_order
+        self.description = description
+        self.node_id = node_id
+
+    def readable_group(self):
+        parts = []
+        for c in self.by_cols:
+            if isinstance(c, Expr):
+                parts.append(str(c))
+            elif isinstance(c, str):
+                parts.append(f'''"{c}"''')
+            else:
+                parts.append(f'''"{str(c)}"''')
+        return ", ".join(parts)
+
+    def len(self) -> 'FlowFrame':
+        """
+        Count number of rows per group. Output column is named 'len'.
+        """
+        # Uses direct code generation as per user's example
+        return self._generate_direct_polars_code("len")
+
+    def count(self) -> 'FlowFrame':
+        """
+        Count number of rows per group. Output column is named 'count'.
+        """
+        # Uses direct code generation as per user's example
+        return self._generate_direct_polars_code("count")
+
+    def agg(self, *agg_exprs, **named_agg_exprs) -> FlowFrame:
+        """
+        Apply EXPLICIT aggregations to grouped data using expressions.
+        """
+        agg_expressions = _parse_inputs_as_iterable(agg_exprs)
+        can_be_converted: bool = not self.maintain_order
+        agg_cols: list[transform_schema.AggColl] = []
+        if can_be_converted:
+            can_be_converted = self._process_group_columns(agg_cols)
+        if can_be_converted:
+            can_be_converted = self._process_agg_expressions(agg_cols, agg_expressions)
+        if can_be_converted:
+            can_be_converted = self._process_named_agg_expressions(agg_cols, named_agg_exprs)
+        node_desc = self.description or f"Aggregate after grouping by {self.readable_group()}"
+        return self._create_agg_node(self.node_id, can_be_converted, agg_cols, agg_expressions, named_agg_exprs, node_desc)
+
+    def _process_group_columns(self, agg_cols: list[transform_schema.AggColl]) -> bool:
+        # (Implementation unchanged from user input)
+        for col_expr in self.by_cols:
+            if isinstance(col_expr, str):
+                agg_cols.append(transform_schema.AggColl(old_name=col_expr, agg="groupby"))
+            elif isinstance(col_expr, Expr):
+                agg_cols.append(transform_schema.AggColl(old_name=col_expr.name, agg="groupby"))
+            elif isinstance(col_expr, Selector):
+                return False
+            else:
+                 return False
+        return True
+
+    @staticmethod
+    def _process_agg_expressions(agg_cols: list[transform_schema.AggColl], agg_expressions) -> bool:
+        # (Implementation unchanged from user input)
+        for expr in agg_expressions:
+            if isinstance(expr, Expr):
+                agg_func = getattr(expr, "agg_func", None)
+                old_name = getattr(expr, "_initial_column_name", expr.name) or expr.name
+                if agg_func:
+                    agg_cols.append(
+                        transform_schema.AggColl(old_name=old_name, agg=agg_func, new_name=expr.name)
+                    )
+                else:
+                    agg_cols.append(transform_schema.AggColl(old_name=expr.name, agg="first"))
+            elif isinstance(expr, str):
+                agg_cols.append(transform_schema.AggColl(old_name=expr, agg="first"))
+            elif isinstance(expr, Selector):
+                return False
+            else:
+                return False
+        return True
+
+    @staticmethod
+    def _process_named_agg_expressions(agg_cols: list[transform_schema.AggColl], named_agg_exprs: dict) -> bool:
+        for name, expr in named_agg_exprs.items():
+            if expr.is_complex:
+                return False
+            if isinstance(expr, Expr):
+                agg_func = getattr(expr, "agg_func", "first")
+                old_name = getattr(expr, "_initial_column_name", expr.name) or expr.name
+                agg_cols.append(transform_schema.AggColl(old_name=old_name, agg=agg_func, new_name=name))
+            elif isinstance(expr, str):
+                agg_cols.append(transform_schema.AggColl(old_name=expr, agg="first", new_name=name))
+            elif isinstance(expr, tuple) and len(expr) == 2:
+                col_spec, agg_func_str = expr
+                if isinstance(col_spec, Expr):
+                    old_name = getattr(col_spec, "_initial_column_name", col_spec.name) or col_spec.name
+                elif isinstance(col_spec, str):
+                    old_name = col_spec
+                else:
+                    return False
+                if not isinstance(agg_func_str, str):
+                    return False
+                agg_cols.append(transform_schema.AggColl(old_name=old_name, agg=agg_func_str, new_name=name))
+            else:
+                return False
+        return True
+
+    def _create_agg_node(self, node_id_to_use: int, can_be_converted: bool, agg_cols: list, agg_expressions, named_agg_exprs, description: str):
+        """Creates node for explicit aggregations via self.agg()"""
+        # (Implementation unchanged from user input, passes description)
+        if can_be_converted:
+            group_by_settings = input_schema.NodeGroupBy(
+                flow_id=self.parent.flow_graph.flow_id,
+                node_id=node_id_to_use,
+                groupby_input=transform_schema.GroupByInput(agg_cols=agg_cols),
+                pos_x=200, pos_y=200, is_setup=True,
+                depending_on_id=self.parent.node_id,
+                description=description
+            )
+            self.parent.flow_graph.add_group_by(group_by_settings)
+        else:
+            code = self._generate_polars_agg_code(agg_expressions, named_agg_exprs)
+            self.parent._add_polars_code(new_node_id=node_id_to_use, code=code, description=description)
+        return self.parent._create_child_frame(node_id_to_use)
+
+    def _generate_direct_polars_code(self, method_name: str) -> "FlowFrame":
+        """
+        Generates Polars code for simple GroupBy methods like sum(), mean(), len(), count()
+        which operate implicitly or have a standard Polars counterpart.
+        Always uses the Polars code path.
+        """
+        readable_group_str = self.readable_group()
+        code = f"input_df.group_by([{readable_group_str}], maintain_order={self.maintain_order}).{method_name}()"
+        node_description = self.description or f"{method_name.capitalize()} after grouping by {readable_group_str}"
+        self.parent._add_polars_code(new_node_id=self.node_id, code=code, description=node_description)
+        return self.parent._create_child_frame(self.node_id)
+
+    def _generate_polars_agg_code(self, agg_expressions, named_agg_exprs) -> str:
+        """Generate Polars code specifically for explicit .agg() calls."""
+        # (Implementation unchanged from user input)
+        readable_group_str = self.readable_group()
+        agg_strs = [str(expr) for expr in agg_expressions]
+        named_agg_strs = [f"{name}={str(expr)}" for name, expr in named_agg_exprs.items()]
+        all_agg_strs = agg_strs + named_agg_strs
+        agg_combined = ", ".join(all_agg_strs)
+        # Assuming input dataframe is 'input_df' in execution context
+        return f"input_df.group_by([{readable_group_str}], maintain_order={self.maintain_order}).agg({agg_combined})"
+
+    # --- Convenience Methods (No Column Args - Use Direct Code Gen) ---
+
+    def sum(self):
+        """Calculate sum for all non-grouping columns."""
+        return self._generate_direct_polars_code("sum")
+
+    def mean(self):
+        """Calculate mean for all non-grouping columns."""
+        return self._generate_direct_polars_code("mean")
+
+    def median(self):
+        """Calculate median for all non-grouping columns."""
+        return self._generate_direct_polars_code("median")
+
+    def min(self):
+        """Calculate minimum for all non-grouping columns."""
+        # Remove *columns argument
+        return self._generate_direct_polars_code("min")
+
+    def max(self):
+        """Calculate maximum for all non-grouping columns."""
+        # Remove *columns argument
+        return self._generate_direct_polars_code("max")
+
+    def first(self):
+        """Get first value for all non-grouping columns."""
+        # Remove *columns argument
+        return self._generate_direct_polars_code("first")
+
+    def last(self):
+        """Get last value for all non-grouping columns."""
+        # Remove *columns argument
+        return self._generate_direct_polars_code("last")

flowfile_frame/join.py

@@ -0,0 +1,75 @@
+# Assume these imports are correct from your original context
+from flowfile_frame.expr import Column
+from flowfile_core.schemas import transform_schema
+
+
+def _normalize_columns_to_list(columns):
+    """Convert a column specification to a list format.
+
+    Args:
+        columns: Column name or list of column names
+
+    Returns:
+        List of column names/expressions
+    """
+    if columns is None:
+        return []
+    elif isinstance(columns, str):
+        return [columns]
+    elif isinstance(columns, (list, tuple)):
+        return list(columns)
+    else:
+        return [columns]  # Single non-string item
+
+
+def _extract_column_name(col_expr):
+    """Extract a simple column name from various column representations.
+
+    Args:
+        col_expr: Column expression (string, Column object, etc.)
+
+    Returns:
+        tuple: (column_name, needs_polars_code)
+        - column_name is the string name if possible
+        - needs_polars_code is True if this expression requires polars code generation
+    """
+    if isinstance(col_expr, str):
+        return col_expr, False
+
+    if isinstance(col_expr, Column):
+        # If it's a simple unaltered column, use its name
+        if not col_expr._select_input.is_altered:
+            return col_expr.name, False
+        # Otherwise, this requires polars code
+        return col_expr, True
+
+    # Any other expression type needs polars code
+    return col_expr, True
+
+
+def _create_join_mappings(left_columns, right_columns):
+    """Create join mappings between left and right columns.
+
+    Args:
+        left_columns: List of left join columns
+        right_columns: List of right join columns
+
+    Returns:
+        tuple: (join_mappings, needs_polars_code)
+        - join_mappings is a list of JoinMap objects
+        - needs_polars_code is True if any column requires polars code generation
+    """
+    join_mappings = []
+    needs_polars_code = False
+
+    for left_col, right_col in zip(left_columns, right_columns):
+        left_name, left_needs_code = _extract_column_name(left_col)
+        right_name, right_needs_code = _extract_column_name(right_col)
+
+        needs_polars_code = needs_polars_code or left_needs_code or right_needs_code
+
+        # Only create standard join mappings if both columns are simple strings
+        if not left_needs_code and not right_needs_code:
+            join_mappings.append(transform_schema.JoinMap(left_col=left_name, right_col=right_name))
+
+    return join_mappings, needs_polars_code

flowfile_frame/selectors.py

@@ -0,0 +1,242 @@
+from typing import List, Union, TYPE_CHECKING  # Added TYPE_CHECKING
+
+import polars as pl
+
+# --- TYPE CHECKING IMPORTS ---
+# if TYPE_CHECKING:
+#     Import Expr only for type hints
+from flowfile_frame.expr import Expr
+
+
+# --- Selector Base Classes (Compound, Complement) ---
+
+class Selector:
+    """Base class for column selectors, inspired by polars.selectors"""
+    def __init__(self):
+        self._repr_str = self._get_repr_str()  # Use base repr calculation method
+        # No agg_func state stored here anymore
+
+    @property
+    def repr_str(self):
+        return self._repr_str
+
+    def _get_repr_str(self) -> str:
+        """Get representation string for the selector itself."""
+        # Default implementation, specific selectors override this
+        return f"pl.selectors.{self.__class__.__name__}()"
+
+    def __repr__(self) -> str:
+        return self._repr_str
+
+    def __or__(self, other: 'Selector') -> 'CompoundSelector':
+        return CompoundSelector(self, other, 'union')
+
+    def __and__(self, other: 'Selector') -> 'CompoundSelector':
+        return CompoundSelector(self, other, 'intersection')
+
+    def __sub__(self, other: 'Selector') -> 'CompoundSelector':
+        return CompoundSelector(self, other, 'difference')
+
+    def __xor__(self, other: 'Selector') -> 'CompoundSelector':
+        return CompoundSelector(self, other, 'symmetric_difference')
+
+    def __invert__(self) -> 'ComplementSelector':
+        return ComplementSelector(self)
+
+    # --- Aggregation Methods ---
+    # These methods now return Expr objects, importing Expr locally
+
+    def sum(self) -> 'Expr':
+        """Create an expression to sum columns selected by this selector."""
+        # Expr init will handle creating the 'pl.sum(selector)' repr
+        return Expr(expr=None, selector=self, agg_func="sum")
+
+    def mean(self) -> 'Expr':
+        """Create an expression to average columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="mean")
+
+    def median(self) -> 'Expr':
+        """Create an expression to find the median of columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="median")
+
+    def min(self) -> 'Expr':
+        """Create an expression to find the minimum of columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="min")
+
+    def max(self) -> 'Expr':
+        """Create an expression to find the maximum of columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="max")
+
+    def std(self, ddof: int = 1) -> 'Expr':
+        """Create an expression to find the standard deviation of columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="std", ddof=ddof)
+
+    def var(self, ddof: int = 1) -> 'Expr':
+        """Create an expression to find the variance of columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="var", ddof=ddof)
+
+    def first(self) -> 'Expr':
+        """Create an expression to get the first element of columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="first")
+
+    def last(self) -> 'Expr':
+        """Create an expression to get the last element of columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="last")
+
+    def count(self) -> 'Expr':
+        """Create an expression to count elements in columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="count")
+
+    def n_unique(self) -> 'Expr':
+        """Create an expression to count unique elements in columns selected by this selector."""
+        return Expr(expr=None, selector=self, agg_func="n_unique")
+
+    # Removed alias method - belongs on Expr
+
+
+class CompoundSelector(Selector):
+    """Selector representing a compound operation between two selectors"""
+    def __init__(self, left: Selector, right: Selector, operation: str):
+        self.left = left
+        self.right = right
+        self.operation = operation
+        super().__init__()
+
+    def _get_repr_str(self) -> str:
+        op_map = {'union': '|', 'intersection': '&', 'difference': '-', 'symmetric_difference': '^'}
+        op_symbol = op_map.get(self.operation, '|')
+        # Use base repr (_repr_str) of operands
+        left_repr = f"({self.left._repr_str})" if isinstance(self.left, CompoundSelector) else self.left._repr_str
+        right_repr = f"({self.right._repr_str})" if isinstance(self.right, CompoundSelector) else self.right._repr_str
+        return f"{left_repr} {op_symbol} {right_repr}"
+
+
+class ComplementSelector(Selector):
+    """Selector representing the complement (NOT) of another selector"""
+    def __init__(self, selector: Selector):
+        self.selector = selector
+        super().__init__()
+
+    def _get_repr_str(self) -> str:
+        selector_repr = f"({self.selector._repr_str})" if isinstance(self.selector, CompoundSelector) else self.selector._repr_str
+        return f"~{selector_repr}"
+
+
+class NumericSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.numeric()"
+
+
+class FloatSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.float()"
+
+
+class IntegerSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.integer()"
+
+
+class StringSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.string()"
+
+
+class TemporalSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.temporal()"
+
+
+class DatetimeSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.datetime()"
+
+
+class DateSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.date()"
+
+
+class TimeSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.time()"
+
+
+class DurationSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.duration()"
+
+
+class BooleanSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.boolean()"
+
+
+class CategoricalSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.categorical()"
+
+
+class ObjectSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.object()"
+
+
+class ListSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.list()"
+
+
+class StructSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.struct()"
+
+
+class AllSelector(Selector):
+    def _get_repr_str(self) -> str: return "pl.selectors.all()"
+
+
+class DtypeSelector(Selector):
+    def __init__(self, dtypes: Union[pl.DataType, List[pl.DataType]]):
+        self.dtypes = dtypes if isinstance(dtypes, list) else [dtypes]
+        super().__init__()
+
+    def _get_repr_str(self) -> str:
+        dtype_strs = []
+        for dt in self.dtypes:
+            dt_repr = repr(dt)
+            if dt_repr.startswith("DataType"):
+                dt_repr = str(dt).capitalize()
+            dtype_strs.append(f"pl.{dt_repr}")
+        dtype_repr_arg = dtype_strs[0] if len(dtype_strs) == 1 else f"[{', '.join(dtype_strs)}]"
+        return f"pl.selectors.by_dtype({dtype_repr_arg})"
+
+
+class PatternSelector(Selector):
+    def __init__(self, pattern: str):
+        self.pattern = pattern
+        super().__init__()
+
+
+class ContainsSelector(PatternSelector):
+    def _get_repr_str(self) -> str: return f"pl.selectors.contains({self.pattern!r})"
+
+
+class StartsWithSelector(PatternSelector):
+    def _get_repr_str(self) -> str: return f"pl.selectors.starts_with({self.pattern!r})"
+
+
+class EndsWithSelector(PatternSelector):
+    def _get_repr_str(self) -> str: return f"pl.selectors.ends_with({self.pattern!r})"
+
+
+class MatchesSelector(PatternSelector):
+    def _get_repr_str(self) -> str: return f"pl.selectors.matches({self.pattern!r})"
+
+
+def numeric() -> NumericSelector: return NumericSelector()
+def float_() -> FloatSelector: return FloatSelector()
+def integer() -> IntegerSelector: return IntegerSelector()
+def string() -> StringSelector: return StringSelector()
+def temporal() -> TemporalSelector: return TemporalSelector()
+def datetime() -> DatetimeSelector: return DatetimeSelector()
+def date() -> DateSelector: return DateSelector()
+def time() -> TimeSelector: return TimeSelector()
+def duration() -> DurationSelector: return DurationSelector()
+def boolean() -> BooleanSelector: return BooleanSelector()
+def categorical() -> CategoricalSelector: return CategoricalSelector()
+def object_() -> ObjectSelector: return ObjectSelector()
+def list_() -> ListSelector: return ListSelector()
+def struct() -> StructSelector: return StructSelector()
+def all_() -> AllSelector: return AllSelector()
+def by_dtype(dtypes: Union[pl.DataType, List[pl.DataType]]) -> DtypeSelector: return DtypeSelector(dtypes)
+def contains(pattern: str) -> ContainsSelector: return ContainsSelector(pattern)
+def starts_with(pattern: str) -> StartsWithSelector: return StartsWithSelector(pattern)
+def ends_with(pattern: str) -> EndsWithSelector: return EndsWithSelector(pattern)
+def matches(pattern: str) -> MatchesSelector: return MatchesSelector(pattern)