diff options
Diffstat (limited to 'tests/dataframe/integration/dataframe_validator.py')
-rw-r--r-- | tests/dataframe/integration/dataframe_validator.py | 149 |
1 files changed, 149 insertions, 0 deletions
diff --git a/tests/dataframe/integration/dataframe_validator.py b/tests/dataframe/integration/dataframe_validator.py new file mode 100644 index 0000000..6c4642f --- /dev/null +++ b/tests/dataframe/integration/dataframe_validator.py @@ -0,0 +1,149 @@ +import typing as t +import unittest +import warnings + +import sqlglot +from tests.helpers import SKIP_INTEGRATION + +if t.TYPE_CHECKING: + from pyspark.sql import DataFrame as SparkDataFrame + + +@unittest.skipIf(SKIP_INTEGRATION, "Skipping Integration Tests since `SKIP_INTEGRATION` is set") +class DataFrameValidator(unittest.TestCase): + spark = None + sqlglot = None + df_employee = None + df_store = None + df_district = None + spark_employee_schema = None + sqlglot_employee_schema = None + spark_store_schema = None + sqlglot_store_schema = None + spark_district_schema = None + sqlglot_district_schema = None + + @classmethod + def setUpClass(cls): + from pyspark import SparkConf + from pyspark.sql import SparkSession, types + + from sqlglot.dataframe.sql import types as sqlglotSparkTypes + from sqlglot.dataframe.sql.session import SparkSession as SqlglotSparkSession + + # This is for test `test_branching_root_dataframes` + config = SparkConf().setAll([("spark.sql.analyzer.failAmbiguousSelfJoin", "false")]) + cls.spark = SparkSession.builder.master("local[*]").appName("Unit-tests").config(conf=config).getOrCreate() + cls.spark.sparkContext.setLogLevel("ERROR") + cls.sqlglot = SqlglotSparkSession() + cls.spark_employee_schema = types.StructType( + [ + types.StructField("employee_id", types.IntegerType(), False), + types.StructField("fname", types.StringType(), False), + types.StructField("lname", types.StringType(), False), + types.StructField("age", types.IntegerType(), False), + types.StructField("store_id", types.IntegerType(), False), + ] + ) + cls.sqlglot_employee_schema = sqlglotSparkTypes.StructType( + [ + sqlglotSparkTypes.StructField("employee_id", sqlglotSparkTypes.IntegerType(), False), + sqlglotSparkTypes.StructField("fname", sqlglotSparkTypes.StringType(), False), + sqlglotSparkTypes.StructField("lname", sqlglotSparkTypes.StringType(), False), + sqlglotSparkTypes.StructField("age", sqlglotSparkTypes.IntegerType(), False), + sqlglotSparkTypes.StructField("store_id", sqlglotSparkTypes.IntegerType(), False), + ] + ) + employee_data = [ + (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), + ] + cls.df_employee = cls.spark.createDataFrame(data=employee_data, schema=cls.spark_employee_schema) + cls.dfs_employee = cls.sqlglot.createDataFrame(data=employee_data, schema=cls.sqlglot_employee_schema) + cls.df_employee.createOrReplaceTempView("employee") + + cls.spark_store_schema = types.StructType( + [ + types.StructField("store_id", types.IntegerType(), False), + types.StructField("store_name", types.StringType(), False), + types.StructField("district_id", types.IntegerType(), False), + types.StructField("num_sales", types.IntegerType(), False), + ] + ) + cls.sqlglot_store_schema = sqlglotSparkTypes.StructType( + [ + sqlglotSparkTypes.StructField("store_id", sqlglotSparkTypes.IntegerType(), False), + sqlglotSparkTypes.StructField("store_name", sqlglotSparkTypes.StringType(), False), + sqlglotSparkTypes.StructField("district_id", sqlglotSparkTypes.IntegerType(), False), + sqlglotSparkTypes.StructField("num_sales", sqlglotSparkTypes.IntegerType(), False), + ] + ) + store_data = [ + (1, "Hydra", 1, 37), + (2, "Arrow", 2, 2000), + ] + cls.df_store = cls.spark.createDataFrame(data=store_data, schema=cls.spark_store_schema) + cls.dfs_store = cls.sqlglot.createDataFrame(data=store_data, schema=cls.sqlglot_store_schema) + cls.df_store.createOrReplaceTempView("store") + + cls.spark_district_schema = types.StructType( + [ + types.StructField("district_id", types.IntegerType(), False), + types.StructField("district_name", types.StringType(), False), + types.StructField("manager_name", types.StringType(), False), + ] + ) + cls.sqlglot_district_schema = sqlglotSparkTypes.StructType( + [ + sqlglotSparkTypes.StructField("district_id", sqlglotSparkTypes.IntegerType(), False), + sqlglotSparkTypes.StructField("district_name", sqlglotSparkTypes.StringType(), False), + sqlglotSparkTypes.StructField("manager_name", sqlglotSparkTypes.StringType(), False), + ] + ) + district_data = [ + (1, "Temple", "Dogen"), + (2, "Lighthouse", "Jacob"), + ] + cls.df_district = cls.spark.createDataFrame(data=district_data, schema=cls.spark_district_schema) + cls.dfs_district = cls.sqlglot.createDataFrame(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) + + def setUp(self) -> None: + warnings.filterwarnings("ignore", category=ResourceWarning) + self.df_spark_store = self.df_store.alias("df_store") # type: ignore + self.df_spark_employee = self.df_employee.alias("df_employee") # type: ignore + self.df_spark_district = self.df_district.alias("df_district") # type: ignore + self.df_sqlglot_store = self.dfs_store.alias("store") # type: ignore + self.df_sqlglot_employee = self.dfs_employee.alias("employee") # type: ignore + self.df_sqlglot_district = self.dfs_district.alias("district") # type: ignore + + def compare_spark_with_sqlglot( + self, df_spark, df_sqlglot, no_empty=True, skip_schema_compare=False + ) -> t.Tuple["SparkDataFrame", "SparkDataFrame"]: + def compare_schemas(schema_1, schema_2): + for schema in [schema_1, schema_2]: + for struct_field in schema.fields: + struct_field.metadata = {} + self.assertEqual(schema_1, schema_2) + + for statement in df_sqlglot.sql(): + actual_df_sqlglot = self.spark.sql(statement) # type: ignore + df_sqlglot_results = actual_df_sqlglot.collect() + df_spark_results = df_spark.collect() + if not skip_schema_compare: + compare_schemas(df_spark.schema, actual_df_sqlglot.schema) + self.assertEqual(df_spark_results, df_sqlglot_results) + if no_empty: + self.assertNotEqual(len(df_spark_results), 0) + self.assertNotEqual(len(df_sqlglot_results), 0) + return df_spark, actual_df_sqlglot + + @classmethod + def get_explain_plan(cls, df: "SparkDataFrame", mode: str = "extended") -> str: + return df._sc._jvm.PythonSQLUtils.explainString(df._jdf.queryExecution(), mode) # type: ignore |