Skip to content

Commit

Permalink
Fix telemetry test when sql simplifier is disabled (#1925)
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-aalam committed Jul 17, 2024
1 parent eb652e2 commit 7e05572
Showing 1 changed file with 15 additions and 12 deletions.
27 changes: 15 additions & 12 deletions tests/integ/test_telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -586,6 +586,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
df.collect()
# API calls don't change after query is executed
query_plan_height = 2 if sql_simplifier_enabled else 3
filter = 1 if sql_simplifier_enabled else 2
low_impact = 3 if sql_simplifier_enabled else 2

assert df._plan.api_calls == [
{
Expand All @@ -594,8 +596,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
"query_plan_height": query_plan_height,
"query_plan_num_duplicate_nodes": 0,
"query_plan_complexity": {
"filter": 1,
"low_impact": 3,
"filter": filter,
"low_impact": low_impact,
"function": 3,
"column": 3,
"literal": 5,
Expand All @@ -616,8 +618,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
"query_plan_height": query_plan_height,
"query_plan_num_duplicate_nodes": 0,
"query_plan_complexity": {
"filter": 1,
"low_impact": 3,
"filter": filter,
"low_impact": low_impact,
"function": 3,
"column": 3,
"literal": 5,
Expand All @@ -638,8 +640,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
"query_plan_height": query_plan_height,
"query_plan_num_duplicate_nodes": 0,
"query_plan_complexity": {
"filter": 1,
"low_impact": 3,
"filter": filter,
"low_impact": low_impact,
"function": 3,
"column": 3,
"literal": 5,
Expand All @@ -660,8 +662,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
"query_plan_height": query_plan_height,
"query_plan_num_duplicate_nodes": 0,
"query_plan_complexity": {
"filter": 1,
"low_impact": 3,
"filter": filter,
"low_impact": low_impact,
"function": 3,
"column": 3,
"literal": 5,
Expand All @@ -682,8 +684,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
"query_plan_height": query_plan_height,
"query_plan_num_duplicate_nodes": 0,
"query_plan_complexity": {
"filter": 1,
"low_impact": 3,
"filter": filter,
"low_impact": low_impact,
"function": 3,
"column": 3,
"literal": 5,
Expand Down Expand Up @@ -819,14 +821,15 @@ def test_dataframe_stat_functions_api_calls(session):
# check to make sure that the original DF is unchanged
assert df._plan.api_calls == [{"name": "Session.create_dataframe[values]"}]

column = 6 if session.sql_simplifier_enabled else 9
crosstab = df.stat.crosstab("empid", "month")
assert crosstab._plan.api_calls == [
{
"name": "Session.create_dataframe[values]",
"sql_simplifier_enabled": session.sql_simplifier_enabled,
"query_plan_height": 4,
"query_plan_num_duplicate_nodes": 0,
"query_plan_complexity": {"group_by": 1, "column": 6, "literal": 48},
"query_plan_complexity": {"group_by": 1, "column": column, "literal": 48},
},
{
"name": "DataFrameStatFunctions.crosstab",
Expand All @@ -844,7 +847,7 @@ def test_dataframe_stat_functions_api_calls(session):
"sql_simplifier_enabled": session.sql_simplifier_enabled,
"query_plan_height": 4,
"query_plan_num_duplicate_nodes": 0,
"query_plan_complexity": {"group_by": 1, "column": 6, "literal": 48},
"query_plan_complexity": {"group_by": 1, "column": column, "literal": 48},
}
]

Expand Down

0 comments on commit 7e05572

Please sign in to comment.