perf(sqla): avoid unnecessary type check on adhoc column (#23491)

This commit is contained in:
Ville Brofeldt 2023-03-31 18:18:08 +03:00 committed by GitHub
parent bc2ec044b8
commit ee9ef24509
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 26 additions and 18 deletions

View File

@ -1006,15 +1006,19 @@ class SqlaTable(Model, BaseDatasource): # pylint: disable=too-many-public-metho
return self.make_sqla_column_compatible(sqla_metric, label) return self.make_sqla_column_compatible(sqla_metric, label)
def adhoc_column_to_sqla( def adhoc_column_to_sqla( # pylint: disable=too-many-locals
self, self,
col: AdhocColumn, col: AdhocColumn,
force_type_check: bool = False,
template_processor: Optional[BaseTemplateProcessor] = None, template_processor: Optional[BaseTemplateProcessor] = None,
) -> ColumnElement: ) -> ColumnElement:
""" """
Turn an adhoc column into a sqlalchemy column. Turn an adhoc column into a sqlalchemy column.
:param col: Adhoc column definition :param col: Adhoc column definition
:param force_type_check: Should the column type be checked in the db.
This is needed to validate if a filter with an adhoc column
is applicable.
:param template_processor: template_processor instance :param template_processor: template_processor instance
:returns: The metric defined as a sqlalchemy column :returns: The metric defined as a sqlalchemy column
:rtype: sqlalchemy.sql.column :rtype: sqlalchemy.sql.column
@ -1027,28 +1031,28 @@ class SqlaTable(Model, BaseDatasource): # pylint: disable=too-many-public-metho
template_processor=template_processor, template_processor=template_processor,
) )
col_in_metadata = self.get_column(expression) col_in_metadata = self.get_column(expression)
time_grain = col.get("timeGrain")
has_timegrain = col.get("columnType") == "BASE_AXIS" and time_grain
is_dttm = False
if col_in_metadata: if col_in_metadata:
sqla_column = col_in_metadata.get_sqla_col( sqla_column = col_in_metadata.get_sqla_col(
template_processor=template_processor template_processor=template_processor
) )
is_dttm = col_in_metadata.is_temporal is_dttm = col_in_metadata.is_temporal
else: else:
try: sqla_column = literal_column(expression)
sqla_column = literal_column(expression) if has_timegrain or force_type_check:
# probe adhoc column type try:
tbl, _ = self.get_from_clause(template_processor) # probe adhoc column type
qry = sa.select([sqla_column]).limit(1).select_from(tbl) tbl, _ = self.get_from_clause(template_processor)
sql = self.database.compile_sqla_query(qry) qry = sa.select([sqla_column]).limit(1).select_from(tbl)
col_desc = get_columns_description(self.database, sql) sql = self.database.compile_sqla_query(qry)
is_dttm = col_desc[0]["is_dttm"] col_desc = get_columns_description(self.database, sql)
except SupersetGenericDBErrorException as ex: is_dttm = col_desc[0]["is_dttm"]
raise ColumnNotFoundException(message=str(ex)) from ex except SupersetGenericDBErrorException as ex:
raise ColumnNotFoundException(message=str(ex)) from ex
if ( if is_dttm and has_timegrain:
is_dttm
and col.get("columnType") == "BASE_AXIS"
and (time_grain := col.get("timeGrain"))
):
sqla_column = self.db_engine_spec.get_timestamp_expr( sqla_column = self.db_engine_spec.get_timestamp_expr(
col=sqla_column, col=sqla_column,
pdf=None, pdf=None,
@ -1458,7 +1462,11 @@ class SqlaTable(Model, BaseDatasource): # pylint: disable=too-many-public-metho
col_obj = dttm_col col_obj = dttm_col
elif is_adhoc_column(flt_col): elif is_adhoc_column(flt_col):
try: try:
sqla_col = self.adhoc_column_to_sqla(flt_col) sqla_col = self.adhoc_column_to_sqla(
col=flt_col,
force_type_check=True,
template_processor=template_processor,
)
applied_adhoc_filters_columns.append(flt_col) applied_adhoc_filters_columns.append(flt_col)
except ColumnNotFoundException: except ColumnNotFoundException:
rejected_adhoc_filters_columns.append(flt_col) rejected_adhoc_filters_columns.append(flt_col)

View File

@ -28,7 +28,7 @@ from superset.utils.pandas_postprocessing.utils import (
@validate_column_args("index", "columns") @validate_column_args("index", "columns")
def pivot( # pylint: disable=too-many-arguments,too-many-locals def pivot( # pylint: disable=too-many-arguments
df: DataFrame, df: DataFrame,
index: List[str], index: List[str],
aggregates: Dict[str, Dict[str, Any]], aggregates: Dict[str, Dict[str, Any]],