From f446f7f051c40f81a22f40868eba9ec74e9bdf82 Mon Sep 17 00:00:00 2001 From: Joaquin Hui Gomez <132194176+joaquinhuigomez@users.noreply.github.com> Date: Tue, 15 Sep 2026 12:33:41 +0100 Subject: [PATCH] Add row_groups_per_file to to_parquet/write_parquet The Parquet COPY option ROW_GROUPS_PER_FILE is supported by the engine but was not exposed on DuckDBPyRelation.to_parquet/write_parquet, so callers had to drop down to raw SQL to rotate output files by row group count. Pass the value straight through to the COPY options map. Validation mirrors the neighbouring row_group_size block: non-integer input raises InvalidInputException rather than reaching the engine. Fixes #386 --- _duckdb-stubs/__init__.pyi | 2 ++ src/include/duckdb_python/pyrelation.hpp | 10 +++++----- src/pyrelation.cpp | 17 ++++++++++++---- src/pyrelation/initialize.cpp | 9 +++++---- tests/fast/api/test_to_parquet.py | 25 ++++++++++++++++++++++++ 5 files changed, 50 insertions(+), 13 deletions(-) diff --git a/_duckdb-stubs/__init__.pyi b/_duckdb-stubs/__init__.pyi index 8770483f..ea7eefac 100644 --- a/_duckdb-stubs/__init__.pyi +++ b/_duckdb-stubs/__init__.pyi @@ -734,6 +734,7 @@ class DuckDBPyRelation: field_ids: ParquetFieldsOptions | None = None, row_group_size_bytes: int | str | None = None, row_group_size: int | None = None, + row_groups_per_file: int | None = None, overwrite: bool | None = None, per_thread_output: bool | None = None, use_tmp_file: bool | None = None, @@ -790,6 +791,7 @@ class DuckDBPyRelation: field_ids: ParquetFieldsOptions | None = None, row_group_size_bytes: str | int | None = None, row_group_size: int | None = None, + row_groups_per_file: int | None = None, overwrite: bool | None = None, per_thread_output: bool | None = None, use_tmp_file: bool | None = None, diff --git a/src/include/duckdb_python/pyrelation.hpp b/src/include/duckdb_python/pyrelation.hpp index f71a6327..ee94fd45 100644 --- a/src/include/duckdb_python/pyrelation.hpp +++ b/src/include/duckdb_python/pyrelation.hpp @@ -213,11 +213,11 @@ struct DuckDBPyRelation { void ToParquet(const string &filename, const nb::object &compression = nb::none(), const nb::object &field_ids = nb::none(), const nb::object &row_group_size_bytes = nb::none(), - const nb::object &row_group_size = nb::none(), const nb::object &overwrite = nb::none(), - const nb::object &per_thread_output = nb::none(), const nb::object &use_tmp_file = nb::none(), - const nb::object &partition_by = nb::none(), const nb::object &write_partition_columns = nb::none(), - const nb::object &append = nb::none(), const nb::object &filename_pattern = nb::none(), - const nb::object &file_size_bytes = nb::none()); + const nb::object &row_group_size = nb::none(), const nb::object &row_groups_per_file = nb::none(), + const nb::object &overwrite = nb::none(), const nb::object &per_thread_output = nb::none(), + const nb::object &use_tmp_file = nb::none(), const nb::object &partition_by = nb::none(), + const nb::object &write_partition_columns = nb::none(), const nb::object &append = nb::none(), + const nb::object &filename_pattern = nb::none(), const nb::object &file_size_bytes = nb::none()); void ToCSV(const string &filename, const nb::object &sep = nb::none(), const nb::object &na_rep = nb::none(), const nb::object &header = nb::none(), const nb::object "echar = nb::none(), diff --git a/src/pyrelation.cpp b/src/pyrelation.cpp index 24dda20e..f109548d 100644 --- a/src/pyrelation.cpp +++ b/src/pyrelation.cpp @@ -1262,10 +1262,11 @@ static Value NestedDictToStruct(const nb::object &dictionary) { void DuckDBPyRelation::ToParquet(const string &filename, const nb::object &compression, const nb::object &field_ids, const nb::object &row_group_size_bytes, const nb::object &row_group_size, - const nb::object &overwrite, const nb::object &per_thread_output, - const nb::object &use_tmp_file, const nb::object &partition_by, - const nb::object &write_partition_columns, const nb::object &append, - const nb::object &filename_pattern, const nb::object &file_size_bytes) { + const nb::object &row_groups_per_file, const nb::object &overwrite, + const nb::object &per_thread_output, const nb::object &use_tmp_file, + const nb::object &partition_by, const nb::object &write_partition_columns, + const nb::object &append, const nb::object &filename_pattern, + const nb::object &file_size_bytes) { identifier_map_t> options; if (!nb::none().is(compression)) { @@ -1306,6 +1307,14 @@ void DuckDBPyRelation::ToParquet(const string &filename, const nb::object &compr options["row_group_size"] = {Value(row_group_size_int)}; } + if (!nb::none().is(row_groups_per_file)) { + if (!nb::isinstance(row_groups_per_file)) { + throw InvalidInputException("to_parquet only accepts 'row_groups_per_file' as an integer"); + } + int64_t row_groups_per_file_int = (int64_t)nb::int_(row_groups_per_file); + options["row_groups_per_file"] = {Value(row_groups_per_file_int)}; + } + if (!nb::none().is(partition_by)) { if (!nb::isinstance(partition_by)) { throw InvalidInputException("to_parquet only accepts 'partition_by' as a list of strings"); diff --git a/src/pyrelation/initialize.cpp b/src/pyrelation/initialize.cpp index 9c5a562b..22808479 100644 --- a/src/pyrelation/initialize.cpp +++ b/src/pyrelation/initialize.cpp @@ -39,10 +39,11 @@ static void InitializeConsumers(nb::class_ &m) { "Write the relation object to a Parquet file in 'file_name'", nb::arg("file_name"), nb::kw_only(), nb::arg("compression") = nb::none(), nb::arg("field_ids") = nb::none(), nb::arg("row_group_size_bytes") = nb::none(), nb::arg("row_group_size") = nb::none(), - nb::arg("overwrite") = nb::none(), nb::arg("per_thread_output") = nb::none(), - nb::arg("use_tmp_file") = nb::none(), nb::arg("partition_by") = nb::none(), - nb::arg("write_partition_columns") = nb::none(), nb::arg("append") = nb::none(), - nb::arg("filename_pattern") = nb::none(), nb::arg("file_size_bytes") = nb::none()); + nb::arg("row_groups_per_file") = nb::none(), nb::arg("overwrite") = nb::none(), + nb::arg("per_thread_output") = nb::none(), nb::arg("use_tmp_file") = nb::none(), + nb::arg("partition_by") = nb::none(), nb::arg("write_partition_columns") = nb::none(), + nb::arg("append") = nb::none(), nb::arg("filename_pattern") = nb::none(), + nb::arg("file_size_bytes") = nb::none()); DefineMethod( {"to_csv", "write_csv"}, m, &DuckDBPyRelation::ToCSV, "Write the relation object to a CSV file in 'file_name'", diff --git a/tests/fast/api/test_to_parquet.py b/tests/fast/api/test_to_parquet.py index 71d5e00e..643b40be 100644 --- a/tests/fast/api/test_to_parquet.py +++ b/tests/fast/api/test_to_parquet.py @@ -68,6 +68,31 @@ def test_row_group_size(self): parquet_rel = duckdb.read_parquet(temp_file_name) assert rel.execute().fetchall() == parquet_rel.execute().fetchall() + def test_row_groups_per_file(self): + temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names())) # noqa: PTH118 + + # use same test data as external/duckdb/test/sql/copy/row_groups_per_file.test, which also + # pins threads to 1: file rotation is best-effort and only exact on a single thread + con = duckdb.connect(config={"threads": 1}) + rel = con.from_query("SELECT i AS col_a, i AS col_b FROM range(0,10000) tbl(i);") + rel.to_parquet(temp_file_name, row_group_size=2000, row_groups_per_file=1) + + # 5 row groups of 2000 rows, one row group per file + files = list(pathlib.Path(temp_file_name).iterdir()) + assert len(files) == 5, f"Expected 5 files, got {len(files)}" + + # Verify data integrity + result = con.read_parquet(f"{temp_file_name}/*.parquet") + assert len(result.execute().fetchall()) == 10000 + + def test_row_groups_per_file_invalid_type(self): + temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names())) # noqa: PTH118 + rel = duckdb.sql("SELECT 1 AS i") + with pytest.raises( + duckdb.InvalidInputException, match="to_parquet only accepts 'row_groups_per_file' as an integer" + ): + rel.to_parquet(temp_file_name, row_groups_per_file="x") + @pytest.mark.parametrize("write_columns", [None, True, False]) def test_partition(self, write_columns): temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names())) # noqa: PTH118