james-willis commented on code in PR #1159: URL: https://github.com/apache/sedona-db/pull/1159#discussion_r3883331244
########## python/sedonadb/python/sedonadb/testing_spark.py: ########## @@ -0,0 +1,176 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Sedona Spark as a :class:`~sedonadb.testing.DBEngine`. + +The compatibility target for SedonaDB's SQL surface is Sedona Spark, so parity +tests broadcast one shared SQL string to both engines and compare strictly — +the same pattern the geometry suite uses for SedonaDB vs PostGIS. This lives in +its own module because the engine needs pyspark, a JVM, and network access, +none of which the core :mod:`sedonadb.testing` module should pull in. + +Bootstrapping a ``SparkSession`` (downloading the Sedona jars from Maven, JVM +startup) costs tens of seconds, so these tests are opt-in: :meth:`create_or_skip` +skips — even under ``SEDONADB_PYTHON_NO_SKIP_TESTS`` — unless +``SEDONADB_RUN_SPARK_TESTS`` is set, in which case construction failures +propagate. + +Requires Spark 4.0+ (results travel out of the JVM through ``DataFrame.toArrow``, +which preserves nulls Arrow-natively — unlike ``toPandas``, which would coerce a +nodata ``None`` to ``NaN`` and mask exactly the value under test). +""" + +import os + +import pyarrow as pa + +from sedonadb.testing import DBEngine + +# Sedona jar coordinates. Override with SEDONADB_SEDONA_SPARK_PACKAGES (full +# Maven coordinates) when testing against a different Sedona release. +SEDONA_SPARK_VERSION = "1.9.0" +GEOTOOLS_WRAPPER_VERSION = "1.9.0-33.5" + + +class SedonaSpark(DBEngine): + """Runs Sedona Spark SQL — the compatibility-target dialect. + + One local ``SparkSession`` is bootstrapped per process and shared across + engine instances. Rasters are read from GeoTIFF files with + ``RS_FromGeoTiff`` over a ``binaryFile`` scan (see :meth:`create_raster_view`). + """ + + _spark = None + + def __init__(self): + self._session = self._ensure_session() + + @classmethod + def name(cls) -> str: + return "sedona-spark" + + @classmethod + def install_hint(cls) -> str: + return ( + "- Run `pip install 'pyspark>=4.0' apache-sedona` (needs a JVM; the " + "first run downloads the Sedona jars from Maven)\n" + "- Set SEDONADB_RUN_SPARK_TESTS=true to opt in" + ) + + @classmethod + def create_or_skip(cls, *args, **kwargs) -> "SedonaSpark": + import pytest + + if os.environ.get("SEDONADB_RUN_SPARK_TESTS", "false") not in ("true", "1"): + pytest.skip("Sedona Spark parity tests are opt-in:\n" + cls.install_hint()) + return cls(*args, **kwargs) + + @classmethod + def _ensure_session(cls): + if SedonaSpark._spark is None: + from sedona.spark import SedonaContext + + config = ( + SedonaContext.builder() + .master("local[2]") + .appName("sedonadb-spark-parity") + .config("spark.jars.packages", cls._packages()) + .config("spark.jars.ivy", cls._ivy_dir()) + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + SedonaSpark._spark = SedonaContext.create(config) + return SedonaSpark._spark + + @staticmethod + def _packages() -> str: + env = os.environ.get("SEDONADB_SEDONA_SPARK_PACKAGES") + if env: + return env + import pyspark + + major, minor = (int(part) for part in pyspark.__version__.split(".")[:2]) + # Sedona publishes per-Spark-minor artifacts. A pyspark older than every + # published artifact is an error (a newer jar on an older runtime fails + # at class load); a pyspark newer than the newest artifact tries the + # newest, which usually loads — override the coordinates if it doesn't. + known = ("3.5", "4.0") Review Comment: done -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
