fix(table): enforce declared schema for file registration

mplang/backends/table_impl.py

@@ -22,7 +22,7 @@
 import base64
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, ClassVar, Protocol, Self, runtime_checkable
+from typing import Any, ClassVar, Protocol, Self, cast, runtime_checkable
 
 import duckdb
 import pandas as pd
@@ -383,6 +383,41 @@ def close(self) -> None:
 }
 
 
+_duckdb_type_mapping = {
+    pa.bool_(): "BOOLEAN",
+    pa.int8(): "TINYINT",
+    pa.int16(): "SMALLINT",
+    pa.int32(): "INTEGER",
+    pa.int64(): "BIGINT",
+    pa.uint8(): "UTINYINT",
+    pa.uint16(): "USMALLINT",
+    pa.uint32(): "UINTEGER",
+    pa.uint64(): "UBIGINT",
+    pa.float32(): "FLOAT",
+    pa.float64(): "DOUBLE",
+    pa.string(): "VARCHAR",
+    pa.date64(): "DATE",
+    pa.time32("ms"): "TIME",
+    pa.timestamp("ms"): "TIMESTAMP_MS",
+    pa.binary(): "BLOB",
+    pa.json_(): "JSON",
+}
+
+
+def _duckdb_type_name(data_type: pa.DataType) -> str:
+    """Return the DuckDB SQL type corresponding to an Arrow schema type."""
+    if pa.types.is_decimal(data_type):
+        return f"DECIMAL({data_type.precision}, {data_type.scale})"
+    if data_type in _duckdb_type_mapping:
+        return _duckdb_type_mapping[data_type]
+    raise ValueError(f"Schema type {data_type} is not supported by DuckDB SQL")
+
+
+def _quote_identifier(identifier: str) -> str:
+    """Quote a SQL identifier using DuckDB's double-quote escaping rules."""
+    return f'"{identifier.replace(chr(34), chr(34) * 2)}"'
+
+
 def _pa_schema(s: elt.TableType) -> pa.Schema:
     fields = []
     for k, v in s.schema.items():
@@ -404,25 +439,37 @@ class FileTableSource(TableSource):
     def register(
         self, conn: duckdb.DuckDBPyConnection, name: str, replace: bool = True
     ) -> None:
-        """Register the file as a view in DuckDB."""
-        func_name = ""
-        match self.format:
+        """Register the file as a view while enforcing its declared schema."""
+        match self.format.lower():
             case "parquet":
-                func_name = "read_parquet"
+                relation = conn.read_parquet([self.path])
             case "csv":
-                func_name = "read_csv_auto"
+                relation = conn.from_csv_auto(self.path)
             case "json":
-                func_name = "read_json_auto"
+                relation = conn.read_json(self.path)
             case _:
                 raise ValueError(f"Unsupported format: {self.format}")
 
-        safe_path = self.path.replace("'", "''")
-        base_query = f"SELECT * FROM {func_name}('{safe_path}')"
-        if replace:
-            query = f"CREATE OR REPLACE VIEW {name} AS {base_query}"
-        else:
-            query = f"CREATE VIEW {name} AS {base_query}"
-        conn.execute(query)
+        if self.schema:
+            missing_columns = [
+                column for column in self.schema.names if column not in relation.columns
+            ]
+            if missing_columns:
+                raise ValueError(
+                    f"Cannot apply schema to {self.path}: missing columns "
+                    f"{missing_columns}; available columns are {relation.columns}"
+                )
+
+            projections = []
+            for field in self.schema:
+                identifier = _quote_identifier(field.name)
+                projections.append(
+                    f"CAST({identifier} AS {_duckdb_type_name(field.type)}) "
+                    f"AS {identifier}"
+                )
+            relation = relation.project(", ".join(projections))
+
+        relation.create_view(name, replace=replace)
 
     def open(self, batch_size: int = DEFAULT_BATCH_SIZE) -> TableReader:
         """Create a streaming reader for the file."""
@@ -681,6 +728,22 @@ def table2tensor_impl(interpreter: Interpreter, op: Operation, table_val: Any) -
         )
 
     arr = table_to_numpy(tbl)
+
+    output_type = op.outputs[0].type
+    if not isinstance(output_type, elt.TensorType) or not isinstance(
+        output_type.element_type, elt.ScalarType
+    ):
+        raise TypeError(f"table2tensor: expected scalar TensorType, got {output_type}")
+
+    from mplang.dialects import dtypes
+
+    expected_dtype = dtypes.to_numpy(cast(elt.ScalarType, output_type.element_type))
+    if arr.dtype != expected_dtype:
+        raise TypeError(
+            f"table2tensor: runtime dtype {arr.dtype} does not match "
+            f"declared tensor dtype {expected_dtype}"
+        )
+
     return TensorValue.wrap(arr)
 
 

tests/backends/test_table_impl.py

@@ -18,6 +18,7 @@
 from types import SimpleNamespace
 
 import duckdb
+import numpy as np
 import pyarrow as pa
 import pytest
 
@@ -588,6 +589,141 @@ def test_file_table_source():
             pass  # This should raise ValueError during open
 
 
+@pytest.mark.parametrize("format_name", ["parquet", "csv", "json"])
+def test_file_table_source_register_applies_declared_schema(tmp_path, format_name):
+    """SQL registration and direct reads must enforce the same schema."""
+    import pyarrow.csv as pa_csv
+    import pyarrow.parquet as pq
+
+    from mplang.backends.table_impl import FileTableSource
+
+    path = tmp_path / f"labels.{format_name}"
+    inferred_table = pa.table({"id": [1, 2], "y": [0, 1]})
+    match format_name:
+        case "parquet":
+            pq.write_table(inferred_table, path)
+        case "csv":
+            pa_csv.write_csv(inferred_table, path)
+        case "json":
+            inferred_table.to_pandas().to_json(path, orient="records", lines=True)
+
+    schema = pa.schema([("id", pa.int64()), ("y", pa.float64())])
+    expected = pa.table({
+        "id": pa.array([1, 2], type=pa.int64()),
+        "y": pa.array([0.0, 1.0], type=pa.float64()),
+    })
+    source = FileTableSource(str(path), format_name, schema)
+
+    with source.open() as reader:
+        assert reader.read_all() == expected
+
+    with duckdb.connect() as conn:
+        source.register(conn, "input_table")
+        assert conn.table("input_table").fetch_arrow_table() == expected
+
+
+def test_file_table_source_register_escapes_path_and_identifiers(tmp_path):
+    """File paths and schema identifiers are not interpolated as raw SQL."""
+    import pyarrow.csv as pa_csv
+
+    from mplang.backends.table_impl import FileTableSource
+
+    column_name = 'label"value'
+    view_name = 'input"table'
+    path = tmp_path / "labels'quoted.csv"
+    pa_csv.write_csv(pa.table({column_name: [0, 1]}), path)
+    source = FileTableSource(str(path), "csv", pa.schema([(column_name, pa.float64())]))
+
+    with duckdb.connect() as conn:
+        source.register(conn, view_name)
+        result = conn.sql('SELECT * FROM "input""table"').fetch_arrow_table()
+
+    assert result == pa.table({column_name: pa.array([0.0, 1.0], type=pa.float64())})
+
+
+def test_file_table_source_register_rejects_missing_schema_column(tmp_path):
+    import pyarrow.csv as pa_csv
+
+    from mplang.backends.table_impl import FileTableSource
+
+    path = tmp_path / "missing.csv"
+    pa_csv.write_csv(pa.table({"y": [0, 1]}), path)
+    source = FileTableSource(
+        str(path),
+        "csv",
+        pa.schema([("y", pa.float64()), ("missing", pa.float64())]),
+    )
+
+    with duckdb.connect() as conn:
+        with pytest.raises(ValueError, match=r"missing columns.*missing"):
+            source.register(conn, "input_table")
+
+
+def test_file_table_source_register_reports_invalid_schema_conversion(tmp_path):
+    import pyarrow.csv as pa_csv
+
+    from mplang.backends.table_impl import FileTableSource
+
+    path = tmp_path / "invalid.csv"
+    pa_csv.write_csv(pa.table({"y": ["0", "invalid"]}), path)
+    source = FileTableSource(str(path), "csv", pa.schema([("y", pa.float64())]))
+
+    with duckdb.connect() as conn:
+        source.register(conn, "input_table")
+        with pytest.raises(duckdb.ConversionException, match=r"invalid.*DOUBLE"):
+            conn.table("input_table").fetch_arrow_table()
+
+
+def test_csv_schema_survives_sql_table2tensor_compile_evaluate(tmp_path):
+    """Regression for catalog f64 labels whose CSV values look integral."""
+    import pyarrow.csv as pa_csv
+
+    import mplang as mp
+    from mplang.backends.tensor_impl import TensorValue
+    from mplang.runtime.interpreter import Interpreter
+
+    path = tmp_path / "labels.csv"
+    pa_csv.write_csv(pa.table({"feature": [10, 20], "y": [0, 1]}), path)
+    input_type = elt.TableType({"feature": elt.i64, "y": elt.f64})
+    output_type = elt.TableType({"y": elt.f64})
+
+    def workload():
+        source = table.read(str(path), schema=input_type)
+        selected = table.run_sql(
+            "SELECT y FROM source", out_type=output_type, source=source
+        )
+        return table.table2tensor(selected, number_rows=-1)
+
+    interpreter = Interpreter()
+    program = mp.compile(workload, context=interpreter)
+    result = mp.evaluate(program, context=interpreter)
+
+    assert program.graph.outputs[0].type == elt.TensorType(elt.f64, (-1, 1))
+    assert isinstance(result.runtime_obj, TensorValue)
+    np.testing.assert_array_equal(
+        result.runtime_obj.unwrap(),
+        np.array([[0.0], [1.0]], dtype=np.float64),
+    )
+
+
+def test_table2tensor_rejects_runtime_dtype_different_from_declared_type():
+    """Backend values cannot silently violate table2tensor's IR output type."""
+
+    def workload():
+        source = table.constant({"y": [0, 1]})
+        declared_float = table.run_sql(
+            "SELECT y FROM source",
+            out_type=elt.TableType({"y": elt.f64}),
+            source=source,
+        )
+        return table.table2tensor(declared_float, number_rows=-1)
+
+    with pytest.raises(
+        TypeError, match=r"runtime dtype int64.*declared tensor dtype float64"
+    ):
+        workload()
+
+
 def test_write_basic(tmp_path):
     filename = tmp_path / "test_write_basic"
     data = {"x": [10, 20], "y": ["foo", "bar"]}