diff --git a/mplang/backends/table_impl.py b/mplang/backends/table_impl.py index d5eae462..49aa9145 100644 --- a/mplang/backends/table_impl.py +++ b/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) diff --git a/tests/backends/test_table_impl.py b/tests/backends/test_table_impl.py index c3176163..5cc618a5 100644 --- a/tests/backends/test_table_impl.py +++ b/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"]}