e6data-spark-compatibility 1.0.0__py3-none-any.whl

e6_spark_compat/__init__.py

@@ -0,0 +1,50 @@
+"""
+e6-spark-compat: PySpark and Sedona compatibility layer for e6data
+"""
+
+# Fix multiprocessing issue on macOS/Windows before any imports
+import platform
+import multiprocessing
+import os
+
+# Only apply fix if:
+# 1. We're on macOS (Darwin) or Windows
+# 2. Not explicitly disabled via environment variable
+# 3. The current method is 'spawn' (problematic)
+if os.environ.get('E6_SPARK_COMPAT_NO_MP_FIX') != '1':
+    try:
+        current_method = multiprocessing.get_start_method(allow_none=True)
+        system = platform.system()
+        
+        # macOS defaults to 'spawn' which causes re-imports
+        # current_method can be None initially, which defaults to 'spawn' on macOS
+        if system == 'Darwin' and (current_method == 'spawn' or current_method is None):
+            multiprocessing.set_start_method('fork', force=True)
+        # Windows only supports 'spawn', so we can't fix it this way
+        elif system == 'Windows' and current_method == 'spawn':
+            # Log a warning but continue
+            import warnings
+            warnings.warn(
+                "e6-spark-compat: Windows detected. Scripts may execute multiple times. "
+                "Use 'if __name__ == \"__main__\":' guard to prevent issues.",
+                RuntimeWarning
+            )
+    except (RuntimeError, ValueError):
+        # Already set or platform doesn't support the change
+        pass
+
+# Initialize vendored dependencies first (patches sys.modules for sqlglot)
+from . import _vendor
+
+# Initialize logging configuration before importing other modules
+from .core.logging_config import configure_logging
+configure_logging()
+
+from .core.session import SparkSession
+from .sql.dataframe import DataFrame
+from .sql.column import Column
+from .sql.row import Row
+from .core.sql_generator import set_use_cte
+
+__version__ = "1.0.0"
+__all__ = ["SparkSession", "DataFrame", "Column", "Row", "set_use_cte"]

e6_spark_compat/_vendor/__init__.py

@@ -0,0 +1,35 @@
+"""
+Vendored dependencies for e6-spark-compat.
+
+This directory contains vendored (bundled) copies of dependencies
+to avoid external package requirements for PyPI publishing.
+"""
+import sys
+import os
+import importlib.util
+
+# Pre-register sqlglot in sys.modules BEFORE importing it
+# This is necessary because the vendored sqlglot uses internal imports like 'from sqlglot import ...'
+def _setup_vendored_sqlglot():
+    """Setup vendored sqlglot by pre-registering it in sys.modules."""
+    vendor_dir = os.path.dirname(__file__)
+    sqlglot_dir = os.path.join(vendor_dir, 'sqlglot')
+
+    # Create a module spec for the vendored sqlglot
+    spec = importlib.util.spec_from_file_location(
+        'sqlglot',
+        os.path.join(sqlglot_dir, '__init__.py'),
+        submodule_search_locations=[sqlglot_dir]
+    )
+
+    # Create the module and add to sys.modules BEFORE executing it
+    sqlglot_module = importlib.util.module_from_spec(spec)
+    sys.modules['sqlglot'] = sqlglot_module
+
+    # Now execute the module (this allows internal imports to work)
+    spec.loader.exec_module(sqlglot_module)
+
+    # Also make it available as e6_spark_compat._vendor.sqlglot
+    sys.modules['e6_spark_compat._vendor.sqlglot'] = sqlglot_module
+
+_setup_vendored_sqlglot()

e6_spark_compat/_vendor/sqlglot/__init__.py

@@ -0,0 +1,180 @@
+# ruff: noqa: F401
+"""
+.. include:: ../README.md
+
+----
+"""
+
+from __future__ import annotations
+
+import logging
+import typing as t
+
+from sqlglot import expressions as exp
+from sqlglot.dialects.dialect import Dialect as Dialect, Dialects as Dialects
+from sqlglot.diff import diff as diff
+from sqlglot.errors import (
+    ErrorLevel as ErrorLevel,
+    ParseError as ParseError,
+    TokenError as TokenError,
+    UnsupportedError as UnsupportedError,
+)
+from sqlglot.expressions import (
+    Expression as Expression,
+    alias_ as alias,
+    and_ as and_,
+    case as case,
+    cast as cast,
+    column as column,
+    condition as condition,
+    delete as delete,
+    except_ as except_,
+    from_ as from_,
+    func as func,
+    insert as insert,
+    intersect as intersect,
+    maybe_parse as maybe_parse,
+    merge as merge,
+    not_ as not_,
+    or_ as or_,
+    select as select,
+    subquery as subquery,
+    table_ as table,
+    to_column as to_column,
+    to_identifier as to_identifier,
+    to_table as to_table,
+    union as union,
+)
+from sqlglot.generator import Generator as Generator
+from sqlglot.parser import Parser as Parser
+from sqlglot.schema import MappingSchema as MappingSchema, Schema as Schema
+from sqlglot.tokens import (
+    Token as Token,
+    Tokenizer as Tokenizer,
+    TokenType as TokenType,
+)
+
+if t.TYPE_CHECKING:
+    from sqlglot._typing import E
+    from sqlglot.dialects.dialect import DialectType as DialectType
+
+logger = logging.getLogger("sqlglot")
+
+
+try:
+    from sqlglot._version import __version__, __version_tuple__
+except ImportError:
+    logger.error(
+        "Unable to set __version__, run `pip install -e .` or `python setup.py develop` first."
+    )
+
+
+pretty = False
+"""Whether to format generated SQL by default."""
+
+
+def tokenize(sql: str, read: DialectType = None, dialect: DialectType = None) -> t.List[Token]:
+    """
+    Tokenizes the given SQL string.
+
+    Args:
+        sql: the SQL code string to tokenize.
+        read: the SQL dialect to apply during tokenizing (eg. "spark", "hive", "presto", "mysql").
+        dialect: the SQL dialect (alias for read).
+
+    Returns:
+        The resulting list of tokens.
+    """
+    return Dialect.get_or_raise(read or dialect).tokenize(sql)
+
+
+def parse(
+    sql: str, read: DialectType = None, dialect: DialectType = None, **opts
+) -> t.List[t.Optional[Expression]]:
+    """
+    Parses the given SQL string into a collection of syntax trees, one per parsed SQL statement.
+
+    Args:
+        sql: the SQL code string to parse.
+        read: the SQL dialect to apply during parsing (eg. "spark", "hive", "presto", "mysql").
+        dialect: the SQL dialect (alias for read).
+        **opts: other `sqlglot.parser.Parser` options.
+
+    Returns:
+        The resulting syntax tree collection.
+    """
+    return Dialect.get_or_raise(read or dialect).parse(sql, **opts)
+
+
+@t.overload
+def parse_one(sql: str, *, into: t.Type[E], **opts) -> E: ...
+
+
+@t.overload
+def parse_one(sql: str, **opts) -> Expression: ...
+
+
+def parse_one(
+    sql: str,
+    read: DialectType = None,
+    dialect: DialectType = None,
+    into: t.Optional[exp.IntoType] = None,
+    **opts,
+) -> Expression:
+    """
+    Parses the given SQL string and returns a syntax tree for the first parsed SQL statement.
+
+    Args:
+        sql: the SQL code string to parse.
+        read: the SQL dialect to apply during parsing (eg. "spark", "hive", "presto", "mysql").
+        dialect: the SQL dialect (alias for read)
+        into: the SQLGlot Expression to parse into.
+        **opts: other `sqlglot.parser.Parser` options.
+
+    Returns:
+        The syntax tree for the first parsed statement.
+    """
+
+    dialect = Dialect.get_or_raise(read or dialect)
+    if into:
+        result = dialect.parse_into(into, sql, **opts)
+    else:
+        result = dialect.parse(sql, **opts)
+    for expression in result:
+        if not expression:
+            raise ParseError(f"No expression was parsed from '{sql}'")
+        return expression
+    else:
+        raise ParseError(f"No expression was parsed from '{sql}'")
+
+
+def transpile(
+    sql: str,
+    read: DialectType = None,
+    write: DialectType = None,
+    identity: bool = True,
+    error_level: t.Optional[ErrorLevel] = None,
+    **opts,
+) -> t.List[str]:
+    """
+    Parses the given SQL string in accordance with the source dialect and returns a list of SQL strings transformed
+    to conform to the target dialect. Each string in the returned list represents a single transformed SQL statement.
+
+    Args:
+        sql: the SQL code string to transpile.
+        read: the source dialect used to parse the input string (eg. "spark", "hive", "presto", "mysql").
+        write: the target dialect into which the input should be transformed (eg. "spark", "hive", "presto", "mysql").
+        identity: if set to `True` and if the target dialect is not specified the source dialect will be used as both:
+            the source and the target dialect.
+        error_level: the desired error level of the parser.
+        **opts: other `sqlglot.generator.Generator` options.
+
+    Returns:
+        The list of transpiled SQL statements.
+    """
+    write = (read if write is None else write) if identity else write
+    write = Dialect.get_or_raise(write)
+    return [
+        write.generate(expression, copy=False, **opts) if expression else ""
+        for expression in parse(sql, read, error_level=error_level)
+    ]

e6_spark_compat/_vendor/sqlglot/__main__.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import argparse
+import sys
+import typing as t
+
+import sqlglot
+
+parser = argparse.ArgumentParser(description="Transpile SQL")
+parser.add_argument(
+    "sql",
+    metavar="sql",
+    type=str,
+    help="SQL statement(s) to transpile, or - to parse stdin.",
+)
+parser.add_argument(
+    "--read",
+    dest="read",
+    type=str,
+    default=None,
+    help="Dialect to read default is generic",
+)
+parser.add_argument(
+    "--write",
+    dest="write",
+    type=str,
+    default=None,
+    help="Dialect to write default is generic",
+)
+parser.add_argument(
+    "--no-identify",
+    dest="identify",
+    action="store_false",
+    help="Don't auto identify fields",
+)
+parser.add_argument(
+    "--no-pretty",
+    dest="pretty",
+    action="store_false",
+    help="Compress sql",
+)
+parser.add_argument(
+    "--parse",
+    dest="parse",
+    action="store_true",
+    help="Parse and return the expression tree",
+)
+parser.add_argument(
+    "--tokenize",
+    dest="tokenize",
+    action="store_true",
+    help="Tokenize and return the tokens list",
+)
+parser.add_argument(
+    "--error-level",
+    dest="error_level",
+    type=str,
+    default="IMMEDIATE",
+    help="IGNORE, WARN, RAISE, IMMEDIATE (default)",
+)
+parser.add_argument(
+    "--version",
+    action="version",
+    version=sqlglot.__version__,
+    help="Display the SQLGlot version",
+)
+
+
+args = parser.parse_args()
+error_level = sqlglot.ErrorLevel[args.error_level.upper()]
+
+sql = sys.stdin.read() if args.sql == "-" else args.sql
+
+if args.parse:
+    objs: t.Union[t.List[str], t.List[sqlglot.tokens.Token]] = [
+        repr(expression)
+        for expression in sqlglot.parse(
+            sql,
+            read=args.read,
+            error_level=error_level,
+        )
+    ]
+elif args.tokenize:
+    objs = sqlglot.Dialect.get_or_raise(args.read).tokenize(sql)
+else:
+    objs = sqlglot.transpile(
+        sql,
+        read=args.read,
+        write=args.write,
+        identify=args.identify,
+        pretty=args.pretty,
+        error_level=error_level,
+    )
+
+for obj in objs:
+    print(obj)

e6_spark_compat/_vendor/sqlglot/_typing.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import typing as t
+
+import sqlglot
+
+if t.TYPE_CHECKING:
+    from typing_extensions import Literal as Lit  # noqa
+
+# A little hack for backwards compatibility with Python 3.7.
+# For example, we might want a TypeVar for objects that support comparison e.g. SupportsRichComparisonT from typeshed.
+# But Python 3.7 doesn't support Protocols, so we'd also need typing_extensions, which we don't want as a dependency.
+A = t.TypeVar("A", bound=t.Any)
+B = t.TypeVar("B", bound="sqlglot.exp.Binary")
+E = t.TypeVar("E", bound="sqlglot.exp.Expression")
+F = t.TypeVar("F", bound="sqlglot.exp.Func")
+T = t.TypeVar("T")

e6_spark_compat/_vendor/sqlglot/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.0.1.dev6817'
+__version_tuple__ = version_tuple = (0, 0, 1, 'dev6817')
+
+__commit_id__ = commit_id = 'g8a0cb2080'

e6_spark_compat/_vendor/sqlglot/dialects/__init__.py

@@ -0,0 +1,126 @@
+# ruff: noqa: F401
+"""
+## Dialects
+
+While there is a SQL standard, most SQL engines support a variation of that standard. This makes it difficult
+to write portable SQL code. SQLGlot bridges all the different variations, called "dialects", with an extensible
+SQL transpilation framework.
+
+The base `sqlglot.dialects.dialect.Dialect` class implements a generic dialect that aims to be as universal as possible.
+
+Each SQL variation has its own `Dialect` subclass, extending the corresponding `Tokenizer`, `Parser` and `Generator`
+classes as needed.
+
+### Implementing a custom Dialect
+
+Creating a new SQL dialect may seem complicated at first, but it is actually quite simple in SQLGlot:
+
+```python
+from sqlglot import exp
+from sqlglot.dialects.dialect import Dialect
+from sqlglot.generator import Generator
+from sqlglot.tokens import Tokenizer, TokenType
+
+
+class Custom(Dialect):
+    class Tokenizer(Tokenizer):
+        QUOTES = ["'", '"']  # Strings can be delimited by either single or double quotes
+        IDENTIFIERS = ["`"]  # Identifiers can be delimited by backticks
+
+        # Associates certain meaningful words with tokens that capture their intent
+        KEYWORDS = {
+            **Tokenizer.KEYWORDS,
+            "INT64": TokenType.BIGINT,
+            "FLOAT64": TokenType.DOUBLE,
+        }
+
+    class Generator(Generator):
+        # Specifies how AST nodes, i.e. subclasses of exp.Expression, should be converted into SQL
+        TRANSFORMS = {
+            exp.Array: lambda self, e: f"[{self.expressions(e)}]",
+        }
+
+        # Specifies how AST nodes representing data types should be converted into SQL
+        TYPE_MAPPING = {
+            exp.DataType.Type.TINYINT: "INT64",
+            exp.DataType.Type.SMALLINT: "INT64",
+            exp.DataType.Type.INT: "INT64",
+            exp.DataType.Type.BIGINT: "INT64",
+            exp.DataType.Type.DECIMAL: "NUMERIC",
+            exp.DataType.Type.FLOAT: "FLOAT64",
+            exp.DataType.Type.DOUBLE: "FLOAT64",
+            exp.DataType.Type.BOOLEAN: "BOOL",
+            exp.DataType.Type.TEXT: "STRING",
+        }
+```
+
+The above example demonstrates how certain parts of the base `Dialect` class can be overridden to match a different
+specification. Even though it is a fairly realistic starting point, we strongly encourage the reader to study existing
+dialect implementations in order to understand how their various components can be modified, depending on the use-case.
+
+----
+"""
+
+import importlib
+import threading
+
+DIALECTS = [
+    "Athena",
+    "BigQuery",
+    "ClickHouse",
+    "Databricks",
+    "Doris",
+    "Dremio",
+    "Drill",
+    "Druid",
+    "DuckDB",
+    "Dune",
+    "Fabric",
+    "Hive",
+    "Materialize",
+    "MySQL",
+    "Oracle",
+    "Postgres",
+    "Presto",
+    "PRQL",
+    "Redshift",
+    "RisingWave",
+    "SingleStore",
+    "Snowflake",
+    "Spark",
+    "Spark2",
+    "SQLite",
+    "StarRocks",
+    "Tableau",
+    "Teradata",
+    "Trino",
+    "TSQL",
+    "Exasol",
+    "E6",
+]
+
+MODULE_BY_DIALECT = {name: name.lower() for name in DIALECTS}
+DIALECT_MODULE_NAMES = MODULE_BY_DIALECT.values()
+
+MODULE_BY_ATTRIBUTE = {
+    **MODULE_BY_DIALECT,
+    "Dialect": "dialect",
+    "Dialects": "dialect",
+}
+
+__all__ = list(MODULE_BY_ATTRIBUTE)
+
+# We use a reentrant lock because a dialect may depend on (i.e., import) other dialects.
+# Without it, the first dialect import would never be completed, because subsequent
+# imports would be blocked on the lock held by the first import.
+_import_lock = threading.RLock()
+
+
+def __getattr__(name):
+    module_name = MODULE_BY_ATTRIBUTE.get(name)
+    if module_name:
+        with _import_lock:
+            module = importlib.import_module(f"sqlglot.dialects.{module_name}")
+        return getattr(module, name)
+
+    raise AttributeError(f"module {__name__} has no attribute {name}")