Skip to main content

PySpark and Sedona compatibility layer for e6data

Project description

e6-spark-compat

A PySpark and Sedona compatibility layer for e6data, providing seamless migration from PySpark to e6data with minimal code changes.

Overview

e6-spark-compat is a production-ready compatibility library that translates PySpark and Apache Sedona operations into optimized e6data SQL queries. The library implements lazy evaluation, professional SQL generation using SQLGlot, and comprehensive PySpark API compatibility.

Key Features

  • Lazy Evaluation: Build query plans that execute only when actions are called
  • Professional SQL Generation: SQLGlot-based AST optimization for high-performance queries
  • Complete PySpark API: DataFrame operations, SQL functions, window functions, and aggregations
  • Spatial Operations: Full Apache Sedona compatibility with ST_* functions
  • Window Functions: Complete Window specification API for advanced analytics
  • Multiple File Formats: Parquet, ORC, CSV, JSON, and GeoParquet support
  • Dual-Mode Test Suite: 1270 tests run offline (SQL validation) and live (cluster execution)
  • CI/CD: GitHub Actions with offline (automatic) and live (manual) test runs

What this is (and isn't)

This is a PySpark-shaped façade over the e6data SQL engine. It re-implements the public pyspark.sql.* API surface as a parallel set of classes that build a lazy query plan, emit SQL via a vendored SQLGlot fork, and ship that SQL to e6data over gRPC. There is no JVM, no Spark cluster, no executors, no shuffle — only SQL.

This is NOT a wrapper around PySpark. PySpark is not a runtime dependency (see setup.py). In particular:

  • e6_spark_compat.sql.DataFrame is not a subclass of pyspark.sql.DataFrame. Any third-party code that does isinstance(df, pyspark.sql.DataFrame) will return False.
  • Databricks notebook builtins are not supported. The Databricks runtime injects display(df), dbutils.fs.*, MLflow log_input(df), the Spark UI, spark.sparkContext.*, etc. These integrate with the DBR-native Spark DataFrame type and will reject (or silently mis-handle) an e6_spark_compat DataFrame. Use df.show() instead of display(df); avoid sparkContext.*.
  • Anything outside the public PySpark DataFrame API is on the wrong side of the seam. If a Databricks integration calls something that isn't documented in the PySpark Python API, assume it won't work and test before shipping.
  • Python @udfs are AST-translated to SQL CASE/expressions by sql/udf_translator.py — not executed as Python on workers. Only UDFs that the translator can express in SQL will work.

Installation

For Users

# Install from PyPI
pip install e6data-spark-compatibility

# Install from GitHub
pip install git+https://github.com/e6data/e6-spark-compat.git

# Install from local clone
git clone https://github.com/e6data/e6-spark-compat.git
cd e6-spark-compat
pip install -e .

For Developers

# Clone the repository
git clone https://github.com/e6data/e6-spark-compat.git
cd e6-spark-compat

# Install with development dependencies
pip install -e ".[dev]"

# Run tests
pytest tests/

Quick Start

Migration from PySpark

Simply change your import statements - everything else stays the same:

# Before: PySpark imports
# from pyspark.sql import SparkSession
# from pyspark.sql.functions import col, upper, count, sum, row_number
# from pyspark.sql.window import Window

# After: e6-spark-compat imports
from e6_spark_compat import SparkSession
from e6_spark_compat.sql.functions import col, upper, count, sum, row_number
from e6_spark_compat.sql.window import Window

# Connection configuration
spark = (SparkSession.builder
    .appName("E6DataExample")
    .config("spark.e6data.host", "your-host")
    .config("spark.e6data.username", "your-username")
    .config("spark.e6data.password", "your-password")
    .config("spark.e6data.database", "your-database")
    .config("spark.e6data.catalog", "your-catalog")
    .config("spark.e6data.cluster", "your-cluster-name")
    .config("spark.e6data.secure", True)
    .getOrCreate())

# All PySpark operations work identically
df = spark.read.parquet("s3://bucket/path/to/data.parquet")

# Transformations (lazy evaluation)
result = (df.filter(col("age") > 21)
    .select("name", upper(col("city")).alias("city_upper"), "salary")
    .groupBy("city_upper")
    .agg(
        count("*").alias("total_people"),
        sum("salary").alias("total_salary")
    )
    .orderBy(col("total_people").desc()))

# Action triggers execution
result.show()

Window Functions

from e6_spark_compat.sql.window import Window
from e6_spark_compat.sql.functions import row_number, rank, lag

# Window specifications
window_spec = Window.partitionBy("department").orderBy("salary")
frame_spec = (Window.partitionBy("team").orderBy("hire_date")
    .rowsBetween(Window.UNBOUNDED_PRECEDING, Window.CURRENT_ROW))

# Window functions
df_with_analytics = df.select(
    "*",
    row_number().over(window_spec).alias("rank"),
    lag("salary", 1).over(window_spec).alias("prev_salary"),
    sum("salary").over(frame_spec).alias("running_total")
)

Spatial Support

For spatial operations with Sedona compatibility:

from e6_spark_compat.sedona import SedonaRegistrator
from e6_spark_compat.sql.functions import expr

# Register Sedona functions
SedonaRegistrator.registerAll(spark)

# Use spatial functions
points_df = spark.read.parquet("s3://bucket/points.parquet")
polygons_df = spark.read.parquet("s3://bucket/polygons.parquet")

# Spatial join
result = points_df.join(
    polygons_df,
    expr("ST_Contains(polygons_df.geometry, ST_Point(points_df.lon, points_df.lat))"),
    "inner"
)

Comprehensive Feature Set

DataFrame Operations

  • Transformations: select(), filter()/where(), join(), groupBy(), orderBy()/sort(), limit(), distinct()
  • Set Operations: union(), intersect(), exceptAll()
  • Column Operations: withColumn(), withColumnRenamed(), drop(), cast()
  • Caching: cache()/persist(), unpersist()

Action Methods

  • Data Retrieval: collect(), count(), first(), take(), head()
  • Display: show(), explain(), describe()
  • Export: toPandas(), write.parquet(), write.orc(), write.csv(), write.json()
  • Views: createOrReplaceTempView(), createGlobalTempView()

SQL Functions (135+)

  • String Functions: upper, lower, concat, substring, trim, split, regexp_replace, length
  • Math Functions: abs, round, floor, ceil, sqrt, pow, sin, cos, log
  • Aggregate Functions: sum, avg, count, min, max, first, last, collect_list
  • Date/Time Functions: year, month, dayofmonth, hour, minute, date_add, date_sub, datediff
  • Window Functions: row_number, rank, dense_rank, percent_rank, ntile, lag, lead
  • Conditional Functions: when, coalesce, isnull, isnan, greatest, least
  • Hash Functions: md5, sha1, sha2, hash, xxhash64
  • Array Functions: array, array_contains, explode, size, array_sort, and 16 more
  • Map Functions: map_keys, map_values, map_from_arrays, map_concat
  • JSON Functions: get_json_object, json_extract, from_json, to_json

Spatial Operations (Sedona Compatible)

  • Geometry Functions: All ST_* functions including ST_Point, ST_Polygon, ST_Buffer
  • Spatial Relationships: ST_Contains, ST_Intersects, ST_Distance, ST_Within
  • Transformations: ST_Transform, ST_Centroid, ST_ConvexHull
  • Indexing: H3 and GeoHash spatial indexing support

Testing

Test Suite Overview

The test suite uses dual-mode execution: tests run offline by default (mock connection, SQL validation) and can also execute on a live e6data cluster when credentials are provided.

A third axis — the tests/oracle/ suite — uses real Apache PySpark (local[*] mode, no cluster) as a ground-truth reference. The same query runs on both engines; results must match row-for-row. This catches semantic/emission gaps that pure SQL-generation tests cannot — e.g. e6 may emit "valid-looking" SQL that produces different rows than PySpark, or emit non-portable tokens that strict-ANSI engines reject. The suite skips cleanly when pyspark or Java are unavailable (so default CI is unaffected). See tests/oracle/README.md for the design.

Category Tests Description
E2E SQL Generation 360 Full pipeline: PySpark API → SQL string validation
TPC-DS Translations 37 Standard benchmark queries (37 of 99 TPC-DS queries)
TPC-DS Analytics 22 Complex analytics workloads (star schema, RFM, ETL pipelines)
SQLGlot Expressions 288 SQL expression generation via SQLGlot
Integration 227 Query plan node composition
Unit 200 Individual class/method tests
Writer 38 DataFrameWriter operations
Workloads 94 Customer workload pipeline tests (Kroger DQ + DFM analytics)
Total 1270 1263 pass, 7 known failures

Customer Scripts

Standalone scripts that exercise real customer pipelines end-to-end:

Script Description Dataset
kroger_original_test.py Kroger pre-DQ validation pipeline mars_full
e6_dfm_on_mars_full.py DFM pipeline patterns (39 tests) mars_full
e6_dfm_test.py DFM feature coverage (42 tests) TPC-DS
*_rust.py variants Same scripts with FORMAT_NUMBER workaround for rust engine
# Run customer script on Java engine
E6DATA_HOST=... S3_READ_BASE=... python tests/customer_scripts/kroger_original_test.py

# Run on rust engine
E6DATA_HOST=... S3_READ_BASE=... python tests/customer_scripts/kroger_original_test_rust.py

Running Tests

# Offline (default) — fast, no cluster needed (~5 seconds)
./run_tests.sh

# Specific file or directory
./run_tests.sh offline 1 tests/workloads/
./run_tests.sh offline 1 tests/workloads/test_dfm_pipeline.py
./run_tests.sh offline 1 tests/workloads/test_dfm_pipeline.py::TestDFMMissingFunctions

# By keyword or marker
./run_tests.sh offline 1 -k "dfm"
./run_tests.sh offline 1 -m e2e

# Or use pytest directly
pytest tests/e2e/                    # E2E tests
pytest tests/tpcds/                  # TPC-DS queries
pytest tests/e2e/ -m spatial         # Spatial tests

Live Mode (cluster execution)

Set environment variables to run the same tests against a real e6data cluster:

export E6DATA_HOST="your-cluster-host"
export E6DATA_PORT="443"
export E6DATA_USERNAME="your-user"
export E6DATA_PASSWORD="your-token"
export E6DATA_DATABASE="tpcds_1000_delta"
export E6DATA_CATALOG="your-catalog"
export E6DATA_CLUSTER="your-cluster"

# Sequential with full HTML report
./run_tests.sh live

# Parallel with N workers
./run_tests.sh live 4
./run_tests.sh live 8

In live mode, each test generates SQL (same as offline) and executes it on the cluster with LIMIT 5. The HTML report captures query IDs, row counts, execution times, and error details.

Test Reports

Reports are auto-generated on every run:

  • HTML: reports/test_report.html — browsable report with expandable details per test:
    • PySpark input source code
    • Generated SQL
    • Query ID (live mode)
    • Engine error details (live mode, failures only)
  • JUnit XML: reports/test_results.xml — for CI/CD integration

CI / GitHub Actions

Offline tests run automatically on every push/PR — no setup needed.

Live tests are triggered manually (requires GitHub secrets configured):

From GitHub UI:

  1. Go to Actions tab → e6-spark-compat CI
  2. Click Run workflow
  3. Select branch, cluster (general / rust-azure-benchmark), and worker count
  4. Click Run workflow

From CLI:

gh workflow run workflow.yaml --ref main -f run_live=true -f cluster=general -f concurrency=4

The CI summary includes category breakdown tables, failure details, and the dorny/test-reporter renders an interactive test results tab.

Architecture

Data flow

┌─────────────────────────────────────────────────────────────────────┐
│  User code:   from e6_spark_compat import SparkSession              │
│               spark.read.format("delta").load(...).groupBy(...)     │
└─────────────────────────┬───────────────────────────────────────────┘
                          │ PySpark-shaped API surface (parallel impl)
                          ▼
┌─────────────────────────────────────────────────────────────────────┐
│  e6_spark_compat/sql/         (façade — what the user touches)      │
│    dataframe.py, column.py, functions.py, reader.py, writer.py,     │
│    window.py, group.py, row.py, types.py, udf_translator.py         │
└─────────────────────────┬───────────────────────────────────────────┘
                          │ Every op builds an AST node, no execution
                          ▼
┌─────────────────────────────────────────────────────────────────────┐
│  e6_spark_compat/core/        (engine — the real work)              │
│    query_plan.py   ─→ lazy AST: TableScan / Filter / Project /      │
│                       Join / Aggregate / Sort / Limit / …           │
│    sql_generator.py─→ walks AST → SQLGlot Expression → SQL string   │
│    session.py      ─→ SparkSession.builder().config(...).build()    │
│    connection.py   ─→ gRPC ship to e6data                           │
└─────────────────────────┬───────────────────────────────────────────┘
                          │ SQL string (e6 dialect)
                          ▼
        ┌──────────────────────────────────────────────┐
        │  _vendor/sqlglot/  (vendored e6 fork, pinned │
        │                     to a specific commit)    │
        └──────────────────────────────────────────────┘
                          │
                          ▼ gRPC via e6data-python-connector
                  ┌─────────────────┐
                  │  e6data engine  │
                  └─────────────────┘

Module map

Module Lines Role
core/sql_generator.py 2,880 AST → SQL emission, projection rewrites, column-ambiguity fixes, CTE flattening
sql/udf_translator.py 1,812 AST-translates Python @udf bodies into SQL CASE/expressions
sql/dataframe.py 1,274 DataFrame, DataFrameNaFunctions, DataFrameStatFunctions
sql/functions.py 1,171 135+ SQL functions (col, lit, sum, when, …)
core/query_plan.py 849 Lazy AST node types and CTE assembly
sql/column.py 779 Column + operator overloads (+, -, &, `
sql/writer.py 544 DataFrameWriter (Parquet/ORC/CSV/JSON/noop)
sql/reader.py 368 spark.read.* for Parquet/Delta/CSV/JSON/Iceberg
core/session.py 302 SparkSession.builder() + connection plumbing
sql/window.py 281 Window.partitionBy/orderBy/rowsBetween

Core components

  1. Query Plan Tree: AST nodes representing SQL operations (Filter, Project, Join, Aggregate, etc.) — core/query_plan.py
  2. SQLGlot Integration: SQL generation through a vendored e6 fork of SQLGlot at e6_spark_compat/_vendor/sqlglot/, pinned to a specific commit (not pip-installed). Provides the e6 dialect.
  3. Lazy Evaluation: Operations build query plans without immediate execution
  4. Connection Management: gRPC transport via e6data-python-connector (core/connection.py). See the silent-fallback caveat below — connection errors do not raise by default.

Silent connection-error fallback. When a gRPC call from the connector raises (e.g. INTERNAL, network drop, malformed engine response), several action paths in sql/dataframe.py (notably count(), collect(), toPandas()) catch the exception, log "Connection error bypassed: <error>" at WARNING level, and return a safe default0 for count(), [] for collect(), an empty DataFrame for toPandas() — instead of propagating. This was added to "keep notebooks running" but means an engine failure can silently corrupt downstream logic (a check that depends on df.count() == 0 will succeed even though the query never ran). If you need strict failures, grep for "Connection error bypassed" in your run logs after every batch — it is the only signal — or wrap calls in an explicit retry/raise harness.

Execution flow

  1. Build Phase: DataFrame operations create query plan nodes (no I/O)
  2. Generation Phase: An action (show, collect, count, toPandas, write.save) walks the AST through SQLGenerator to produce a SQL string. set_use_cte(bool) toggles CTE vs. flattened-subquery emission.
  3. Execution Phase: SQL ships via gRPC to e6data
  4. Result Phase: Results returned as Row lists or pandas.DataFrame (toPandas)

What triggers a network roundtrip

Everything else is lazy — only these calls hit the engine:

Call What's sent Returns
df.show(n) / df.display(n) Generated SQL with LIMIT n prints to stdout
df.collect() Generated SQL (no limit) List[Row]
df.count() SELECT COUNT(*) FROM (<generated SQL>) int
df.first() / df.take(n) / df.head(n) Generated SQL with LIMIT 1 / LIMIT n Row or List[Row]
df.toPandas() Generated SQL (no limit) pandas.DataFrame
df.isEmpty() Generated SQL with LIMIT 1 bool
df.write.format(...).save(...) INSERT INTO … / CREATE TABLE … nothing (engine-side)
df.write.format("noop")... Generated SQL (acts as a dry-run materialization) nothing
spark.read.format(...).load(path) Schema-discovery call (DESCRIBE / metadata fetch) then lazy plan a lazy DataFrame — but the schema-discovery roundtrip happens at construction time, not at action time

That last row is a common surprise: even though spark.read.…load(...) is described as lazy in PySpark, this library issues a metadata roundtrip now so subsequent column-resolution can validate references locally. If your script reads 50 Delta paths in a loop, you've already issued 50 engine calls before any action runs.

What lives outside this seam

The architecture above ends at the public PySpark Python API. Anything that isn't on that surface is not implemented here:

  • Databricks notebook builtins (display(df), dbutils.*, displayHTML)
  • JVM-level Spark APIs (spark.sparkContext, RDD, Accumulator, broadcast)
  • The Spark UI / event log / SparkListener
  • MLflow-Spark integrations that consume a pyspark.sql.DataFrame directly
  • Anything that uses isinstance(df, pyspark.sql.DataFrame) for dispatch

If your customer pipeline touches any of these, either replace the call (display(df)df.show()) or run that step on a real Spark / DBR runtime.

Known semantic gaps (library ⇄ engine)

The compatibility seam has two failure modes worth distinguishing — both matter because the README's "everything works" framing oversells the second.

Library-side gaps (we accept the method, then misbehave):

  • fillna(scalar) is a no-op — no COALESCE emitted
  • dropna() default emits WHERE * IS NOT NULL instead of expanding columns
  • dropna(how='all') without subset emits NOT (* IS NULL)
  • col("x")[y] subscript — Column.__getitem__ not implemented (breaks melt patterns)
  • createDataFrame(pandas_df) — pandas truth-value ambiguity raises
  • spark.read.csv(path, sep="|", header=True, ...)all options are silently dropped since 2026-04-14 (only the path makes it to the engine). Non-default delimiters will silently produce wrong data. Track at sql/reader.py:_read_csv.
  • Three SQL-emission bugs that produce non-portable SQL (e6 happens to accept it; DBR / BigQuery / Snowflake / DuckDB all reject it):
    • join(on=[list], how="outer") emits ON l.k=r.k instead of USING(k)
    • chained withColumn collapses into one SELECT with forward-referenced columns
    • join(how='leftsemi' | 'leftanti' | 'leftouter' | 'rightouter' | 'fullouter') emits a single token (LEFTSEMI JOIN) instead of two (LEFT SEMI JOIN)

Engine-side gaps (library emits valid SQL; the engine rejects it):

  • IGNORE NULLS syntax (first/last with ignorenulls=True)
  • ISNAN, BROUND, PERCENTILE_APPROX, ARRAYS_ZIP + EXPLODE
  • FORMAT_NUMBER (rust engine only — Java engine supports it)
  • LOG (rust engine only)
  • WRITE INTO (rust engine only)
  • DELTA format writes (both engines)
  • LAST_VALUE as scalar (DuckDB-side test-harness gap, not the engine)

Why this matters: offline tests catch only library-side gaps; engine-side gaps surface only in live mode and only on the cluster build you're testing against. The current known-failure count is ~7 tests — kept as hard failures rather than xfail so they stay visible in CI output.

For the live status of these gaps, see CLAUDE.md ("Known Bugs & Missing Features") which is updated as the library evolves.

Development

Project Structure

e6_spark_compat/
├── core/           # Session, connection, query plan, SQL generator
├── sql/            # DataFrame, Column, functions, window, reader, writer, types
├── spatial/        # Sedona-compatible ST_* functions
└── sedona/         # SedonaRegistrator compatibility shim

tests/
├── e2e/              # E2E SQL generation tests (dual-mode: offline + live)
├── tpcds/            # TPC-DS benchmark query translations (37 queries)
├── integration/      # Query plan composition tests
├── sqlglot_tests/    # SQLGlot expression tests
├── unit/             # Unit tests
├── workloads/        # Customer workload pipeline tests
├── customer_scripts/ # Standalone customer pipeline scripts
└── live/             # Live cluster-only execution tests

Code Quality

# Format code
black e6_spark_compat/

# Run linting
flake8 e6_spark_compat/

# Type checking
mypy e6_spark_compat/

Migration Guide

From PySpark

  1. Update Imports: Change pyspark imports to e6_spark_compat
  2. Configuration: Update SparkSession configuration for e6data connection
  3. Test: Run your existing PySpark code - it should work unchanged

From Sedona

  1. Spatial Functions: Replace Sedona imports with e6_spark_compat.spatial.functions
  2. Registration: Use SedonaRegistrator.registerAll(spark) for compatibility
  3. Geometry Types: Spatial operations work identically to Sedona

Contributing

We welcome contributions! Here's how to get started:

  1. Fork the repository
  2. Create a feature branch
  3. Add tests for new functionality
  4. Ensure all tests pass (pytest tests/)
  5. Submit a pull request

Support

  • Documentation: See docs/pyspark-compatibility/ for detailed guides
  • Issues: Report bugs and feature requests on GitHub
  • Contact: support@e6data.com for enterprise support

License

Apache License 2.0 - see LICENSE for details.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

e6data_spark_compatibility-1.0.7.tar.gz (623.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

e6data_spark_compatibility-1.0.7-py3-none-any.whl (673.2 kB view details)

Uploaded Python 3

File details

Details for the file e6data_spark_compatibility-1.0.7.tar.gz.

File metadata

File hashes

Hashes for e6data_spark_compatibility-1.0.7.tar.gz
Algorithm Hash digest
SHA256 3c5bd5856f961beabc4caf60f8213f233174c63e7803fcbb5d027c520efa9cfc
MD5 7f2a04ce21618adfc25bfc32287f2686
BLAKE2b-256 cc72f388ad0b629903d132e6a9a609c1c1a4f064657557518849a487c00907d0

See more details on using hashes here.

Provenance

The following attestation bundles were made for e6data_spark_compatibility-1.0.7.tar.gz:

Publisher: workflow.yaml on e6data/e6-spark-compat

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file e6data_spark_compatibility-1.0.7-py3-none-any.whl.

File metadata

File hashes

Hashes for e6data_spark_compatibility-1.0.7-py3-none-any.whl
Algorithm Hash digest
SHA256 467d2fcef0baf39b88350c3da1355008e51fb4a1fdc96becc7fa9cc135e44c75
MD5 f787cd5d5d89df0230ad257f07111ff8
BLAKE2b-256 1a1619f3fc60a99a0f5af7778ed8b2e519807fdb6ae193bb884c5f370abb3914

See more details on using hashes here.

Provenance

The following attestation bundles were made for e6data_spark_compatibility-1.0.7-py3-none-any.whl:

Publisher: workflow.yaml on e6data/e6-spark-compat

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page