Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Compilation of a labeled custom FunctionElement fails when used in grouping #1150

Open
snapiri opened this issue Dec 24, 2024 · 0 comments · May be fixed by #1155
Open

Compilation of a labeled custom FunctionElement fails when used in grouping #1150

snapiri opened this issue Dec 24, 2024 · 0 comments · May be fixed by #1155
Assignees
Labels
api: bigquery Issues related to the googleapis/python-bigquery-sqlalchemy API.

Comments

@snapiri
Copy link
Contributor

snapiri commented Dec 24, 2024

Python 3.10.15
sqlalchemy_bigquery 1.12.0 (issue introduced in 1.11.0).

When creating a new FunctionalElement to compile some custom function that does not support the default dialect, and grouping by that element, we get an error as we are using the default dialect instead of compiler dialect.

Steps to reproduce

  1. Create a new FunctionElement
  2. Create a function that compiles it. Make sure it does not support the default dialect.
  3. Create a query that groups by a field using this function and adding a label to it

Code example

import sqlalchemy

class CustomLower(sqlalchemy.sql.functions.FunctionElement):
    name = "custom_lower"

@sqlalchemy.ext.compiler.compiles(CustomLower)
def compile_custom_intersect(element, compiler, **kwargs):
    if compiler.dialect.name != "bigquery":
        raise sqlalchemy.exc.CompileError(
            f"custom_lower is not supported for dialect '{compiler.dialect.name}'"
        )

    clauses = list(element.clauses)
    field = compiler.process(clauses[0], **kwargs)
    return f"LOWER({field})"

db_metadata = MetaData()
engine = create_engine("bigquery://", future=True)

table1 = setup_table(
    engine,
    "table1",
    metadata,
    sqlalchemy.Column("foo", sqlalchemy.String),
    sqlalchemy.Column("bar", sqlalchemy.Integer),
)

lower_foo = CustomLower(table1.c.foo).label("some_label")
q = (
    sqlalchemy.select(lower_foo, sqlalchemy.func.max(table1.c.bar))
    .select_from(table1)
    .group_by(lower_foo)
)

print(q.compile(engine).string)

Error received

sqlalchemy.exc.CompileError: custom_lower is not supported for dialect 'default'

This is because when trying to understand if we have the grouping_ops in the element we use str() to stringify it instead of compiling it using the dialect.

The solution should be in visit_label:
We should not use str(column_label), but rather column_label.compile(dialect=self.dialect).string

Example fix

    def visit_label(self, *args, within_group_by=False, **kwargs):
        # Use labels in GROUP BY clause.
        #
        # Flag set in the group_by_clause method. Works around missing
        # equivalent to supports_simple_order_by_label for group by.
        if within_group_by:
            column_label = args[0]
            sql_keywords = {"GROUPING SETS", "ROLLUP", "CUBE"}
            label_str = column_label.compile(dialect=self.dialect).string
            if not any(keyword in label_str for keyword in sql_keywords):
                kwargs["render_label_as_label"] = column_label

        return super(BigQueryCompiler, self).visit_label(*args, **kwargs)
@product-auto-label product-auto-label bot added the api: bigquery Issues related to the googleapis/python-bigquery-sqlalchemy API. label Dec 24, 2024
@snapiri snapiri linked a pull request Jan 19, 2025 that will close this issue
3 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
api: bigquery Issues related to the googleapis/python-bigquery-sqlalchemy API.
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants