westonpace commented on code in PR #34651: URL: https://github.com/apache/arrow/pull/34651#discussion_r1198949600
##########
cpp/src/arrow/engine/substrait/serde_test.cc:
##########
@@ -5145,6 +5145,128 @@ TEST(Substrait, CompoundEmitWithFilter) {
CheckRoundTripResult(std::move(expected_table), buf, {}, conversion_options);
}
+TEST(Substrait, SortAndFetch) {
+ // Sort by A, ascending, take items [2, 5), then sort by B descending
+ std::string substrait_json = R"({
+ "version": {
+ "major_number": 9999,
+ "minor_number": 9999,
+ "patch_number": 9999
+ },
+ "relations": [
+ {
+ "rel": {
+ "sort": {
+ "input": {
+ "fetch": {
+ "input": {
+ "sort": {
+ "input": {
+ "read": {
+ "base_schema": {
+ "names": [
+ "A",
+ "B"
+ ],
+ "struct": {
+ "types": [
+ {
+ "i32": {}
+ },
+ {
+ "i32": {}
+ }
+ ]
+ }
+ },
+ "namedTable": {
+ "names": [
+ "table"
+ ]
+ }
+ }
+ },
+ "sorts": [
+ {
+ "expr": {
+ "selection": {
+ "directReference": {
+ "structField": {
+ "field": 0
+ }
+ },
+ "rootReference": {}
+ }
+ },
+ "direction":
"SORT_DIRECTION_ASC_NULLS_FIRST"
+ }
+ ]
+ }
+ },
+ "offset": 2,
+ "count": 3
+ }
+ },
+ "sorts": [
+ {
+ "expr": {
+ "selection": {
+ "directReference": {
+ "structField": {
+ "field": 1
+ }
+ },
+ "rootReference": {}
+ }
+ },
+ "direction": "SORT_DIRECTION_DESC_NULLS_LAST"
+ }
+ ]
+ }
+ }
+ }
+ ],
+ "extension_uris": [],
+ "extensions": []
+})";
+
+ ASSERT_OK_AND_ASSIGN(auto buf, internal::SubstraitFromJSON("Plan",
substrait_json));
+ auto test_schema = schema({field("A", int32()), field("B", int32())});
+
+ auto input_table = TableFromJSON(test_schema, {R"([
+ [null, null],
+ [5, 8],
+ [null, null],
+ [null, null],
+ [3, 4],
+ [9, 6],
+ [4, 5]
+ ])"});
+
+ // First sort by A, ascending, nulls first to yield rows:
+ // 0, 2, 3, 4, 6, 1, 5
+ // Apply fetch to grab rows 3, 4, 6
+ // Then sort by B, descending, to yield rows 6, 4, 3
+
+ auto output_table = TableFromJSON(test_schema, {R"([
+ [4, 5],
+ [3, 4],
+ [null, null]
+ ])"});
+
+ NamedTableProvider table_provider = [&](const std::vector<std::string>&
names,
+ const Schema&) {
+ std::shared_ptr<acero::ExecNodeOptions> options =
+ std::make_shared<acero::TableSourceNodeOptions>(input_table);
+ return acero::Declaration("table_source", {}, options, "mock_source");
+ };
Review Comment:
Good cleanup, thanks. I've switched to this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
