GH-34252: [Java] Support ScannerBuilder::Project or ScannerBuilder::Filter as a Substrait proto extended expression

cpp/cmake_modules/ThirdpartyToolchain.cmake

@@ -1849,7 +1849,7 @@ macro(build_substrait)
 
   # Note: not all protos in Substrait actually matter to plan
   # consumption. No need to build the ones we don't need.
-  set(SUBSTRAIT_PROTOS algebra extensions/extensions plan type)
+  set(SUBSTRAIT_PROTOS algebra extended_expression extensions/extensions plan type)
   set(ARROW_SUBSTRAIT_PROTOS extension_rels)
   set(ARROW_SUBSTRAIT_PROTOS_DIR "${CMAKE_SOURCE_DIR}/proto")
 

cpp/src/arrow/dataset/file_parquet.cc

@@ -224,6 +224,41 @@ Status ResolveOneFieldRef(
   return Status::OK();
 }
 
+bool IsValidFieldRef(const FieldRef& ref) {
+  if (ref.IsName()) return true;
+  if (const auto* nested_refs = ref.nested_refs()) {
+    for (const auto& nested_ref : *nested_refs) {
+      if (!nested_ref.IsName()) return false;
+    }
+    return true;
+  }
+  return false;
+}
+
+// Converts a field ref into a position-independent ref (containing only a sequence of
+// names) based on the dataset schema. Returns `false` if no conversion was needed.
+Result<bool> ValidateFieldRef(const FieldRef& ref, const Schema& dataset_schema,
+                              FieldRef* out) {
+  if (ARROW_PREDICT_TRUE(IsValidFieldRef(ref))) {
+    return false;
+  }
+
+  ARROW_ASSIGN_OR_RAISE(auto path, ref.FindOne(dataset_schema));
+  std::vector<FieldRef> named_refs;
+  named_refs.reserve(path.indices().size());
+
+  const FieldVector* child_fields = &dataset_schema.fields();
+  for (auto index : path) {
+    const auto& child_field = *(*child_fields)[index];
+    named_refs.emplace_back(child_field.name());
+    child_fields = &child_field.type()->fields();
+  }
+
+  *out =
+      named_refs.size() == 1 ? std::move(named_refs[0]) : FieldRef(std::move(named_refs));
+  return true;
+}
+
 // Compute the column projection based on the scan options
 Result<std::vector<int>> InferColumnProjection(const parquet::arrow::FileReader& reader,
                                                const ScanOptions& options) {
@@ -248,7 +283,12 @@ Result<std::vector<int>> InferColumnProjection(const parquet::arrow::FileReader&
   }
 
   std::vector<int> columns_selection;
-  for (const auto& ref : field_refs) {
+  for (auto& ref : field_refs) {
+    // In the (unlikely) absence of a known dataset schema, we require that all
+    // materialized refs are named.
+    if (options.dataset_schema) {
+      ARROW_RETURN_NOT_OK(ValidateFieldRef(ref, *options.dataset_schema, &ref));
+    }
     RETURN_NOT_OK(ResolveOneFieldRef(manifest, ref, field_lookup, duplicate_fields,
                                      &columns_selection));
   }

cpp/src/arrow/dataset/file_parquet_test.cc

@@ -694,5 +694,55 @@ TEST(TestParquetStatistics, NullMax) {
   EXPECT_EQ(stat_expression->ToString(), "(x >= 1)");
 }
 
+// Tests round-trip projection with nested/indexed FieldRefs
+// https://github.com/apache/arrow/issues/35579
+TEST(TestRoundTrip, ProjectedFieldRefs) {
+  auto test_schema = schema(
+      {field("id", uint32()),
+       field("info", struct_({field("name", utf8()),
+                              field("data", struct_({field("amount", float64()),
+                                                     field("percent", float32())}))}))});
+  auto test_table = TableFromJSON(test_schema, {R"([
+    {"id": 1, "info": {"name": "a", "data": {"amount": 10.3, "percent": 0.1}}},
+    {"id": 2, "info": {"name": "b", "data": {"amount": 11.6, "percent": 0.2}}},
+    {"id": 3, "info": {"name": "c", "data": {"amount": 12.9, "percent": 0.3}}},
+    {"id": 4, "info": {"name": "d", "data": {"amount": 14.2, "percent": 0.4}}},
+    {"id": 5, "info": {"name": "e", "data": {"amount": 15.5, "percent": 0.5}}},
+    {"id": 6, "info": {"name": "f", "data": {"amount": 16.8, "percent": 0.6}}}])"});
+  ASSERT_OK(test_table->ValidateFull());
+
+  ASSERT_OK_AND_ASSIGN(auto fs, fs::internal::MockFileSystem::Make(fs::kNoTime, {}));
+  ASSERT_OK_AND_ASSIGN(auto sink, fs->OpenOutputStream("test.parquet"));
+  ASSERT_OK(parquet::arrow::WriteTable(*test_table, arrow::default_memory_pool(), sink));
+  ASSERT_OK(sink->Close());
+  auto format = std::make_shared<ParquetFileFormat>();
+  ASSERT_OK_AND_ASSIGN(auto factory,
+                       FileSystemDatasetFactory::Make(fs, fs::FileSelector(), format,
+                                                      FileSystemFactoryOptions{}));
+  ASSERT_OK_AND_ASSIGN(auto dataset, factory->Finish());
+  AssertSchemaEqual(test_schema, dataset->schema());
+
+  auto expected_schema = schema({field("value", float32())});
+  auto expected_table = TableFromJSON(expected_schema, {R"([
+    {"value": 0.1},{"value": 0.2},{"value": 0.3},
+    {"value": 0.4},{"value": 0.5},{"value": 0.6}])"});
+
+  std::vector<FieldRef> equivalent_refs = {
+      FieldRef("info", "data", "percent"), FieldRef("info", 1, 1),
+      FieldRef(1, 1, "percent"),           FieldRef(1, 1, 1),
+      FieldRef(1, FieldRef("data", 1)),    FieldRef(FieldRef(1), FieldRef(1, 1)),
+  };
+  for (const auto& ref : equivalent_refs) {
+    ARROW_SCOPED_TRACE("ref = ", ref.ToString());
+
+    ScannerBuilder builder(dataset);
+    ASSERT_OK(builder.Project({field_ref(ref)}, {"value"}));
+    ASSERT_OK_AND_ASSIGN(auto scanner, builder.Finish());
+    ASSERT_OK_AND_ASSIGN(auto actual_table, scanner->ToTable());
+
+    AssertTablesEqual(*expected_table, *actual_table);
+  }
+}
+
 }  // namespace dataset
 }  // namespace arrow

cpp/src/arrow/engine/CMakeLists.txt

@@ -21,6 +21,7 @@ arrow_install_all_headers("arrow/engine")
 
 set(ARROW_SUBSTRAIT_SRCS
     substrait/expression_internal.cc
+    substrait/extended_expression_internal.cc
     substrait/extension_set.cc
     substrait/extension_types.cc
     substrait/options.cc

cpp/src/arrow/engine/substrait/expression_internal.cc

@@ -1167,12 +1167,36 @@ Result<std::unique_ptr<substrait::Expression>> ToProto(
   }
 
   // other expression types dive into extensions immediately
-  ARROW_ASSIGN_OR_RAISE(
-      ExtensionIdRegistry::ArrowToSubstraitCall converter,
-      ext_set->registry()->GetArrowToSubstraitCall(call->function_name));
-  ARROW_ASSIGN_OR_RAISE(SubstraitCall substrait_call, converter(*call));
-  ARROW_ASSIGN_OR_RAISE(std::unique_ptr<substrait::Expression::ScalarFunction> scalar_fn,
-                        EncodeSubstraitCall(substrait_call, ext_set, conversion_options));
+  Result<ExtensionIdRegistry::ArrowToSubstraitCall> maybe_converter =
+      ext_set->registry()->GetArrowToSubstraitCall(call->function_name);
+
+  ExtensionIdRegistry::ArrowToSubstraitCall converter;
+  std::unique_ptr<substrait::Expression::ScalarFunction> scalar_fn;
+  if (maybe_converter.ok()) {
+    converter = *maybe_converter;
+    ARROW_ASSIGN_OR_RAISE(SubstraitCall substrait_call, converter(*call));
+    ARROW_ASSIGN_OR_RAISE(
+        scalar_fn, EncodeSubstraitCall(substrait_call, ext_set, conversion_options));
+  } else if (maybe_converter.status().IsNotImplemented() &&
+             conversion_options.allow_arrow_extensions) {
+    if (call->options) {
+      return Status::NotImplemented(
+          "The function ", call->function_name,
+          " has no Substrait mapping.  Arrow extensions are enabled but the call "
+          "contains function options and there is no current mechanism to encode those.");
+    }
+    SubstraitCall substrait_call(
+        Id{kArrowSimpleExtensionFunctionsUri, call->function_name},
+        call->type.GetSharedPtr(),
+        /*nullable=*/true);
+    for (int i = 0; i < static_cast<int>(call->arguments.size()); i++) {
+      substrait_call.SetValueArg(i, call->arguments[i]);
+    }
+    ARROW_ASSIGN_OR_RAISE(
+        scalar_fn, EncodeSubstraitCall(substrait_call, ext_set, conversion_options));
+  } else {
+    return maybe_converter.status();
+  }
   out->set_allocated_scalar_function(scalar_fn.release());
   return std::move(out);
 }

cpp/src/arrow/engine/substrait/extended_expression_internal.cc

@@ -0,0 +1,196 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// This API is EXPERIMENTAL.
+
+#include "arrow/engine/substrait/extended_expression_internal.h"
+
+#include "arrow/engine/substrait/expression_internal.h"
+#include "arrow/engine/substrait/relation_internal.h"
+#include "arrow/engine/substrait/type_internal.h"
+#include "arrow/engine/substrait/util.h"
+#include "arrow/engine/substrait/util_internal.h"
+#include "arrow/util/iterator.h"
+#include "arrow/util/string.h"
+
+namespace arrow {
+namespace engine {
+
+namespace {
+Result<ExtensionSet> GetExtensionSetFromExtendedExpression(
+    const substrait::ExtendedExpression& expr,
+    const ConversionOptions& conversion_options, const ExtensionIdRegistry* registry) {
+  return GetExtensionSetFromMessage(expr, conversion_options, registry);
+}
+
+Status AddExtensionSetToExtendedExpression(const ExtensionSet& ext_set,
+                                           substrait::ExtendedExpression* expr) {
+  return AddExtensionSetToMessage(ext_set, expr);
+}
+
+Status VisitNestedFields(const DataType& type,
+                         std::function<Status(const Field&)> visitor) {
+  if (!is_nested(type.id())) {
+    return Status::OK();
+  }
+  for (const auto& field : type.fields()) {
+    ARROW_RETURN_NOT_OK(VisitNestedFields(*field->type(), visitor));
+    ARROW_RETURN_NOT_OK(visitor(*field));
+  }
+  return Status::OK();
+}
+
+Result<NamedExpression> ExpressionFromProto(
+    const substrait::ExpressionReference& expression, const Schema& input_schema,
+    const ExtensionSet& ext_set, const ConversionOptions& conversion_options,
+    const ExtensionIdRegistry* registry) {
+  NamedExpression named_expr;
+  switch (expression.expr_type_case()) {
+    case substrait::ExpressionReference::ExprTypeCase::kExpression: {
+      ARROW_ASSIGN_OR_RAISE(
+          named_expr.expression,
+          FromProto(expression.expression(), ext_set, conversion_options));
+      break;
+    }
+    case substrait::ExpressionReference::ExprTypeCase::kMeasure: {
+      return Status::NotImplemented("ExtendedExpression containing aggregate functions");
+    }
+    default: {
+      return Status::Invalid(
+          "Unrecognized substrait::ExpressionReference::ExprTypeCase: ",
+          expression.expr_type_case());
+    }
+  }
+
+  ARROW_ASSIGN_OR_RAISE(named_expr.expression, named_expr.expression.Bind(input_schema));
+  const DataType& output_type = *named_expr.expression.type();
+
+  // An expression reference has multiple output names which is redundant given the type
+  // should contain any nested names.  We only need the final field name.  However, we
+  // should verify that the output names match the schema
+  int output_name_idx = 0;
+  ARROW_RETURN_NOT_OK(VisitNestedFields(output_type, [&](const Field& field) {
+    if (output_name_idx >= expression.output_names_size()) {
+      return Status::Invalid("Ambiguous plan.  Expression had ",
+                             expression.output_names_size(),
+                             " output names but the field in base_schema had type ",
+                             output_type.ToString(), " which needs more output names");
+    }
+    if (field.name() != expression.output_names(output_name_idx)) {
+      return Status::Invalid("Ambiguous plan.  Expression had output type ",
+                             output_type.ToString(),
+                             " which contains a nested field named ", field.name(),
+                             " but the output_names in the Substrait message contained ",
+                             expression.output_names(output_name_idx));
+    }
+    output_name_idx++;
+    return Status::OK();
+  }));
+  // The last name is the actual field name that we can't verify but there should only
+  // be one extra name.
+  if (output_name_idx < expression.output_names_size() - 1) {
+    return Status::Invalid("Ambiguous plan.  Expression had ",
+                           expression.output_names_size(),
+                           " output names but the field in base_schema had type ",
+                           output_type.ToString(), " which doesn't have enough fields");
+  }
+  if (expression.output_names_size() == 0) {
+    // This is potentially invalid substrait but we can handle it
+    named_expr.name = "";
+  } else {
+    named_expr.name = expression.output_names(expression.output_names_size() - 1);
+  }
+  return named_expr;
+}
+
+Result<std::unique_ptr<substrait::ExpressionReference>> CreateExpressionReference(
+    const std::string& name, const Expression& expr, ExtensionSet* ext_set,
+    const ConversionOptions& conversion_options) {
+  auto expr_ref = std::make_unique<substrait::ExpressionReference>();
+  ARROW_RETURN_NOT_OK(VisitNestedFields(*expr.type(), [&](const Field& field) {
+    expr_ref->add_output_names(field.name());
+    return Status::OK();
+  }));
+  expr_ref->add_output_names(name);
+  ARROW_ASSIGN_OR_RAISE(std::unique_ptr<substrait::Expression> expression,
+                        ToProto(expr, ext_set, conversion_options));
+  expr_ref->set_allocated_expression(expression.release());
+  return std::move(expr_ref);
+}
+
+}  // namespace
+
+Result<BoundExpressions> FromProto(const substrait::ExtendedExpression& expression,
+                                   ExtensionSet* ext_set_out,
+                                   const ConversionOptions& conversion_options,
+                                   const ExtensionIdRegistry* registry) {
+  BoundExpressions bound_expressions;
+  ARROW_RETURN_NOT_OK(CheckVersion(expression.version().major_number(),
+                                   expression.version().minor_number()));
+  if (expression.has_advanced_extensions()) {
+    return Status::NotImplemented("Advanced extensions in ExtendedExpression");
+  }
+  ARROW_ASSIGN_OR_RAISE(
+      ExtensionSet ext_set,
+      GetExtensionSetFromExtendedExpression(expression, conversion_options, registry));
+
+  ARROW_ASSIGN_OR_RAISE(bound_expressions.schema,
+                        FromProto(expression.base_schema(), ext_set, conversion_options));
+
+  bound_expressions.named_expressions.reserve(expression.referred_expr_size());
+
+  for (const auto& referred_expr : expression.referred_expr()) {
+    ARROW_ASSIGN_OR_RAISE(NamedExpression named_expr,
+                          ExpressionFromProto(referred_expr, *bound_expressions.schema,
+                                              ext_set, conversion_options, registry));
+    bound_expressions.named_expressions.push_back(std::move(named_expr));
+  }
+
+  if (ext_set_out) {
+    *ext_set_out = std::move(ext_set);
+  }
+
+  return std::move(bound_expressions);
+}
+
+Result<std::unique_ptr<substrait::ExtendedExpression>> ToProto(
+    const BoundExpressions& bound_expressions, ExtensionSet* ext_set,
+    const ConversionOptions& conversion_options) {
+  auto expression = std::make_unique<substrait::ExtendedExpression>();
+  expression->set_allocated_version(CreateVersion().release());
+  ARROW_ASSIGN_OR_RAISE(std::unique_ptr<substrait::NamedStruct> base_schema,
+                        ToProto(*bound_expressions.schema, ext_set, conversion_options));
+  expression->set_allocated_base_schema(base_schema.release());
+  for (const auto& named_expression : bound_expressions.named_expressions) {
+    Expression bound_expr = named_expression.expression;
+    if (!bound_expr.IsBound()) {
+      // This will use the default function registry.  Most of the time that will be fine.
+      // In the cases where this is not what the user wants then the user should make sure
+      // to pass in bound expressions.
+      ARROW_ASSIGN_OR_RAISE(bound_expr, bound_expr.Bind(*bound_expressions.schema));
+    }
+    ARROW_ASSIGN_OR_RAISE(std::unique_ptr<substrait::ExpressionReference> expr_ref,
+                          CreateExpressionReference(named_expression.name, bound_expr,
+                                                    ext_set, conversion_options));
+    expression->mutable_referred_expr()->AddAllocated(expr_ref.release());
+  }
+  RETURN_NOT_OK(AddExtensionSetToExtendedExpression(*ext_set, expression.get()));
+  return std::move(expression);
+}
+
+}  // namespace engine
+}  // namespace arrow