diff options
Diffstat (limited to 'sqlglot/executor')
-rw-r--r-- | sqlglot/executor/__init__.py | 61 | ||||
-rw-r--r-- | sqlglot/executor/env.py | 1 | ||||
-rw-r--r-- | sqlglot/executor/table.py | 7 |
3 files changed, 51 insertions, 18 deletions
diff --git a/sqlglot/executor/__init__.py b/sqlglot/executor/__init__.py index 04621b5..67b4b00 100644 --- a/sqlglot/executor/__init__.py +++ b/sqlglot/executor/__init__.py @@ -1,5 +1,13 @@ +""" +.. include:: ../../posts/python_sql_engine.md +---- +""" + +from __future__ import annotations + import logging import time +import typing as t from sqlglot import maybe_parse from sqlglot.errors import ExecuteError @@ -11,42 +19,63 @@ from sqlglot.schema import ensure_schema logger = logging.getLogger("sqlglot") +if t.TYPE_CHECKING: + from sqlglot.dialects.dialect import DialectType + from sqlglot.executor.table import Tables + from sqlglot.expressions import Expression + from sqlglot.schema import Schema -def execute(sql, schema=None, read=None, tables=None): + +def execute( + sql: str | Expression, + schema: t.Optional[t.Dict | Schema] = None, + read: DialectType = None, + tables: t.Optional[t.Dict] = None, +) -> Table: """ Run a sql query against data. Args: - sql (str|sqlglot.Expression): a sql statement - schema (dict|sqlglot.optimizer.Schema): database schema. - This can either be an instance of `sqlglot.optimizer.Schema` or a mapping in one of - the following forms: - 1. {table: {col: type}} - 2. {db: {table: {col: type}}} - 3. {catalog: {db: {table: {col: type}}}} - read (str): the SQL dialect to apply during parsing - (eg. "spark", "hive", "presto", "mysql"). - tables (dict): additional tables to register. + sql: a sql statement. + schema: database schema. + This can either be an instance of `Schema` or a mapping in one of the following forms: + 1. {table: {col: type}} + 2. {db: {table: {col: type}}} + 3. {catalog: {db: {table: {col: type}}}} + read: the SQL dialect to apply during parsing (eg. "spark", "hive", "presto", "mysql"). + tables: additional tables to register. + Returns: - sqlglot.executor.Table: Simple columnar data structure. + Simple columnar data structure. """ - tables = ensure_tables(tables) + tables_ = ensure_tables(tables) + if not schema: schema = { name: {column: type(table[0][column]).__name__ for column in table.columns} - for name, table in tables.mapping.items() + for name, table in tables_.mapping.items() } + schema = ensure_schema(schema) - if tables.supported_table_args and tables.supported_table_args != schema.supported_table_args: + + if tables_.supported_table_args and tables_.supported_table_args != schema.supported_table_args: raise ExecuteError("Tables must support the same table args as schema") + expression = maybe_parse(sql, dialect=read) + now = time.time() expression = optimize(expression, schema, leave_tables_isolated=True) + logger.debug("Optimization finished: %f", time.time() - now) logger.debug("Optimized SQL: %s", expression.sql(pretty=True)) + plan = Plan(expression) + logger.debug("Logical Plan: %s", plan) + now = time.time() - result = PythonExecutor(tables=tables).execute(plan) + result = PythonExecutor(tables=tables_).execute(plan) + logger.debug("Query finished: %f", time.time() - now) + return result diff --git a/sqlglot/executor/env.py b/sqlglot/executor/env.py index 04dc938..ba9cbbd 100644 --- a/sqlglot/executor/env.py +++ b/sqlglot/executor/env.py @@ -171,5 +171,6 @@ ENV = { "STRPOSITION": str_position, "SUB": null_if_any(lambda e, this: e - this), "SUBSTRING": substring, + "TIMESTRTOTIME": null_if_any(lambda arg: datetime.datetime.fromisoformat(arg)), "UPPER": null_if_any(lambda arg: arg.upper()), } diff --git a/sqlglot/executor/table.py b/sqlglot/executor/table.py index f1b5b54..27e3e5e 100644 --- a/sqlglot/executor/table.py +++ b/sqlglot/executor/table.py @@ -1,5 +1,7 @@ from __future__ import annotations +import typing as t + from sqlglot.helper import dict_depth from sqlglot.schema import AbstractMappingSchema @@ -106,11 +108,11 @@ class Tables(AbstractMappingSchema[Table]): pass -def ensure_tables(d: dict | None) -> Tables: +def ensure_tables(d: t.Optional[t.Dict]) -> Tables: return Tables(_ensure_tables(d)) -def _ensure_tables(d: dict | None) -> dict: +def _ensure_tables(d: t.Optional[t.Dict]) -> t.Dict: if not d: return {} @@ -127,4 +129,5 @@ def _ensure_tables(d: dict | None) -> dict: columns = tuple(table[0]) if table else () rows = [tuple(row[c] for c in columns) for row in table] result[name] = Table(columns=columns, rows=rows) + return result |