summaryrefslogtreecommitdiffstats
path: root/tests/test_executor.py
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2022-09-15 16:46:17 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2022-09-15 16:46:17 +0000
commit28cc22419e32a65fea2d1678400265b8cabc3aff (patch)
treeff9ac1991fd48490b21ef6aa9015a347a165e2d9 /tests/test_executor.py
parentInitial commit. (diff)
downloadsqlglot-28cc22419e32a65fea2d1678400265b8cabc3aff.tar.xz
sqlglot-28cc22419e32a65fea2d1678400265b8cabc3aff.zip
Adding upstream version 6.0.4.upstream/6.0.4
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'tests/test_executor.py')
-rw-r--r--tests/test_executor.py72
1 files changed, 72 insertions, 0 deletions
diff --git a/tests/test_executor.py b/tests/test_executor.py
new file mode 100644
index 0000000..9afa225
--- /dev/null
+++ b/tests/test_executor.py
@@ -0,0 +1,72 @@
+import unittest
+
+import duckdb
+import pandas as pd
+from pandas.testing import assert_frame_equal
+
+from sqlglot import exp, parse_one
+from sqlglot.executor import execute
+from sqlglot.executor.python import Python
+from tests.helpers import FIXTURES_DIR, TPCH_SCHEMA, load_sql_fixture_pairs
+
+DIR = FIXTURES_DIR + "/optimizer/tpc-h/"
+
+
+class TestExecutor(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ cls.conn = duckdb.connect()
+
+ for table in TPCH_SCHEMA:
+ cls.conn.execute(
+ f"""
+ CREATE VIEW {table} AS
+ SELECT *
+ FROM READ_CSV_AUTO('{DIR}{table}.csv.gz')
+ """
+ )
+
+ cls.cache = {}
+ cls.sqls = [
+ (sql, expected)
+ for _, sql, expected in load_sql_fixture_pairs("optimizer/tpc-h/tpc-h.sql")
+ ]
+
+ @classmethod
+ def tearDownClass(cls):
+ cls.conn.close()
+
+ def cached_execute(self, sql):
+ if sql not in self.cache:
+ self.cache[sql] = self.conn.execute(sql).fetchdf()
+ return self.cache[sql]
+
+ def rename_anonymous(self, source, target):
+ for i, column in enumerate(source.columns):
+ if "_col_" in column:
+ source.rename(columns={column: target.columns[i]}, inplace=True)
+
+ def test_py_dialect(self):
+ self.assertEqual(Python().generate(parse_one("'x '''")), r"'x \''")
+
+ def test_optimized_tpch(self):
+ for sql, optimized in self.sqls[0:20]:
+ a = self.cached_execute(sql)
+ b = self.conn.execute(optimized).fetchdf()
+ self.rename_anonymous(b, a)
+ assert_frame_equal(a, b)
+
+ def test_execute_tpch(self):
+ def to_csv(expression):
+ if isinstance(expression, exp.Table):
+ return parse_one(
+ f"READ_CSV('{DIR}{expression.name}.csv.gz', 'delimiter', '|') AS {expression.name}"
+ )
+ return expression
+
+ for sql, _ in self.sqls[0:3]:
+ a = self.cached_execute(sql)
+ sql = parse_one(sql).transform(to_csv).sql(pretty=True)
+ table = execute(sql, TPCH_SCHEMA)
+ b = pd.DataFrame(table.rows, columns=table.columns)
+ assert_frame_equal(a, b, check_dtype=False)