Support any scalar type in PivotLonger features

Author: Rafał Hibner

cpp/src/arrow/acero/options.h

@@ -780,15 +780,15 @@ class ARROW_ACERO_EXPORT TableSinkNodeOptions : public ExecNodeOptions {
 
 /// \brief a row template that describes one row that will be generated for each input row
 struct ARROW_ACERO_EXPORT PivotLongerRowTemplate {
-  PivotLongerRowTemplate(std::vector<std::string> feature_values,
+  PivotLongerRowTemplate(std::vector<std::shared_ptr<Scalar>> feature_values,
                          std::vector<std::optional<FieldRef>> measurement_values)
       : feature_values(std::move(feature_values)),
         measurement_values(std::move(measurement_values)) {}
   /// A (typically unique) set of feature values for the template, usually derived from a
   /// column name
   ///
   /// These will be used to populate the feature columns
-  std::vector<std::string> feature_values;
+  std::vector<std::shared_ptr<Scalar>> feature_values;
   /// The fields containing the measurements to use for this row
   ///
   /// These will be used to populate the measurement columns.  If nullopt then nulls

cpp/src/arrow/acero/pivot_longer_node.cc

@@ -42,7 +42,7 @@ namespace {
 
 // A row template that's been bound to a schema
 struct BoundRowTemplate {
-  std::vector<std::string> feature_values;
+  std::vector<std::shared_ptr<Scalar>> feature_values;
   std::vector<std::optional<FieldPath>> measurement_paths;
 
   static Result<BoundRowTemplate> Make(const PivotLongerRowTemplate& unbound,
@@ -65,7 +65,7 @@ struct BoundRowTemplate {
   }
 
  private:
-  BoundRowTemplate(std::vector<std::string> feature_values,
+  BoundRowTemplate(std::vector<std::shared_ptr<Scalar>> feature_values,
                    std::vector<std::optional<FieldPath>> measurement_paths)
       : feature_values(std::move(feature_values)),
         measurement_paths(std::move(measurement_paths)) {}
@@ -89,6 +89,8 @@ class PivotLongerNode : public ExecNode, public TracedNode {
           "have names");
     }
 
+    std::vector<std::shared_ptr<DataType>> feature_types(
+        options.feature_field_names.size());
     for (const auto& row_template : options.row_templates) {
       if (row_template.feature_values.size() != options.feature_field_names.size()) {
         return Status::Invalid("There were names given for ",
@@ -103,11 +105,28 @@ class PivotLongerNode : public ExecNode, public TracedNode {
             " measurement columns but one of the row templates only had ",
             row_template.measurement_values.size(), " field references");
       }
+
+      for (std::size_t i = 0; i < row_template.feature_values.size(); i++) {
+        if (feature_types[i]) {
+          if (!feature_types[i]->Equals(row_template.feature_values[i]->type)) {
+            return Status::Invalid(
+                "Mixed feature types at column ", options.feature_field_names[i],
+                ".  Some row templates had the type ", feature_types[i]->ToString(),
+                " but later row templates had the type ",
+                row_template.feature_values[i]->type->ToString(),
+                ".  All row templates must have same type for each feature "
+                "column.");
+          }
+        } else {
+          feature_types[i] = row_template.feature_values[i]->type;
+        }
+      }
     }
 
     std::vector<std::shared_ptr<Field>> fields(input_schema->fields());
-    for (const auto& name : options.feature_field_names) {
-      fields.push_back(field(name, utf8()));
+    for (std::size_t i = 0; i < options.feature_field_names.size(); i++) {
+      fields.push_back(
+          field(options.feature_field_names[i], std::move(feature_types[i])));
     }
     std::vector<std::shared_ptr<DataType>> measurement_types(
         options.measurement_field_names.size());

cpp/src/arrow/acero/pivot_longer_node_test.cc

@@ -43,9 +43,15 @@ TEST(PivotLongerNode, Basic) {
           ->Table(kRowsPerBatch, kNumBatches);
 
   PivotLongerNodeOptions options;
-  options.feature_field_names = {"feature1", "feature2"};
+  options.feature_field_names = {"feature1", "feature2", "feature3"};
   options.measurement_field_names = {"meas1", "meas2"};
-  options.row_templates = {{{"a", "x"}, {{1}, {3}}}, {{"b", "y"}, {{2}, std::nullopt}}};
+  options.row_templates = {
+      {{std::make_shared<StringScalar>("a"), std::make_shared<StringScalar>("x"),
+        std::make_shared<UInt32Scalar>(12)},
+       {{1}, {3}}},
+      {{std::make_shared<StringScalar>("b"), std::make_shared<StringScalar>("y"),
+        std::make_shared<UInt32Scalar>(13)},
+       {{2}, std::nullopt}}};
 
   Declaration plan = Declaration::Sequence({
       {"table_source", TableSourceNodeOptions(std::move(input))},
@@ -62,6 +68,7 @@ TEST(PivotLongerNode, Basic) {
       field("f3", uint32()),
       field("feature1", utf8()),
       field("feature2", utf8()),
+      field("feature3", uint32()),
       field("meas1", uint32()),
       field("meas2", uint32()),
   });
@@ -98,19 +105,27 @@ TEST(PivotLongerNode, Error) {
              "There were names given for 1 feature columns but one of the row templates "
              "only had 0 feature values");
 
-  options.row_templates = {{{"x"}, {}}};
+  options.row_templates = {{{std::make_shared<StringScalar>("x")}, {}}};
   CheckError(
       options,
       "There were names given for 1 measurement columns but one of the row templates "
       "only had 0 field references");
 
-  options.row_templates = {{{"x"}, {{0}}}, {{"y"}, {{1}}}};
+  options.row_templates = {{{std::make_shared<StringScalar>("x")}, {{0}}},
+                           {{std::make_shared<StringScalar>("y")}, {{1}}}};
   CheckError(
       options,
       "Some row templates had the type uint32 but later row templates had the type bool");
 
-  options.row_templates = {{{"x"}, {std::nullopt}}, {{"y"}, {std::nullopt}}};
+  options.row_templates = {{{std::make_shared<StringScalar>("x")}, {std::nullopt}},
+                           {{std::make_shared<StringScalar>("y")}, {std::nullopt}}};
   CheckError(options, "All row templates had nullopt");
+
+  options.row_templates = {{{std::make_shared<StringScalar>("x")}, {{0}}},
+                           {{std::make_shared<UInt32Scalar>(1)}, {{0}}}};
+  CheckError(options,
+             "Some row templates had the type string but later row templates had the "
+             "type uint32");
 }
 
 // The following examples are smaller versions of examples taken from
@@ -130,11 +145,11 @@ TEST(PivotLongerNode, ExamplesFromTidyr1) {
   PivotLongerNodeOptions options;
   options.feature_field_names = {"income"};
   options.measurement_field_names = {"count"};
-  options.row_templates = {{{"<$10k"}, {{1}}},
-                           {{"$10k-20k"}, {{2}}},
-                           {{"$20k-30k"}, {{3}}},
-                           {{"$30k-40k"}, {{4}}},
-                           {{"$40k-50k"}, {{5}}}};
+  options.row_templates = {{{std::make_shared<StringScalar>("<$10k")}, {{1}}},
+                           {{std::make_shared<StringScalar>("$10k-20k")}, {{2}}},
+                           {{std::make_shared<StringScalar>("$20k-30k")}, {{3}}},
+                           {{std::make_shared<StringScalar>("$30k-40k")}, {{4}}},
+                           {{std::make_shared<StringScalar>("$40k-50k")}, {{5}}}};
 
   Declaration plan = Declaration::Sequence(
       {{"table_source", TableSourceNodeOptions(std::move(input))},
@@ -183,7 +198,8 @@ TEST(PivotLongerNode, ExamplesFromTidyr2) {
   PivotLongerNodeOptions options;
   options.feature_field_names = {"week"};
   options.measurement_field_names = {"rank"};
-  options.row_templates = {{{"1"}, {{2}}}, {{"2"}, {{3}}}};
+  options.row_templates = {{{std::make_shared<StringScalar>("1")}, {{2}}},
+                           {{std::make_shared<StringScalar>("2")}, {{3}}}};
 
   Declaration plan = Declaration::Sequence(
       {{"table_source", TableSourceNodeOptions(std::move(input))},
@@ -222,7 +238,13 @@ TEST(PivotLongerNode, ExamplesFromTidyr3) {
   PivotLongerNodeOptions options;
   options.feature_field_names = {"diagnosis", "gender", "age"};
   options.measurement_field_names = {"count"};
-  options.row_templates = {{{"sp", "m", "014"}, {{1}}}, {{"sp", "m", "1524"}, {{2}}}};
+  options.row_templates = {
+      {{std::make_shared<StringScalar>("sp"), std::make_shared<StringScalar>("m"),
+        std::make_shared<StringScalar>("014")},
+       {{1}}},
+      {{std::make_shared<StringScalar>("sp"), std::make_shared<StringScalar>("m"),
+        std::make_shared<StringScalar>("1524")},
+       {{2}}}};
 
   Declaration plan = Declaration::Sequence(
       {{"table_source", TableSourceNodeOptions(std::move(input))},
@@ -261,7 +283,8 @@ TEST(PivotLongerNode, ExamplesFromTidyr4) {
   PivotLongerNodeOptions options;
   options.feature_field_names = {"set"};
   options.measurement_field_names = {"x", "y"};
-  options.row_templates = {{{"1"}, {{0}, {2}}}, {{"2"}, {{1}, {3}}}};
+  options.row_templates = {{{std::make_shared<StringScalar>("1")}, {{0}, {2}}},
+                           {{std::make_shared<StringScalar>("2")}, {{1}, {3}}}};
 
   Declaration plan = Declaration::Sequence(
       {{"table_source", TableSourceNodeOptions(std::move(input))},