Support Python->Rust Struct/Table type bindings. #19

python/cocoindex/op.py

@@ -1,10 +1,10 @@
 """
 Facilities for defining cocoindex operations.
 """
+import dataclasses
 import inspect
 
 from typing import get_type_hints, Protocol, Any, Callable, dataclass_transform
-from dataclasses import dataclass
 from enum import Enum
 from threading import Lock
 
@@ -28,7 +28,7 @@ def __new__(mcs, name, bases, attrs, category: OpCategory | None = None):
             setattr(cls, '_op_category', category)
         else:
             # It's the specific class providing specific fields.
-            cls = dataclass(cls)
+            cls = dataclasses.dataclass(cls)
         return cls
 
 class SourceSpec(metaclass=SpecMeta, category=OpCategory.SOURCE): # pylint: disable=too-few-public-methods
@@ -59,6 +59,14 @@ def __call__(self, spec: dict[str, Any], *args, **kwargs):
         result_type = executor.analyze(*args, **kwargs)
         return (dump_type(result_type), executor)
 
+def to_engine_value(value: Any) -> Any:
+    """Convert a Python value to an engine value."""
+    if dataclasses.is_dataclass(value):
+        return [to_engine_value(getattr(value, f.name)) for f in dataclasses.fields(value)]
+    elif isinstance(value, list) or isinstance(value, tuple):
+        return [to_engine_value(v) for v in value]
+    return value
+
 _gpu_dispatch_lock = Lock()
 
 def executor_class(gpu: bool = False, cache: bool = False, behavior_version: int | None = None) -> Callable[[type], type]:
@@ -162,9 +170,10 @@ def __call__(self, *args, **kwargs):
                     # For now, we use a lock to ensure only one task is executed at a time.
                     # TODO: Implement multi-processing dispatching.
                     with _gpu_dispatch_lock:
-                        return super().__call__(*args, **kwargs)
+                        output = super().__call__(*args, **kwargs)
                 else:
-                    return super().__call__(*args, **kwargs)
+                    output = super().__call__(*args, **kwargs)
+                return to_engine_value(output)
 
         _WrappedClass.__name__ = cls.__name__
 

src/ops/py_factory.rs

@@ -1,18 +1,19 @@
-use std::sync::Arc;
+use std::{collections::BTreeMap, sync::Arc};
 
 use axum::async_trait;
 use blocking::unblock;
 use futures::FutureExt;
 use pyo3::{
     exceptions::PyException,
     pyclass, pymethods,
-    types::{IntoPyDict, PyAnyMethods, PyString, PyTuple},
+    types::{IntoPyDict, PyAnyMethods, PyList, PyString, PyTuple},
     Bound, IntoPyObjectExt, Py, PyAny, PyResult, Python,
 };
 
 use crate::{
     base::{schema, value},
     builder::plan,
+    py::IntoPyResult,
 };
 use anyhow::Result;
 
@@ -89,6 +90,28 @@ fn basic_value_from_py_object<'py>(
     Ok(result)
 }
 
+fn field_values_from_py_object<'py>(
+    schema: &schema::StructSchema,
+    v: &Bound<'py, PyAny>,
+) -> PyResult<value::FieldValues> {
+    let list = v.extract::<Vec<Bound<'py, PyAny>>>()?;
+    if list.len() != schema.fields.len() {
+        return Err(PyException::new_err(format!(
+            "struct field number mismatch, expected {}, got {}",
+            schema.fields.len(),
+            list.len()
+        )));
+    }
+    Ok(value::FieldValues {
+        fields: schema
+            .fields
+            .iter()
+            .zip(list.into_iter())
+            .map(|(f, v)| value_from_py_object(&f.value_type.typ, &v))
+            .collect::<PyResult<Vec<_>>>()?,
+    })
+}
+
 fn value_from_py_object<'py>(
     typ: &schema::ValueType,
     v: &Bound<'py, PyAny>,
@@ -100,6 +123,39 @@ fn value_from_py_object<'py>(
             schema::ValueType::Basic(typ) => {
                 value::Value::Basic(basic_value_from_py_object(typ, v)?)
             }
+            schema::ValueType::Struct(schema) => {
+                value::Value::Struct(field_values_from_py_object(schema, v)?)
+            }
+            schema::ValueType::Collection(schema) => {
+                let list = v.extract::<Vec<Bound<'py, PyAny>>>()?;
+                let values = list
+                    .into_iter()
+                    .map(|v| field_values_from_py_object(&schema.row, &v))
+                    .collect::<PyResult<Vec<_>>>()?;
+                match schema.kind {
+                    schema::CollectionKind::Collection => {
+                        value::Value::Collection(values.into_iter().map(|v| v.into()).collect())
+                    }
+                    schema::CollectionKind::List => {
+                        value::Value::List(values.into_iter().map(|v| v.into()).collect())
+                    }
+                    schema::CollectionKind::Table => value::Value::Table(
+                        values
+                            .into_iter()
+                            .map(|v| {
+                                let mut iter = v.fields.into_iter();
+                                let key = iter.next().unwrap().to_key().into_py_result()?;
+                                Ok((
+                                    key,
+                                    value::ScopeValue(value::FieldValues {
+                                        fields: iter.collect::<Vec<_>>(),
+                                    }),
+                                ))
+                            })
+                            .collect::<PyResult<BTreeMap<_, _>>>()?,
+                    ),
+                }
+            }
             _ => {
                 return Err(PyException::new_err(format!(
                     "unsupported value type: {}",