Add support for LISTOF function, array return, and EXPLODE operator in PyDough

documentation/dsl.md

@@ -19,6 +19,7 @@ This page describes the specification of the PyDough DSL. The specification incl
    * [SINGULAR](#singular)
    * [BEST](#best)
    * [CROSS](#cross)
+   * [EXPLODE](#explode)
 - [User Generated Collections](#user-generated-collections)
     * [range_collection](#range_collection)
     * [dataframe_collection](#dataframe_collection)
@@ -36,7 +37,7 @@ This page describes the specification of the PyDough DSL. The specification incl
 ## Example Graph
 
 The examples in this document use a metadata graph (named `GRAPH`) with the following collections:
-- `People`: records of every known person. Scalar properties: `first_name`, `middle_name`, `last_name`, `ssn`, `birth_date`, `email`, `current_address_id`.
+- `People`: records of every known person. Scalar properties: `first_name`, `middle_name`, `last_name`, `ssn`, `birth_date`, `email`, `current_address_id`, `phone_numbers`.
 - `Addresses`: records of every known address. Scalar properties: `address_id`, `street_number`, `street_name`, `apartment`, `zip_code`, `city`, `state`.
 - `Packages`: records of every known package. Scalar properties: `package_id`, `customer_ssn`, `shipping_address_id`, `billing_address_id`, `order_date`, `arrival_date`, `package_cost`.
 
@@ -1536,6 +1537,147 @@ People.CALCULATE(Packages=COUNT(People.packages)).CROSS(Packages)
 People.CROSS(Addresses).current_address
 ```
 
+<!-- TOC --><a name="explode"></a>
+### EXPLODE
+
+A PyDough operation that explodes each row from a collection into multiple rows, i.e. from flattening a column of array data, or by splitting up a string column on a delimiter. The outputted collection will be a sub-collection of the original context containing the exploded data, and an optional indexed column keeping track of the indices of each value of the exploded data within a single row. This newly generated sub-collection has no other sub-collections from the original collection. The syntax for this operation is `collection.EXPLODE(...)`. `EXPLODE` has the following arguments:
+- `data` (required): the expression from the current context being exploded (either an array or string). This expression cannot reference any sub-collections of the current context.
+- `name` (required): the name of the collection created from the explosion operation (similar to `PARTITION`).
+- `value_name` (required): a string literal declaring the name of the new column that will be used to store the exploded data.
+- `index_name` (optional): a string literal declaring the name of the new column that will be used to store the indices of the exploded data.  If not provided, this column is not generated. The `index_name` is required if `is_distinct` is False. The indices are 0-indexed.
+- `version` (optional, default=`"array"`): either `"array"` or `"string"`, stating whether the data to explode is an array being flattened or a string being split on a delimiter.
+- `delimiter` (optional): a string literal indicating the delimiter that should be used to split up the string if `version="string"`. If `delimiter` is an empty string, the string will be split into individual characters.
+- `filtering` (optional, default=`True`): `True` if it is possible for not every row in the original collection to be preserved in the exploded sub-collection (i.e. if one of the arrays is empty), and `False` otherwise.
+- `is_distinct` (optional, default=`False`): `True` if each row of the exploded data is unique within the set of all other values from that same original row of unexploded data, and `False` otherwise. An `index_name` can only be omitted if `is_distinct` is `True`.
+
+> [!IMPORTANT]
+> This feature is only supported in certain dialects. It is currently supported for Snowflake, DataBricks, DuckDB, Postgres and Trino.
+
+**Good Example #1**: List out every phone number had by every person (assuming `phone_numbers` is an array of strings).
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, "numbers", value_name='phone_number', is_distinct=True)
+```
+
+**Good Example #2**: For each person, find the first phone number they have (assuming `phone_numbers` is an array of strings).
+
+```py
+%%pydough
+People.CALCULATE(first_name, last_name)
+      .EXPLODE(phone_numbers, "numbers", value_name='phone_number', index_name='idx', is_distinct=True)
+      .WHERE(idx == 0)
+      .CALCULATE(first_name, last_name, phone_number)
+```
+
+**Good Example #3**: For each person, count how many phone numbers they have (assuming `phone_numbers` is an array of strings).
+```py
+%%pydough
+exploded_numbers = EXPLODE(phone_numbers, "numbers", value_name='phone_number', is_distinct=True)
+People.CALCULATE(first_name, last_name, n_phone_numbers=COUNT(exploded_numbers))
+```
+
+**Good Example #4**: List out every phone number had by every person (assuming `phone_numbers` is a string of comma separated values).
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, "numbers", value_name='phone_number', version="string", delimiter=",", is_distinct=True)
+```
+
+**Good Example #5**: List the number of times each character of the alphabet used within first names of people.
+
+```py
+%%pydough
+People.EXPLODE(LOWER(first_name), "characters", value_name='char', index_name="idx", version="string", delimiter="")
+      .PARTITION(name='letters', by=char)
+      .CALCULATE(char, n_uses=COUNT(characters))
+```
+
+**Data Example**:
+
+Suppose we have the following collection of data `thesaurus` containing words and some of their synonyms in an array:
+| word    | synonyms                       |
+|---------|--------------------------------|
+| 'wise'  | ['sage', 'insightful', 'keen'] |
+| 'old'   | ['elderly', 'ancient']         |
+| 'my'    | []                             |
+| 'large' | ['big']                        |
+
+Now suppose the following PyDough code is used to transform `thesaurus` using the `EXPLODE` operator (note: `filtering=True` because one of the rows is an empty array):
+```py
+%%pydough
+thesaurus.CALCULATE(word)
+         .EXPLODE(synonyms, "words", value_name='synonym', index_name='syn_idx', filtering=True)
+         .CALCULATE(word, syn_idx, synonym)
+```
+
+The result would be the following table:
+| word    | syn_idx | synonym      |
+|---------|---------|--------------|
+| 'wise'  | 0       | 'sage'       |
+| 'wise'  | 1       | 'insightful' |
+| 'wise'  | 2       | 'keen'       |
+| 'old'   | 0       | 'elderly'    |
+| 'old'   | 1       | 'ancient'    |
+| 'large' | 0       | 'big'        |
+
+**Bad Example #1**: Missing the `name`.
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, value_name='phone_number', index_name="idx")
+```
+
+**Bad Example #2**: Missing the `value_name`.
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, "numbers", index_name="idx")
+```
+
+**Bad Example #3**: Missing the `index_name` when `is_distinct=True` is not provided.
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, "numbers", value_name='phone_number')
+```
+
+**Bad Example #4**: Providing an invalid `version`
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, "numbers", value_name='phone_number', index_name="idx", version="party")
+```
+
+**Bad Example #5**: Missing the `delimiter` when `version="string"`
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, "numbers", value_name='phone_number', index_name="idx", version="string")
+```
+
+**Bad Example #6**: Not providing a string literal for the delimiter.
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, "numbers", value_name='phone_number', index_name="idx", version="string", delimiter=first_name)
+```
+
+**Bad Example #7**: Attempting to access a sub-collection after exploding.
+
+```py
+%%pydough
+People.EXPLODE(phone_numbers, "numbers", value_name='phone_number', index_name="idx")
+      .packages
+```
+
+**Bad Example #8**: Accessing a sub-collection when declaring the data to explode.
+
+```py
+%%pydough
+People.EXPLODE(LISTOF(packages.package_cost), "costs", value_name='cost', index_name="idx")
+```
+
 <!-- TOC --><a name="user-generated-collections"></a>
 ## User Generated Collections
 
@@ -1605,12 +1747,17 @@ The supported PyDough types for `dataframe_collection` are:
 - `NumericType`: includes float, integer, infinity, Nan.
 - `BooleanType`: True or False.
 - `StringType`: alphanumeric characters.
-- `Datetype`: date and datetime.
+- `DateType`: date and datetime.
+- `ArrayType`: arrays of data.
 - `UnknownType`: used for all `None` columns.
 
 Note: MySQL by default does not support infinity values. When PyDough detects 
 infinity value with `DatabaseDiatect.MYSQL` an error will be raised.
 
+> [!IMPORTANT]
+> `ArrayType` is only supported for certain dialects: Trino, Postgres, DuckDB, Databricks.
+> Postgres has a limited ability to support rows with an empty array, depending on the type of the column. These sorts of array literals are only supported when the overall column type is an array of booleans, numbers, strings, or datetime values.
+
 #### Example 1
 
 ```python

documentation/functions.md

@@ -72,6 +72,7 @@ Below is the list of every function/operator currently supported in PyDough as a
    * [HASNOT](#hasnot)
    * [VAR](#var)
    * [STD](#std)
+   * [LISTOF](#listof)
 - [Window Functions](#window-functions)
    * [RANKING](#ranking)
    * [PERCENTILE](#percentile)
@@ -1089,6 +1090,34 @@ Parts.CALCULATE(std = STD(supply_records.supply_cost))
 Parts.CALCULATE(std = STD(supply_records.supply_cost, type="sample"))
 ```
 
+<!-- TOC --><a name="std"></a>
+
+### LISTOF
+
+The `LISTOF` function collects a set of values into an array.
+
+> [!IMPORTANT]
+> This function is only supported in certain dialects. It is currently supported for Snowflake, DataBricks, DuckDB, Postgres and Trino.
+
+```py
+# For each region, list the names of all nations inside that region
+Regions.CALCULATE(region_name=name, nation_names=LISTOF(nations.name))
+
+# For each customer, list the three largest quantities purchased made by that
+# customer.
+selected_lines = orders.lines.best(by=quantity.DESC(), per='Customers', n_best=3)
+Customers.CALCULATE(customer_name=name, quantities=LISTOF(selected_lines.quantity))
+```
+
+The first of these examples would produce the following output:
+| region_name   | nation_names                                                  |
+|---------------|---------------------------------------------------------------|
+| "AFRICA"      | ["ALGERIA",  "ETHIOPIA",  "KENYA",  "MOROCCO",  "MOZAMBIQUE"] |
+| "AMERICA"     | ["ARGENTINA", "BRAZIL", "CANADA", "PERU", "UNITED STATES"]    |
+| "ASIA"        | ["INDIA", "INDONESIA", "JAPAN", "CHINA", "VIETNAM"]           |
+| "EUROPE"      | ["FRANCE", "GERMANY", "ROMANIA", "RUSSIA", "UNITED KINGDOM"]  |
+| "MIDDLE EAST" | ["EGYPT", "IRAN", "IRAQ", "JORDAN", "SAUDI ARABIA"]           |
+
 <!-- TOC --><a name="window-functions"></a>
 
 ## Window Functions

documentation/metadata.md

@@ -466,12 +466,12 @@ Below are several examples the JSON for such verifiers:
 
 The JSON for a function deducer, used in the `output signature` field of a function definition, specifies the rules for determining the output type of a call to the function in terms of its input arguments. If a verifier is not provided, the default assumption is that the function call outputs an expression of type `"unknown"`.
 
-Each deducer has a mandatory string field `type` specifying what kind of verifier it is. The currently supported values are `"constant"` and `"select argument"`.
+Each deducer has a mandatory string field `type` specifying what kind of verifier it is. The currently supported values are `"constant"`, `"select argument"` and `"array of"`.
 
 <!-- TOC --><a name="function-deducer-type-constant"></a>
 ### Function Deducer Type: Constant
 
-Function deducers of this type have a type string of `"constant"` and correspond to a function call that always returns the same type. Verifiers of this type have the following additional key-value pairs in their metadata JSON object:
+Function deducers of this type have a type string of `"constant"` and correspond to a function call that always returns the same type. Deducers of this type have the following additional key-value pairs in their metadata JSON object:
 
 - `value` (required): a type string ([see here for more information](#pydough-type-strings)) indicating what type the function always returns.
 
@@ -484,7 +484,7 @@ Below are several examples the JSON for such deducers:
 <!-- TOC --><a name="function-deducer-type-select-argument"></a>
 ### Function Deducer Type: Select Argument
 
-Function deducers of this type have a type string of `"select argument"` and correspond to a function call that always returns a value of the same type as a specific argument. Verifiers of this type have the following additional key-value pairs in their metadata JSON object:
+Function deducers of this type have a type string of `"select argument"` and correspond to a function call that always returns a value of the same type as a specific argument. Deducers of this type have the following additional key-value pairs in their metadata JSON object:
 
 - `value` (required): a non-negative integer indicating which input argument to the function call should determine the output type of the function when called.
 
@@ -493,6 +493,18 @@ Below are several examples the JSON for such deducers:
 - Returns the type of the first argument: `{"type": "select argument", "value": 0}`
 - Returns the type of the second argument: `{"type": "select argument", "value": 1}`
 
+<!-- TOC --><a name="function-deducer-type-array-of"></a>
+### Function Deducer Type: Array Of
+
+Function deducers of this type have a type string of `"array of"` and correspond to a function call that returns an array type. Deducers of this type have the following additional key-value pairs in their metadata JSON object:
+
+- `element type` (required): a JSON object containing the specification for another function deducer containing the element type of the array.
+
+Below are several examples the JSON for such deducers:
+
+- Returns an array of strings: `{"type": "array of", "element type": {"type": "constant", "value": "string"}}`
+- Returns an array where the elements have the same type as the first argument: `{"type": "array of", "element type": {"type": "select argument", "value": 0}}`
+
 <!-- TOC --><a name="pydough-type-strings"></a>
 ## PyDough Type Strings
 

pydough/conversion/agg_removal.py

@@ -11,6 +11,7 @@
     Aggregate,
     CallExpression,
     EmptySingleton,
+    Explode,
     Filter,
     GeneratedTable,
     Join,
@@ -218,6 +219,36 @@ def deduce_join_uniqueness(
     return result
 
 
+def deduce_explode_uniqueness(
+    unique_terms: set[frozenset[str]], explode: Explode
+) -> set[frozenset[str]]:
+    """
+    Helper function to transforms the uniqueness sets after an Explode
+    operation duplicates rows.
+
+    Args:
+        `unique_terms`: the uniqueness sets of the input to the explode.
+        `explode`: the explode node.
+
+    Returns:
+        The uniqueness sets of the output of the explode.
+    """
+    # Build up the list of column names that imply uniqueness within one of the
+    # expanded row sets (either the index column, or the expanded data, or
+    # both)
+    cross_terms: list[str] = []
+    if explode.explode_spec.index_name is not None:
+        cross_terms.append(explode.explode_spec.index_name)
+    if explode.explode_spec.is_distinct:
+        cross_terms.append(explode.explode_spec.value_name)
+    # Add each of the cross terms to all of the uniqueness sets
+    result: set[frozenset[str]] = set()
+    for term_set in unique_terms:
+        for cross_term in cross_terms:
+            result.add(term_set | frozenset([cross_term]))
+    return result
+
+
 def aggregation_uniqueness_helper(
     node: RelationalNode,
 ) -> tuple[RelationalNode, set[frozenset[str]]]:
@@ -286,6 +317,10 @@ def aggregation_uniqueness_helper(
                 node, unique_sets
             )
             return node, final_uniqueness
+        case Explode():
+            node._input, input_uniqueness = aggregation_uniqueness_helper(node.input)
+            input_uniqueness = bubble_uniqueness(input_uniqueness, node.columns, None)
+            return node, deduce_explode_uniqueness(input_uniqueness, node)
         # Empty singletons don't have uniqueness information.
         case EmptySingleton():
             return node, set()

pydough/conversion/filter_pushdown.py

@@ -12,6 +12,7 @@
     CallExpression,
     ColumnReference,
     EmptySingleton,
+    Explode,
     Filter,
     GeneratedTable,
     Join,
@@ -319,6 +320,27 @@ def visit_generated_table(self, generated_table: GeneratedTable) -> RelationalNo
         # cannot be pushed down any further.
         return self.flush_remaining_filters(generated_table, self.filters, set())
 
+    def visit_explode(self, explode: Explode) -> RelationalNode:
+        pushable_filters: set[RelationalExpression]
+        remaining_filters: set[RelationalExpression]
+        # Push all filters that only depend on columns that pass through
+        # from the input, as opposed to being generated outputs.
+        allowed_cols: set[str] = set()
+        for name, expr in explode.columns.items():
+            if not (
+                isinstance(expr, ColumnReference)
+                and expr.name
+                in (explode.explode_spec.value_name, explode.explode_spec.index_name)
+            ):
+                allowed_cols.add(name)
+        pushable_filters, remaining_filters = partition_expressions(
+            self.filters,
+            lambda expr: only_references_columns(expr, allowed_cols),
+        )
+        return self.flush_remaining_filters(
+            explode, remaining_filters, pushable_filters
+        )
+
 
 def push_filters(node: RelationalNode, session: PyDoughSession) -> RelationalNode:
     """

pydough/conversion/hybrid_decorrelater.py

@@ -27,6 +27,7 @@
 from .hybrid_operations import (
     HybridCalculate,
     HybridChildPullUp,
+    HybridExplode,
     HybridFilter,
     HybridNoop,
     HybridPartition,
@@ -307,6 +308,14 @@ def correl_ref_purge(
                         correl_level,
                         new_parent_uni_keys,
                     )
+                if isinstance(operation, HybridExplode):
+                    operation.explode_data = self.remove_correl_refs(
+                        operation.explode_data,
+                        old_parent,
+                        child_height,
+                        correl_level,
+                        new_parent_uni_keys,
+                    )
             # Repeat the process on the ancestor until either loop guard
             # condition is no longer True. Only update the child height if we
             # are still making steps from the original tree, as opposed to from
@@ -554,6 +563,10 @@ def find_correlated_children(self, hybrid: HybridTree) -> None:
                 correl_levels = max(
                     correl_levels, operation.condition.count_correlated_levels()
                 )
+            if isinstance(operation, HybridExplode):
+                correl_levels = max(
+                    correl_levels, operation.explode_data.count_correlated_levels()
+                )
 
         assert correl_levels <= len(self.stack)
         for i in range(-1, -correl_levels - 1, -1):