如何用Pyspark/Python识别含Z值的GeoJSON几何并截断Z值

需求回顾 #

• DataFrame包含geometry列，存储MultiPolygon或Polygon类型的GeoJSON格式数据
• 目标1：识别出坐标包含[X,Y,Z]的行
• 目标2：移除几何数据中所有Z值，保留X,Y
• 目标3：记录每行中首个带Z值的坐标

方案一：Python（Pandas）实现 #

步骤1：定义几何处理函数 #

针对Polygon和MultiPolygon的层级结构差异，写递归逻辑处理坐标：

import pandas as pd

def process_geometry(geo):
    has_z = False
    first_z_coord = None
    processed_coords = []
    
    # 处理MultiPolygon结构：外层数组包含多个Polygon的坐标集合
    if geo['type'] == 'MultiPolygon':
        for polygon in geo['coordinates']:
            processed_polygon = []
            for ring in polygon:
                processed_ring = []
                for coord in ring:
                    # 检查坐标是否含Z值
                    if len(coord) >= 3:
                        if not has_z:
                            has_z = True
                            first_z_coord = coord
                        # 仅保留X,Y坐标
                        processed_ring.append(coord[:2])
                    else:
                        processed_ring.append(coord)
                processed_polygon.append(processed_ring)
            processed_coords.append(processed_polygon)
    
    # 处理Polygon结构：外层数组是坐标环集合
    elif geo['type'] == 'Polygon':
        processed_polygon = []
        for ring in geo['coordinates']:
            processed_ring = []
            for coord in ring:
                if len(coord) >= 3:
                    if not has_z:
                        has_z = True
                        first_z_coord = coord
                    processed_ring.append(coord[:2])
                else:
                    processed_ring.append(coord)
            processed_polygon.append(processed_ring)
        processed_coords.append(processed_polygon)
    
    # 组装处理后的2D几何对象
    processed_geo = {
        'type': geo['type'],
        'coordinates': processed_coords
    }
    
    return has_z, first_z_coord, processed_geo

步骤2：应用函数到DataFrame #

生成新列并筛选目标行：

# 批量处理geometry列，生成三个新字段
source_df[['has_z', 'first_z_coord', 'geometry_2d']] = source_df['geometry'].apply(
    lambda x: pd.Series(process_geometry(x))
)

# 筛选出包含Z值的行
z_rows = source_df[source_df['has_z'] == True]

方案二：PySpark实现 #

PySpark需用UDF处理分布式数据，避免collect遍历：

步骤1：定义UDF并注册 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import StructType, StructField, BooleanType, ArrayType, DoubleType, StringType

# 初始化Spark会话
spark = SparkSession.builder.appName("GeometryZProcessing").getOrCreate()

# 定义UDF返回值的Schema
result_schema = StructType([
    StructField("has_z", BooleanType(), nullable=True),
    StructField("first_z_coord", ArrayType(DoubleType()), nullable=True),
    StructField("geometry_2d", StructType([
        StructField("type", StringType(), nullable=True),
        StructField("coordinates", ArrayType(ArrayType(ArrayType(DoubleType()))), nullable=True)
    ]))
])

def process_geo_spark(geo):
    has_z = False
    first_z_coord = None
    processed_coords = []
    
    if geo['type'] == 'MultiPolygon':
        for polygon in geo['coordinates']:
            processed_polygon = []
            for ring in polygon:
                processed_ring = []
                for coord in ring:
                    if len(coord) >=3:
                        if not has_z:
                            has_z = True
                            first_z_coord = coord
                        processed_ring.append(coord[:2])
                    else:
                        processed_ring.append(coord)
                processed_polygon.append(processed_ring)
            processed_coords.append(processed_polygon)
    elif geo['type'] == 'Polygon':
        processed_polygon = []
        for ring in geo['coordinates']:
            processed_ring = []
            for coord in ring:
                if len(coord) >=3:
                    if not has_z:
                        has_z = True
                        first_z_coord = coord
                    processed_ring.append(coord[:2])
                else:
                    processed_ring.append(coord)
            processed_polygon.append(processed_ring)
        processed_coords.append(processed_polygon)
    
    processed_geo = {
        'type': geo['type'],
        'coordinates': processed_coords
    }
    
    return (has_z, first_z_coord, processed_geo)

# 注册UDF
geo_process_udf = udf(process_geo_spark, result_schema)

步骤2：处理DataFrame并筛选 #

# 应用UDF并展开结果字段
processed_df = source_df.withColumn("geo_result", geo_process_udf(source_df['geometry'])) \
    .select(
        "*",
        "geo_result.has_z",
        "geo_result.first_z_coord",
        "geo_result.geometry_2d"
    )

# 筛选含Z值的行
z_rows_spark = processed_df.filter(processed_df.has_z == True)

关键说明 #

• 两种方案均适配Polygon和MultiPolygon的嵌套结构，确保所有坐标环都被处理
• first_z_coord仅返回该行中第一个出现的带Z坐标，无Z值则为None
• 处理后的geometry_2d为纯2D几何数据，可直接用于空间分析

内容的提问来源于stack exchange，提问作者SPena