From d9ada72a229d7c985b21b15163417d54b831fec2 Mon Sep 17 00:00:00 2001 From: Dylan Whitaker Date: Tue, 6 Oct 2026 11:13:04 +0100 Subject: [PATCH 1/5] Read partitioned parquet directories as single tables A subdirectory of parquet files (e.g. Hive-partitioned yr=2023/p00001.parquet) is now one table named after the directory, instead of one table per part file. All files in a directory must share columns, types and partition keys, otherwise make-tables exits with an error recommending the data be reorganised. Partition keys become TEXT columns. DuckDB sees each dataset through a temporary view so source databases are not modified. Co-Authored-By: Claude Sonnet 5.5 --- datafaker/db_utils.py | 50 ++++--- datafaker/parquet2orm.py | 197 +++++++++++++++++++++++---- docs/source/tutorial_parquet.rst | 37 +++++ tests/test_parquet2orm.py | 225 ++++++++++++++++++++++++++++++- 4 files changed, 468 insertions(+), 41 deletions(-) diff --git a/datafaker/db_utils.py b/datafaker/db_utils.py index 7ed5f17c..c6ee6c4f 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, @@ -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..da35dcce 100644 --- a/datafaker/parquet2orm.py +++ b/datafaker/parquet2orm.py @@ -1,36 +1,187 @@ """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) -> 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_difference(expected: Mapping[str, Any], actual: Mapping[str, Any]) -> str: + """Describe how the columns 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] + changed = [ + f"{c} ({expected[c]} vs {actual[c]})" + for c in expected + if c in actual and str(expected[c]) != str(actual[c]) + ] + if missing: + problems.append(f"missing columns {missing}") + if extra: + problems.append(f"extra columns {extra}") + if changed: + problems.append(f"different types {changed}") + return "; ".join(problems) + + +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 columns with the + same types, and the same partition keys. Directories that contain + unrelated tables are rejected with an error message. + + :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 schema. + """ + first = table.files[0] + dtypes = _file_dtypes(first) + if not table.is_dataset: + return dtypes + keys = _partition_keys(table, first) + for path in table.files[1:]: + problem = _describe_difference(dtypes, _file_dtypes(path)) + 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 and types. 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) + 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 +198,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 +240,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 +271,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", @@ -178,4 +327,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..b65f1e90 100644 --- a/docs/source/tutorial_parquet.rst +++ b/docs/source/tutorial_parquet.rst @@ -82,6 +82,43 @@ 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 columns and types. +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. + +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..0882a9d1 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,216 @@ 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 loose files are tables and directories are tables.""" + self.write_files( + { + "fruit.parquet": {"fruit_id": [1]}, + "visit/yr=2023/p1.parquet": {"visit_id": [1]}, + "empty/readme.txt": {}, + } + ) + (self.parquet_dir / "empty" / "readme.txt").unlink() + (self.parquet_dir / "empty").rmdir() + 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_rejected(self, mock_logger: MagicMock) -> None: + """Test the same column name with a different type is rejected.""" + self.write_files( + { + "visit/yr=2022/p1.parquet": {"visit_id": [1]}, + "visit/yr=2023/p1.parquet": {"visit_id": ["one"]}, + } + ) + with self.assertRaises(Exit): + get_parquet_orm(self.parquet_dir) + message = ( + mock_logger.error.call_args.args[0] % mock_logger.error.call_args.args[1:] + ) + self.assertIn("different types", message) + self.assertIn("visit_id", message) + + 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): + """Tests that DuckDB can query partitioned parquet directories.""" + + def setUp(self) -> None: + """Make a partitioned dataset and a loose file.""" + 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) # type: ignore[arg-type] + ) + ) + with engine.connect() as conn: + count = conn.execute(text('SELECT count(*) FROM "visit"')).scalar() + self.assertEqual(count, 4) From fdf1fa6d89855663680398a7bfe10e0ec0f5d5c0 Mon Sep 17 00:00:00 2001 From: Dylan Whitaker Date: Tue, 6 Oct 2026 14:13:58 +0100 Subject: [PATCH 2/5] Allow column type differences between partitions Differing column names or partition keys remain a hard error. Differing column types are now a single warning per column, using the first file's type, and the comparison ignores case so int32 equals Int32. Pandas nullable types such as Int32 now map to the SQL type of their numpy dtype. Co-Authored-By: Claude Sonnet 5.5 --- datafaker/parquet2orm.py | 51 ++++++++++++++++++++---------- docs/source/tutorial_parquet.rst | 15 ++++++--- tests/test_parquet2orm.py | 53 ++++++++++++++++++++++++++------ 3 files changed, 89 insertions(+), 30 deletions(-) diff --git a/datafaker/parquet2orm.py b/datafaker/parquet2orm.py index da35dcce..4dfaa47c 100644 --- a/datafaker/parquet2orm.py +++ b/datafaker/parquet2orm.py @@ -104,51 +104,55 @@ def _partition_keys(table: ParquetTable, path: Path) -> list[str]: return keys -def _describe_difference(expected: Mapping[str, Any], actual: Mapping[str, Any]) -> str: - """Describe how the columns of ``actual`` differ from ``expected``.""" +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] - changed = [ - f"{c} ({expected[c]} vs {actual[c]})" - for c in expected - if c in actual and str(expected[c]) != str(actual[c]) - ] if missing: problems.append(f"missing columns {missing}") if extra: problems.append(f"extra columns {extra}") - if changed: - problems.append(f"different types {changed}") 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 columns with the - same types, and the same partition keys. Directories that contain - unrelated tables are rejected with an error message. + 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 schema. + :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:]: - problem = _describe_difference(dtypes, _file_dtypes(path)) + 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 and types. If the " + "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.", @@ -159,6 +163,19 @@ def get_table_dtypes(table: ParquetTable) -> dict[str, Any]: 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}} @@ -319,7 +336,9 @@ def _get_table_orm( 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 diff --git a/docs/source/tutorial_parquet.rst b/docs/source/tutorial_parquet.rst index b65f1e90..31ef9fd8 100644 --- a/docs/source/tutorial_parquet.rst +++ b/docs/source/tutorial_parquet.rst @@ -109,11 +109,16 @@ Large tables are often saved as a directory of Parquet files, partitioned 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 columns and types. -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. +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 diff --git a/tests/test_parquet2orm.py b/tests/test_parquet2orm.py index 0882a9d1..dad61237 100644 --- a/tests/test_parquet2orm.py +++ b/tests/test_parquet2orm.py @@ -281,21 +281,56 @@ def test_differing_columns_are_rejected(self, mock_logger: MagicMock) -> None: self.assertIn("move each table into its own directory", message) @patch("datafaker.parquet2orm.logger") - def test_differing_types_are_rejected(self, mock_logger: MagicMock) -> None: - """Test the same column name with a different type is rejected.""" + 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": ["one"]}, + "visit/yr=2023/p1.parquet": {"visit_id": [1.5]}, + "visit/yr=2023/p2.parquet": {"visit_id": [2.5]}, } ) - with self.assertRaises(Exit): - get_parquet_orm(self.parquet_dir) - message = ( - mock_logger.error.call_args.args[0] % mock_logger.error.call_args.args[1:] + 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" ) - self.assertIn("different types", message) - self.assertIn("visit_id", message) + 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.""" From 5ce844e3acb3fdb74da66b1c4860db607f1a1635 Mon Sep 17 00:00:00 2001 From: Dylan Whitaker Date: Tue, 6 Oct 2026 14:43:18 +0100 Subject: [PATCH 3/5] Accept str for parquet_dir to satisfy ty check parquet_dir is read back from orm.yaml as text, so create_db_engine and find_parquet_tables now declare Path | str. Drops the mypy-style type: ignore in the test, which ty does not honour. Co-Authored-By: Claude Sonnet 5.5 --- datafaker/db_utils.py | 2 +- datafaker/parquet2orm.py | 2 +- tests/test_parquet2orm.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/datafaker/db_utils.py b/datafaker/db_utils.py index c6ee6c4f..f317a0da 100644 --- a/datafaker/db_utils.py +++ b/datafaker/db_utils.py @@ -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.""" diff --git a/datafaker/parquet2orm.py b/datafaker/parquet2orm.py index 4dfaa47c..a8d5030e 100644 --- a/datafaker/parquet2orm.py +++ b/datafaker/parquet2orm.py @@ -45,7 +45,7 @@ def _parquet_files_under(directory: Path) -> tuple[Path, ...]: ) -def find_parquet_tables(directory: Path) -> list[ParquetTable]: +def find_parquet_tables(directory: Path | str) -> list[ParquetTable]: """ Find the tables represented by the parquet files in a directory. diff --git a/tests/test_parquet2orm.py b/tests/test_parquet2orm.py index dad61237..0bf9fc0a 100644 --- a/tests/test_parquet2orm.py +++ b/tests/test_parquet2orm.py @@ -394,7 +394,7 @@ 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) # type: ignore[arg-type] + "duckdb:///:memory:", parquet_dir=str(self.parquet_dir) ) ) with engine.connect() as conn: From e42cf7be7fe0307a94b1b9e23d9ed04054aa7787 Mon Sep 17 00:00:00 2001 From: Dylan Whitaker Date: Wed, 7 Oct 2026 17:24:36 +0100 Subject: [PATCH 4/5] Running Black and updating tests and docstrings with PR comments --- tests/test_parquet2orm.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/tests/test_parquet2orm.py b/tests/test_parquet2orm.py index 0bf9fc0a..42aece68 100644 --- a/tests/test_parquet2orm.py +++ b/tests/test_parquet2orm.py @@ -209,16 +209,15 @@ def test_same_part_names_in_different_tables_do_not_collide(self) -> None: self.assertSetEqual(set(orm["visit"]["columns"].keys()), {"visit_id", "yr"}) def test_files_and_directories_mix(self) -> None: - """Test loose files are tables and directories are tables.""" + """Test a mixture of loose parquet files and parquet directories 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]}, - "empty/readme.txt": {}, } ) - (self.parquet_dir / "empty" / "readme.txt").unlink() - (self.parquet_dir / "empty").rmdir() + (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"}) @@ -345,10 +344,10 @@ def test_differing_partition_keys_are_rejected(self) -> None: class PartitionedParquetEngine(TestCase): - """Tests that DuckDB can query partitioned parquet directories.""" + """Integration tests for querying Parquet data through the parquet_dir option.""" def setUp(self) -> None: - """Make a partitioned dataset and a loose file.""" + """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) @@ -393,9 +392,7 @@ def test_dataset_views_are_not_written_to_the_source_database(self) -> None: 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) - ) + 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() From 5d82a51dc37071a2824fdfd0558f281982d553c8 Mon Sep 17 00:00:00 2001 From: Dylan Whitaker Date: Wed, 7 Oct 2026 17:24:36 +0100 Subject: [PATCH 5/5] Running Black and updating tests and docstrings with PR comments --- tests/test_parquet2orm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_parquet2orm.py b/tests/test_parquet2orm.py index 42aece68..f4564a34 100644 --- a/tests/test_parquet2orm.py +++ b/tests/test_parquet2orm.py @@ -209,7 +209,8 @@ def test_same_part_names_in_different_tables_do_not_collide(self) -> None: 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 directories are read as separate tables. Checks empty directory are not treated as partitioned tables.""" + """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]},