如何让Arrow Flight服务器接收多Schema的GeoDataFrame与DataFrame？

针对你需要发送geopandas.GeoDataFrame和pandas.DataFrame两种不同Schema数据到Arrow Flight服务器的需求，这里提供三种可行的实现方式：

方案一：添加类型标记列合并为单表发送 #

给两个数据集分别添加一个用于区分的标记列，合并成一个Arrow表后发送，服务器接收后再根据标记拆分出原始数据集。

客户端代码示例 #

import pyarrow as pa
import pandas as pd
import geopandas as gpd
from pyarrow.flight import FlightClient, FlightDescriptor

# 构造测试数据
geo_df = gpd.GeoDataFrame({'name': ['A', 'B'], 'geometry': gpd.points_from_xy([1,2], [3,4])})
coords_df = pd.DataFrame({'x': [10,20], 'y': [30,40]})

# 添加类型标记列
geo_df['dataset_type'] = 'geo'
coords_df['dataset_type'] = 'coords'

# 转换为Arrow表并合并
geo_table = pa.Table.from_pandas(geo_df)
coords_table = pa.Table.from_pandas(coords_df)
combined_table = pa.concat_tables([geo_table, coords_table])

# 发送到服务器
client = FlightClient("grpc://localhost:50051")
descriptor = FlightDescriptor.for_path("combined_data")
writer, _ = client.do_put(descriptor, combined_table.schema)
writer.write_table(combined_table)
writer.close()

服务器代码示例 #

import pyarrow as pa
from pyarrow.flight import FlightServerBase, FlightDescriptor, FlightInfo, SchemaResult, FlightData

class CustomFlightServer(FlightServerBase):
    def __init__(self, location):
        super().__init__(location)
        self.geo_data = None
        self.coords_data = None

    def do_put(self, context, descriptor, reader, writer):
        # 读取合并后的表
        combined_table = reader.read_all()
        # 根据标记列拆分
        self.geo_data = combined_table.filter(pa.field('dataset_type') == 'geo').drop(['dataset_type'])
        self.coords_data = combined_table.filter(pa.field('dataset_type') == 'coords').drop(['dataset_type'])
        # 执行后续计算逻辑
        self._run_calculation()

    def _run_calculation(self):
        # 这里写你的内存计算逻辑
        if self.geo_data is not None and self.coords_data is not None:
            print("Geo数据行数:", self.geo_data.num_rows)
            print("坐标数据行数:", self.coords_data.num_rows)

if __name__ == "__main__":
    server = CustomFlightServer("grpc://localhost:50051")
    print("Server starting on port 50051...")
    server.serve()

方案二：分两次调用do_put，用FlightDescriptor区分 #

通过两次独立的do_put请求，用不同的FlightDescriptor（比如path参数）标记不同的数据集，服务器根据descriptor将数据存储到不同的内存变量中，待两个数据集都接收完成后执行计算。

客户端代码示例 #

import pyarrow as pa
import pandas as pd
import geopandas as gpd
from pyarrow.flight import FlightClient, FlightDescriptor

# 构造测试数据
geo_df = gpd.GeoDataFrame({'name': ['A', 'B'], 'geometry': gpd.points_from_xy([1,2], [3,4])})
coords_df = pd.DataFrame({'x': [10,20], 'y': [30,40]})

# 转换为Arrow表
geo_table = pa.Table.from_pandas(geo_df)
coords_table = pa.Table.from_pandas(coords_df)

# 发送Geo数据
client = FlightClient("grpc://localhost:50051")
geo_descriptor = FlightDescriptor.for_path("geo_data")
writer, _ = client.do_put(geo_descriptor, geo_table.schema)
writer.write_table(geo_table)
writer.close()

# 发送坐标数据
coords_descriptor = FlightDescriptor.for_path("coords_data")
writer, _ = client.do_put(coords_descriptor, coords_table.schema)
writer.write_table(coords_table)
writer.close()

服务器代码示例 #

import pyarrow as pa
from pyarrow.flight import FlightServerBase, FlightDescriptor

class CustomFlightServer(FlightServerBase):
    def __init__(self, location):
        super().__init__(location)
        self.geo_data = None
        self.coords_data = None

    def do_put(self, context, descriptor, reader, writer):
        # 根据descriptor的path判断数据类型
        path = descriptor.path[0]
        if path == "geo_data":
            self.geo_data = reader.read_all()
        elif path == "coords_data":
            self.coords_data = reader.read_all()
        
        # 检查是否两个数据都已接收，执行计算
        if self.geo_data is not None and self.coords_data is not None:
            self._run_calculation()

    def _run_calculation(self):
        # 内存计算逻辑
        print("开始计算...")
        # 示例：关联Geo数据和坐标数据的逻辑
        # ...

if __name__ == "__main__":
    server = CustomFlightServer("grpc://localhost:50051")
    print("Server starting on port 50051...")
    server.serve()

方案三：用Arrow Struct类型封装两个数据集 #

将两个数据集封装到一个Arrow Struct类型的字段中，作为单个表发送，服务器接收后提取Struct字段中的两个子表。

客户端代码示例 #

import pyarrow as pa
import pandas as pd
import geopandas as gpd
from pyarrow.flight import FlightClient, FlightDescriptor

# 构造测试数据
geo_df = gpd.GeoDataFrame({'name': ['A', 'B'], 'geometry': gpd.points_from_xy([1,2], [3,4])})
coords_df = pd.DataFrame({'x': [10,20], 'y': [30,40]})

# 转换为Arrow RecordBatch
geo_batch = pa.RecordBatch.from_pandas(geo_df)
coords_batch = pa.RecordBatch.from_pandas(coords_df)

# 构造Struct类型的字段
struct_type = pa.struct([
    ('geo_data', geo_batch.schema),
    ('coords_data', coords_batch.schema)
])
# 创建包含单个Struct值的RecordBatch
batch = pa.RecordBatch.from_arrays(
    [pa.array([{'geo_data': geo_batch.to_pylist()[0], 'coords_data': coords_batch.to_pylist()[0]},
               {'geo_data': geo_batch.to_pylist()[1], 'coords_data': coords_batch.to_pylist()[1]}], type=struct_type)],
    names=['combined']
)
table = pa.Table.from_batches([batch])

# 发送到服务器
client = FlightClient("grpc://localhost:50051")
descriptor = FlightDescriptor.for_path("struct_combined")
writer, _ = client.do_put(descriptor, table.schema)
writer.write_table(table)
writer.close()

服务器代码示例 #

import pyarrow as pa
from pyarrow.flight import FlightServerBase, FlightDescriptor

class CustomFlightServer(FlightServerBase):
    def __init__(self, location):
        super().__init__(location)
        self.geo_data = None
        self.coords_data = None

    def do_put(self, context, descriptor, reader, writer):
        table = reader.read_all()
        # 提取Struct字段中的数据
        geo_arrays = table['combined'].field('geo_data').flatten()
        coords_arrays = table['combined'].field('coords_data').flatten()
        
        # 重建原始表
        self.geo_data = pa.Table.from_arrays(geo_arrays, schema=geo_arrays.type.schema)
        self.coords_data = pa.Table.from_arrays(coords_arrays, schema=coords_arrays.type.schema)
        
        # 执行计算
        self._run_calculation()

    def _run_calculation(self):
        print("Geo数据Schema:", self.geo_data.schema)
        print("坐标数据Schema:", self.coords_data.schema)

if __name__ == "__main__":
    server = CustomFlightServer("grpc://localhost:50051")
    print("Server starting on port 50051...")
    server.serve()

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