Adding support for user defined functions in PyDough (#380)

Resolves #382. Adds ability to define additional function operators in the PyDough metadata, with three formats at first:
- SQL Alias: calls to the function get translated directly 1:1 to a function call in the database's SQL dialect (scalar or aggregation)
- SQL Window Alias: Same as SQL Alias but only for window functions
- SQL Macro: calls to the function inject their arguments' SQL texts into a Python format string

Also updated docstrings throughout the codebase to avoid using types of the inputs/outputs since those should be part of the type annotations (not the docstrings).

Author: knassre-bodo

documentation/metadata.md

@@ -54,7 +54,7 @@ def metadata(self) -> GraphMetadata | None:
         Get the active metadata graph.
 
         Returns:
-            GraphMetadata: The active metadata graph.
+            The active metadata graph.
         """
         return self._metadata
 
@@ -64,7 +64,7 @@ def metadata(self, graph: GraphMetadata | None) -> None:
         Set the active metadata graph.
 
         Args:
-            graph (GraphMetadata | None): The metadata graph to set.
+            graph: The metadata graph to set.
         """
         self._metadata = graph
 
@@ -74,7 +74,7 @@ def config(self) -> PyDoughConfigs:
         Get the active PyDough configuration.
 
         Returns:
-            PyDoughConfigs: The active PyDough configuration.
+            The active PyDough configuration.
         """
         return self._config
 
@@ -84,7 +84,7 @@ def config(self, config: PyDoughConfigs) -> None:
         Set the active PyDough configuration.
 
         Args:
-            config (PyDoughConfigs): The PyDough configuration to set.
+            `config`: The PyDough configuration to set.
         """
         self._config = config
 
@@ -94,7 +94,7 @@ def database(self) -> DatabaseContext:
         Get the active database context.
 
         Returns:
-            DatabaseContext: The active database context.
+            The active database context.
         """
         return self._database
 
@@ -104,7 +104,7 @@ def database(self, context: DatabaseContext) -> None:
         Set the active database context.
 
         Args:
-            context (DatabaseContext): The database context to set.
+            `context`: The database context to set.
         """
         self._database = context
 
@@ -114,13 +114,13 @@ def connect_database(self, database_name: str, **kwargs) -> DatabaseContext:
         the corresponding context in case the user wants/needs to modify it.
 
         Args:
-            database_name (str): The name of the database to connect to.
+            `database_name`: The name of the database to connect to.
             **kwargs: Additional keyword arguments to pass to the connection.
                 All arguments must be accepted using the supported connect API
                 for the dialect. Most likely the database path will be required.
 
         Returns:
-            DatabaseContext: The newly created database context.
+            The newly created database context.
         """
         context: DatabaseContext = load_database_context(database_name, **kwargs)
         self.database = context
@@ -135,13 +135,13 @@ def load_metadata_graph(self, graph_path: str, graph_name: str) -> GraphMetadata
         property directly later.
 
         Args:
-            graph_path (str): The path to load the graph. At this time this must be on
+            `graph_path`: The path to load the graph. At this time this must be on
                 the user's local file system.
-            graph_name (str): The name under which to load the graph from the file. This
+            `graph_name`: The name under which to load the graph from the file. This
                 is to allow loading multiple graphs from the same json file.
 
         Returns:
-            GraphMetadata: The loaded metadata graph.
+            The loaded metadata graph.
         """
         graph: GraphMetadata = parse_json_metadata_from_file(graph_path, graph_name)
         self.metadata = graph

pydough/configs/session.py

@@ -1142,19 +1142,22 @@ def translate_child_pullup(self, node: HybridChildPullUp) -> TranslationOutput:
         return TranslationOutput(child_result.relational_node, new_expressions)
 
     def translate_hybridroot(self, context: TranslationOutput) -> TranslationOutput:
-        """Converts a HybridRoot node into a relational tree.
-        This method shifts all expressions in the given context back by one level,
-        effectively removing the HybridRoot from the context (re-aligning them to the parent context's scope).
-        This is needed when stepping out of a nested context.
-        The HybridRoot itself does not introduce a new relational operation but serves as a logical boundary.
-        This method prepares the context so that subsequent operations refer to the correct expression depth.
+        """
+        Converts a HybridRoot node into a relational tree. This method shifts
+        all expressions in the given context back by one level, effectively
+        removing the HybridRoot from the context (re-aligning them to the
+        parent context's scope). This is needed when stepping out of a nested
+        context. The HybridRoot itself does not introduce a new relational
+        operation but serves as a logical boundary. This method prepares the
+        context so that subsequent operations refer to the correct expression
+        depth.
 
         Args:
-            context (TranslationOutput): The current translation context
-            associated with the HybridRoot. Must not be None.
+            `context`: The current translation context associated with the
+            HybridRoot. Must not be None.
 
         Returns:
-            TranslationOutput: The translated output payload.
+            The translated output payload.
         """
         new_expressions: dict[HybridExpr, ColumnReference] = {}
         for expr, column_ref in context.expressions.items():
@@ -1324,14 +1327,12 @@ def make_relational_ordering(
     Converts a list of collation expressions into a list of ExpressionSortInfo.
 
     Args:
-        collation (list[CollationExpression]): The list of collation
-            expressions to convert.
-        expressions (dict[HybridExpr, ColumnReference]): The dictionary of
-            expressions to use for the relational ordering.
+        `collation`: The list of collation expressions to convert.
+        `expressions`: The dictionary of expressions to use for the relational
+        ordering.
 
     Returns:
-        list[ExpressionSortInfo]: The ordering expressions converted into
-        ExpressionSortInfo.
+        The ordering expressions converted into `ExpressionSortInfo`.
     """
     orderings: list[ExpressionSortInfo] = []
     for col_expr in collation:

pydough/conversion/relational_converter.py

@@ -15,13 +15,13 @@ def load_database_context(database_name: str, **kwargs) -> DatabaseContext:
     Load the database context with the appropriate connection and dialect.
 
     Args:
-        database (str): The name of the database to connect to.
-        **kwargs: Additional keyword arguments to pass to the connection.
+        `database`: The name of the database to connect to.
+        `**kwargs`: Additional keyword arguments to pass to the connection.
             All arguments must be accepted using the supported connect API
             for the dialect.
 
     Returns:
-        DatabaseContext: The database context object.
+        The database context object.
     """
     supported_databases = {"sqlite"}
     connection: DatabaseConnection
@@ -44,7 +44,7 @@ def load_sqlite_connection(**kwargs) -> DatabaseConnection:
     around the DB 2.0 connect API.
 
     Returns:
-        DatabaseConnection: A database connection object for SQLite.
+        A database connection object for SQLite.
     """
     if "database" not in kwargs:
         raise ValueError("SQLite connection requires a database path.")

pydough/database_connectors/builtin_databases.py

@@ -38,7 +38,7 @@ def execute_query_df(self, sql: str) -> pd.DataFrame:
         types are in scope and how we need to test them.
 
         Args:
-            sql (str): The SQL query to execute.
+            `sql`: The SQL query to execute.
 
         Returns:
             list[pt.Any]: A list of rows returned by the query.
@@ -83,10 +83,10 @@ def from_string(dialect: str) -> "DatabaseDialect":
         """Convert a string to a DatabaseDialect enum.
 
         Args:
-            dialect (str): The string representation of the dialect.
+            `dialect`: The string representation of the dialect.
 
         Returns:
-            DatabaseDialect: The dialect enum.
+            The dialect enum.
         """
         if dialect == "ansi":
             return DatabaseDialect.ANSI

pydough/database_connectors/database_connector.py

@@ -107,15 +107,15 @@ def to_sql(node: UnqualifiedNode, **kwargs) -> str:
     Convert the given unqualified tree to a SQL string.
 
     Args:
-        node (UnqualifiedNode): The node to convert to SQL.
-        **kwargs: Additional arguments to pass to the conversion for testing.
+        `node`: The node to convert to SQL.
+        `**kwargs`: Additional arguments to pass to the conversion for testing.
             From a user perspective these values should always be derived from
             the active session, but to allow a simple + extensible testing
             infrastructure in the future, any of these can be passed in using
             the name of the field in session.py.
 
     Returns:
-        str: The SQL string corresponding to the unqualified query.
+        The SQL string corresponding to the unqualified query.
     """
     graph: GraphMetadata
     config: PyDoughConfigs
@@ -139,15 +139,15 @@ def to_df(node: UnqualifiedNode, **kwargs) -> pd.DataFrame:
     DataFrame.
 
     Args:
-        node (UnqualifiedNode): The node to convert to a DataFrame.
-        **kwargs: Additional arguments to pass to the conversion for testing.
+        `node`: The node to convert to a DataFrame.
+        `**kwargs`: Additional arguments to pass to the conversion for testing.
             From a user perspective these values should always be derived from
             the active session, but to allow a simple + extensible testing
             infrastructure in the future, any of these can be passed in using
             the name of the field in session.py.
 
     Returns:
-        pd.DataFrame: The DataFrame corresponding to the unqualified query.
+        The DataFrame corresponding to the unqualified query.
     """
     graph: GraphMetadata
     config: PyDoughConfigs

pydough/evaluation/evaluate_unqualified.py

@@ -17,6 +17,7 @@
     "PyDoughPredicate",
     "extract_array",
     "extract_bool",
+    "extract_integer",
     "extract_object",
     "extract_string",
     "is_bool",
@@ -375,6 +376,31 @@ def extract_bool(json_obj: dict, key_name: str, obj_name: str) -> bool:
     return value
 
 
+def extract_integer(json_obj: dict, key_name: str, obj_name: str) -> int:
+    """
+    Extracts an integer field from a JSON object, returning the integer field
+    and verifying that the field exists and is well formed.
+
+    Args:
+        `json_obj`: the JSON object to extract the string from.
+        `key_name`: the name of the key in the JSON object that
+        contains the string.
+        `obj_name`: the name of the object being extracted from, to be used
+        in error messages.
+
+    Returns:
+        The integer value of the field.
+
+    Raises:
+        `PyDoughMetadataException` if the JSON object does not contain a key
+        with the name `key_name`, or if the value of the key is not an integer.
+    """
+    HasPropertyWith(key_name, is_integer).verify(json_obj, obj_name)
+    value = json_obj[key_name]
+    assert isinstance(value, int)
+    return value
+
+
 def extract_array(json_obj: dict, key_name: str, obj_name: str) -> list:
     """
     Extracts an array field from a JSON object, returning the string field

pydough/metadata/errors.py

@@ -2,9 +2,16 @@
 Definition of PyDough metadata for a graph.
 """
 
+from typing import TYPE_CHECKING
+
 from pydough.metadata.abstract_metadata import AbstractMetadata
 from pydough.metadata.errors import HasType, PyDoughMetadataException, is_valid_name
 
+if TYPE_CHECKING:
+    from pydough.pydough_operators import (
+        ExpressionFunctionOperator,
+    )
+
 
 class GraphMetadata(AbstractMetadata):
     """
@@ -17,6 +24,7 @@ class GraphMetadata(AbstractMetadata):
         "version",
         "collections",
         "relationships",
+        "functions",
         "additional definitions",
         "verified pydough analysis",
         "extra semantic info",
@@ -39,6 +47,7 @@ def __init__(
         self._verified_pydough_analysis: list[dict] | None = verified_pydough_analysis
         self._name: str = name
         self._collections: dict[str, AbstractMetadata] = {}
+        self._functions: dict[str, ExpressionFunctionOperator] = {}
         super().__init__(description, synonyms, extra_semantic_info)
 
     @property
@@ -55,6 +64,13 @@ def collections(self) -> dict[str, AbstractMetadata]:
         """
         return self._collections
 
+    @property
+    def functions(self) -> dict[str, "ExpressionFunctionOperator"]:
+        """
+        The user defined functions contained within the graph.
+        """
+        return self._functions
+
     @property
     def error_name(self) -> str:
         return f"graph {self.name!r}"
@@ -131,3 +147,46 @@ def get_collection(self, collection_name: str) -> AbstractMetadata:
 
     def __getitem__(self, key: str):
         return self.get_collection(key)
+
+    def get_function_names(self) -> list[str]:
+        """
+        Fetches all of the names of user defined functions in the graph.
+        """
+        return list(self.functions)
+
+    def get_function(self, function_name: str) -> "ExpressionFunctionOperator":
+        """
+        Fetches a specific function's metadata from within the graph by name.
+        """
+        if function_name not in self.functions:
+            raise PyDoughMetadataException(
+                f"{self.error_name} does not have a function named {function_name!r}"
+            )
+        return self.functions[function_name]
+
+    def add_function(self, name: str, function: "ExpressionFunctionOperator") -> None:
+        """
+        Adds a new user defined function to the graph.
+
+        Args:
+            `name`: the name of the function.
+            `function`: the function operator being inserted into the graph.
+
+        Raises:
+            `PyDoughMetadataException`: if `function` cannot be inserted
+            into the graph because of a name collision.
+        """
+        is_valid_name.verify(name, "function name")
+        if name == self.name:
+            raise PyDoughMetadataException(
+                f"Function name {name!r} cannot be the same as the graph name {self.name!r}"
+            )
+        if name in self.get_collection_names():
+            raise PyDoughMetadataException(
+                f"Function name {name!r} cannot be the same as a collection name in {self.error_name}"
+            )
+        if name in self.functions:
+            raise PyDoughMetadataException(
+                f"Function {name!r} already exists in {self.error_name}"
+            )
+        self.functions[name] = function