datablade 0.0.0__py3-none-any.whl → 0.0.6__py3-none-any.whl

datablade/sql/ddl_pyarrow.py

@@ -0,0 +1,411 @@
+"""Parquet schema-driven DDL generation using PyArrow."""
+
+from __future__ import annotations
+
+import logging
+import pathlib
+from dataclasses import dataclass
+from typing import Any, List, Mapping, Optional, Union
+
+from ..utils.messages import print_verbose
+from ..utils.strings import coerce_path
+from .ddl import _qualify_name
+from .dialects import Dialect
+from .quoting import quote_identifier
+from .schema_spec import resolve_column_spec
+
+logger = logging.getLogger("datablade")
+
+
+@dataclass(frozen=True)
+class DroppedColumn:
+    """Metadata about a dropped column during Parquet DDL generation."""
+
+    name: str
+    arrow_type: str
+    reason: str
+
+
+@dataclass(frozen=True)
+class FallbackColumn:
+    """Metadata about a column handled via JSON fallback."""
+
+    name: str
+    arrow_type: str
+    sql_type: str
+
+
+@dataclass(frozen=True)
+class ParquetDDLMetadata:
+    """Details about columns dropped or handled via fallback."""
+
+    dropped_columns: List[DroppedColumn]
+    fallback_columns: List[FallbackColumn]
+
+
+def _require_pyarrow():
+    """Import pyarrow lazily to keep core dependencies light."""
+    try:
+        import pyarrow as pa  # type: ignore
+        import pyarrow.parquet as pq  # type: ignore
+    except ImportError as exc:  # pragma: no cover
+        raise ImportError(
+            "Parquet DDL generation requires 'pyarrow'. Install with: pip install pyarrow"
+        ) from exc
+
+    return pa, pq
+
+
+def _is_complex_arrow_type(data_type) -> bool:
+    pa, _ = _require_pyarrow()
+    return (
+        pa.types.is_struct(data_type)
+        or pa.types.is_list(data_type)
+        or pa.types.is_large_list(data_type)
+        or pa.types.is_fixed_size_list(data_type)
+        or pa.types.is_map(data_type)
+        or pa.types.is_union(data_type)
+    )
+
+
+def _json_fallback_sql_type(dialect: Dialect) -> str:
+    if dialect == Dialect.SQLSERVER:
+        return "nvarchar(max)"
+    if dialect == Dialect.POSTGRES:
+        return "text"
+    if dialect == Dialect.MYSQL:
+        return "TEXT"
+    if dialect == Dialect.DUCKDB:
+        return "VARCHAR"
+    raise NotImplementedError(f"Dialect not supported: {dialect}")
+
+
+def _sql_type_from_arrow(data_type, dialect: Dialect) -> Optional[str]:  # noqa: C901
+    """Map a pyarrow.DataType to a SQL type string.
+
+    Returns None when there is no clean mapping and the caller should drop the column.
+    """
+
+    pa, _ = _require_pyarrow()
+
+    # Dictionary-encoded columns behave like their value type for DDL purposes.
+    if pa.types.is_dictionary(data_type):
+        return _sql_type_from_arrow(data_type.value_type, dialect)
+
+    # Nested/complex types: no clean general mapping across dialects.
+    if (
+        pa.types.is_struct(data_type)
+        or pa.types.is_list(data_type)
+        or pa.types.is_large_list(data_type)
+        or pa.types.is_fixed_size_list(data_type)
+        or pa.types.is_map(data_type)
+        or pa.types.is_union(data_type)
+    ):
+        return None
+
+    if dialect == Dialect.SQLSERVER:
+        if pa.types.is_boolean(data_type):
+            return "bit"
+        if pa.types.is_int8(data_type) or pa.types.is_int16(data_type):
+            return "smallint"
+        if pa.types.is_int32(data_type):
+            return "int"
+        if pa.types.is_int64(data_type):
+            return "bigint"
+        if pa.types.is_uint8(data_type) or pa.types.is_uint16(data_type):
+            return "int"
+        if pa.types.is_uint32(data_type):
+            return "bigint"
+        if pa.types.is_uint64(data_type):
+            return "decimal(20, 0)"
+        if pa.types.is_float16(data_type) or pa.types.is_float32(data_type):
+            return "real"
+        if pa.types.is_float64(data_type):
+            return "float"
+        if pa.types.is_decimal(data_type):
+            precision = min(int(data_type.precision), 38)
+            scale = int(data_type.scale)
+            return f"decimal({precision}, {scale})"
+        if pa.types.is_date(data_type):
+            return "date"
+        if pa.types.is_time(data_type):
+            return "time"
+        if pa.types.is_timestamp(data_type):
+            # SQL Server has datetimeoffset for tz-aware values.
+            return "datetimeoffset" if data_type.tz is not None else "datetime2"
+        if pa.types.is_binary(data_type) or pa.types.is_large_binary(data_type):
+            return "varbinary(max)"
+        if pa.types.is_fixed_size_binary(data_type):
+            return (
+                f"varbinary({int(data_type.byte_width)})"
+                if int(data_type.byte_width) <= 8000
+                else "varbinary(max)"
+            )
+        if pa.types.is_string(data_type) or pa.types.is_large_string(data_type):
+            return "nvarchar(max)"
+
+        # Anything else (including null) is not reliably representable.
+        return None
+
+    if dialect == Dialect.POSTGRES:
+        if pa.types.is_boolean(data_type):
+            return "boolean"
+        if pa.types.is_int8(data_type) or pa.types.is_int16(data_type):
+            return "smallint"
+        if pa.types.is_int32(data_type):
+            return "integer"
+        if pa.types.is_int64(data_type):
+            return "bigint"
+        if pa.types.is_unsigned_integer(data_type):
+            # Postgres has no unsigned ints; use a wider signed or numeric.
+            if pa.types.is_uint8(data_type) or pa.types.is_uint16(data_type):
+                return "integer"
+            if pa.types.is_uint32(data_type):
+                return "bigint"
+            if pa.types.is_uint64(data_type):
+                return "numeric(20, 0)"
+        if pa.types.is_float16(data_type) or pa.types.is_float32(data_type):
+            return "real"
+        if pa.types.is_float64(data_type):
+            return "double precision"
+        if pa.types.is_decimal(data_type):
+            precision = int(data_type.precision)
+            scale = int(data_type.scale)
+            return f"numeric({precision}, {scale})"
+        if pa.types.is_date(data_type):
+            return "date"
+        if pa.types.is_time(data_type):
+            return "time"
+        if pa.types.is_timestamp(data_type):
+            return "timestamptz" if data_type.tz is not None else "timestamp"
+        if pa.types.is_binary(data_type) or pa.types.is_large_binary(data_type):
+            return "bytea"
+        if pa.types.is_fixed_size_binary(data_type):
+            return "bytea"
+        if pa.types.is_string(data_type) or pa.types.is_large_string(data_type):
+            return "text"
+
+        return None
+
+    if dialect == Dialect.MYSQL:
+        if pa.types.is_boolean(data_type):
+            return "TINYINT(1)"
+        if pa.types.is_int8(data_type) or pa.types.is_int16(data_type):
+            return "SMALLINT"
+        if pa.types.is_int32(data_type):
+            return "INT"
+        if pa.types.is_int64(data_type):
+            return "BIGINT"
+        if pa.types.is_unsigned_integer(data_type):
+            # MySQL supports UNSIGNED, but we keep mappings consistent with the existing
+            # pandas-based DDL generator (signed types).
+            if pa.types.is_uint8(data_type) or pa.types.is_uint16(data_type):
+                return "INT"
+            if pa.types.is_uint32(data_type):
+                return "BIGINT"
+            if pa.types.is_uint64(data_type):
+                return "DECIMAL(20, 0)"
+        if pa.types.is_float16(data_type) or pa.types.is_float32(data_type):
+            return "FLOAT"
+        if pa.types.is_float64(data_type):
+            return "DOUBLE"
+        if pa.types.is_decimal(data_type):
+            precision = min(int(data_type.precision), 65)
+            scale = min(int(data_type.scale), 30, precision)
+            return f"DECIMAL({precision}, {scale})"
+        if pa.types.is_date(data_type):
+            return "DATE"
+        if pa.types.is_time(data_type):
+            return "TIME"
+        if pa.types.is_timestamp(data_type):
+            return "DATETIME"
+        if pa.types.is_binary(data_type) or pa.types.is_large_binary(data_type):
+            return "LONGBLOB"
+        if pa.types.is_fixed_size_binary(data_type):
+            width = int(data_type.byte_width)
+            return f"VARBINARY({width})" if width <= 65535 else "LONGBLOB"
+        if pa.types.is_string(data_type) or pa.types.is_large_string(data_type):
+            return "TEXT"
+
+        return None
+
+    if dialect == Dialect.DUCKDB:
+        if pa.types.is_boolean(data_type):
+            return "BOOLEAN"
+        if pa.types.is_signed_integer(data_type):
+            return "BIGINT"
+        if pa.types.is_unsigned_integer(data_type):
+            return "UBIGINT"
+        if pa.types.is_floating(data_type):
+            return "DOUBLE"
+        if pa.types.is_decimal(data_type):
+            precision = int(data_type.precision)
+            scale = int(data_type.scale)
+            return f"DECIMAL({precision}, {scale})"
+        if pa.types.is_date(data_type):
+            return "DATE"
+        if pa.types.is_time(data_type):
+            return "TIME"
+        if pa.types.is_timestamp(data_type):
+            return "TIMESTAMPTZ" if data_type.tz is not None else "TIMESTAMP"
+        if pa.types.is_binary(data_type) or pa.types.is_large_binary(data_type):
+            return "BLOB"
+        if pa.types.is_fixed_size_binary(data_type):
+            return "BLOB"
+        if pa.types.is_string(data_type) or pa.types.is_large_string(data_type):
+            return "VARCHAR"
+
+        return None
+
+    raise NotImplementedError(f"Dialect not supported: {dialect}")
+
+
+def generate_create_table_from_parquet(
+    parquet_path: str | pathlib.Path,
+    catalog: Optional[str] = None,
+    schema: Optional[str] = None,
+    table: str = "table",
+    drop_existing: bool = True,
+    dialect: Dialect = Dialect.SQLSERVER,
+    use_go: bool = False,
+    schema_spec: Optional[Mapping[str, Any]] = None,
+    verbose: bool = False,
+    fallback_to_json: bool = False,
+    return_metadata: bool = False,
+) -> Union[str, tuple[str, ParquetDDLMetadata]]:
+    """Generate a CREATE TABLE statement from a Parquet file schema.
+
+    This reads the Parquet schema only (via PyArrow) and does not materialize data.
+
+    Columns whose Parquet types have no clean mapping for the chosen dialect are
+    dropped, and a warning is logged under logger name 'datablade'. If
+    fallback_to_json is enabled, complex types are instead mapped to a text
+    column intended to store JSON-encoded values. Use return_metadata to receive
+    details about dropped and fallback-mapped columns.
+
+    When dialect is SQL Server and use_go is True, a GO batch separator is
+    inserted after a USE statement when a catalog is provided.
+
+    schema_spec may provide per-column sql_type/nullable overrides.
+    """
+
+    path_obj = coerce_path(
+        parquet_path,
+        must_exist=True,
+        verbose=verbose,
+        label="parquet_path",
+    )
+    if not isinstance(table, str) or not table.strip():
+        raise ValueError("table must be a non-empty string")
+    if catalog is not None and (not isinstance(catalog, str) or not catalog.strip()):
+        raise ValueError("catalog, if provided, must be a non-empty string")
+    if schema is not None and (not isinstance(schema, str) or not schema.strip()):
+        raise ValueError("schema, if provided, must be a non-empty string")
+    if not isinstance(use_go, bool):
+        raise TypeError("use_go must be a boolean")
+
+    _, pq = _require_pyarrow()
+
+    # Read Parquet metadata only; this does not load row data.
+    arrow_schema = pq.ParquetFile(path_obj).schema_arrow
+
+    qualified_name = _qualify_name(catalog, schema, table, dialect)
+    lines: List[str] = []
+    dropped_columns: List[DroppedColumn] = []
+    fallback_columns: List[FallbackColumn] = []
+
+    for field in arrow_schema:
+        column_name = str(field.name)
+        defaults, column_spec = resolve_column_spec(column_name, schema_spec)
+        sql_type_override = column_spec.get("sql_type")
+        if sql_type_override is not None:
+            if not isinstance(sql_type_override, str) or not sql_type_override.strip():
+                raise ValueError(
+                    f"schema_spec.columns['{column_name}'].sql_type must be a non-empty string"
+                )
+            sql_type = sql_type_override.strip()
+        else:
+            sql_type = _sql_type_from_arrow(field.type, dialect)
+
+        if sql_type is None:
+            if fallback_to_json and _is_complex_arrow_type(field.type):
+                fallback_sql_type = _json_fallback_sql_type(dialect)
+                fallback_columns.append(
+                    FallbackColumn(
+                        name=str(field.name),
+                        arrow_type=str(field.type),
+                        sql_type=fallback_sql_type,
+                    )
+                )
+                sql_type = fallback_sql_type
+            else:
+                dropped_columns.append(
+                    DroppedColumn(
+                        name=str(field.name),
+                        arrow_type=str(field.type),
+                        reason="unsupported type",
+                    )
+                )
+                logger.warning(
+                    "Dropping Parquet column %r (type=%s) for dialect=%s: unsupported type",
+                    field.name,
+                    str(field.type),
+                    dialect.value,
+                )
+                continue
+
+        nullable = field.nullable
+        for label, value in (
+            ("nullable", column_spec.get("nullable")),
+            ("allow_null", column_spec.get("allow_null")),
+            ("defaults.nullable", defaults.get("nullable")),
+            ("defaults.allow_null", defaults.get("allow_null")),
+        ):
+            if value is None:
+                continue
+            if not isinstance(value, bool):
+                raise TypeError(f"{label} must be a boolean")
+            nullable = value
+            break
+
+        null_str = "NULL" if nullable else "NOT NULL"
+        lines.append(
+            f"    {quote_identifier(column_name, dialect)} {sql_type} {null_str}"
+        )
+
+    if not lines:
+        raise ValueError(
+            "No supported columns found in Parquet schema for the selected dialect"
+        )
+
+    body = ",\n".join(lines)
+
+    drop_clause = ""
+    if drop_existing:
+        if dialect == Dialect.SQLSERVER:
+            object_id_name = qualified_name.replace("'", "''")
+            if catalog:
+                batch_sep = "GO\n" if use_go else ""
+                drop_clause = (
+                    f"USE {quote_identifier(catalog, dialect)};\n"
+                    f"{batch_sep}IF OBJECT_ID('{object_id_name}') IS NOT NULL "
+                    f"DROP TABLE {qualified_name};\n"
+                )
+            else:
+                drop_clause = (
+                    f"IF OBJECT_ID('{object_id_name}') IS NOT NULL "
+                    f"DROP TABLE {qualified_name};\n"
+                )
+        else:
+            drop_clause = f"DROP TABLE IF EXISTS {qualified_name};\n"
+
+    statement = f"{drop_clause}CREATE TABLE {qualified_name} (\n{body}\n);"
+    print_verbose(
+        f"Generated CREATE TABLE from Parquet schema for {qualified_name}", verbose
+    )
+    if return_metadata:
+        metadata = ParquetDDLMetadata(
+            dropped_columns=dropped_columns, fallback_columns=fallback_columns
+        )
+        return statement, metadata
+    return statement

datablade/sql/dialects.py

@@ -0,0 +1,12 @@
+"""Enumeration of SQL dialects supported by datablade."""
+
+from enum import Enum
+
+
+class Dialect(str, Enum):
+    """Supported SQL dialects for datablade DDL helpers."""
+
+    SQLSERVER = "sqlserver"
+    POSTGRES = "postgres"
+    MYSQL = "mysql"
+    DUCKDB = "duckdb"

datablade/sql/quoting.py

@@ -0,0 +1,44 @@
+"""Identifier quoting for supported SQL dialects."""
+
+from typing import Optional
+
+from .dialects import Dialect
+
+
+def quote_identifier(name: Optional[str], dialect: Dialect = Dialect.SQLSERVER) -> str:
+    """
+    Quote an identifier for the given SQL dialect.
+
+    Args:
+        name: Identifier to quote; must be non-empty string.
+        dialect: Target SQL dialect.
+
+    Returns:
+        Quoted identifier string.
+
+    Raises:
+        ValueError: If name is missing/empty.
+        TypeError: If name is not a string.
+        NotImplementedError: If dialect is unsupported.
+    """
+    if name is None:
+        raise ValueError("name must be provided")
+    if not isinstance(name, str):
+        raise TypeError("name must be a string")
+    cleaned = name.strip()
+    if not cleaned:
+        raise ValueError("name must be a non-empty string")
+
+    if dialect == Dialect.SQLSERVER:
+        return f"[{cleaned.replace('[', '').replace(']', '')}]"
+    if dialect == Dialect.POSTGRES:
+        escaped = cleaned.replace('"', '""')
+        return f'"{escaped}"'
+    if dialect == Dialect.MYSQL:
+        escaped = cleaned.replace("`", "``")
+        return f"`{escaped}`"
+    if dialect == Dialect.DUCKDB:
+        escaped = cleaned.replace('"', '""')
+        return f'"{escaped}"'
+
+    raise NotImplementedError(f"Dialect not supported: {dialect}")

datablade/sql/schema_spec.py

@@ -0,0 +1,65 @@
+"""Schema specification helpers for DDL generation."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, Optional, Tuple
+
+
+def _as_mapping(value: Any, label: str) -> dict:
+    if value is None:
+        return {}
+    if not isinstance(value, Mapping):
+        raise TypeError(f"{label} must be a mapping")
+    return dict(value)
+
+
+def resolve_schema_spec(
+    schema_spec: Optional[Mapping[str, Any]],
+) -> Tuple[dict, dict]:
+    """Return (defaults, columns) mappings for a schema spec."""
+    if schema_spec is None:
+        return {}, {}
+    if not isinstance(schema_spec, Mapping):
+        raise TypeError("schema_spec must be a mapping")
+
+    defaults = _as_mapping(schema_spec.get("defaults"), "schema_spec.defaults")
+    columns = _as_mapping(schema_spec.get("columns"), "schema_spec.columns")
+    return defaults, columns
+
+
+def resolve_column_spec(
+    column_name: str,
+    schema_spec: Optional[Mapping[str, Any]],
+) -> Tuple[dict, dict]:
+    """Return (defaults, column_spec) for a column name."""
+    defaults, columns = resolve_schema_spec(schema_spec)
+    if not columns:
+        return defaults, {}
+
+    column_spec = columns.get(column_name)
+    if column_spec is None:
+        column_spec = columns.get(str(column_name))
+
+    if column_spec is None:
+        return defaults, {}
+    if not isinstance(column_spec, Mapping):
+        raise TypeError(f"schema_spec.columns['{column_name}'] must be a mapping")
+    return defaults, dict(column_spec)
+
+
+def resolve_string_policy(
+    column_name: str,
+    defaults: dict,
+    column_spec: dict,
+) -> dict:
+    """Merge defaults + column string policy overrides."""
+    string_defaults = _as_mapping(defaults.get("string"), "schema_spec.defaults.string")
+    string_overrides = _as_mapping(
+        column_spec.get("string"),
+        f"schema_spec.columns['{column_name}'].string",
+    )
+    policy = {**string_defaults, **string_overrides}
+    if "defined_pad" in policy and "pad" not in policy:
+        policy["pad"] = policy["defined_pad"]
+    return policy