From 362d3e74af0b755ab30e85a56f86b2c2d70e5d6c Mon Sep 17 00:00:00 2001 From: Jia Yu Date: Thu, 3 Sep 2026 00:14:54 -0700 Subject: [PATCH 1/2] fix(python/sedonadb): resolve GeoArrow scalars with their edge type, and common scalar values The literal resolver rebuilt a GeoArrow scalar's type from its CRS alone, so a spherical (geography) scalar came back as planar geometry, and it registered only the WKB scalar class, so WKT and native-encoded scalars fell through to pa.array([obj]) and failed. Carry the edge type through and route every GeoArrow scalar class through the same WKB path. Also resolve values that pa.array([obj]) rejects or handles lossily: pandas.NA and numpy.ma.masked become NULL, pandas.NaT a timestamp NULL, pandas Timestamp and Timedelta keep their full resolution and time zone instead of being truncated to microseconds, NumPy datetime64/timedelta64 values in any unit convert at a lossless Arrow resolution (ambiguous or sub-nanosecond units are rejected the way pandas rejects them, with an overflow check), 0-d arrays resolve as their typed scalar, structured numpy.void values become structs, and null Arrow scalars of nested or extension types resolve through their typed one-element array. pandas 3 class names are registered alongside the pandas 2 ones. Closes #1214. --- python/sedonadb/python/sedonadb/context.py | 8 + .../sedonadb/python/sedonadb/expr/literal.py | 149 ++++++++++++++- python/sedonadb/tests/expr/test_literal.py | 170 ++++++++++++++++++ 3 files changed, 325 insertions(+), 2 deletions(-) diff --git a/python/sedonadb/python/sedonadb/context.py b/python/sedonadb/python/sedonadb/context.py index 26173a4116..4dd046668d 100644 --- a/python/sedonadb/python/sedonadb/context.py +++ b/python/sedonadb/python/sedonadb/context.py @@ -558,6 +558,8 @@ def lit(self, value: Any) -> LiteralExpr: is accepted by `pyarrow.array([...])` is supported in addition to: - Shapely geometries become SedonaDB geometry objects. + - GeoArrow scalars (WKB, WKT, or native encodings) become SedonaDB + geometries with CRS and edge type (planar or spherical) preserved. - GeoSeries objects of length 1 become SedonaDB geometries with CRS preserved. - GeoDataFrame objects with a single column and single row become @@ -569,6 +571,12 @@ def lit(self, value: Any) -> LiteralExpr: value. - pyproj CRS objects become PROJJSON strings (e.g., so they may be used in `ST_SetCRS()`, `ST_Point()`, or `ST_GeomFromWKT()`). + - pandas `Timestamp` and `Timedelta` values keep their full resolution + (and time zone); `pandas.NA` and `numpy.ma.masked` become NULL and + `pandas.NaT` a timestamp NULL. + - NumPy `datetime64`/`timedelta64` values in any unit convert at a + lossless Arrow resolution, 0-d arrays resolve as their scalar, and + structured `numpy.void` values become structs. """ return lit_expr(value, ctx=self) diff --git a/python/sedonadb/python/sedonadb/expr/literal.py b/python/sedonadb/python/sedonadb/expr/literal.py index bda5bf55bc..84442d8d6f 100644 --- a/python/sedonadb/python/sedonadb/expr/literal.py +++ b/python/sedonadb/python/sedonadb/expr/literal.py @@ -134,6 +134,11 @@ def _resolve_arrow_lit(obj: Any): import pyarrow as pa + # A null Arrow scalar of a nested or extension type is not accepted by + # pa.array([obj]); its one-element array spelling carries the same type. + if isinstance(obj, pa.Scalar) and not obj.is_valid: + return pa.array([None], type=obj.type) + try: return pa.array([obj]) except Exception as e: @@ -143,8 +148,11 @@ def _resolve_arrow_lit(obj: Any): def _lit_from_geoarrow_scalar(obj): + # Every GeoArrow scalar (WKB, WKT, native point/linestring/...) exposes + # its WKB, so one path serves them all. The edge type travels with the + # CRS: a spherical (geography) scalar must not come back planar. wkb_value = None if obj.value is None else obj.wkb - return _lit_from_wkb_and_crs(wkb_value, obj.type.crs) + return _lit_from_wkb_and_crs(wkb_value, obj.type.crs, obj.type.edge_type) def _lit_from_dataframe(obj): @@ -191,15 +199,137 @@ def _lit_from_shapely(obj): return _lit_from_wkb_and_crs(obj.wkb, None) -def _lit_from_wkb_and_crs(wkb, crs): +def _lit_from_wkb_and_crs(wkb, crs, edge_type=None): import geoarrow.pyarrow as ga import pyarrow as pa type = ga.wkb().with_crs(crs) + if edge_type is not None: + type = type.with_edge_type(edge_type) storage = pa.array([wkb], type.storage_type) return type.wrap_array(storage) +def _lit_from_missing(obj): + # pandas.NA and numpy.ma.masked both mean "no value". + import pyarrow as pa + + return pa.array([None]) + + +def _lit_from_nat(obj): + # NaT is a datetime missing value in pandas (assigning it yields a + # datetime64 column), so it resolves to a typed timestamp null rather + # than an untyped NULL. + import pyarrow as pa + + return pa.array([None], pa.timestamp("ns")) + + +def _lit_from_pandas_timestamp(obj): + # pa.array([Timestamp]) treats it as a datetime and resolves at + # microseconds, silently dropping nanoseconds. .asm8 is the instant as a + # numpy datetime64 at the Timestamp's own unit (UTC for a zone-aware + # value), which pyarrow converts exactly; the zone is then re-attached + # by a cast, which reinterprets the naive values as UTC without shifting + # them. + import pyarrow as pa + + resolved = pa.array([obj.asm8]) + if obj.tz is None: + return resolved + return resolved.cast(pa.timestamp(resolved.type.unit, pa.scalar(obj).type.tz)) + + +def _lit_from_pandas_timedelta(obj): + # Same nanosecond concern as Timestamp: .asm8 keeps the unit. + import pyarrow as pa + + return pa.array([obj.asm8]) + + +_NUMPY_TEMPORAL_UNITS = { + # Arrow-native units keep their resolution. + "s": "s", + "ms": "ms", + "us": "us", + "ns": "ns", + # Whole multiples of seconds convert exactly to seconds. + "W": "s", + "D": "s", + "h": "s", + "m": "s", +} + + +def _lit_from_numpy_temporal(obj): + # pyarrow only understands the four Arrow units, so a datetime64[D] (the + # default unit for a bare date string) or a timedelta64[W] fails outright. + # Convert at a lossless unit instead of forcing nanoseconds: the ns range + # covers only 1677-2262, so an unchecked astype would silently wrap a + # coarse value centuries away. + import numpy as np + import pyarrow as pa + + is_datetime = isinstance(obj, np.datetime64) + kind = "datetime64" if is_datetime else "timedelta64" + if np.isnat(obj): + return pa.array( + [None], pa.timestamp("ns") if is_datetime else pa.duration("ns") + ) + + unit = np.datetime_data(obj.dtype)[0] + if unit in _NUMPY_TEMPORAL_UNITS: + target = _NUMPY_TEMPORAL_UNITS[unit] + elif is_datetime and unit in ("Y", "M"): + # Calendar year/month positions are exact instants for a datetime + # (a timedelta in months or years has no fixed length). + target = "s" + elif is_datetime: + # Sub-nanosecond datetimes narrow to nanoseconds; the round-trip + # check below rejects the ones that lose precision. + target = "ns" + else: + raise ValueError( + f"Can't create SedonaDB literal from a {kind}[{unit}] value: use an " + f"unambiguous unit no finer than nanoseconds" + ) + + converted = obj.astype(f"{kind}[{target}]") + if converted.astype(obj.dtype) != obj: + # A same-unit conversion is the identity, so a mismatch is either a + # sub-nanosecond value with no exact ns form or a coarse value whose + # seconds form overflows int64. + if unit in _NUMPY_TEMPORAL_UNITS or unit in ("Y", "M"): + raise OverflowError(f"{obj!r} does not fit the Arrow {target!r} resolution") + raise ValueError(f"{obj!r} loses precision at the Arrow 'ns' resolution") + return pa.array([converted]) + + +def _lit_from_numpy_array(obj): + # A 0-d array is one value: unwrap to the typed NumPy scalar (which keeps + # the dtype, unlike .item()) and resolve that. Anything with dimensions + # keeps the generic behavior of becoming a single list value. + import pyarrow as pa + + if obj.ndim == 0: + return _resolve_arrow_lit(obj[()]) + return pa.array([obj]) + + +def _lit_from_numpy_void(obj): + # A plain void's payload is its bytes. A structured scalar becomes a + # typed Arrow struct so field names and dtypes survive (flattened to a + # tuple it would lose both). + import pyarrow as pa + + if obj.dtype.fields is None: + return pa.array([obj.item()]) + fields = [(name, pa.from_numpy_dtype(obj.dtype[name])) for name in obj.dtype.names] + payload = {name: obj[name].item() for name in obj.dtype.names} + return pa.array([payload], pa.struct(fields)) + + def _lit_from_crs(crs): return _resolve_arrow_lit(crs.to_json()) @@ -216,9 +346,22 @@ def _qualified_type_name(obj): # pandas < 3.0 "pandas.core.frame.DataFrame": _lit_from_dataframe, "pandas.core.series.Series": _lit_from_series, + "pandas._libs.missing.NAType": _lit_from_missing, + "pandas._libs.tslibs.nattype.NaTType": _lit_from_nat, + "pandas._libs.tslibs.timestamps.Timestamp": _lit_from_pandas_timestamp, + "pandas._libs.tslibs.timedeltas.Timedelta": _lit_from_pandas_timedelta, # pandas >= 3.0 "pandas.DataFrame": _lit_from_dataframe, "pandas.Series": _lit_from_series, + "pandas.api.typing.NAType": _lit_from_missing, + "pandas.api.typing.NaTType": _lit_from_nat, + "pandas.Timestamp": _lit_from_pandas_timestamp, + "pandas.Timedelta": _lit_from_pandas_timedelta, + "numpy.datetime64": _lit_from_numpy_temporal, + "numpy.timedelta64": _lit_from_numpy_temporal, + "numpy.ma.core.MaskedConstant": _lit_from_missing, + "numpy.ndarray": _lit_from_numpy_array, + "numpy.void": _lit_from_numpy_void, "pyproj.crs.crs.CRS": _lit_from_crs, "sedonadb.dataframe.DataFrame": _lit_from_sedonadb, "shapely.geometry.point.Point": _lit_from_shapely, @@ -230,4 +373,6 @@ def _qualified_type_name(obj): "shapely.geometry.multipolygon.MultiPolygon": _lit_from_shapely, "shapely.geometry.collection.GeometryCollection": _lit_from_shapely, "geoarrow.pyarrow._scalar.WkbScalar": _lit_from_geoarrow_scalar, + "geoarrow.pyarrow._scalar.WktScalar": _lit_from_geoarrow_scalar, + "geoarrow.pyarrow._scalar.GeometryExtensionScalar": _lit_from_geoarrow_scalar, } diff --git a/python/sedonadb/tests/expr/test_literal.py b/python/sedonadb/tests/expr/test_literal.py index 42cd43642d..4ba5a3fc09 100644 --- a/python/sedonadb/tests/expr/test_literal.py +++ b/python/sedonadb/tests/expr/test_literal.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import numpy as np import pyarrow as pa import shapely import geopandas @@ -185,3 +186,172 @@ def test_contextless_literal(): with pytest.raises(ValueError, match="Can't pipe Literal"): literal.funcs + + +def test_geoarrow_scalar_keeps_edge_type(): + # The scalar path rebuilt the type from CRS alone, so a spherical + # (geography) scalar came back as planar geometry. + spherical = ga.wkb().with_edge_type(ga.EdgeType.SPHERICAL).with_crs(ga.OGC_CRS84) + for payload in (shapely.Point(1, 1).wkb, None): + lit_array = pa.array(lit(pa.scalar(payload, spherical))) + assert lit_array.type.edge_type == ga.EdgeType.SPHERICAL + assert lit_array.type.crs.to_json_dict()["id"] == { + "authority": "OGC", + "code": "CRS84", + } + assert lit_array.null_count == (1 if payload is None else 0) + + +@pytest.mark.parametrize( + "make", + [ + lambda wkt: ga.as_wkt([wkt])[0], + lambda wkt: ga.as_geoarrow([wkt])[0], + lambda wkt: pa.scalar(shapely.from_wkt(wkt).wkb, ga.large_wkb()), + ], + ids=["wkt", "native", "large_wkb"], +) +def test_geoarrow_non_wkb_scalar_literal(make): + # Only the WKB scalar class was registered; other GeoArrow scalars fell + # through to pa.array([obj]) and failed. They all expose their WKB. + scalar = make("POINT (1 2)") + assert pa.array(lit(scalar)) == ga.as_wkb(["POINT (1 2)"]) + + +@pytest.mark.parametrize( + "type", + [ga.wkt(), ga.point(), ga.wkb().with_crs(ga.OGC_CRS84)], + ids=["wkt", "native", "wkb_with_crs"], +) +def test_geoarrow_null_scalar_literal(type): + lit_array = pa.array(lit(pa.scalar(None, type))) + assert lit_array.null_count == 1 + assert lit_array.type.extension_name == "geoarrow.wkb" + if type.crs is not None: + assert lit_array.type.crs.to_json_dict()["id"] == { + "authority": "OGC", + "code": "CRS84", + } + + +def test_null_arrow_scalar_literal(): + # A null nested scalar is not accepted by pa.array([obj]); the typed + # one-element array is. + for type in ( + pa.list_(pa.int64()), + pa.map_(pa.string(), pa.int64()), + pa.struct([("a", pa.int32())]), + ): + lit_array = pa.array(lit(pa.scalar(None, type))) + assert lit_array.type == type + assert lit_array.null_count == 1 + # Valid nested scalars were already fine and still are. + assert pa.array(lit(pa.scalar([1, 2]))).to_pylist() == [[1, 2]] + + +def test_pandas_missing_literal(): + assert pa.array(lit(pd.NA)) == pa.array([None]) + nat = pa.array(lit(pd.NaT)) + assert nat.type == pa.timestamp("ns") + assert nat.null_count == 1 + + +def test_pandas_timestamp_literal(): + # pa.array([Timestamp]) resolves at microseconds and silently drops + # nanoseconds. + stamp = pd.Timestamp("2026-01-01 00:00:00.000000001") + lit_array = pa.array(lit(stamp)) + assert lit_array.type == pa.timestamp("ns") + assert lit_array[0].as_py() == stamp + + aware = pd.Timestamp("2026-01-01 00:00:00.000000001", tz="US/Pacific") + lit_array = pa.array(lit(aware)) + assert lit_array.type == pa.timestamp("ns", "US/Pacific") + assert lit_array[0].as_py() == aware + + # A coarser-unit Timestamp keeps its unit rather than being forced to ns + # (2500 is outside the ns range). + coarse = pd.Timestamp("2500-01-01").as_unit("s") + lit_array = pa.array(lit(coarse)) + assert lit_array.type == pa.timestamp("s") + assert lit_array[0].as_py() == coarse + + +def test_pandas_timedelta_literal(): + lit_array = pa.array(lit(pd.Timedelta(1))) + assert lit_array.type == pa.duration("ns") + assert lit_array[0].as_py() == pd.Timedelta(1) + + +@pytest.mark.parametrize( + "value,expected", + [ + ( + np.datetime64("2500-01-01", "D"), + pa.array([np.datetime64("2500-01-01T00:00:00", "s")]), + ), + ( + np.datetime64("2500", "Y"), + pa.array([np.datetime64("2500-01-01T00:00:00", "s")]), + ), + ( + np.datetime64("2026-01-01T00:00:00.000000001", "ns"), + pa.array([np.datetime64("2026-01-01T00:00:00.000000001", "ns")]), + ), + (np.datetime64(10**6, "fs"), pa.array([np.datetime64(1, "ns")])), + (np.timedelta64(2, "D"), pa.array([np.timedelta64(2 * 86400, "s")])), + (np.timedelta64(3, "W"), pa.array([np.timedelta64(3 * 7 * 86400, "s")])), + (np.timedelta64(5, "ms"), pa.array([np.timedelta64(5, "ms")])), + ], + ids=["day", "year", "ns", "fs_exact", "td_day", "td_week", "td_ms"], +) +def test_numpy_temporal_literal(value, expected): + # pyarrow only understands the four Arrow units; other units convert at a + # lossless resolution rather than being rejected or forced to ns. + assert pa.array(lit(value)) == expected + + +def test_numpy_temporal_nat_literal(): + # Explicit units: newer NumPy deprecates the unit-less ("generic") NaT. + for value, type in ( + (np.datetime64("NaT", "ns"), pa.timestamp("ns")), + (np.timedelta64("NaT", "ns"), pa.duration("ns")), + (np.datetime64("NaT", "D"), pa.timestamp("ns")), + ): + lit_array = pa.array(lit(value)) + assert lit_array.type == type + assert lit_array.null_count == 1 + + +def test_numpy_temporal_literal_errors(): + # A timedelta in months or years has no fixed length. + with pytest.raises(ValueError, match="unambiguous unit"): + pa.array(lit(np.timedelta64(1, "M"))) + # A sub-nanosecond value with no exact nanosecond form. + with pytest.raises(ValueError, match="loses precision"): + pa.array(lit(np.datetime64(1, "fs"))) + # A coarse value whose seconds form overflows int64. + with pytest.raises(OverflowError): + pa.array(lit(np.timedelta64(2**62, "D"))) + + +def test_numpy_masked_literal(): + assert pa.array(lit(np.ma.masked)) == pa.array([None]) + + +def test_numpy_zero_dim_array_literal(): + # A 0-d array is one value, resolved as its typed scalar (dtype kept). + lit_array = pa.array(lit(np.array(np.int32(5)))) + assert lit_array == pa.array([5], pa.int32()) + lit_array = pa.array(lit(np.array(np.datetime64("2500-01-01")))) + assert lit_array == pa.array([np.datetime64("2500-01-01T00:00:00", "s")]) + # Arrays with dimensions keep resolving as a single list value. + assert pa.array(lit(np.array([1, 2]))).to_pylist() == [[1, 2]] + + +def test_numpy_void_literal(): + assert pa.array(lit(np.void(b"ab"))) == pa.array([b"ab"]) + record = np.array([(3, 1.5)], dtype=[("c", "int16"), ("r", "float32")])[0] + lit_array = pa.array(lit(record)) + assert lit_array.type == pa.struct([("c", pa.int16()), ("r", pa.float32())]) + assert lit_array.to_pylist() == [{"c": 3, "r": 1.5}] From 4da6c5c3c24ceb486e924b71f7d419ac37e04e66 Mon Sep 17 00:00:00 2001 From: Jia Yu Date: Thu, 3 Sep 2026 13:21:44 -0700 Subject: [PATCH 2/2] refactor(python/sedonadb): rename the WKB literal helper, explain the NaT unit The helper now takes an edge type as well as a CRS, so its name names the value it builds from rather than its arguments. Note why NaT resolves at nanoseconds: it carries no unit of its own, nanoseconds is where pandas stores it, and as a null it coerces to the unit and zone the surrounding expression needs. --- python/sedonadb/python/sedonadb/expr/literal.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/python/sedonadb/python/sedonadb/expr/literal.py b/python/sedonadb/python/sedonadb/expr/literal.py index 84442d8d6f..675c3fba82 100644 --- a/python/sedonadb/python/sedonadb/expr/literal.py +++ b/python/sedonadb/python/sedonadb/expr/literal.py @@ -152,7 +152,7 @@ def _lit_from_geoarrow_scalar(obj): # its WKB, so one path serves them all. The edge type travels with the # CRS: a spherical (geography) scalar must not come back planar. wkb_value = None if obj.value is None else obj.wkb - return _lit_from_wkb_and_crs(wkb_value, obj.type.crs, obj.type.edge_type) + return _lit_from_wkb(wkb_value, obj.type.crs, obj.type.edge_type) def _lit_from_dataframe(obj): @@ -173,7 +173,7 @@ def _lit_from_series(obj): if obj.dtype.name == "geometry": first_value = obj.array[0] first_wkb = None if first_value is None else first_value.wkb - return _lit_from_wkb_and_crs(first_wkb, obj.array.crs) + return _lit_from_wkb(first_wkb, obj.array.crs) else: import pyarrow as pa @@ -196,10 +196,10 @@ def _lit_from_sedonadb(obj): def _lit_from_shapely(obj): - return _lit_from_wkb_and_crs(obj.wkb, None) + return _lit_from_wkb(obj.wkb, None) -def _lit_from_wkb_and_crs(wkb, crs, edge_type=None): +def _lit_from_wkb(wkb, crs, edge_type=None): import geoarrow.pyarrow as ga import pyarrow as pa @@ -220,7 +220,9 @@ def _lit_from_missing(obj): def _lit_from_nat(obj): # NaT is a datetime missing value in pandas (assigning it yields a # datetime64 column), so it resolves to a typed timestamp null rather - # than an untyped NULL. + # than an untyped NULL. NaT itself carries no unit; nanoseconds is the + # unit pandas stores it in, and as a null it coerces to whatever unit + # and time zone the surrounding expression needs. import pyarrow as pa return pa.array([None], pa.timestamp("ns"))