work in process

Author: rustyconover

src/query_farm_server_base/action_decoders.py

@@ -1,5 +1,6 @@
 from typing import Any, Literal, TypeVar, get_args, get_origin  # noqa: UP035
 
+import io
 import msgpack
 import pyarrow as pa
 import pyarrow.flight as flight
@@ -17,6 +18,12 @@ def serialize_record_batch(value: pa.RecordBatch, _info: Any) -> bytes | None:
     return sink.getvalue().to_pybytes()
 
 
+def serialize_schema(value: pa.Schema, _info: Any) -> bytes | None:
+    sink = io.BytesIO()
+    pa.ipc.write_schema(value, sink)
+    return sink.getvalue()
+
+
 def serialize_flight_descriptor(value: flight.FlightDescriptor, _info: Any) -> bytes:
     return value.serialize()
 
@@ -195,11 +202,16 @@ class EndpointsParametersParameters(BaseModel):
     column_ids: list[int]
 
     table_function_parameters: pa.RecordBatch | None
+    table_function_input_schema: pa.Schema | None
 
     _validate_table_function_parameters = field_validator(
         "table_function_parameters", mode="before"
     )(deserialize_record_batch_or_none)
 
+    _validate_table_function_input_schema = field_validator(
+        "table_function_input_schema", mode="before"
+    )(deserialize_schema_or_none)
+
 
 class EndpointsParameters(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)  # for Pydantic v2