fix(pinot): dialect date truncation (#35420)

Co-authored-by: bito-code-review[bot] <188872107+bito-code-review[bot]@users.noreply.github.com>
This commit is contained in:
Beto Dealmeida
2025-10-01 13:16:46 -04:00
committed by GitHub
parent 28389de93e
commit aa97d2fe03
2 changed files with 88 additions and 0 deletions

View File

@@ -78,6 +78,12 @@ class Pinot(MySQL):
exp.DataType.Type.UBIGINT: "UNSIGNED",
}
TRANSFORMS = {
**MySQL.Generator.TRANSFORMS,
}
# Remove DATE_TRUNC transformation - Pinot supports standard SQL DATE_TRUNC
TRANSFORMS.pop(exp.DateTrunc, None)
def datatype_sql(self, expression: exp.DataType) -> str:
# Don't use MySQL's VARCHAR size requirement logic
# Just use TYPE_MAPPING for all types
@@ -95,3 +101,8 @@ class Pinot(MySQL):
return f"{type_sql} UNSIGNED{nested}"
return f"{type_sql}{nested}"
def cast_sql(self, expression: exp.Cast, safe_prefix: str | None = None) -> str:
# Pinot doesn't support MySQL's TIMESTAMP() function
# Use standard CAST syntax instead
return super(MySQL.Generator, self).cast_sql(expression, safe_prefix)

View File

@@ -421,3 +421,80 @@ def test_unsigned_type() -> None:
assert "UNSIGNED" in result
assert "BIGINT" in result
def test_date_trunc_preserved() -> None:
"""
Test that DATE_TRUNC is preserved and not converted to MySQL's DATE() function.
"""
sql = "SELECT DATE_TRUNC('day', dt_column) FROM table"
result = sqlglot.parse_one(sql, Pinot).sql(Pinot)
assert "DATE_TRUNC" in result
assert "date_trunc('day'" in result.lower()
# Should not be converted to MySQL's DATE() function
assert result != "SELECT DATE(dt_column) FROM table"
def test_cast_timestamp_preserved() -> None:
"""
Test that CAST AS TIMESTAMP is preserved and not converted to TIMESTAMP() function.
"""
sql = "SELECT CAST(dt_column AS TIMESTAMP) FROM table"
result = sqlglot.parse_one(sql, Pinot).sql(Pinot)
assert "CAST" in result
assert "AS TIMESTAMP" in result
# Should not be converted to MySQL's TIMESTAMP() function
assert "TIMESTAMP(dt_column)" not in result
def test_date_trunc_with_cast_timestamp() -> None:
"""
Test the original complex query with DATE_TRUNC and CAST AS TIMESTAMP.
Verifies that both are preserved in parse/generate round-trip.
"""
sql = """
SELECT
CAST(
DATE_TRUNC(
'day',
CAST(
DATETIMECONVERT(
dt_epoch_ms, '1:MILLISECONDS:EPOCH',
'1:MILLISECONDS:EPOCH', '1:MILLISECONDS'
) AS TIMESTAMP
)
) AS TIMESTAMP
),
SUM(a) + SUM(b)
FROM
"default".c
WHERE
dt_epoch_ms >= 1735690800000
AND dt_epoch_ms < 1759328588000
AND locality != 'US'
GROUP BY
CAST(
DATE_TRUNC(
'day',
CAST(
DATETIMECONVERT(
dt_epoch_ms, '1:MILLISECONDS:EPOCH',
'1:MILLISECONDS:EPOCH', '1:MILLISECONDS'
) AS TIMESTAMP
)
) AS TIMESTAMP
)
LIMIT
10000
"""
result = sqlglot.parse_one(sql, Pinot).sql(Pinot)
# Verify DATE_TRUNC and CAST are preserved
assert "DATE_TRUNC" in result
assert "CAST" in result
# Verify these are NOT converted to MySQL functions
assert "TIMESTAMP(DATETIMECONVERT" not in result
assert result.count("DATE_TRUNC") == 2 # Should appear twice (SELECT and GROUP BY)