diff --git a/datafaker/db_utils.py b/datafaker/db_utils.py index 7ed5f17c..f317a0da 100644 --- a/datafaker/db_utils.py +++ b/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: diff --git a/datafaker/parquet2orm.py b/datafaker/parquet2orm.py index ecabc253..a8d5030e 100644 --- a/datafaker/parquet2orm.py +++ b/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,7 +336,9 @@ 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 @@ -178,4 +346,6 @@ def _dtype_to_sql(dtype: Any) -> str | None: if isinstance(dtype, str): if dtype.startswith("datetime"): return "DATETIME" + if dtype == PARTITION_KEY_DTYPE: + return "TEXT" return None diff --git a/docs/source/tutorial_parquet.rst b/docs/source/tutorial_parquet.rst index d28d9515..31ef9fd8 100644 --- a/docs/source/tutorial_parquet.rst +++ b/docs/source/tutorial_parquet.rst @@ -82,6 +82,48 @@ Example directory structure: ├── artist.parquet └── artwork.parquet +Partitioned Parquet Datasets +"""""""""""""""""""""""""""" + +Large tables are often saved as a directory of Parquet files, partitioned +(for example by year) and split into several part files: + +.. code-block:: text + + input_data/ + ├── artist.parquet # a single file: the table ``artist.parquet`` + └── artwork/ # a directory: the table ``artwork`` + ├── yr=2023/ + │ ├── p00001.parquet + │ └── p00002.parquet + └── yr=__MISSING__/ + └── p00001.parquet + +* A ``.parquet`` file directly inside the input directory is one table, + named after the file, as before. +* A subdirectory is one table, named after the directory, made from all + the Parquet files inside it at any depth. +* Partition directories named ``key=value`` become a ``TEXT`` column + (``yr`` above), because partition keys are not stored in the files. +* If ``--parquet-dir`` points directly at a partitioned dataset (a directory + containing only ``key=value`` subdirectories) the whole directory is one + table, named after that directory. + +All the files of one partitioned table must have the same column names +and the same partition keys. If they do not, ``make-tables`` stops with an +error naming the differing file. This usually means the subdirectory holds +several different tables; move each table into its own directory directly +inside the input directory and run ``make-tables`` again. + +Column types may differ between files (for example when one partition has no +values for a column and so is stored with a different type). This is logged as +a warning and the type from the first file is used. Differences in case only, +such as ``int32`` and ``Int32``, are ignored. + +The synthetic data is written as one file per table (for example +``artwork.parquet``) by ``dump-data --parquet``; the partitioning of the +input is not reproduced. + Building the ORM Configuration ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/tests/test_parquet2orm.py b/tests/test_parquet2orm.py index 84ebb197..f4564a34 100644 --- a/tests/test_parquet2orm.py +++ b/tests/test_parquet2orm.py @@ -1,6 +1,7 @@ """Test the parquet-file-dir to ``orm.yaml`` functionality.""" import datetime import os +import shutil import tempfile from collections.abc import Iterable from pathlib import Path @@ -9,7 +10,10 @@ from unittest.mock import MagicMock, patch import pandas as pd +from sqlalchemy import text +from typer import Exit +from datafaker.db_utils import create_db_engine, get_sync_engine from datafaker.parquet2orm import get_parquet_orm @@ -33,8 +37,8 @@ def __repr__(self) -> str: return f"HasValues<{self.value_set}>" -class Parquet2Orm(TestCase): - """Tests the ``parquet2orm`` function.""" +class ParquetDirTestCase(TestCase): + """Base class providing a temporary directory for parquet files.""" def setUp(self) -> None: """Make a temporary directory for the parquet files.""" @@ -60,6 +64,10 @@ def write_parquet(self, data: dict[str, dict[str, list[Any]]]) -> None: for fn, table_data in data.items(): pd.DataFrame.from_dict(table_data).to_parquet(self.parquet_dir / fn) + +class Parquet2Orm(ParquetDirTestCase): + """Tests the ``parquet2orm`` function.""" + def test_can_infer_column_type(self) -> None: """Test that we can guess obvious column types in parquet files.""" data: dict[str, dict[str, list[Any]]] = { @@ -144,3 +152,249 @@ def test_can_infer_foreign_key(self) -> None: orm["fruit.parquet"]["columns"]["seed_id"]["foreign_keys"], ["seed.parquet.seed_id"], ) + + +class PartitionedParquet2Orm(ParquetDirTestCase): + """Tests reading directories of parquet files as single tables.""" + + def write_files(self, data: dict[str, dict[str, list[Any]]]) -> None: + """ + Write parquet files at the relative paths given. + + :param data: dict of relative file paths to dict of column names to data. + """ + for fn, table_data in data.items(): + path = self.parquet_dir / fn + path.parent.mkdir(parents=True, exist_ok=True) + pd.DataFrame.from_dict(table_data).to_parquet(path) + + def test_partitioned_directory_is_one_table(self) -> None: + """Test a hive-partitioned directory is one table named after it.""" + self.write_files( + { + "visit/yr=2023/p00001.parquet": {"visit_id": [1], "ward": ["a"]}, + "visit/yr=2023/p00002.parquet": {"visit_id": [2], "ward": ["b"]}, + "visit/yr=__MISSING__/p00001.parquet": {"visit_id": [3], "ward": ["c"]}, + } + ) + orm = get_parquet_orm(self.parquet_dir) + assert orm is not None + self.assertSetEqual(set(orm.keys()), {"visit"}) + cols = orm["visit"]["columns"] + self.assertSetEqual(set(cols.keys()), {"visit_id", "ward", "yr"}) + self.assertEqual(cols["yr"]["type"], "TEXT") + self.assertEqual(cols["visit_id"]["type"], "INTEGER") + + @patch("datafaker.parquet2orm.logger") + def test_partition_key_type_is_not_a_warning(self, mock_logger: MagicMock) -> None: + """Test the partition key does not log 'Could not determine type'.""" + self.write_files({"visit/yr=2023/p1.parquet": {"visit_id": [1]}}) + get_parquet_orm(self.parquet_dir) + mock_logger.warning.assert_not_called() + + def test_same_part_names_in_different_tables_do_not_collide(self) -> None: + """Test part files of the same name in different datasets stay separate.""" + self.write_files( + { + "visit/yr=2023/p00001.parquet": {"visit_id": [1]}, + "drug/yr=2023/p00001.parquet": {"drug_id": [1], "dose": [1.5]}, + } + ) + orm = get_parquet_orm(self.parquet_dir) + assert orm is not None + self.assertSetEqual(set(orm.keys()), {"visit", "drug"}) + self.assertSetEqual( + set(orm["drug"]["columns"].keys()), {"drug_id", "dose", "yr"} + ) + self.assertSetEqual(set(orm["visit"]["columns"].keys()), {"visit_id", "yr"}) + + def test_files_and_directories_mix(self) -> None: + """Test a mixture of loose parquet files and parquet dirs are read as separate tables. + Checks empty directory are not treated as partitioned tables.""" + self.write_files( + { + "fruit.parquet": {"fruit_id": [1]}, + "visit/yr=2023/p1.parquet": {"visit_id": [1]}, + } + ) + (self.parquet_dir / "empty").mkdir() + + orm = get_parquet_orm(self.parquet_dir) + assert orm is not None + self.assertSetEqual(set(orm.keys()), {"fruit.parquet", "visit"}) + + def test_nested_plain_directories_are_one_table(self) -> None: + """Test a subdirectory of unpartitioned files is one table.""" + self.write_files( + { + "sub/a.parquet": {"x": [1]}, + "sub/b.parquet": {"x": [2]}, + } + ) + orm = get_parquet_orm(self.parquet_dir) + assert orm is not None + self.assertSetEqual(set(orm.keys()), {"sub"}) + + def test_parquet_dir_that_is_a_dataset(self) -> None: + """Test the directory given can itself be a partitioned table.""" + self.write_files( + { + "yr=2022/p1.parquet": {"visit_id": [1]}, + "yr=2023/p1.parquet": {"visit_id": [2]}, + } + ) + orm = get_parquet_orm(self.parquet_dir) + assert orm is not None + self.assertSetEqual(set(orm.keys()), {self.parquet_dir.resolve().name}) + only = next(iter(orm.values())) + self.assertSetEqual(set(only["columns"].keys()), {"visit_id", "yr"}) + + def test_parquet_dir_dataset_must_be_consistent(self) -> None: + """Test a directory of unrelated hive-named tables is still rejected.""" + self.write_files( + { + "yr=2022/p1.parquet": {"visit_id": [1]}, + "yr=2023/p1.parquet": {"other": ["x"]}, + } + ) + with self.assertRaises(Exit): + get_parquet_orm(self.parquet_dir) + + @patch("datafaker.parquet2orm.logger") + def test_differing_columns_are_rejected(self, mock_logger: MagicMock) -> None: + """Test different tables nested under one directory give a clear error.""" + self.write_files( + { + "sub/a.parquet": {"visit_id": [1], "ward": ["a"]}, + "sub/b.parquet": {"drug_id": [1]}, + } + ) + with self.assertRaises(Exit): + get_parquet_orm(self.parquet_dir) + mock_logger.error.assert_called_once() + message = ( + mock_logger.error.call_args.args[0] % mock_logger.error.call_args.args[1:] + ) + self.assertIn("sub", message) + self.assertIn("missing columns ['visit_id', 'ward']", message) + self.assertIn("extra columns ['drug_id']", message) + self.assertIn("move each table into its own directory", message) + + @patch("datafaker.parquet2orm.logger") + def test_differing_types_are_a_warning(self, mock_logger: MagicMock) -> None: + """Test the same column with a different type warns but is accepted.""" + self.write_files( + { + "visit/yr=2022/p1.parquet": {"visit_id": [1]}, + "visit/yr=2023/p1.parquet": {"visit_id": [1.5]}, + "visit/yr=2023/p2.parquet": {"visit_id": [2.5]}, + } + ) + orm = get_parquet_orm(self.parquet_dir) + assert orm is not None + self.assertEqual(orm["visit"]["columns"]["visit_id"]["type"], "INTEGER") + mock_logger.error.assert_not_called() + mock_logger.warning.assert_called_once() + args = mock_logger.warning.call_args.args + self.assertEqual(args[1:3], ("visit", "visit_id")) + self.assertEqual(args[3], ["float64", "int64"]) + + @patch("datafaker.parquet2orm.logger") + def test_int32_and_nullable_int32_are_the_same_type( + self, mock_logger: MagicMock + ) -> None: + """Test ``int32`` and ``Int32`` do not even warn.""" + self.write_files( + { + "person/yr=2022/p1.parquet": {"age": [1]}, + "person/yr=2023/p1.parquet": {"age": [2]}, + } + ) + pd.DataFrame({"age": pd.array([1], dtype="int32")}).to_parquet( + self.parquet_dir / "person/yr=2022/p1.parquet" + ) + pd.DataFrame({"age": pd.array([None, 2], dtype="Int32")}).to_parquet( + self.parquet_dir / "person/yr=2023/p1.parquet" + ) + orm = get_parquet_orm(self.parquet_dir) + assert orm is not None + self.assertEqual(orm["person"]["columns"]["age"]["type"], "INTEGER") + for call in mock_logger.warning.call_args_list: + self.assertNotIn("different types", call.args[0]) + + def test_nullable_int_first_file_is_still_integer(self) -> None: + """Test a pandas ``Int32`` column is typed ``INTEGER``.""" + self.write_files({"person/yr=2022/p1.parquet": {"age": [1]}}) + pd.DataFrame({"age": pd.array([None, 2], dtype="Int32")}).to_parquet( + self.parquet_dir / "person/yr=2022/p1.parquet" + ) + orm = get_parquet_orm(self.parquet_dir) + assert orm is not None + self.assertEqual(orm["person"]["columns"]["age"]["type"], "INTEGER") + + def test_differing_partition_keys_are_rejected(self) -> None: + """Test files partitioned differently are rejected.""" + self.write_files( + { + "visit/yr=2022/p1.parquet": {"visit_id": [1]}, + "visit/p2.parquet": {"visit_id": [2]}, + } + ) + with self.assertRaises(Exit): + get_parquet_orm(self.parquet_dir) + + +class PartitionedParquetEngine(TestCase): + """Integration tests for querying Parquet data through the parquet_dir option.""" + + def setUp(self) -> None: + """Create partitioned and loose Parquet data for parquet_dir integration tests.""" + super().setUp() + self.parquet_dir = Path(tempfile.mkdtemp(prefix="parq")) + self.addCleanup(shutil.rmtree, self.parquet_dir) + for path, data in { + "visit/yr=2022/p1.parquet": {"visit_id": [1, 2]}, + "visit/yr=2023/p1.parquet": {"visit_id": [3]}, + "visit/yr=2023/p2.parquet": {"visit_id": [4]}, + "fruit.parquet": {"fruit_id": [9]}, + }.items(): + (self.parquet_dir / path).parent.mkdir(parents=True, exist_ok=True) + pd.DataFrame.from_dict(data).to_parquet(self.parquet_dir / path) + + def test_dataset_can_be_queried_by_directory_name(self) -> None: + """Test every partition appears in the table named after the directory.""" + engine = get_sync_engine( + create_db_engine("duckdb:///:memory:", parquet_dir=self.parquet_dir) + ) + with engine.connect() as conn: + rows = conn.execute( + text('SELECT a.visit_id, a.yr FROM "visit" AS a ORDER BY a.visit_id') + ).fetchall() + fruit = conn.execute(text('SELECT * FROM "fruit.parquet"')).fetchall() + self.assertListEqual( + [tuple(r) for r in rows], [(1, 2022), (2, 2022), (3, 2023), (4, 2023)] + ) + self.assertListEqual([tuple(r) for r in fruit], [(9,)]) + + def test_dataset_views_are_not_written_to_the_source_database(self) -> None: + """Test a file-backed source database is not modified.""" + db = self.parquet_dir / "src.duckdb" + engine = get_sync_engine( + create_db_engine(f"duckdb:///{db}", parquet_dir=self.parquet_dir) + ) + with engine.connect() as conn: + conn.execute(text('SELECT count(*) FROM "visit"')) + engine.dispose() + engine = get_sync_engine(create_db_engine(f"duckdb:///{db}")) + with engine.connect() as conn: + with self.assertRaises(Exception): + conn.execute(text('SELECT count(*) FROM "visit"')) + + def test_parquet_dir_may_be_a_string(self) -> None: + """Test the directory read back from ``orm.yaml`` as text still works.""" + engine = get_sync_engine( + create_db_engine("duckdb:///:memory:", parquet_dir=str(self.parquet_dir)) + ) + with engine.connect() as conn: + count = conn.execute(text('SELECT count(*) FROM "visit"')).scalar() + self.assertEqual(count, 4)