Dataframe Collection (User Collections) by john-sanchez31 · Pull Request #469 · bodo-ai/PyDough

john-sanchez31 · 2026-01-08T14:43:17Z

Resolves #162

… SQLGlot bug discovered along the way [RUN CI]

… into John/df_collection

…run sf][run postgres][run s3]

john-sanchez31

Some comments and guides for the review

john-sanchez31 · 2026-01-27T21:09:18Z

pydough/unqualified/qualification.py

                new_ancestry, new_child, ancestry_names = self.split_partition_ancestry(
                    parent, partition_ancestor
                )
+            case UnqualifiedGeneratedCollection():


This changes are related to supporting Partition with GeneratedCollection

john-sanchez31 · 2026-01-27T21:11:55Z

pydough/sqlglot/transform_bindings/base_transform_bindings.py

+
+        Note: This covers the NOT standard representations
+        (Datetime, NumericType, UnknownType) meaning that the expression
+        changes in different dialects.


This method is implemented in all transform bindings

john-sanchez31 · 2026-01-27T21:18:36Z

pyproject.toml

 # Note: sqlite is included in the standard library, so it doesn't need to be listed here.
 # There is a bug in some unit tests when run with sqlglot>=26.8.0
-dependencies = ["pytz", "sqlglot==26.7.0", "pandas>=2.0.0", "jupyterlab"]
+dependencies = ["pytz", "sqlglot==26.7.0", "pandas>=2.0.0", "jupyterlab", "pyarrow"]


pyarrow is used to map types in a dataframe

It should be pinned to the version you used to implement to avoid issues later on.

john-sanchez31 · 2026-01-27T21:20:26Z

tests/conftest.py



+@pytest.fixture(scope="session")
+def get_postgres_defog_graphs() -> graph_fetcher:


Now postgres have its own defog graph. This is needed because a UDF implementation was required for dealership_adv13 which is different for each dialect. Before postgres was sharing graph with sqlite, but no more

john-sanchez31 · 2026-01-27T21:24:07Z

tests/conftest.py



+@pytest.fixture(scope="session")
+def get_dialect_defog_graphs(


For test_defog_until_sql was used the standard graph. This made the sql refsol inconsistent to what was actually being executed. For example in the query I could see the UDF sqlite implementation in all dialect (which is wrong). This allows the usage for each dialect its graph. Actually with this new fixture most of the defog refsol file were changed.

john-sanchez31 · 2026-01-27T21:24:33Z

tests/test_metadata/defog_graphs.json

        "synonyms": ["purchase record", "vehicle sale", "car purchase"]
      }
    ],
+    "functions": [


UDF used for dealership_adv13

john-sanchez31 · 2026-01-27T21:26:22Z

tests/test_metadata/postgres_defog_graphs.json

+        "synonyms": ["purchase record", "vehicle sale", "car purchase"]
+      }
+    ],
+    "functions": [


This file is the same than defog_graphs.json except for this UDF which is necessary

john-sanchez31 · 2026-01-27T21:27:59Z

tests/test_pipeline_defog.py

@@ -756,9 +756,6 @@ def test_graph_structure_defog(defog_graphs: graph_fetcher, graph_name: str) ->
                order_sensitive=True,
            ),
            id="dealership_adv8",


These test needed GeneratedCollection but now are executed.

john-sanchez31 · 2026-01-27T21:28:44Z

tests/test_pipeline_mysql.py

    Test executing the TPC-H custom queries from the original code generation on
    MySQL.
    """
+


MySQL doesn't support infinity values, so we skip the test

john-sanchez31 · 2026-01-27T21:31:38Z

tests/test_sql_refsols/defog_academic_gen10_mysql.sql

@@ -1,5 +1,5 @@
 SELECT
  title COLLATE utf8mb4_bin AS title
-FROM main.publication


All this changes are related to the usage of the new graph for defog

… mysql][run sf][run postgres][run s3]" This reverts commit 03d37d6.

… sf][run postgres]

knassre-bodo · 2026-01-27T21:06:14Z

documentation/dsl.md

+Supported Signatures:
+- `dataframe_collection(dataframe, name)`: generates collection with the given datafram and name.


This is unnecessary here since there is only 1 signature, as opposed to range_collection which has several.

knassre-bodo · 2026-01-27T21:06:52Z

documentation/dsl.md

+It takes in the following arguments:
+
+- `name`: The name of the dataframe collection.
+- `dataframe`: The panda dataframe containing the corresponding data


We explicitly should note which types are supported in the DataFrame currently.

+1
Also, let's document the infinity behavior

Oh, and also stuff like None vs NaN, and NaT

Shoot speaking of... did we test NaT (not-a-time?)

knassre-bodo · 2026-01-28T18:42:50Z

tests/test_pydough_functions/user_collections.py

+
+
+# Windows functions
+def dataframe_collection_window_functions():


Let's also add one where the per=... refers to the generated collection.

hadia206

Good work John. Please see my comment below.

hadia206 · 2026-01-29T21:48:53Z

documentation/dsl.md

+<!-- TOC --><a name="dataframe_collection"></a>
+### `pydough.dataframe_collection`
+
+The `dataframe_collection` creates a collection within a specified pandas dataframe. This is useful for building datasets dynamically.


Suggested change

The `dataframe_collection` creates a collection within a specified pandas dataframe. This is useful for building datasets dynamically.

The `dataframe_collection` creates a collection from a specified Pandas DataFrame. This is useful for building datasets dynamically.

hadia206 · 2026-01-29T21:49:16Z

documentation/dsl.md

+It takes in the following arguments:
+
+- `name`: The name of the dataframe collection.
+- `dataframe`: The panda dataframe containing the corresponding data


Suggested change

- `dataframe`: The panda dataframe containing the corresponding data

- `dataframe`: Pandas DataFrame containing the corresponding data.

hadia206 · 2026-01-29T21:49:43Z

documentation/dsl.md

+Supported Signatures:
+- `dataframe_collection(dataframe, name)`: generates collection with the given datafram and name.


pydough/conversion/hybrid_tree.py

hadia206 · 2026-01-29T21:59:30Z

pydough/user_collections/user_collection_apis.py

@@ -4,7 +4,10 @@

 __all__ = ["range_collection"]


Missing dataframe_collection?

hadia206 · 2026-01-29T22:49:29Z

pydough/user_collections/README.md

+
+  - `DataframeGeneratedCollection`: Class used to create a dataframe collection using the given dataframe and name.
+    - `name`: The name of the dataframe collection.
+    - `dataframe`: The pandas dataframe containing all data (rows and columns).


Suggested change

- `dataframe`: The pandas dataframe containing all data (rows and columns).

- `dataframe`: The Pandas DataFrame containing all data (rows and columns).

hadia206 · 2026-01-29T22:49:44Z

pydough/user_collections/README.md

    - Returns: An instance of `RangeGeneratedCollection`.
+  - `dataframe_collection`: Function to create a dataframe collection with the specified parameters.
+    - `name`: The name of the dataframe collection.
+    - `dataframe`: The pandas dataframe.


Suggested change

- `dataframe`: The pandas dataframe.

- `dataframe`: The Pandas DataFrame.

hadia206 · 2026-01-29T22:50:07Z

pydough/user_collections/README.md

 The user collections module provides a way to create collections that are not part of the static metadata graph but can be generated dynamically based on user input or code. The most common user collection are integer range collections and Pandas DataFrame collections.
-The range collection, generates a sequence of numbers. The `RangeGeneratedCollection` class allows users to define a range collection by specifying the start, end, and step values. The `range_collection` function is a convenient API to create instances of `RangeGeneratedCollection`.
+The range collection, generates a sequence of numbers. The `RangeGeneratedCollection` class allows users to define a range collection by specifying the start, end, and step values. The `range_collection` function is a convenient API to create instances of `RangeGeneratedCollection`.
+The dataframe collection, generates a collection based on the given pandas dataframe. The `DataframeGeneratedCollection` class


Suggested change

The dataframe collection, generates a collection based on the given pandas dataframe. The `DataframeGeneratedCollection` class

The dataframe collection, generates a collection based on the given Pandas Dataframe. The `DataframeGeneratedCollection` class

hadia206 · 2026-01-29T22:50:53Z

pydough/user_collections/README.md

-The range collection, generates a sequence of numbers. The `RangeGeneratedCollection` class allows users to define a range collection by specifying the start, end, and step values. The `range_collection` function is a convenient API to create instances of `RangeGeneratedCollection`.
+The range collection, generates a sequence of numbers. The `RangeGeneratedCollection` class allows users to define a range collection by specifying the start, end, and step values. The `range_collection` function is a convenient API to create instances of `RangeGeneratedCollection`.
+The dataframe collection, generates a collection based on the given pandas dataframe. The `DataframeGeneratedCollection` class
+allows user to create a collection by specifying the dataframe and name. The `dataframe_collection` function is an effective API


for consistency with range description

Suggested change

allows user to create a collection by specifying the dataframe and name. The `dataframe_collection` function is an effective API

allows user to create a collection by specifying the dataframe and name. The `dataframe_collection` function is a convenient API

hadia206 · 2026-01-29T22:51:20Z

pydough/user_collections/user_collection_apis.py

+) -> UnqualifiedGeneratedCollection:
+    """
+    Implementation of the `pydough.dataframe_collection` function, which provides
+    a way to create a collection of pandas dataframe in PyDough.


Suggested change

a way to create a collection of pandas dataframe in PyDough.

a way to create a collection of Pandas DataFrame in PyDough.

knassre-bodo

A few more comments based on some of Hadia's feedback + a few things I spotted while taking second/third looks in some of those files.

knassre-bodo · 2026-01-31T09:28:49Z

documentation/dsl.md

+Supported Signatures:
+- `dataframe_collection(dataframe, name)`: generates collection with the given datafram and name.
+
+#### Example


Should we perhaps also have a more involved example that does more things with the collection?

Reminder to update the example here

knassre-bodo · 2026-01-31T09:31:04Z

pydough/sqlglot/override_qualify.py

-                        table_alias.append("columns", exp.to_identifier(f"_col_{i}"))
+                        # PyDough change: adjust the formula to match the
+                        # dialect's
+                        if isinstance(dialect, SQLite):
+                            table_alias.append(
+                                "columns", exp.to_identifier(f"column{i + 1}")
+                            )
+                        else:
+                            table_alias.append(
+                                "columns", exp.to_identifier(f"_col_{i}")
+                            )


For context, this naming convention matters a lot when handling SQLite, and SQLGlot uses the wrong one so it glitches out, but if we use the SQLite convention for other dialects it causes SQLGlot to glitch out in other places. This is not an ideal fix, but its much more straightforward than taking SQLGlot apart to actually fix the issue.

knassre-bodo · 2026-01-31T09:32:51Z

pydough/user_collections/dataframe_collection.py

+
+    @property
+    def unique_column_names(self) -> list[str]:
+        return list(dict.fromkeys(self.columns))


Aww heck I knew I forgot something in the spec: we need the option to specify subsets of the columns that are unique (just like we do in metadata). It can be an optional argument in pydough.dataframe_collection which has type list[str | list[str]] where the strings are the column names. For instance, in your dsl.md exmaple, it would be ["color", "idx"] since both columns are unique.

We should then have tests on the behavior of window function / correlation tests with DataFrames that use this argument.

Reminder to do this^

knassre-bodo · 2026-01-31T09:35:09Z

tests/test_pipeline_tpch_custom.py

+                "r = pydough.range_collection('tbl', 'v', 0, 500, 13).CALCULATE(first_digit=INTEGER(STRING(v)[:1]))\n"
+                "result = r.PARTITION(name='digits', by=first_digit).CALCULATE(first_digit, n=COUNT(tbl))",
+                "TPCH",


These new tests cover some categories of edge cases John discovered while testing DataFrames which we realized hadn't been covered/handled yet for simple range (hence all the weird HybridTree changes).

knassre-bodo · 2026-01-31T09:37:19Z

pydough/user_collections/dataframe_collection.py

+            elif len(dataframe[col]) == 0:
+                raise TypeError(
+                    f"Column '{col}' is empty. All columns must have at least one value."
+                )


Oh shoot that reminds me... we should have some tests for DataFrames with 1+ columns but no rows. There's a lot of edge cases to consider there, e.g. what happens if you cross join it with something else.

knassre-bodo · 2026-01-31T09:38:44Z

pydough/user_collections/dataframe_collection.py

+        match field_type:
+            case _ if pa.types.is_null(field_type):


If all of these are going to be case _ if ...:, we should just be doing a a bunch of if ... elif ... elif ... instead.

knassre-bodo · 2026-01-31T09:39:27Z

pydough/user_collections/dataframe_collection.py

+            case _ if pa.types.is_binary(field_type) or pa.types.is_large_binary(
+                field_type
+            ):
+                return StringType()


Shouldn't this be rejected for now?

knassre-bodo · 2026-01-31T09:40:00Z

pydough/user_collections/dataframe_collection.py

+            case _:
+                return UnknownType()


Are there cases where this happens that are valid? If so what are there? (should be documented in comments). Otherwise, we should explicitly ban things with clear error messages.

Same question. Let's be clear what are those unknowns

…[run sf][run postgres][run mysql][run s3]

hadia206

You addressed a lot of things. Almost there :)

hadia206 · 2026-02-06T20:13:39Z

documentation/dsl.md

+Supported Signatures:
+- `dataframe_collection(dataframe, name)`: generates collection with the given datafram and name.
+
+#### Example


Reminder to update the example here

hadia206 · 2026-02-06T20:20:45Z

pydough/user_collections/dataframe_collection.py

+            and self.name == other.name
+            and self.columns == other.columns


Do we need these? If we already check self.dataframe.equals in the line below it?

hadia206 · 2026-02-06T20:29:02Z

pydough/user_collections/dataframe_collection.py

+
+            except pa.ArrowInvalid:
+                raise ValueError(
+                    f"Mixed types in column '{col}'. All values in a column must be of the same type."


Let's generalize error message here. I believe this error is not just for mixed types.
It can appear with other type conversion issues

Example I have seen online
ArrowInvalid: Invalid null value in PyArrow means data doesn't match the expected type

hadia206 · 2026-02-06T20:29:41Z

pydough/user_collections/dataframe_collection.py

+            raise ValueError("Structs are not supported for dataframe collections")
+
+        elif pa_types.is_dictionary(field_type):
+            raise ValueError("Dictionaries are not supported for dataframe collections")


what about tuples? Do we support that?

hadia206 · 2026-02-06T20:36:10Z

documentation/dsl.md

+
+The supported PyDough types for `dataframe_collection` are:
+- `NumericType`: includes float, integer, infinity, Nan.
+- `BooleanType`: includes classic true and false.


Suggested change

- `BooleanType`: includes classic true and false.

- `BooleanType`: True and False.

hadia206 · 2026-02-06T20:40:59Z

documentation/dsl.md

+The supported PyDough types for `dataframe_collection` are:
+- `NumericType`: includes float, integer, infinity, Nan.
+- `BooleanType`: includes classic true and false.
+- `StringType`: includes all aphanumeric caracters.


Removed and fixed the typo

Suggested change

- `StringType`: includes all aphanumeric caracters.

- `StringType`: alphanumeric characters.

hadia206 · 2026-02-06T20:42:11Z

documentation/dsl.md

+- `NumericType`: includes float, integer, infinity, Nan.
+- `BooleanType`: includes classic true and false.
+- `StringType`: includes all aphanumeric caracters.
+- `Datetype`: includes date and datetime.


Timestamp? NaT?

Suggested change

- `Datetype`: includes date and datetime.

- `Datetype`: date and datetime.

hadia206 · 2026-02-06T20:44:54Z

pydough/user_collections/dataframe_collection.py

+- float64
+- decimal
+- bool
+- datetime64


date?
what about timestamp?

hadia206 · 2026-02-06T20:54:50Z

pydough/user_collections/user_collection_apis.py

+
+
+def dataframe_collection(
+    name: str, dataframe: pd.DataFrame


This is missing Kian's request here

hadia206 · 2026-02-06T20:56:32Z

pydough/user_collections/dataframe_collection.py

+        else:
+            return UnknownType()


Let's clarify what are these if any are still missing.
See my comment in the dsl.md as well.

Initial documentation

dfad733

john-sanchez31 marked this pull request as draft January 8, 2026 14:43

john-sanchez31 self-assigned this Jan 8, 2026

john-sanchez31 added enhancement New feature or request extensibility Increasing situations in which PyDough works effort - high major issue that will require multiple steps or complex design testing Alters the testing/CI process for PyDough labels Jan 8, 2026

john-sanchez31 and others added 23 commits January 12, 2026 15:05

base df collection implementation for sqlite, ansi and mysql

3d10cc0

types fixed

4ea44f8

ref sql added

486d785

implementation df collections for postgres and snowflake

0eef40d

datatypes fixed

036f39a

datatypes, numbers and inf test added

c882800

string and cross df collection tests (no fix for partition yet)

4e84394

WIP: patition with user generated collections

60b9703

Adding more range colleciton partition tests and fixing qualification…

4d6ca40

… SQLGlot bug discovered along the way [RUN CI]

fixing comments and deleting unneccesary case

63fb854

Merge branch 'John/df_collection' of https://github.com/bodo-ai/PyDough…

b0d942e

… into John/df_collection

partition 2 df collection test

f3f22ae

df collection where date test added

7a49b16

df collection top_k test added

0318e25

dataframe_collection_best test added

a2d3dd6

bad test and window function test added

c500b54

docstring and refactored code [run all]

2038159

pyarrow dependency added [run all]

cb77bd8

pyarrow dependency changed [run all]

78ec168

testing [run all]

3c2d350

testing [run all]

f15d3a1

reverting [run all]

eb53764

connectors version locked [run all]

29276b1

john-sanchez31 added 8 commits January 22, 2026 12:49

range test updated [run all]

6ec7c96

conflicts solved

3c9c571

dataframe test changed to tpch_custom file

62c32ce

defog deealership_adv8 test fixed

01f94d2

refsol sql added

c246449

defog dealership_adv13 added

4ffe44f

testing [run all]

b60aa13

dataframe bad 2 fixed and sf_masked test updated [run ci][run mysql][…

03d37d6

…run sf][run postgres][run s3]

john-sanchez31 marked this pull request as ready for review January 27, 2026 18:57

john-sanchez31 added 4 commits January 27, 2026 15:04

dataframe docstring added

cceebc1

Merge branch 'main' into John/df_collection

5e5653a

testing [run ci][run s3][run mysql][run sf][run postgres]

b8b4ce1

minor fixes [run ci][run s3][run mysql][run postgres][run sf]

0923d64

john-sanchez31 commented Jan 27, 2026

View reviewed changes

john-sanchez31 requested review from a team, hadia206, juankx-bodo and knassre-bodo and removed request for a team January 27, 2026 21:34

john-sanchez31 added 2 commits January 27, 2026 16:00

Revert "dataframe bad 2 fixed and sf_masked test updated [run ci][run…

56baf54

… mysql][run sf][run postgres][run s3]" This reverts commit 03d37d6.

test dataframe_collection_bad_2 fixed [run ci][run s3][run mysql][run…

454c6e1

… sf][run postgres]

knassre-bodo reviewed Jan 28, 2026

View reviewed changes

hadia206 reviewed Jan 29, 2026

View reviewed changes

knassre-bodo reviewed Jan 31, 2026

View reviewed changes

documentation fixed, df validation and generation factorized [run ci]…

2d7514e

…[run sf][run postgres][run mysql][run s3]

john-sanchez31 requested review from hadia206 and knassre-bodo February 2, 2026 16:45

hadia206 reviewed Feb 6, 2026

View reviewed changes



		@pytest.fixture(scope="session")
		def get_postgres_defog_graphs() -> graph_fetcher:



		@pytest.fixture(scope="session")
		def get_dialect_defog_graphs(

		Supported Signatures:
		- `dataframe_collection(dataframe, name)`: generates collection with the given datafram and name.



		# Windows functions
		def dataframe_collection_window_functions():

	The `dataframe_collection` creates a collection within a specified pandas dataframe. This is useful for building datasets dynamically.
	The `dataframe_collection` creates a collection from a specified Pandas DataFrame. This is useful for building datasets dynamically.

	- `dataframe`: The panda dataframe containing the corresponding data
	- `dataframe`: Pandas DataFrame containing the corresponding data.

	- `dataframe`: The pandas dataframe containing all data (rows and columns).
	- `dataframe`: The Pandas DataFrame containing all data (rows and columns).

	- `dataframe`: The pandas dataframe.
	- `dataframe`: The Pandas DataFrame.

	The dataframe collection, generates a collection based on the given pandas dataframe. The `DataframeGeneratedCollection` class
	The dataframe collection, generates a collection based on the given Pandas Dataframe. The `DataframeGeneratedCollection` class

	allows user to create a collection by specifying the dataframe and name. The `dataframe_collection` function is an effective API
	allows user to create a collection by specifying the dataframe and name. The `dataframe_collection` function is a convenient API

	a way to create a collection of pandas dataframe in PyDough.
	a way to create a collection of Pandas DataFrame in PyDough.

		and self.name == other.name
		and self.columns == other.columns

	- `BooleanType`: includes classic true and false.
	- `BooleanType`: True and False.

	- `StringType`: includes all aphanumeric caracters.
	- `StringType`: alphanumeric characters.

	- `Datetype`: includes date and datetime.
	- `Datetype`: date and datetime.

Conversation

john-sanchez31 commented Jan 8, 2026

Uh oh!

john-sanchez31 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hadia206 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

knassre-bodo left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment