diff --git a/cpp/src/arrow/acero/options.cc b/cpp/src/arrow/acero/options.cc index 8bb1e10f3cb5..40f7726ea76f 100644 --- a/cpp/src/arrow/acero/options.cc +++ b/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 batches) { } } // namespace +PivotLongerRowTemplate::PivotLongerRowTemplate( + std::vector feature_values, + std::vector> 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(std::move(feature_value))); + } +} + ExecBatchSourceNodeOptions::ExecBatchSourceNodeOptions( std::shared_ptr schema, std::vector batches, ::arrow::internal::Executor* io_executor) diff --git a/cpp/src/arrow/acero/options.h b/cpp/src/arrow/acero/options.h index 827e9ea775d7..8420793b9929 100644 --- a/cpp/src/arrow/acero/options.h +++ b/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 feature_values, + PivotLongerRowTemplate(std::vector> feature_values, std::vector> measurement_values) : feature_values(std::move(feature_values)), measurement_values(std::move(measurement_values)) {} + PivotLongerRowTemplate(std::vector feature_values, + std::vector> 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 feature_values; + std::vector> 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 diff --git a/cpp/src/arrow/acero/pivot_longer_node.cc b/cpp/src/arrow/acero/pivot_longer_node.cc index c8f2a5c7b06a..04b96dc8dcd1 100644 --- a/cpp/src/arrow/acero/pivot_longer_node.cc +++ b/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 feature_values; + std::vector> feature_values; std::vector> measurement_paths; static Result Make(const PivotLongerRowTemplate& unbound, @@ -65,7 +65,7 @@ struct BoundRowTemplate { } private: - BoundRowTemplate(std::vector feature_values, + BoundRowTemplate(std::vector> feature_values, std::vector> 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> 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> 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> measurement_types( options.measurement_field_names.size()); diff --git a/cpp/src/arrow/acero/pivot_longer_node_test.cc b/cpp/src/arrow/acero/pivot_longer_node_test.cc index 9c548a2f23f3..7b71f60be055 100644 --- a/cpp/src/arrow/acero/pivot_longer_node_test.cc +++ b/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("a"), std::make_shared("x"), + std::make_shared(12)}, + {{1}, {3}}}, + {{std::make_shared("b"), std::make_shared("y"), + std::make_shared(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 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{}, {{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{}, {{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("x")}, {{0}}}, + {{std::make_shared(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(1)}, {{2}}}, + {{std::make_shared(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 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
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);