Redis 按距离处理地理位置

redis_lab=$(mktemp -d)
curl --fail --location --output "$redis_lab/redis.tar.gz" https://download.redis.io/releases/redis-8.10.1.tar.gz
tar -xzf "$redis_lab/redis.tar.gz" -C "$redis_lab"
make -C "$redis_lab/redis-8.10.1" build redis -j2 MALLOC=libc
export REDIS_SERVER="$redis_lab/redis-8.10.1/src/redis-server"
python3 -m venv "$redis_lab/venv"
"$redis_lab/venv/bin/python" -m pip install redis==8.1.0

import math

from contextlib import contextmanager
import os
from pathlib import Path
import subprocess
from tempfile import TemporaryDirectory
import time

import redis
from redis.backoff import NoBackoff
from redis.retry import Retry


@contextmanager
def instance():
    with TemporaryDirectory(prefix="redis-lab-", dir="/tmp") as directory:
        endpoint = str(Path(directory) / "redis.sock")
        process = subprocess.Popen([
            os.environ["REDIS_SERVER"], "--port", "0",
            "--unixsocket", endpoint, "--unixsocketperm", "700",
            "--save", "", "--appendonly", "no", "--dir", directory,
            "--maxmemory-policy", "noeviction",
        ], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
        client = redis.Redis(unix_socket_path=endpoint, decode_responses=False,
                             socket_timeout=2, retry=Retry(NoBackoff(), 0))
        try:
            deadline = time.monotonic() + 5
            while True:
                assert process.poll() is None, "Redis exited during startup"
                try:
                    if client.ping():
                        break
                except redis.ConnectionError:
                    if time.monotonic() >= deadline:
                        raise
                    time.sleep(0.02)
            yield client
        finally:
            client.close()
            if process.poll() is None:
                process.terminate()
            try:
                _, errors = process.communicate(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
                process.communicate()
                raise
            assert process.returncode == 0, errors.decode()


with instance() as client:
    palermo = (13.361389, 38.115556)
    catania = (15.087269, 37.502669)
    assert client.geoadd("places", [
        *palermo, b"Palermo", *catania, b"Catania",
    ]) == 2
    assert client.type("places") == b"zset"
    position = client.geopos("places", "Palermo")[0]
    assert all(abs(actual - expected) < 0.00001
               for actual, expected in zip(position, palermo))
    distance = client.geodist("places", "Palermo", "Catania", unit="km")
    assert math.isclose(distance, 166.2742, abs_tol=0.001)
    assert client.geodist("places", "Palermo", "missing") is None
    print(f"longitude first; decoded coordinates are approximate; distance {distance:.4f} km")

    query = dict(longitude=15, latitude=37, radius=200, unit="km", sort="ASC")
    results = client.geosearch("places", **query, withdist=True)
    assert [item[0] for item in results] == [b"Catania", b"Palermo"]
    assert results[0][1] < results[1][1] < 200
    assert client.geosearch("places", **query, count=1) == [b"Catania"]
    assert client.geosearch(
        "places", member="Palermo", radius=1, unit="km", sort="ASC",
    ) == [b"Palermo"]
    print("ASC COUNT 1 selects Catania; FROMMEMBER also includes the center member")

    assert client.geosearchstore("subset", "places", **query) == 2
    assert client.zscore("subset", "Catania") == client.zscore("places", "Catania")
    assert client.geosearchstore("distances", "places", **query, storedist=True) == 2
    assert math.isclose(client.zscore("distances", "Catania"),
                        results[0][1], abs_tol=0.001)
    assert client.zscore("distances", "Catania") != client.zscore("places", "Catania")
    print("default GEOSEARCHSTORE preserves position scores; STOREDIST stores kilometers")

    try:
        client.geoadd("places", [0, 90, b"pole"])
    except redis.ResponseError as error:
        assert "invalid longitude,latitude pair" in str(error)
    else:
        raise AssertionError("latitude 90 must be rejected")
    assert client.zcard("places") == 2
    assert client.geoadd("swapped", [
        *palermo, b"correct", palermo[1], palermo[0], b"reversed",
    ]) == 2
    assert client.geodist("swapped", "correct", "reversed", unit="km") > 1000
    print("polar coordinates are rejected; swapped but valid coordinates are not detected")

    assert client.geoadd("places", [13.4, 38.1, b"Palermo"]) == 0
    assert client.zcard("places") == 2
    assert client.geoadd("places", [13.5, 38.2, b"Palermo"], ch=True) == 1
    assert client.zrem("places", b"Catania") == 1
    assert client.geopos("places", "Catania") == [None]
    assert client.zcard("subset") == 2
    print("updates reuse members; ZREM removes locations; stored subsets are snapshots")