Adding upstream version 9.0.1.upstream/9.0.1

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 28 insertions, 0 deletions

diff --git a/tests/dataframe/integration/test_session.py b/tests/dataframe/integration/test_session.py
new file mode 100644
index 0000000..ff1477b
--- /dev/null
+++ b/tests/dataframe/integration/test_session.py
@@ -0,0 +1,28 @@
+from pyspark.sql import functions as F
+
+from sqlglot.dataframe.sql import functions as SF
+from tests.dataframe.integration.dataframe_validator import DataFrameValidator
+
+
+class TestSessionFunc(DataFrameValidator):
+    def test_sql_simple_select(self):
+        query = "SELECT fname, lname FROM employee"
+        df = self.spark.sql(query)
+        dfs = self.sqlglot.sql(query)
+        self.compare_spark_with_sqlglot(df, dfs)
+
+    def test_sql_with_join(self):
+        query = """
+        SELECT
+            e.employee_id
+            , s.store_id    
+        FROM
+            employee e
+            INNER JOIN
+            store s
+            ON
+                e.store_id = s.store_id
+        """
+        df = self.spark.sql(query).groupBy(F.col("store_id")).agg(F.countDistinct(F.col("employee_id")))
+        dfs = self.sqlglot.sql(query).groupBy(SF.col("store_id")).agg(SF.countDistinct(SF.col("employee_id")))
+        self.compare_spark_with_sqlglot(df, dfs, skip_schema_compare=True)