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