Add parameters to control page size in Parquet writer

cpp/include/cudf/io/parquet.hpp

@@ -37,8 +37,10 @@ namespace io {
  * @file
  */
 
-constexpr size_t default_row_group_size_bytes   = 128 * 1024 * 1024;  // 128MB
-constexpr size_type default_row_group_size_rows = 1000000;
+constexpr size_t default_row_group_size_bytes     = 128 * 1024 * 1024;  // 128MB
+constexpr size_type default_row_group_size_rows   = 1000000;
+constexpr size_t default_target_page_size_bytes   = 512 * 1024;
+constexpr size_type default_target_page_size_rows = 20000;
 
 /**
  * @brief Builds parquet_reader_options to use for `read_parquet()`.
@@ -382,6 +384,10 @@ class parquet_writer_options {
   size_t _row_group_size_bytes = default_row_group_size_bytes;
   // Maximum number of rows in row group (unless smaller than a single page)
   size_type _row_group_size_rows = default_row_group_size_rows;
+  // Maximum size of each page (uncompressed)
+  size_t _target_page_size_bytes = default_target_page_size_bytes;
+  // Maximum number of rows in a page
+  size_type _target_page_size_rows = default_target_page_size_rows;
 
   /**
    * @brief Constructor from sink and table.
@@ -482,6 +488,16 @@ class parquet_writer_options {
    */
   auto get_row_group_size_rows() const { return _row_group_size_rows; }
 
+  /**
+   * @brief Returns target page size, in bytes.
+   */
+  auto get_target_page_size_bytes() const { return _target_page_size_bytes; }
+
+  /**
+   * @brief Returns target page size, in rows.
+   */
+  auto get_target_page_size_rows() const { return _target_page_size_rows; }
+
   /**
    * @brief Sets partitions.
    *
@@ -552,24 +568,22 @@ class parquet_writer_options {
   /**
    * @brief Sets the maximum row group size, in bytes.
    */
-  void set_row_group_size_bytes(size_t size_bytes)
-  {
-    CUDF_EXPECTS(
-      size_bytes >= 512 * 1024,
-      "The maximum row group size cannot be smaller than the page size, which is 512KB.");
-    _row_group_size_bytes = size_bytes;
-  }
+  void set_row_group_size_bytes(size_t size_bytes) { _row_group_size_bytes = size_bytes; }
 
   /**
    * @brief Sets the maximum row group size, in rows.
    */
-  void set_row_group_size_rows(size_type size_rows)
-  {
-    CUDF_EXPECTS(
-      size_rows >= 5000,
-      "The maximum row group size cannot be smaller than the page size, which is 5000 rows.");
-    _row_group_size_rows = size_rows;
-  }
+  void set_row_group_size_rows(size_type size_rows) { _row_group_size_rows = size_rows; }
+
+  /**
+   * @brief Sets the maximum page size, in bytes.
+   */
+  void set_target_page_size_bytes(size_t pgsz_bytes) { _target_page_size_bytes = pgsz_bytes; }
+
+  /**
+   * @brief Sets the maximum page size, in rows.
+   */
+  void set_target_page_size_rows(size_type pgsz_rows) { _target_page_size_rows = pgsz_rows; }
 };
 
 class parquet_writer_options_builder {
@@ -699,6 +713,30 @@ class parquet_writer_options_builder {
     return *this;
   }
 
+  /**
+   * @brief Sets the maximum page size, in bytes.
+   *
+   * @param val The page size to use.
+   * @return this for chaining.
+   */
+  parquet_writer_options_builder& target_page_size_bytes(size_t val)
+  {
+    options.set_target_page_size_bytes(val);
+    return *this;
+  }
+
+  /**
+   * @brief Sets the maximum page size, in rows.
+   *
+   * @param val The page size to use.
+   * @return this for chaining.
+   */
+  parquet_writer_options_builder& target_page_size_bytes(size_type val)
+  {
+    options.set_target_page_size_rows(val);
+    return *this;
+  }
+
   /**
    * @brief Sets whether int96 timestamps are written or not in parquet_writer_options.
    *
@@ -783,6 +821,10 @@ class chunked_parquet_writer_options {
   size_t _row_group_size_bytes = default_row_group_size_bytes;
   // Maximum number of rows in row group (unless smaller than a single page)
   size_type _row_group_size_rows = default_row_group_size_rows;
+  // Maximum size of each page (uncompressed)
+  size_t _target_page_size_bytes = default_target_page_size_bytes;
+  // Maximum number of rows in a page
+  size_type _target_page_size_rows = default_target_page_size_rows;
 
   /**
    * @brief Constructor from sink.
@@ -844,6 +886,16 @@ class chunked_parquet_writer_options {
    */
   auto get_row_group_size_rows() const { return _row_group_size_rows; }
 
+  /**
+   * @brief Returns target page size, in bytes.
+   */
+  auto get_target_page_size_bytes() const { return _target_page_size_bytes; }
+
+  /**
+   * @brief Returns target page size, in rows.
+   */
+  auto get_target_page_size_rows() const { return _target_page_size_rows; }
+
   /**
    * @brief Sets metadata.
    *
@@ -888,24 +940,22 @@ class chunked_parquet_writer_options {
   /**
    * @brief Sets the maximum row group size, in bytes.
    */
-  void set_row_group_size_bytes(size_t size_bytes)
-  {
-    CUDF_EXPECTS(
-      size_bytes >= 512 * 1024,
-      "The maximum row group size cannot be smaller than the page size, which is 512KB.");
-    _row_group_size_bytes = size_bytes;
-  }
+  void set_row_group_size_bytes(size_t size_bytes) { _row_group_size_bytes = size_bytes; }
 
   /**
    * @brief Sets the maximum row group size, in rows.
    */
-  void set_row_group_size_rows(size_type size_rows)
-  {
-    CUDF_EXPECTS(
-      size_rows >= 5000,
-      "The maximum row group size cannot be smaller than the page size, which is 5000 rows.");
-    _row_group_size_rows = size_rows;
-  }
+  void set_row_group_size_rows(size_type size_rows) { _row_group_size_rows = size_rows; }
+
+  /**
+   * @brief Sets the maximum page size, in bytes.
+   */
+  void set_target_page_size_bytes(size_t pgsz_bytes) { _target_page_size_bytes = pgsz_bytes; }
+
+  /**
+   * @brief Sets the maximum page size, in rows.
+   */
+  void set_target_page_size_rows(size_type pgsz_rows) { _target_page_size_rows = pgsz_rows; }
 
   /**
    * @brief creates builder to build chunked_parquet_writer_options.
@@ -1025,6 +1075,30 @@ class chunked_parquet_writer_options_builder {
     return *this;
   }
 
+  /**
+   * @brief Sets the maximum page size, in bytes.
+   *
+   * @param val maximum page size
+   * @return this for chaining.
+   */
+  chunked_parquet_writer_options_builder& target_page_size_bytes(size_t val)
+  {
+    options.set_target_page_size_bytes(val);
+    return *this;
+  }
+
+  /**
+   * @brief Sets the maximum page size, in rows.
+   *
+   * @param val maximum page size
+   * @return this for chaining.
+   */
+  chunked_parquet_writer_options_builder& target_page_size_rows(size_type val)
+  {
+    options.set_target_page_size_rows(val);
+    return *this;
+  }
+
   /**
    * @brief move chunked_parquet_writer_options member once it's built.
    */

cpp/src/io/parquet/page_enc.cu

@@ -240,7 +240,9 @@ __global__ void __launch_bounds__(128)
                statistics_merge_group* page_grstats,
                statistics_merge_group* chunk_grstats,
                size_t max_page_comp_data_size,
-               int32_t num_columns)
+               int32_t num_columns,
+               size_t target_page_size_bytes,
+               size_type target_page_size_rows)
 {
   // TODO: All writing seems to be done by thread 0. Could be replaced by thrust foreach
   __shared__ __align__(8) parquet_column_device_view col_g;
@@ -337,8 +339,13 @@ __global__ void __launch_bounds__(128)
       uint32_t max_page_size = (values_in_page * 2 >= ck_g.num_values)   ? 256 * 1024
                                : (values_in_page * 3 >= ck_g.num_values) ? 384 * 1024
                                                                          : 512 * 1024;
+
+      // override max_page_size if target is smaller
+      if (max_page_size > target_page_size_bytes) max_page_size = target_page_size_bytes;
+
       if (num_rows >= ck_g.num_rows ||
-          (values_in_page > 0 && (page_size + fragment_data_size > max_page_size))) {
+          (values_in_page > 0 && (page_size + fragment_data_size > max_page_size)) ||
+          rows_in_page > target_page_size_rows) {
         if (ck_g.use_dictionary) {
           page_size =
             1 + 5 + ((values_in_page * ck_g.dict_rle_bits + 7) >> 3) + (values_in_page >> 8);
@@ -1927,15 +1934,24 @@ void InitEncoderPages(device_2dspan<EncColumnChunk> chunks,
                       device_span<gpu::EncPage> pages,
                       device_span<parquet_column_device_view const> col_desc,
                       int32_t num_columns,
+                      size_t target_page_size_bytes,
+                      size_type target_page_size_rows,
                       statistics_merge_group* page_grstats,
                       statistics_merge_group* chunk_grstats,
                       size_t max_page_comp_data_size,
                       rmm::cuda_stream_view stream)
 {
   auto num_rowgroups = chunks.size().first;
   dim3 dim_grid(num_columns, num_rowgroups);  // 1 threadblock per rowgroup
-  gpuInitPages<<<dim_grid, 128, 0, stream.value()>>>(
-    chunks, pages, col_desc, page_grstats, chunk_grstats, max_page_comp_data_size, num_columns);
+  gpuInitPages<<<dim_grid, 128, 0, stream.value()>>>(chunks,
+                                                     pages,
+                                                     col_desc,
+                                                     page_grstats,
+                                                     chunk_grstats,
+                                                     max_page_comp_data_size,
+                                                     num_columns,
+                                                     target_page_size_bytes,
+                                                     target_page_size_rows);
 }
 
 void EncodePages(device_span<gpu::EncPage> pages,

cpp/src/io/parquet/parquet_gpu.hpp

@@ -575,6 +575,8 @@ void InitEncoderPages(cudf::detail::device_2dspan<EncColumnChunk> chunks,
                       device_span<gpu::EncPage> pages,
                       device_span<parquet_column_device_view const> col_desc,
                       int32_t num_columns,
+                      size_t target_page_size_bytes,
+                      size_type target_page_size_rows,
                       statistics_merge_group* page_grstats,
                       statistics_merge_group* chunk_grstats,
                       size_t max_page_comp_data_size,

cpp/src/io/parquet/writer_impl.cu

@@ -859,7 +859,16 @@ void writer::impl::init_page_sizes(hostdevice_2dvector<gpu::EncColumnChunk>& chu
                                    uint32_t num_columns)
 {
   chunks.host_to_device(stream);
-  gpu::InitEncoderPages(chunks, {}, col_desc, num_columns, nullptr, nullptr, 0, stream);
+  gpu::InitEncoderPages(chunks,
+                        {},
+                        col_desc,
+                        num_columns,
+                        target_page_size_bytes,
+                        target_page_size_rows,
+                        nullptr,
+                        nullptr,
+                        0,
+                        stream);
   chunks.device_to_host(stream, true);
 }
 
@@ -965,6 +974,8 @@ void writer::impl::init_encoder_pages(hostdevice_2dvector<gpu::EncColumnChunk>&
                    pages,
                    col_desc,
                    num_columns,
+                   target_page_size_bytes,
+                   target_page_size_rows,
                    (num_stats_bfr) ? page_stats_mrg.data() : nullptr,
                    (num_stats_bfr > num_pages) ? page_stats_mrg.data() + num_pages : nullptr,
                    max_page_comp_data_size,
@@ -1122,6 +1133,8 @@ writer::impl::impl(std::vector<std::unique_ptr<data_sink>> sinks,
     stream(stream),
     max_row_group_size{options.get_row_group_size_bytes()},
     max_row_group_rows{options.get_row_group_size_rows()},
+    target_page_size_bytes(options.get_target_page_size_bytes()),
+    target_page_size_rows(options.get_target_page_size_rows()),
     compression_(to_parquet_compression(options.get_compression())),
     stats_granularity_(options.get_stats_level()),
     int96_timestamps(options.is_enabled_int96_timestamps()),
@@ -1144,6 +1157,8 @@ writer::impl::impl(std::vector<std::unique_ptr<data_sink>> sinks,
     stream(stream),
     max_row_group_size{options.get_row_group_size_bytes()},
     max_row_group_rows{options.get_row_group_size_rows()},
+    target_page_size_bytes(options.get_target_page_size_bytes()),
+    target_page_size_rows(options.get_target_page_size_rows()),
     compression_(to_parquet_compression(options.get_compression())),
     stats_granularity_(options.get_stats_level()),
     int96_timestamps(options.is_enabled_int96_timestamps()),

cpp/src/io/parquet/writer_impl.hpp

@@ -210,6 +210,8 @@ class writer::impl {
 
   size_t max_row_group_size          = default_row_group_size_bytes;
   size_type max_row_group_rows       = default_row_group_size_rows;
+  size_t target_page_size_bytes      = default_target_page_size_bytes;
+  size_type target_page_size_rows    = default_target_page_size_rows;
   Compression compression_           = Compression::UNCOMPRESSED;
   statistics_freq stats_granularity_ = statistics_freq::STATISTICS_NONE;
   bool int96_timestamps              = false;

cpp/tests/io/parquet_test.cpp

@@ -3184,6 +3184,10 @@ TEST_F(ParquetReaderTest, EmptyOutput)
   CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view());
 }
 
+/* temporarily remove this test.  with page size being tunable, can no longer just test for
+ * row groups being smaller than 512KiB.  ideally one would use the configured page size, but
+ * then that would require setting the page size parameter before the row group size, which
+ * seems odd.  should probably test for consistent parameters when instantiating the writer.
 TEST_F(ParquetWriterTest, RowGroupSizeInvalid)
 {
   const auto unused_table = std::make_unique<table>();
@@ -3205,5 +3209,6 @@ TEST_F(ParquetWriterTest, RowGroupSizeInvalid)
                  .row_group_size_bytes(511 << 10),
                cudf::logic_error);
 }
+*/
 
 CUDF_TEST_PROGRAM_MAIN()