refactor: use DATE_TRUNC for Elasticsearch time grain (#25717)

Co-authored-by: Mikel Vuka <mikel.vuka@zalando.de>
This commit is contained in:
Mikel Vuka 2023-10-20 19:05:05 +02:00 committed by GitHub
parent 8fb0c8da56
commit 9972ac6908
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 20 deletions

View File

@ -40,15 +40,19 @@ class ElasticSearchEngineSpec(BaseEngineSpec): # pylint: disable=abstract-metho
allows_subqueries = True
allows_sql_comments = False
_date_trunc_functions = {
"DATETIME": "DATE_TRUNC",
}
_time_grain_expressions = {
None: "{col}",
TimeGrain.SECOND: "HISTOGRAM({col}, INTERVAL 1 SECOND)",
TimeGrain.MINUTE: "HISTOGRAM({col}, INTERVAL 1 MINUTE)",
TimeGrain.HOUR: "HISTOGRAM({col}, INTERVAL 1 HOUR)",
TimeGrain.DAY: "HISTOGRAM({col}, INTERVAL 1 DAY)",
TimeGrain.WEEK: "DATE_TRUNC('week', {col})",
TimeGrain.MONTH: "HISTOGRAM({col}, INTERVAL 1 MONTH)",
TimeGrain.YEAR: "HISTOGRAM({col}, INTERVAL 1 YEAR)",
TimeGrain.SECOND: "{func}('second', {col})",
TimeGrain.MINUTE: "{func}('minute', {col})",
TimeGrain.HOUR: "{func}('hour', {col})",
TimeGrain.DAY: "{func}('day', {col})",
TimeGrain.WEEK: "{func}('week', {col})",
TimeGrain.MONTH: "{func}('month', {col})",
TimeGrain.YEAR: "{func}('year', {col})",
}
type_code_map: dict[int, str] = {} # loaded from get_datatype only if needed

View File

@ -14,27 +14,30 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from parameterized import parameterized
from sqlalchemy import column
from superset.constants import TimeGrain
from superset.db_engine_specs.elasticsearch import ElasticSearchEngineSpec
from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec
class TestElasticsearchDbEngineSpec(TestDbEngineSpec):
def test_time_grain_week_expression(self):
@parameterized.expand(
[
[TimeGrain.SECOND, "DATE_TRUNC('second', ts)"],
[TimeGrain.MINUTE, "DATE_TRUNC('minute', ts)"],
[TimeGrain.HOUR, "DATE_TRUNC('hour', ts)"],
[TimeGrain.DAY, "DATE_TRUNC('day', ts)"],
[TimeGrain.WEEK, "DATE_TRUNC('week', ts)"],
[TimeGrain.MONTH, "DATE_TRUNC('month', ts)"],
[TimeGrain.YEAR, "DATE_TRUNC('year', ts)"],
]
)
def test_time_grain_expressions(self, time_grain, expected_time_grain_expression):
col = column("ts")
col.type = "datetime"
expected_time_grain_expression = "DATE_TRUNC('week', ts)"
col.type = "DATETIME"
actual = ElasticSearchEngineSpec.get_timestamp_expr(
col=col, pdf=None, time_grain="P1W"
)
self.assertEqual(str(actual), expected_time_grain_expression)
def test_time_grain_hour_expression(self):
col = column("ts")
col.type = "datetime"
expected_time_grain_expression = "HISTOGRAM(ts, INTERVAL 1 HOUR)"
actual = ElasticSearchEngineSpec.get_timestamp_expr(
col=col, pdf=None, time_grain="PT1H"
col=col, pdf=None, time_grain=time_grain
)
self.assertEqual(str(actual), expected_time_grain_expression)