mplang-nightly 0.1.dev334__tar.gz → 0.1.dev336__tar.gz

{mplang_nightly-0.1.dev334 → mplang_nightly-0.1.dev336}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev334
+Version: 0.1.dev336
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev334 → mplang_nightly-0.1.dev336}/mplang/backends/util.py

@@ -101,21 +101,16 @@ def _refine_stablehlo(code: str, args: tuple[TensorLike, ...]) -> str:
         cmd = [
             "stablehlo-opt",
             input_file,
+            "--inline",
             f"--stablehlo-refine-arguments={refine_args}",
             "--stablehlo-refine-shapes",
             "--stablehlo-canonicalize-dynamism",
+            "--stablehlo-check-shape-assertions",
         ]
 
         result = subprocess.run(cmd, capture_output=True, text=True, check=True)
 
-        # Filter out shape_assertion custom calls that cause compilation errors
-        filtered_output = "\n".join(
-            line
-            for line in result.stdout.split("\n")
-            if "stablehlo.custom_call @shape_assertion" not in line
-        )
-
-        return filtered_output
+        return result.stdout
 
     except subprocess.CalledProcessError as e:
         # Re-throw with additional information

{mplang_nightly-0.1.dev334 → mplang_nightly-0.1.dev336}/mplang/edsl/serde.py

@@ -144,6 +144,7 @@ def to_json(obj: Any) -> dict[str, Any]:
     3. Be a list/tuple of serializable objects
     4. Be a dict with string keys and serializable values
     5. Be a numpy ndarray
+    6. Be a pyarrow Table or pandas DataFrame (round-tripped via Arrow IPC)
 
     Args:
         obj: Object to serialize
@@ -196,6 +197,36 @@ def to_json(obj: Any) -> dict[str, Any]:
             "data": base64.b64encode(obj.tobytes()).decode("ascii"),
         }
 
+    # Tabular data with named columns. Both pyarrow Table and pandas
+    # DataFrame expose ``__array__`` and would otherwise fall into the
+    # ndarray branch below, losing column names and dtype info. Use pyarrow
+    # IPC for lossless round-trip. Trigger: any op-attr carrying tabular
+    # data (e.g. ``ibis.memtable({...})`` placed via ``mp.put`` →
+    # ``simp.constant`` → ``table.constant_p``).
+    try:
+        import pyarrow as pa
+    except ImportError:  # pragma: no cover - pyarrow is a hard dep elsewhere
+        pa = None
+    try:
+        import pandas as pd
+    except ImportError:  # pragma: no cover
+        pd = None
+
+    if pa is not None and isinstance(obj, pa.Table):
+        table_obj: Any = obj
+    elif pa is not None and pd is not None and isinstance(obj, pd.DataFrame):
+        table_obj = pa.Table.from_pandas(obj, preserve_index=False)
+    else:
+        table_obj = None
+    if table_obj is not None:
+        sink = pa.BufferOutputStream()
+        with pa.ipc.new_stream(sink, table_obj.schema) as writer:
+            writer.write_table(table_obj)
+        return {
+            "_kind": "_pa_table",
+            "data": base64.b64encode(sink.getvalue().to_pybytes()).decode("ascii"),
+        }
+
     # Array-like (JAX, etc.) - convert to numpy
     if hasattr(obj, "__array__"):
         return to_json(np.asarray(obj))
@@ -298,6 +329,13 @@ def from_json(data: dict[str, Any]) -> Any:
         # Always reshape - empty tuple () means scalar, which requires reshape
         return arr.reshape(shape)
 
+    # Tabular data round-tripped via pyarrow IPC (see to_json companion).
+    if kind == "_pa_table":
+        import pyarrow as pa
+
+        buf = base64.b64decode(data["data"])
+        return pa.ipc.open_stream(pa.BufferReader(buf)).read_all()
+
     # Registered classes
     if kind in _CLASS_REGISTRY:
         cls = _CLASS_REGISTRY[kind]
@@ -485,6 +523,30 @@ def _to_binary_json(obj: Any, ctx: BinaryContext) -> dict[str, Any]:
         seg_idx = ctx.add_segment(obj)
         return {"_kind": "_bytes_bref", "bref": seg_idx}
 
+    # Tabular data (parity with to_json) — emit Arrow IPC bytes as a
+    # binary segment instead of inlining base64. Must come before the
+    # ``__array__`` check, since pa.Table / pd.DataFrame both expose it.
+    try:
+        import pyarrow as pa
+    except ImportError:  # pragma: no cover
+        pa = None
+    try:
+        import pandas as pd
+    except ImportError:  # pragma: no cover
+        pd = None
+
+    table_obj: Any = None
+    if pa is not None and isinstance(obj, pa.Table):
+        table_obj = obj
+    elif pa is not None and pd is not None and isinstance(obj, pd.DataFrame):
+        table_obj = pa.Table.from_pandas(obj, preserve_index=False)
+    if table_obj is not None:
+        sink = pa.BufferOutputStream()
+        with pa.ipc.new_stream(sink, table_obj.schema) as writer:
+            writer.write_table(table_obj)
+        seg_idx = ctx.add_segment(sink.getvalue().to_pybytes())
+        return {"_kind": "_pa_table_bref", "bref": seg_idx}
+
     # Array-like (e.g., JAX arrays)
     if hasattr(obj, "__array__"):
         return _to_binary_json(np.asarray(obj), ctx)
@@ -538,6 +600,14 @@ def _from_binary_json(data: dict[str, Any], segments: list[memoryview]) -> Any:
         arr = np.frombuffer(raw, dtype=np.dtype(data["dtype"]))
         return arr.reshape(tuple(data["shape"])).copy()
 
+    if kind == "_pa_table_bref":
+        import pyarrow as pa
+
+        # ``pa.BufferReader`` accepts any buffer-protocol object, so we pass
+        # the memoryview through to avoid an extra ``bytes(...)`` copy.
+        buf = segments[data["bref"]]
+        return pa.ipc.open_stream(pa.BufferReader(buf)).read_all()
+
     if kind == "_list":
         return [_from_binary_json(item, segments) for item in data["items"]]
 

{mplang_nightly-0.1.dev334 → mplang_nightly-0.1.dev336}/tests/edsl/test_serde.py

@@ -275,6 +275,85 @@ class TestBinaryWireFormat:
         assert result == payload
 
 
+# =============================================================================
+# Tests: Tabular Data (pyarrow.Table / pandas.DataFrame)
+# =============================================================================
+
+
+class TestArrowTableRoundTrip:
+    """Tabular data must round-trip losslessly, preserving column names + dtypes.
+
+    Both ``pa.Table`` and ``pd.DataFrame`` expose ``__array__``. Before the
+    fix they fell into the numpy ndarray branch which collapsed the columns
+    into a single 2-D array, losing names and per-column dtypes.
+    """
+
+    def test_to_json_pa_table_roundtrip(self):
+        import pyarrow as pa
+
+        t = pa.table({
+            "bin_id": [0, 1, 2, 3],
+            "edge": [0.1, 0.5, 0.9, 1.0],
+            "name": ["a", "b", "c", "d"],
+        })
+        payload = serde.to_json(t)
+        assert payload["_kind"] == "_pa_table"
+
+        result = serde.from_json(payload)
+        assert isinstance(result, pa.Table)
+        assert result.schema.equals(t.schema)
+        assert result.equals(t)
+
+    def test_to_json_pandas_dataframe_roundtrip_as_pa_table(self):
+        pd = pytest.importorskip("pandas")
+        import pyarrow as pa
+
+        df = pd.DataFrame({
+            "bin_id": [0, 1, 2, 3],
+            "edge": [0.1, 0.5, 0.9, 1.0],
+        })
+        payload = serde.to_json(df)
+        assert payload["_kind"] == "_pa_table"
+
+        result = serde.from_json(payload)
+        # Deserialized as pa.Table (canonical tabular form). Column names and
+        # numeric dtypes must survive.
+        assert isinstance(result, pa.Table)
+        assert result.column_names == ["bin_id", "edge"]
+        assert pa.Table.from_pandas(df, preserve_index=False).equals(result)
+
+    def test_dumps_binary_pa_table_uses_segment_reference(self):
+        import pyarrow as pa
+
+        t = pa.table({"x": [1, 2, 3], "y": [10.0, 20.0, 30.0]})
+        serialized = serde.dumps_binary(t)
+        meta_len = struct.unpack_from("<I", serialized, 0)[0]
+        meta = json.loads(serialized[4 : 4 + meta_len].decode("utf-8"))
+
+        assert meta["_kind"] == "_pa_table_bref"
+        assert "bref" in meta
+        assert "data" not in meta  # binary path must not base64-inline
+
+        result = serde.loads_binary(serialized)
+        assert isinstance(result, pa.Table)
+        assert result.equals(t)
+
+    def test_dumps_binary_pandas_dataframe_uses_segment_reference(self):
+        pd = pytest.importorskip("pandas")
+        import pyarrow as pa
+
+        df = pd.DataFrame({"a": [1, 2], "b": [3.5, 4.5]})
+        serialized = serde.dumps_binary(df)
+        meta_len = struct.unpack_from("<I", serialized, 0)[0]
+        meta = json.loads(serialized[4 : 4 + meta_len].decode("utf-8"))
+
+        assert meta["_kind"] == "_pa_table_bref"
+
+        result = serde.loads_binary(serialized)
+        assert isinstance(result, pa.Table)
+        assert result.column_names == ["a", "b"]
+
+
 # =============================================================================
 # Tests: Error Handling
 # =============================================================================