Create table_input_metadata from a table_metadata

cpp/include/cudf/io/types.hpp

@@ -230,12 +230,24 @@ struct table_metadata {
     per_file_user_data;  //!< Per file format-dependent metadata as key-values pairs
 };
 
+class table_input_metadata;
+
 /**
  * @brief Table with table metadata used by io readers to return the metadata by value
  */
 struct table_with_metadata {
   std::unique_ptr<table> tbl;  //!< Table
   table_metadata metadata;     //!< Table metadata
+
+  /**
+   * @brief Return a `table_input_metadata` populated with data from this `table_with_metadata`.
+   *
+   * The returned `table_input_metadata` will preserve the hierarchy, naming, and nullability
+   * of the contained table.
+   *
+   * @return `table_input_metadata`
+   */
+  table_input_metadata get_table_input_metadata();
 };
 
 /**
@@ -800,6 +812,16 @@ class table_input_metadata {
    */
   table_input_metadata(table_view const& table);
 
+  /**
+   * @brief Construct a new table_input_metadata from a table_with_metadata.
+   *
+   * The constructed table_input_metadata has the same structure as the passed table_with_metadata,
+   * and also preserves any naming and nullability info from the original table.
+   *
+   * @param table_and_metadata The table_with_metadata to construct metadata for
+   */
+  table_input_metadata(table_with_metadata const& table_and_metadata);
+
   std::vector<column_in_metadata> column_metadata;  //!< List of column metadata
 };
 

cpp/src/io/functions.cpp

@@ -503,6 +503,11 @@ std::unique_ptr<std::vector<uint8_t>> merge_row_group_metadata(
   return detail_parquet::writer::merge_row_group_metadata(metadata_list);
 }
 
+table_input_metadata table_with_metadata::get_table_input_metadata()
+{
+  return table_input_metadata(*this);
+}
+
 table_input_metadata::table_input_metadata(table_view const& table)
 {
   // Create a metadata hierarchy using `table`
@@ -517,6 +522,31 @@ table_input_metadata::table_input_metadata(table_view const& table)
     table.begin(), table.end(), std::back_inserter(this->column_metadata), get_children);
 }
 
+table_input_metadata::table_input_metadata(table_with_metadata const& table_and_metadata)
+{
+  auto const& table = table_and_metadata.tbl->view();
+  auto const& names = table_and_metadata.metadata.schema_info;
+
+  // Create a metadata hierarchy with naming and nullability using `table_and_metadata`
+  std::function<column_in_metadata(column_view const&, column_name_info const&)> get_children =
+    [&](column_view const& col, column_name_info const& name) {
+      auto col_meta = column_in_metadata{name.name};
+      col_meta.set_nullability(col.nullable());
+      std::transform(col.child_begin(),
+                     col.child_end(),
+                     name.children.begin(),
+                     std::back_inserter(col_meta.children),
+                     get_children);
+      return col_meta;
+    };
+
+  std::transform(table.begin(),
+                 table.end(),
+                 names.begin(),
+                 std::back_inserter(this->column_metadata),
+                 get_children);
+}
+
 /**
  * @copydoc cudf::io::write_parquet
  */

cpp/tests/io/parquet_test.cpp

@@ -6599,4 +6599,68 @@ TEST_F(ParquetWriterTest, TimestampMicrosINT96NoOverflow)
   CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view());
 }
 
+TEST_F(ParquetWriterTest, PreserveNullability)
+{
+  constexpr auto num_rows = 100;
+
+  auto const col0_data = random_values<int32_t>(num_rows);
+  auto const col1_data = random_values<int32_t>(num_rows);
+
+  auto const col0_validity = cudf::test::iterators::no_nulls();
+  auto const col1_validity =
+    cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 2 == 0; });
+
+  column_wrapper<int32_t> col0{col0_data.begin(), col0_data.end(), col0_validity};
+  column_wrapper<int32_t> col1{col1_data.begin(), col1_data.end(), col1_validity};
+  auto const col2 = make_parquet_list_list_col<int>(0, num_rows, 5, 8, true);
+
+  auto const expected = table_view{{col0, col1, *col2}};
+
+  cudf::io::table_input_metadata expected_metadata(expected);
+  expected_metadata.column_metadata[0].set_name("mandatory");
+  expected_metadata.column_metadata[0].set_nullability(false);
+  expected_metadata.column_metadata[1].set_name("optional");
+  expected_metadata.column_metadata[1].set_nullability(true);
+  expected_metadata.column_metadata[2].set_name("lists");
+  expected_metadata.column_metadata[2].set_nullability(true);
+  expected_metadata.column_metadata[2].child(0).set_name("offsets");
+  expected_metadata.column_metadata[2].child(0).set_nullability(false);
+  expected_metadata.column_metadata[2].child(1).set_name("element");
+  expected_metadata.column_metadata[2].child(1).set_nullability(false);
+  expected_metadata.column_metadata[2].child(1).child(0).set_name("offsets");
+  expected_metadata.column_metadata[2].child(1).child(0).set_nullability(false);
+  expected_metadata.column_metadata[2].child(1).child(1).set_name("element");
+  expected_metadata.column_metadata[2].child(1).child(1).set_nullability(true);
+
+  auto const filepath = temp_env->get_temp_filepath("PreserveNullability.parquet");
+  cudf::io::parquet_writer_options out_opts =
+    cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, expected)
+      .metadata(expected_metadata);
+
+  cudf::io::write_parquet(out_opts);
+
+  cudf::io::parquet_reader_options const in_opts =
+    cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath});
+  auto const result        = cudf::io::read_parquet(in_opts);
+  auto const read_metadata = cudf::io::table_input_metadata{result};
+
+  // test that expected_metadata matches read_metadata
+  std::function<void(cudf::io::column_in_metadata, cudf::io::column_in_metadata)>
+    compare_names_and_nullability = [&](auto lhs, auto rhs) {
+      EXPECT_EQ(lhs.get_name(), rhs.get_name());
+      EXPECT_EQ(lhs.nullable(), rhs.nullable());
+      EXPECT_EQ(lhs.num_children(), rhs.num_children());
+      for (int i = 0; i < lhs.num_children(); ++i) {
+        compare_names_and_nullability(lhs.child(i), rhs.child(i));
+      }
+    };
+
+  EXPECT_EQ(expected_metadata.column_metadata.size(), read_metadata.column_metadata.size());
+
+  for (size_t i = 0; i < expected_metadata.column_metadata.size(); ++i) {
+    compare_names_and_nullability(expected_metadata.column_metadata[i],
+                                  read_metadata.column_metadata[i]);
+  }
+}
+
 CUDF_TEST_PROGRAM_MAIN()