Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions superset/sql/dialects/pinot.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,12 @@ class Generator(MySQL.Generator):
exp.DataType.Type.UBIGINT: "UNSIGNED",
}

TRANSFORMS = {
**MySQL.Generator.TRANSFORMS,
}
# Remove DATE_TRUNC transformation - Pinot supports standard SQL DATE_TRUNC
TRANSFORMS.pop(exp.DateTrunc, None)
Comment thread
betodealmeida marked this conversation as resolved.

def datatype_sql(self, expression: exp.DataType) -> str:
# Don't use MySQL's VARCHAR size requirement logic
# Just use TYPE_MAPPING for all types
Expand All @@ -95,3 +101,8 @@ def datatype_sql(self, expression: exp.DataType) -> str:
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)
Comment thread
betodealmeida marked this conversation as resolved.
Comment thread
betodealmeida marked this conversation as resolved.
77 changes: 77 additions & 0 deletions tests/unit_tests/sql/dialects/pinot_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 or "DATE_TRUNC('DAY'" in result
Comment thread
betodealmeida marked this conversation as resolved.
Outdated
# 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)
Loading