GH-47390: [C++][Acero] Allow for any type of scalar in Pivot longer features

cpp/src/arrow/acero/options.cc

@@ -18,6 +18,7 @@
 #include "arrow/acero/options.h"
 #include "arrow/acero/exec_plan.h"
 #include "arrow/io/util_internal.h"
+#include "arrow/scalar.h"
 #include "arrow/table.h"
 #include "arrow/util/async_generator.h"
 #include "arrow/util/logging.h"
@@ -62,6 +63,17 @@ ExecBatchIteratorMaker VecToItMaker(std::vector<ExecBatch> batches) {
 }
 }  // namespace
 
+PivotLongerRowTemplate::PivotLongerRowTemplate(
+    std::vector<std::string> feature_values,
+    std::vector<std::optional<FieldRef>> measurement_values)
+    : measurement_values(std::move(measurement_values)) {
+  this->feature_values.reserve(feature_values.size());
+  for (auto& feature_value : feature_values) {
+    this->feature_values.push_back(
+        std::make_shared<StringScalar>(std::move(feature_value)));
+  }
+}
+
 ExecBatchSourceNodeOptions::ExecBatchSourceNodeOptions(
     std::shared_ptr<Schema> schema, std::vector<ExecBatch> batches,
     ::arrow::internal::Executor* io_executor)

cpp/src/arrow/acero/options.h

@@ -780,15 +780,17 @@ 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)) {}
+  PivotLongerRowTemplate(std::vector<std::string> feature_values,
+                         std::vector<std::optional<FieldRef>> 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,32 @@ 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 (!row_template.feature_values[i]) {
+          return Status::Invalid("Feature value at column ",
+                                 options.feature_field_names[i], " must not be null");
+        }
+        if (feature_types[i]) {
+          if (!feature_types[i]->Equals(row_template.feature_values[i]->type)) {
+            return Status::TypeError(
+                "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()),
   });
@@ -70,7 +77,8 @@ TEST(PivotLongerNode, Basic) {
   AssertSchemaEqual(expected_out_schema, output->schema());
 }
 
-void CheckError(const PivotLongerNodeOptions& options, const std::string& message) {
+void CheckError(const PivotLongerNodeOptions& options, const std::string& message,
+                StatusCode code = StatusCode::Invalid) {
   std::shared_ptr<Table> input = gen::Gen({gen::Step(), gen::Random(boolean())})
                                      ->FailOnError()
                                      ->Table(/*rows_per_chunk=*/1, /*num_chunks=*/1);
@@ -81,19 +89,19 @@ void CheckError(const PivotLongerNodeOptions& options, const std::string& messag
   });
 
   ASSERT_THAT(DeclarationToStatus(std::move(plan)),
-              Raises(StatusCode::Invalid, testing::HasSubstr(message)));
+              Raises(code, testing::HasSubstr(message)));
 }
 
 TEST(PivotLongerNode, Error) {
   PivotLongerNodeOptions options;
   CheckError(options, "There must be at least one row template");
 
-  options.row_templates = {{{}, {{0}}}};
+  options.row_templates = {{std::vector<std::string>{}, {{0}}}};
   CheckError(options, "at least one feature column and one measurement column");
 
   options.feature_field_names = {"feat1"};
   options.measurement_field_names = {"meas1"};
-  options.row_templates = {{{}, {{0}}}};
+  options.row_templates = {{std::vector<std::string>{}, {{0}}}};
   CheckError(options,
              "There were names given for 1 feature columns but one of the row templates "
              "only had 0 feature values");
@@ -111,6 +119,13 @@ TEST(PivotLongerNode, Error) {
 
   options.row_templates = {{{"x"}, {std::nullopt}}, {{"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",
+             StatusCode::TypeError);
 }
 
 // The following examples are smaller versions of examples taken from
@@ -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<UInt32Scalar>(1)}, {{2}}},
+                           {{std::make_shared<UInt32Scalar>(2)}, {{3}}}};
 
   Declaration plan = Declaration::Sequence(
       {{"table_source", TableSourceNodeOptions(std::move(input))},
@@ -196,14 +212,14 @@ TEST(PivotLongerNode, ExamplesFromTidyr2) {
                        DeclarationToTable(std::move(plan)));
 
   std::shared_ptr<Schema> expected_schema =
-      schema({field("artist", utf8()), field("track", utf8()), field("week", utf8()),
+      schema({field("artist", utf8()), field("track", utf8()), field("week", uint32()),
               field("rank", float64())});
   std::shared_ptr<Table> expected = TableFromJSON(expected_schema, {{
                                                                        R"([
-        ["2 Pac", "Baby Don't Cry", "1", 87],
-        ["2Ge+her", "The Hardest Part Of", "1", 91],
-        ["2 Pac", "Baby Don't Cry", "2", 82],
-        ["2Ge+her", "The Hardest Part Of", "2", 87]
+        ["2 Pac", "Baby Don't Cry", 1, 87],
+        ["2Ge+her", "The Hardest Part Of", 1, 91],
+        ["2 Pac", "Baby Don't Cry", 2, 82],
+        ["2Ge+her", "The Hardest Part Of", 2, 87]
     ])"}});
 
   AssertTablesEqual(*expected, *output, /*same_chunk_layout=*/false);