[df] Add RNTupleWriteOptions to RSnapshotOptions

N.B., compression settings that have been set directly through the
snapshot options are propagated to the RNTuple write options, provided
that they haven't been set there as well.

Author: enirolf

tree/dataframe/inc/ROOT/RSnapshotOptions.hxx

@@ -11,8 +11,9 @@
 #ifndef ROOT_RSNAPSHOTOPTIONS
 #define ROOT_RSNAPSHOTOPTIONS
 
+#include "ROOT/RNTupleWriteOptions.hxx"
 #include <Compression.h>
-#include <optional>
+
 #include <string_view>
 #include <string>
 
@@ -31,6 +32,7 @@ struct RSnapshotOptions {
    RSnapshotOptions() = default;
    RSnapshotOptions(std::string_view mode, ECAlgo comprAlgo, int comprLevel, int autoFlush, int splitLevel, bool lazy,
                     bool overwriteIfExists = false, bool vector2RVec = true, int basketSize = -1,
+                    ROOT::RNTupleWriteOptions ntupleWriteOpts = ROOT::RNTupleWriteOptions(),
                     ESnapshotOutputFormat outputFormat = ESnapshotOutputFormat::kDefault)
       : fMode(mode),
         fCompressionAlgorithm(comprAlgo),
@@ -41,6 +43,7 @@ struct RSnapshotOptions {
         fOverwriteIfExists(overwriteIfExists),
         fVector2RVec(vector2RVec),
         fBasketSize(basketSize),
+        fNTupleWriteOpts(ntupleWriteOpts),
         fOutputFormat(outputFormat)
    {
    }
@@ -55,7 +58,8 @@ struct RSnapshotOptions {
    bool fVector2RVec = true;         ///< If set to true will convert std::vector columns to RVec when saving to disk
    int fBasketSize = -1;             ///< Set a custom basket size option. For more details, see
                                      ///< https://root.cern/manual/trees/#baskets-clusters-and-the-tree-header
-   ESnapshotOutputFormat fOutputFormat = ESnapshotOutputFormat::kDefault; ///< Which data format to write to
+   ROOT::RNTupleWriteOptions fNTupleWriteOpts = ROOT::RNTupleWriteOptions(); ///< RNTuple-specific write options
+   ESnapshotOutputFormat fOutputFormat = ESnapshotOutputFormat::kDefault;    ///< Which data format to write to
 };
 } // namespace RDF
 } // namespace ROOT

tree/dataframe/src/RDFSnapshotHelpers.cxx

@@ -852,8 +852,14 @@ void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Initialize()
    }
    model->Freeze();
 
-   ROOT::RNTupleWriteOptions writeOptions;
-   writeOptions.SetCompression(fOptions.fCompressionAlgorithm, fOptions.fCompressionLevel);
+   // Propagate fOptions.fCompressionAlgorithm and fOptions.fCompressionLevel to fOptions.fNTupleWriteOpts *only* if the
+   // compression settings in fNTupleWriteOpts have not been changed, and the compression algorithm or level in fOptions
+   // have.
+   if (fOptions.fNTupleWriteOpts.GetCompression() == RCompressionSetting::EDefaults::kUseGeneralPurpose &&
+       (fOptions.fCompressionAlgorithm != RCompressionSetting::EAlgorithm::kZLIB ||
+        fOptions.fCompressionLevel != RCompressionSetting::ELevel::kDefaultZLIB)) {
+      fOptions.fNTupleWriteOpts.SetCompression(fOptions.fCompressionAlgorithm, fOptions.fCompressionLevel);
+   }
 
    fOutputFile.reset(TFile::Open(fFileName.c_str(), fOptions.fMode.c_str()));
    if (!fOutputFile)
@@ -871,7 +877,8 @@ void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Initialize()
 
    // The RNTupleParallelWriter has exclusive access to the underlying TFile, no further synchronization is needed for
    // calls to Fill() (in Exec) and FlushCluster() (in FinalizeTask).
-   fWriter = ROOT::Experimental::RNTupleParallelWriter::Append(std::move(model), fNTupleName, *outputDir, writeOptions);
+   fWriter = ROOT::Experimental::RNTupleParallelWriter::Append(std::move(model), fNTupleName, *outputDir,
+                                                               fOptions.fNTupleWriteOpts);
 }
 
 void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::InitTask(TTreeReader *, unsigned int slot)

tree/dataframe/test/dataframe_snapshot_ntuple.cxx

@@ -101,25 +101,93 @@ TEST(RDFSnapshotRNTuple, LazyNotTriggered)
    EXPECT_TRUE(gSystem->AccessPathName(fileGuard.GetPath().c_str()));
 }
 
-TEST(RDFSnapshotRNTuple, Compression)
+TEST(RDFSnapshotRNTuple, WriteOpts)
 {
-   FileRAII fileGuard{"RDFSnapshotRNTuple_compression.root"};
+   FileRAII fileGuard{"RDFSnapshotRNTuple_write_opts.root"};
    const std::vector<std::string> columns = {"x"};
 
    auto df = ROOT::RDataFrame(25ull).Define("x", [] { return 10; });
 
+   ROOT::RNTupleWriteOptions writeOpts;
+   writeOpts.SetEnablePageChecksums(false);
+
    RSnapshotOptions opts;
    opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
-   opts.fCompressionAlgorithm = ROOT::RCompressionSetting::EAlgorithm::kLZ4;
-   opts.fCompressionLevel = 4;
+   opts.fNTupleWriteOpts = writeOpts;
 
    auto sdf = df.Snapshot("ntuple", fileGuard.GetPath(), "x", opts);
 
    EXPECT_EQ(columns, sdf->GetColumnNames());
 
    auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
-   auto compSettings = *reader->GetDescriptor().GetClusterDescriptor(0).GetColumnRange(0).GetCompressionSettings();
-   EXPECT_EQ(404, compSettings);
+   EXPECT_FALSE(reader->GetDescriptor().GetClusterDescriptor(0).GetPageRange(0).GetPageInfos()[0].HasChecksum());
+}
+
+TEST(RDFSnapshotRNTuple, Compression)
+{
+   FileRAII fileGuard{"RDFSnapshotRNTuple_compression.root"};
+   const std::vector<std::string> columns = {"x"};
+
+   auto df = ROOT::RDataFrame(25ull).Define("x", [] { return 10; });
+
+   // Default should be taken from RNTupleWriteOptions provided by RSnapshotOptions (zstd)
+   {
+      RSnapshotOptions opts;
+      opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
+
+      auto sdf = df.Snapshot("ntuple", fileGuard.GetPath(), "x", opts);
+
+      EXPECT_EQ(columns, sdf->GetColumnNames());
+
+      auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
+      auto compSettings = *reader->GetDescriptor().GetClusterDescriptor(0).GetColumnRange(0).GetCompressionSettings();
+      EXPECT_EQ(505, compSettings);
+   }
+   // Directly through RSnapshotOptions
+   {
+      RSnapshotOptions opts;
+      opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
+      opts.fCompressionAlgorithm = ROOT::RCompressionSetting::EAlgorithm::kLZ4;
+      opts.fCompressionLevel = 4;
+
+      auto sdf = df.Snapshot("ntuple", fileGuard.GetPath(), "x", opts);
+
+      EXPECT_EQ(columns, sdf->GetColumnNames());
+
+      auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
+      auto compSettings = *reader->GetDescriptor().GetClusterDescriptor(0).GetColumnRange(0).GetCompressionSettings();
+      EXPECT_EQ(404, compSettings);
+   }
+   // Through RNTupleWriteOptions provided by RSnapshotOptions
+   {
+      RSnapshotOptions opts;
+      opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
+      opts.fNTupleWriteOpts.SetCompression(404);
+
+      auto sdf = df.Snapshot("ntuple", fileGuard.GetPath(), "x", opts);
+
+      EXPECT_EQ(columns, sdf->GetColumnNames());
+
+      auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
+      auto compSettings = *reader->GetDescriptor().GetClusterDescriptor(0).GetColumnRange(0).GetCompressionSettings();
+      EXPECT_EQ(404, compSettings);
+   }
+   // When both are set, preference is given to RNTupleWriteOptions
+   {
+      RSnapshotOptions opts;
+      opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
+      opts.fCompressionAlgorithm = ROOT::RCompressionSetting::EAlgorithm::kLZ4;
+      opts.fCompressionLevel = 4;
+      opts.fNTupleWriteOpts.SetCompression(207);
+
+      auto sdf = df.Snapshot("ntuple", fileGuard.GetPath(), "x", opts);
+
+      EXPECT_EQ(columns, sdf->GetColumnNames());
+
+      auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
+      auto compSettings = *reader->GetDescriptor().GetClusterDescriptor(0).GetColumnRange(0).GetCompressionSettings();
+      EXPECT_EQ(207, compSettings);
+   }
 }
 
 class RDFSnapshotRNTupleTest : public ::testing::Test {