Help with TDEngine custom dialect

[itepifanio]

Hi, I'm building a custom dialect for TDEngine, they have a PARTITION BY syntax similar to GROUP BY but I wasn't able to make the dialect accept the PARTITION BY , it always returns ParseError: Invalid expression / Unexpected token. Line 6, Col: 13..

The partition by is already an available token at sqlglot and I'm using the same flow available at the parser.py, so I'm not sure what I'm doing wrong.

import sqlglot
from sqlglot.helper import apply_index_offset, seq_get

from sqlglot.errors import ParseError
from sqlglot.parser import Parser
from sqlglot import tokens
from sqlglot.expressions import Expression, Func
from sqlglot import Dialect, Generator, Parser, Tokenizer, exp, generator

class AsUnit(Func):
    arg_types = {"this": True, "target": True}

def _build_as_unit(args: list) -> exp.Expression:
    return AsUnit(this=seq_get(args, 0), target=seq_get(args, 1))

class Interp(Func):
    arg_types = {"this": True}

def _build_interp(args: list) -> exp.Expression:
    return Interp(this=seq_get(args, 0))

class TDEngine(Dialect):
    class Tokenizer(tokens.Tokenizer):
        QUOTES = ["'", '"']  # Strings can be delimited by either single or double quotes
        IDENTIFIERS = ["`"]  # Identifiers can be delimited by backticks


    class Parser(Parser):

        FUNCTIONS = {
            **Parser.FUNCTIONS,
            "AS_UNIT": _build_as_unit,
            "interp": _build_interp,
        }

        EXPRESSION_PARSERS = {
            **Parser.EXPRESSION_PARSERS,
            exp.Partition: lambda self: self._parse_group,
        }

        QUERY_MODIFIER_PARSERS = {
            **Parser.QUERY_MODIFIER_PARSERS,
            tokens.TokenType.PARTITION_BY: lambda self: ("partition", self._parse_group())
        }

        def _parse_group(self):
            group_tokens = [tokens.TokenType.GROUP_BY, tokens.TokenType.PARTITION_BY]

            if not any(self._match(token) for token in group_tokens):
                return None

            return self.expression(
                exp.Group if self._match(tokens.TokenType.GROUP_BY) else exp.Partition,
                expressions=self._parse_csv(self._parse_conjunction),
                grouping_sets=self._parse_grouping_sets(),
                cube=self._match(tokens.TokenType.CUBE) and self._parse_wrapped_id_vars(),
                rollup=self._match(tokens.TokenType.ROLLUP) and self._parse_wrapped_id_vars(),
            )

This is the query I'm trying to parse:

description_1 = '[desc1]'
start = "2024-01-05T00:00:00"
end = "2024-01-07T00:00:00"
description_2 = '[desc2]'
description_3 = '[desc3]'

sql = f"""
     SELECT _irowts as time,
            interp(`value`) as tag_desc_2,
            assets
        FROM sensors 
        WHERE descriptions = '{description_2}' 
        partition by tbname 
        RANGE('{start}', '{end}') EVERY(1h) FILL(prev)
        order by time
"""

sqlglot.parse_one(sql, dialect=TDEngine)

[georgesittas]

There are a couple of issues. First of all, RANGE seems to have syntax that your dialect won't support. Then, this section is problematic:

            return self.expression(
                exp.Group if self._match(tokens.TokenType.GROUP_BY) else exp.Partition,
                expressions=self._parse_csv(self._parse_conjunction),
                grouping_sets=self._parse_grouping_sets(),
                cube=self._match(tokens.TokenType.CUBE) and self._parse_wrapped_id_vars(),
                rollup=self._match(tokens.TokenType.ROLLUP) and self._parse_wrapped_id_vars(),
            )

For the first arg, you've already consumed the group by token, so the match is incorrect there as you'll never have a GROUP BY token to consume. You need to leverage self._prev and peek back to decide what expression you'll create. Then, you're misusing exp.Partition which does not expect the rest of the args you're setting in the self.expression call. Perhaps you need a new expression type?

I'd suggest studying the code a bit more and then move on with the implementation once you have a good understanding of the basics. Look into past PRs as well, that'll help.

[itepifanio]

@georgesittas thank you! I forgot that _match was consuming the token, that's why I couldn't make any progress in my code.

I'll improve this code eventually, but I'm writing the current solution for future reference.

class TDEngine(Dialect):
    class Tokenizer(tokens.Tokenizer):
        QUOTES = ["'", '"']  # Strings can be delimited by either single or double quotes
        IDENTIFIERS = ["`"]  # Identifiers can be delimited by backticks

    class Parser(Parser):

        FUNCTIONS = {
            **Parser.FUNCTIONS,
            "AS_UNIT": _build_as_unit,
            "interp": _build_interp,
        }

        EXPRESSION_PARSERS = {
            **Parser.EXPRESSION_PARSERS,
            exp.Partition: lambda self: self._parse_group,
        }

        QUERY_MODIFIER_PARSERS = {
            **Parser.QUERY_MODIFIER_PARSERS,
            tokens.TokenType.PARTITION_BY: lambda self: ("partition", self._parse_group())
        }

        def _parse_group(self):
            group_tokens = [tokens.TokenType.GROUP_BY, tokens.TokenType.PARTITION_BY]

            if not any(self._match(token, advance=False) for token in group_tokens):
                return None

            if self._match(tokens.TokenType.GROUP_BY, advance=False):
                self._advance()
                return  self.expression(
                    exp.Group,
                    expressions=self._parse_csv(self._parse_conjunction),
                    grouping_sets=self._parse_grouping_sets(),
                    cube=self._match(tokens.TokenType.CUBE) and self._parse_wrapped_id_vars(),
                    rollup=self._match(tokens.TokenType.ROLLUP) and self._parse_wrapped_id_vars(),
                )
            self._advance()
            if self._match_text_seq("ID"):
                expressions: list[exp.Expression] = [
                    self.expression(exp.PartitionId, this=self._parse_string())
                ]
            else:
                expressions = self._parse_expressions()

            return self.expression(
                exp.Partition,
                expressions=expressions,
            )