dbcli/sqlglot
1
0
Fork 0
Code Issues Activity

Merging upstream version 11.7.1.

Browse source 
Signed-off-by: Daniel Baumann <daniel@debian.org>
This commit is contained in:
Daniel Baumann 2025-02-13 15:52:09 +01:00
parent 0c053462ae
commit 8d96084fad
Signed by: daniel
GPG key ID: FBB4F0E80A80222F
144 changed files with 44104 additions and 39367 deletions
Download patch file Download diff file Expand all files Collapse all files

85
tests/dataframe/integration/test_dataframe.py
View file

@ -276,6 +276,7 @@ class TestDataframeFunc(DataFrameValidator):
self.df_spark_employee.store_id,
self.df_spark_store.store_name,
self.df_spark_store["num_sales"],
F.lit("literal_value"),
)
dfs_joined = self.df_sqlglot_employee.join(
self.df_sqlglot_store,
@ -289,6 +290,7 @@ class TestDataframeFunc(DataFrameValidator):
self.df_sqlglot_employee.store_id,
self.df_sqlglot_store.store_name,
self.df_sqlglot_store["num_sales"],
SF.lit("literal_value"),
)
self.compare_spark_with_sqlglot(df_joined, dfs_joined)
@ -330,8 +332,8 @@ class TestDataframeFunc(DataFrameValidator):
def test_join_inner_equality_multiple_bitwise_and(self):
df_joined = self.df_spark_employee.join(
self.df_spark_store,
on=(self.df_spark_employee.store_id == self.df_spark_store.store_id)
& (self.df_spark_employee.age == self.df_spark_store.num_sales),
on=(self.df_spark_store.store_id == self.df_spark_employee.store_id)
& (self.df_spark_store.num_sales == self.df_spark_employee.age),
how="inner",
).select(
self.df_spark_employee.employee_id,
@ -344,8 +346,8 @@ class TestDataframeFunc(DataFrameValidator):
)
dfs_joined = self.df_sqlglot_employee.join(
self.df_sqlglot_store,
on=(self.df_sqlglot_employee.store_id == self.df_sqlglot_store.store_id)
& (self.df_sqlglot_employee.age == self.df_sqlglot_store.num_sales),
on=(self.df_sqlglot_store.store_id == self.df_sqlglot_employee.store_id)
& (self.df_sqlglot_store.num_sales == self.df_sqlglot_employee.age),
how="inner",
).select(
self.df_sqlglot_employee.employee_id,
@ -443,6 +445,81 @@ class TestDataframeFunc(DataFrameValidator):
)
self.compare_spark_with_sqlglot(df, dfs)
def test_triple_join_no_select(self):
df = (
self.df_employee.join(
self.df_store,
on=self.df_employee["employee_id"] == self.df_store["store_id"],
how="left",
)
.join(
self.df_district,
on=self.df_store["store_id"] == self.df_district["district_id"],
how="left",
)
.orderBy(F.col("employee_id"))
)
dfs = (
self.dfs_employee.join(
self.dfs_store,
on=self.dfs_employee["employee_id"] == self.dfs_store["store_id"],
how="left",
)
.join(
self.dfs_district,
on=self.dfs_store["store_id"] == self.dfs_district["district_id"],
how="left",
)
.orderBy(SF.col("employee_id"))
)
self.compare_spark_with_sqlglot(df, dfs)
def test_triple_joins_filter(self):
df = (
self.df_employee.join(
self.df_store,
on=self.df_employee["employee_id"] == self.df_store["store_id"],
how="left",
).join(
self.df_district,
on=self.df_store["store_id"] == self.df_district["district_id"],
how="left",
)
).filter(F.coalesce(self.df_store["num_sales"], F.lit(0)) > 100)
dfs = (
self.dfs_employee.join(
self.dfs_store,
on=self.dfs_employee["employee_id"] == self.dfs_store["store_id"],
how="left",
).join(
self.dfs_district,
on=self.dfs_store["store_id"] == self.dfs_district["district_id"],
how="left",
)
).filter(SF.coalesce(self.dfs_store["num_sales"], SF.lit(0)) > 100)
self.compare_spark_with_sqlglot(df, dfs)
def test_triple_join_column_name_only(self):
df = (
self.df_employee.join(
self.df_store,
on=self.df_employee["employee_id"] == self.df_store["store_id"],
how="left",
)
.join(self.df_district, on="district_id", how="left")
.orderBy(F.col("employee_id"))
)
dfs = (
self.dfs_employee.join(
self.dfs_store,
on=self.dfs_employee["employee_id"] == self.dfs_store["store_id"],
how="left",
)
.join(self.dfs_district, on="district_id", how="left")
.orderBy(SF.col("employee_id"))
)
self.compare_spark_with_sqlglot(df, dfs)
def test_join_select_and_select_start(self):
df = self.df_spark_employee.select(
F.col("fname"), F.col("lname"), F.col("age"), F.col("store_id")

7
tests/dataframe/unit/test_dataframe_writer.py
View file

@ -86,3 +86,10 @@ class TestDataFrameWriter(DataFrameSQLValidator):
"CREATE TABLE table_name AS SELECT `t12441`.`employee_id` AS `employee_id`, `t12441`.`fname` AS `fname`, `t12441`.`lname` AS `lname`, `t12441`.`age` AS `age`, `t12441`.`store_id` AS `store_id` FROM `t12441` AS `t12441`",
]
self.compare_sql(df, expected_statements)
def test_quotes(self):
sqlglot.schema.add_table('"Test"', {'"ID"': "STRING"})
df = self.spark.table('"Test"')
self.compare_sql(
df.select(df['"ID"']), ["SELECT `Test`.`ID` AS `ID` FROM `Test` AS `Test`"]
)

13
tests/dataframe/unit/test_functions.py
View file

@ -807,14 +807,17 @@ class TestFunctions(unittest.TestCase):
self.assertEqual("DATE_ADD(cola, 2)", col.sql())
col_col_for_add = SF.date_add("cola", "colb")
self.assertEqual("DATE_ADD(cola, colb)", col_col_for_add.sql())
current_date_add = SF.date_add(SF.current_date(), 5)
self.assertEqual("DATE_ADD(CURRENT_DATE, 5)", current_date_add.sql())
self.assertEqual("DATEADD(day, 5, CURRENT_DATE)", current_date_add.sql(dialect="snowflake"))
def test_date_sub(self):
col_str = SF.date_sub("cola", 2)
self.assertEqual("DATE_SUB(cola, 2)", col_str.sql())
self.assertEqual("DATE_ADD(cola, -2)", col_str.sql())
col = SF.date_sub(SF.col("cola"), 2)
self.assertEqual("DATE_SUB(cola, 2)", col.sql())
self.assertEqual("DATE_ADD(cola, -2)", col.sql())
col_col_for_add = SF.date_sub("cola", "colb")
self.assertEqual("DATE_SUB(cola, colb)", col_col_for_add.sql())
self.assertEqual("DATE_ADD(cola, colb * -1)", col_col_for_add.sql())
def test_date_diff(self):
col_str = SF.date_diff("cola", "colb")
@ -957,9 +960,9 @@ class TestFunctions(unittest.TestCase):
def test_sha1(self):
col_str = SF.sha1("Spark")
self.assertEqual("SHA1('Spark')", col_str.sql())
self.assertEqual("SHA('Spark')", col_str.sql())
col = SF.sha1(SF.col("cola"))
self.assertEqual("SHA1(cola)", col.sql())
self.assertEqual("SHA(cola)", col.sql())
def test_sha2(self):
col_str = SF.sha2("Spark", 256)