summaryrefslogtreecommitdiffstats
path: root/tests/dataframe
diff options
context:
space:
mode:
Diffstat (limited to 'tests/dataframe')
-rw-r--r--tests/dataframe/integration/dataframe_validator.py6
-rw-r--r--tests/dataframe/unit/test_dataframe_writer.py8
-rw-r--r--tests/dataframe/unit/test_session.py14
3 files changed, 14 insertions, 14 deletions
diff --git a/tests/dataframe/integration/dataframe_validator.py b/tests/dataframe/integration/dataframe_validator.py
index c84a342..22d4982 100644
--- a/tests/dataframe/integration/dataframe_validator.py
+++ b/tests/dataframe/integration/dataframe_validator.py
@@ -135,9 +135,9 @@ class DataFrameValidator(unittest.TestCase):
data=district_data, schema=cls.sqlglot_district_schema
)
cls.df_district.createOrReplaceTempView("district")
- sqlglot.schema.add_table("employee", cls.sqlglot_employee_schema)
- sqlglot.schema.add_table("store", cls.sqlglot_store_schema)
- sqlglot.schema.add_table("district", cls.sqlglot_district_schema)
+ sqlglot.schema.add_table("employee", cls.sqlglot_employee_schema, dialect="spark")
+ sqlglot.schema.add_table("store", cls.sqlglot_store_schema, dialect="spark")
+ sqlglot.schema.add_table("district", cls.sqlglot_district_schema, dialect="spark")
def setUp(self) -> None:
warnings.filterwarnings("ignore", category=ResourceWarning)
diff --git a/tests/dataframe/unit/test_dataframe_writer.py b/tests/dataframe/unit/test_dataframe_writer.py
index 3f45468..303d2f9 100644
--- a/tests/dataframe/unit/test_dataframe_writer.py
+++ b/tests/dataframe/unit/test_dataframe_writer.py
@@ -30,7 +30,7 @@ class TestDataFrameWriter(DataFrameSQLValidator):
@mock.patch("sqlglot.schema", MappingSchema())
def test_insertInto_byName(self):
- sqlglot.schema.add_table("table_name", {"employee_id": "INT"})
+ sqlglot.schema.add_table("table_name", {"employee_id": "INT"}, dialect="spark")
df = self.df_employee.write.byName.insertInto("table_name")
expected = "INSERT INTO table_name SELECT `a1`.`employee_id` AS `employee_id` FROM VALUES (1, 'Jack', 'Shephard', 37, 1), (2, 'John', 'Locke', 65, 1), (3, 'Kate', 'Austen', 37, 2), (4, 'Claire', 'Littleton', 27, 2), (5, 'Hugo', 'Reyes', 29, 100) AS `a1`(`employee_id`, `fname`, `lname`, `age`, `store_id`)"
self.compare_sql(df, expected)
@@ -88,8 +88,8 @@ class TestDataFrameWriter(DataFrameSQLValidator):
self.compare_sql(df, expected_statements)
def test_quotes(self):
- sqlglot.schema.add_table('"Test"', {'"ID"': "STRING"})
- df = self.spark.table('"Test"')
+ sqlglot.schema.add_table("`Test`", {"`ID`": "STRING"}, dialect="spark")
+ df = self.spark.table("`Test`")
self.compare_sql(
- df.select(df['"ID"']), ["SELECT `Test`.`ID` AS `ID` FROM `Test` AS `Test`"]
+ df.select(df["`ID`"]), ["SELECT `test`.`id` AS `id` FROM `test` AS `test`"]
)
diff --git a/tests/dataframe/unit/test_session.py b/tests/dataframe/unit/test_session.py
index 0970a2e..4c275e9 100644
--- a/tests/dataframe/unit/test_session.py
+++ b/tests/dataframe/unit/test_session.py
@@ -71,7 +71,7 @@ class TestDataframeSession(DataFrameSQLValidator):
@mock.patch("sqlglot.schema", MappingSchema())
def test_sql_select_only(self):
query = "SELECT cola, colb FROM table"
- sqlglot.schema.add_table("table", {"cola": "string", "colb": "string"})
+ sqlglot.schema.add_table("table", {"cola": "string", "colb": "string"}, dialect="spark")
df = self.spark.sql(query)
self.assertEqual(
"SELECT `table`.`cola` AS `cola`, `table`.`colb` AS `colb` FROM `table` AS `table`",
@@ -80,17 +80,17 @@ class TestDataframeSession(DataFrameSQLValidator):
@mock.patch("sqlglot.schema", MappingSchema())
def test_select_quoted(self):
- sqlglot.schema.add_table('"TEST"', {"name": "string"})
+ sqlglot.schema.add_table("`TEST`", {"name": "string"}, dialect="spark")
self.assertEqual(
- SparkSession().table('"TEST"').select(F.col("name")).sql(dialect="snowflake")[0],
- '''SELECT "TEST"."name" AS "name" FROM "TEST" AS "TEST"''',
+ SparkSession().table("`TEST`").select(F.col("name")).sql(dialect="snowflake")[0],
+ '''SELECT "test"."name" AS "name" FROM "test" AS "test"''',
)
@mock.patch("sqlglot.schema", MappingSchema())
def test_sql_with_aggs(self):
query = "SELECT cola, colb FROM table"
- sqlglot.schema.add_table("table", {"cola": "string", "colb": "string"})
+ sqlglot.schema.add_table("table", {"cola": "string", "colb": "string"}, dialect="spark")
df = self.spark.sql(query).groupBy(F.col("cola")).agg(F.sum("colb"))
self.assertEqual(
"WITH t38189 AS (SELECT cola, colb FROM table), t42330 AS (SELECT cola, colb FROM t38189) SELECT cola, SUM(colb) FROM t42330 GROUP BY cola",
@@ -100,7 +100,7 @@ class TestDataframeSession(DataFrameSQLValidator):
@mock.patch("sqlglot.schema", MappingSchema())
def test_sql_create(self):
query = "CREATE TABLE new_table AS WITH t1 AS (SELECT cola, colb FROM table) SELECT cola, colb, FROM t1"
- sqlglot.schema.add_table("table", {"cola": "string", "colb": "string"})
+ sqlglot.schema.add_table("table", {"cola": "string", "colb": "string"}, dialect="spark")
df = self.spark.sql(query)
expected = "CREATE TABLE new_table AS SELECT `table`.`cola` AS `cola`, `table`.`colb` AS `colb` FROM `table` AS `table`"
self.compare_sql(df, expected)
@@ -108,7 +108,7 @@ class TestDataframeSession(DataFrameSQLValidator):
@mock.patch("sqlglot.schema", MappingSchema())
def test_sql_insert(self):
query = "WITH t1 AS (SELECT cola, colb FROM table) INSERT INTO new_table SELECT cola, colb FROM t1"
- sqlglot.schema.add_table("table", {"cola": "string", "colb": "string"})
+ sqlglot.schema.add_table("table", {"cola": "string", "colb": "string"}, dialect="spark")
df = self.spark.sql(query)
expected = "INSERT INTO new_table SELECT `table`.`cola` AS `cola`, `table`.`colb` AS `colb` FROM `table` AS `table`"
self.compare_sql(df, expected)