fix: pinot select query logic

superset/db_engine_specs/pinot.py

@@ -79,24 +79,18 @@ def get_timestamp_expr(
         else:
             seconds_or_ms = "MILLISECONDS" if pdf == "epoch_ms" else "SECONDS"
             tf = f"1:{seconds_or_ms}:EPOCH"
-        granularity = cls.get_time_grain_expressions().get(time_grain)
-        if not granularity:
-            raise NotImplementedError("No pinot grain spec for " + str(time_grain))
+        if time_grain:
+            granularity = cls.get_time_grain_expressions().get(time_grain)
+            if not granularity:
+                raise NotImplementedError("No pinot grain spec for " + str(time_grain))
+        else:
+            return TimestampExpression(f"{{col}}", col)
         # In pinot the output is a string since there is no timestamp column like pg
-        time_expr = f'DATETIMECONVERT({{col}}, "{tf}", "{tf}", "{granularity}")'
+        time_expr = f"DATETIMECONVERT({{col}}, '{tf}', '{tf}', '{granularity}')"
         return TimestampExpression(time_expr, col)
 
     @classmethod
     def make_select_compatible(
         cls, groupby_exprs: Dict[str, ColumnElement], select_exprs: List[ColumnElement]
     ) -> List[ColumnElement]:
-        # Pinot does not want the group by expr's to appear in the select clause
-        select_sans_groupby = []
-        # We want identity and not equality, so doing the filtering manually
-        for sel in select_exprs:
-            for gr in groupby_exprs:
-                if sel is gr:
-                    break
-            else:
-                select_sans_groupby.append(sel)
-        return select_sans_groupby
+        return select_exprs

tests/db_engine_specs/pinot_tests.py

@@ -29,5 +29,5 @@ def test_pinot_time_expression_sec_one_1m_grain(self):
         result = str(expr.compile())
         self.assertEqual(
             result,
-            'DATETIMECONVERT(tstamp, "1:SECONDS:EPOCH", "1:SECONDS:EPOCH", "1:MONTHS")',
+            "DATETIMECONVERT(tstamp, '1:SECONDS:EPOCH', '1:SECONDS:EPOCH', '1:MONTHS')",
         )  # noqa