Read partitioned parquet directories as single tables

datafaker/db_utils.py

@@ -1,7 +1,6 @@
 """Utility functions."""
 import gzip
 import io
-import os
 from collections.abc import Callable, Mapping
 from pathlib import Path
 from types import ModuleType
@@ -41,6 +40,7 @@
 )
 from typer import Exit
 
+from datafaker.parquet2orm import ParquetTable, find_parquet_tables
 from datafaker.utils import (
     T,
     get_ignored_table_names,
@@ -145,7 +145,7 @@ def create_db_engine(
     db_dsn: str,
     schema_name: Optional[str] = None,
     use_asyncio: bool = False,
-    parquet_dir: Optional[Path] = None,
+    parquet_dir: Optional[Union[Path, str]] = None,
     **kwargs: Any,
 ) -> MaybeAsyncEngine:
     """Create a SQLAlchemy Engine."""
@@ -174,19 +174,25 @@ def create_db_engine(
             engine = engine.execution_options(schema_translate_map={None: schema_name})
         else:
             settings["search_path"] = schema_name
+    dataset_views: list[str] = []
     if parquet_dir is not None:
-        joined = ",".join(_find_parquet_directories(parquet_dir))
+        tables = find_parquet_tables(parquet_dir)
+        joined = ",".join(
+            sorted({str(t.root.parent) for t in tables if not t.is_dataset})
+        )
         # double up single quotes
         dj = joined.replace("'", "''")
         # enclose in single quotes
         settings["file_search_path"] = f"'{dj}'"
+        dataset_views = [_dataset_view_sql(t) for t in tables if t.is_dataset]
 
-    if settings:
+    if settings or dataset_views:
         event_engine = get_sync_engine(engine)
 
         @event.listens_for(event_engine, "connect", insert=True)
         def connect(dbapi_connection: DBAPIConnection, _: Any) -> None:
             set_db_settings(dbapi_connection, settings)
+            create_dataset_views(dbapi_connection, dataset_views)
 
     return engine
 
@@ -231,21 +237,33 @@ def get_metadata(engine: Engine, schema_name: Optional[str] = None) -> MetaData:
     return md
 
 
-def _find_parquet_directories(parquet_dir: Path) -> list[str]:
-    """Find all the directories under ``parquet_dir`` that contain parquet files."""
-    return [
-        path
-        for path, _, filenames in os.walk(parquet_dir)
-        if _names_include_parquet(Path(path), filenames)
-    ]
+def _sql_quote(text: str, quote: str) -> str:
+    """Enclose ``text`` in ``quote``, doubling any ``quote`` inside it."""
+    return f"{quote}{text.replace(quote, quote * 2)}{quote}"
 
 
-def _names_include_parquet(path: Path, file_names: Iterable[str]) -> bool:
-    for fn in file_names:
-        entry = path / fn
-        if entry.is_file() and entry.suffix in {".parquet", ".parq"}:
-            return True
-    return False
+def _dataset_view_sql(table: ParquetTable) -> str:
+    """Make the SQL that exposes a directory of parquet files as one table."""
+    files = ",".join(_sql_quote(str(f), "'") for f in table.files)
+    return (
+        f"CREATE OR REPLACE TEMP VIEW {_sql_quote(table.name, chr(34))} AS "
+        f"SELECT * FROM read_parquet([{files}], "
+        "hive_partitioning=true, union_by_name=true)"
+    )
+
+
+def create_dataset_views(
+    connection: DBAPIConnection, view_statements: Iterable[str]
+) -> None:
+    """
+    Make each partitioned parquet dataset available as a table.
+
+    The views are temporary so nothing is written into the source database.
+    """
+    cursor = connection.cursor()
+    for statement in view_statements:
+        cursor.execute(statement)
+    cursor.close()
 
 
 def set_db_settings(connection: DBAPIConnection, settings: Mapping[str, str]) -> None:

datafaker/parquet2orm.py

@@ -1,36 +1,204 @@
 """Add to ORM structure based on parquet files."""
-import os
 import re
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 from fastparquet import ParquetFile
+from typer import Exit
 
 from datafaker.utils import logger
 
+PARQUET_SUFFIXES = {".parquet", ".parq"}
+
+# The dtype recorded for a Hive partition key (such as the ``yr`` in ``yr=2023``).
+# Partition keys are not stored in the files themselves.
+PARTITION_KEY_DTYPE = "category"
+
+_hive_dir_re = re.compile(r"^([^=]+)=(.*)$")
+
+
+@dataclass(frozen=True)
+class ParquetTable:
+    """A table found in a directory of parquet files."""
+
+    name: str
+    """The name of the table, as it will appear in ``orm.yaml``."""
+    root: Path
+    """The parquet file for a single-file table, or the directory of a dataset."""
+    files: tuple[Path, ...]
+    """All the parquet files that make up the table, sorted."""
+    is_dataset: bool
+    """True if the table is a (possibly partitioned) directory of parquet files."""
+
+
+def _parquet_files_under(directory: Path) -> tuple[Path, ...]:
+    """Find all the parquet files in ``directory`` and below, sorted."""
+    return tuple(
+        sorted(
+            entry
+            for entry in directory.rglob("*")
+            if entry.suffix in PARQUET_SUFFIXES and entry.is_file()
+        )
+    )
+
+
+def find_parquet_tables(directory: Path | str) -> list[ParquetTable]:
+    """
+    Find the tables represented by the parquet files in a directory.
+
+    * A parquet file directly in ``directory`` is a table named after the file.
+    * A subdirectory of ``directory`` containing parquet files at any depth
+      is a single (partitioned) table named after the subdirectory.
+    * If ``directory`` has no parquet files directly in it, and all the
+      subdirectories containing parquet files are Hive partitions
+      (named ``key=value``), then ``directory`` is itself a single table.
+
+    :param directory: The directory to search for parquet files in.
+    :return: The tables found, files first, then datasets, each sorted by name.
+    """
+    directory = Path(directory)
+    children = sorted(directory.iterdir())
+    file_tables = [
+        ParquetTable(entry.name, entry, (entry,), False)
+        for entry in children
+        if entry.suffix in PARQUET_SUFFIXES and entry.is_file()
+    ]
+    dir_files = [
+        (entry, files)
+        for entry in children
+        if entry.is_dir()
+        for files in [_parquet_files_under(entry)]
+        if files
+    ]
+    if not file_tables and dir_files:
+        if all(_hive_dir_re.match(entry.name) for entry, _ in dir_files):
+            name = directory.resolve().name
+            return [
+                ParquetTable(
+                    name,
+                    directory,
+                    tuple(f for _, files in dir_files for f in files),
+                    True,
+                )
+            ]
+    return file_tables + [
+        ParquetTable(entry.name, entry, files, True) for entry, files in dir_files
+    ]
+
+
+def _file_dtypes(path: Path) -> dict[str, Any]:
+    """Get the column names and dtypes stored in one parquet file."""
+    return dict(ParquetFile(path).dtypes)
+
+
+def _partition_keys(table: ParquetTable, path: Path) -> list[str]:
+    """Get the Hive partition keys of one file of a dataset, in path order."""
+    keys = []
+    for part in path.relative_to(table.root).parts[:-1]:
+        match = _hive_dir_re.match(part)
+        if match:
+            keys.append(match.group(1))
+    return keys
+
+
+def _describe_column_difference(
+    expected: Mapping[str, Any], actual: Mapping[str, Any]
+) -> str:
+    """Describe how the column names of ``actual`` differ from ``expected``."""
+    problems = []
+    missing = [c for c in expected if c not in actual]
+    extra = [c for c in actual if c not in expected]
+    if missing:
+        problems.append(f"missing columns {missing}")
+    if extra:
+        problems.append(f"extra columns {extra}")
+    return "; ".join(problems)
+
+
+def _type_name(dtype: Any) -> str:
+    """Get a name for a dtype that ignores case, so ``int32`` equals ``Int32``."""
+    return str(dtype).lower()
+
+
+def get_table_dtypes(table: ParquetTable) -> dict[str, Any]:
+    """
+    Get the columns and dtypes of a table, checking that its files agree.
+
+    All the files of a partitioned table must have the same column names and
+    the same partition keys. Directories that contain unrelated tables are
+    rejected with an error message. Column types may differ between files
+    (for example when a partition has no values for a column); this is
+    reported as a warning and the type from the first file is used.
+
+    :param table: The table to examine.
+    :return: A dict of column names to dtypes. Partition keys are included.
+    :raises typer.Exit: If the files of a dataset do not all have the same columns.
+    """
+    first = table.files[0]
+    dtypes = _file_dtypes(first)
+    if not table.is_dataset:
+        return dtypes
+    keys = _partition_keys(table, first)
+    type_names: dict[str, set[str]] = {}
+    for path in table.files[1:]:
+        file_dtypes = _file_dtypes(path)
+        problem = _describe_column_difference(dtypes, file_dtypes)
+        if not problem and _partition_keys(table, path) != keys:
+            problem = "different partition directories"
+        if problem:
+            logger.error(
+                "Directory %s cannot be read as a single partitioned table: "
+                "%s differs from %s (%s). All the files in a partitioned "
+                "table must have the same columns. If the "
+                "directories inside %s hold different tables, move each "
+                "table into its own directory directly inside the parquet "
+                "directory and run this command again.",
+                table.name,
+                path,
+                first,
+                problem,
+                table.root,
+            )
+            raise Exit(1)
+        for column, dtype in file_dtypes.items():
+            if _type_name(dtype) != _type_name(dtypes[column]):
+                names = type_names.setdefault(column, {_type_name(dtypes[column])})
+                names.add(_type_name(dtype))
+    for column, names in type_names.items():
+        logger.warning(
+            "Column %s.%s has different types in different files: %s. "
+            "Using the type from %s.",
+            table.name,
+            column,
+            sorted(names),
+            first,
+        )
+    return {**dtypes, **{k: PARTITION_KEY_DTYPE for k in keys if k not in dtypes}}
+
 
 def get_parquet_orm(directory: Path) -> dict[str, Any] | None:
     """
     Read the parquet files, guess a database structure from that.
 
+    A directory of parquet files (possibly partitioned) is read as one table.
+
     :param directory: The directory to search for parquet files in.
     :return: The ORM dictionary on success, None on failure.
+    :raises typer.Exit: If a directory holds files with differing schemas.
     """
     logger.debug("Examining directory %s", directory)
     if not directory.is_dir():
         logger.error("%s is not a directory", directory)
         return None
-    parquets: dict[str, ParquetFile] = {}  # type: ignore[no-any-unimported]
-    for path, _, filenames in os.walk(directory):
-        for filename in filenames:
-            entry = Path(path) / filename
-            if entry.is_file() and entry.suffix in {".parquet", ".parq"}:
-                logger.debug("examining file %s", entry)
-                parquets[entry.name] = ParquetFile(entry)
+    tables = find_parquet_tables(directory)
+    dtypes = {table.name: get_table_dtypes(table) for table in tables}
+    guesser = _ColumnGuesser(dtypes)
     return {
-        name: _get_table_orm(parquet, name, _ColumnGuesser(parquets))
-        for name, parquet in parquets.items()
+        name: _get_table_orm(table_dtypes, name, guesser)
+        for name, table_dtypes in dtypes.items()
     }
 
 
@@ -47,14 +215,14 @@ def _get_words(s: str) -> set[str]:
 class _ColumnGuesser:
     """Guesses foreign key targets."""
 
-    def __init__(self, tables: dict[str, ParquetFile]) -> None:  # type: ignore[no-any-unimported]
-        """Initialize the column guesser."""
+    def __init__(self, tables: Mapping[str, Iterable[str]]) -> None:
+        """Initialize the column guesser from table names to column names."""
         self.tn_cn_words: list[tuple[str, str, set[str]]] = []
         self.tn2words: dict[str, set[str]] = {}
         for name, table in tables.items():
             table_words = _get_words(name)
             self.tn2words[name] = table_words
-            for column in table.columns:
+            for column in table:
                 self.tn_cn_words.append(
                     (name, column, table_words | _get_words(column))
                 )
@@ -89,23 +257,21 @@ def get_likely_foreign_key_target(
         return best_so_far
 
 
-def _get_table_orm(  # type: ignore[no-any-unimported]
-    table: ParquetFile, name: str, column_guesser: _ColumnGuesser
+def _get_table_orm(
+    dtypes: Mapping[str, Any], name: str, column_guesser: _ColumnGuesser
 ) -> dict[str, Any]:
     """
     Guess the ORM configuration of the table passed in.
 
-    :param table: The parquet of the table to guess the config of.
-    :param name: The filename of the parquet file
-    :param tables: A dict of parquet file names to parquet files of all the tables.
+    :param dtypes: The column names and dtypes of the table to guess the config of.
+    :param name: The filename of the parquet file, or directory name of the dataset.
+    :param column_guesser: Guesses foreign key targets across all the tables.
     :return: A reasonable ORM configuration for this table.
     """
-    column_types = {
-        column: _dtype_to_sql(dtype) for column, dtype in table.dtypes.items()
-    }
+    column_types = {column: _dtype_to_sql(dtype) for column, dtype in dtypes.items()}
     # Get everything before the last dot
     # (or the whole thing if there is no dot)
-    name_pref = name[: name.rfind(".")]
+    name_pref = name.rsplit(".", 1)[0]
     name_words = _get_words(name_pref)
     cols_orm = {}
     likely_primaries = []
@@ -122,7 +288,7 @@ def _get_table_orm(  # type: ignore[no-any-unimported]
                 col_orm["nullable"] = True
             fk_pair = column_guesser.get_likely_foreign_key_target(column)
             if fk_pair is not None:
-                (fk_table, fk_column) = fk_pair
+                fk_table, fk_column = fk_pair
                 if fk_table != name:
                     logger.debug(
                         "Column %s.%s guessed as being a foreign key to %s.%s",
@@ -170,12 +336,16 @@ def _get_table_orm(  # type: ignore[no-any-unimported]
 
 
 def _dtype_to_sql(dtype: Any) -> str | None:
-    """Convert a numpy datatype into a SQL type."""
+    """Convert a numpy or Pandas datatype into a SQL type."""
+    # Pandas nullable types such as ``Int32`` wrap a numpy dtype
+    dtype = getattr(dtype, "numpy_dtype", dtype)
     if isinstance(dtype, np.dtype):
         if dtype.shape != () or dtype.kind not in _numpy_dtype_to_sql:
             return None
         return _numpy_dtype_to_sql[dtype.kind]
     if isinstance(dtype, str):
         if dtype.startswith("datetime"):
             return "DATETIME"
+        if dtype == PARTITION_KEY_DTYPE:
+            return "TEXT"
     return None