Skip to content

Commit

Permalink
GH-41923: [C++] Fix ExecuteScalar deduce all_scalar with chunked_array (
Browse files Browse the repository at this point in the history
#41925)

### Rationale for this change

ExecBatch says it's a `array` or `scalar` for execution, however, it could contains the `chunked_array` in some cases. The `is_scalar` in `compute` is wrong when input contains `chunked_array`.

### What changes are included in this PR?

Fix `is_scalar` in compute

### Are these changes tested?

Yes

### Are there any user-facing changes?

bugfix

* GitHub Issue: #41923

Lead-authored-by: mwish <maplewish117@gmail.com>
Co-authored-by: mwish <anmmscs_maple@qq.com>
Signed-off-by: mwish <maplewish117@gmail.com>
  • Loading branch information
mapleFU and mapleFU authored Jun 13, 2024
1 parent 512d245 commit 8ae1edb
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 3 deletions.
4 changes: 1 addition & 3 deletions cpp/src/arrow/compute/expression.cc
Original file line number Diff line number Diff line change
Expand Up @@ -763,9 +763,7 @@ Result<Datum> ExecuteScalarExpression(const Expression& expr, const ExecBatch& i
for (size_t i = 0; i < arguments.size(); ++i) {
ARROW_ASSIGN_OR_RAISE(
arguments[i], ExecuteScalarExpression(call->arguments[i], input, exec_context));
if (arguments[i].is_array()) {
all_scalar = false;
}
all_scalar &= arguments[i].is_scalar();
}

int64_t input_length;
Expand Down
35 changes: 35 additions & 0 deletions cpp/src/arrow/compute/expression_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -909,6 +909,41 @@ TEST(Expression, ExecuteCallWithNoArguments) {
EXPECT_EQ(actual.length(), kCount);
}

TEST(Expression, ExecuteChunkedArray) {
// GH-41923: compute should generate the right result if input
// ExecBatch is `chunked_array`.
auto input_schema = struct_({field("a", struct_({
field("a", float64()),
field("b", float64()),
}))});

auto chunked_array_input = ChunkedArrayFromJSON(input_schema, {R"([
{"a": {"a": 6.125, "b": 3.375}},
{"a": {"a": 0.0, "b": 1}}
])",
R"([
{"a": {"a": -1, "b": 4.75}}
])"});

ASSERT_OK_AND_ASSIGN(auto table_input,
Table::FromChunkedStructArray(chunked_array_input));

auto expr = add(field_ref(FieldRef("a", "a")), field_ref(FieldRef("a", "b")));

ASSERT_OK_AND_ASSIGN(expr, expr.Bind(input_schema));
std::vector<Datum> inputs{table_input->column(0)};
ExecBatch batch{inputs, 3};

ASSERT_OK_AND_ASSIGN(Datum res, ExecuteScalarExpression(expr, batch));

AssertDatumsEqual(res, ArrayFromJSON(float64(),
R"([
9.5,
1,
3.75
])"));
}

TEST(Expression, ExecuteDictionaryTransparent) {
ExpectExecute(
equal(field_ref("a"), field_ref("b")),
Expand Down

0 comments on commit 8ae1edb

Please sign in to comment.