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

docs/source/java/dataset.rst

@@ -132,12 +132,10 @@ within method ``Scanner::schema()``:
 
 .. _java-dataset-projection:
 
-Projection
-==========
+Projection (Subset of Columns)
+==============================
 
-User can specify projections in ScanOptions. For ``FileSystemDataset``, only
-column projection is allowed for now, which means, only column names
-in the projection list will be accepted. For example:
+User can specify projections in ScanOptions. For example:
 
 .. code-block:: Java
 
@@ -159,6 +157,27 @@ Or use shortcut construtor:
 
 Then all columns will be emitted during scanning.
 
+Projection (Produce New Columns) and Filters
+============================================
+
+User can specify projections (new columns) or filters in ScanOptions using Substrait. For example:
+
+.. code-block:: Java
+
+   ByteBuffer substraitExpressionFilter = getSubstraitExpressionFilter();
+   ByteBuffer substraitExpressionProject = getSubstraitExpressionProjection();
+   // Use Substrait APIs to create an Expression and serialize to a ByteBuffer
+   ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768)
+                .columns(Optional.empty())
+                .substraitExpressionFilter(substraitExpressionFilter)
+                .substraitExpressionProjection(getSubstraitExpressionProjection())
+                .build();
+
+.. seealso::
+
+   :doc:`Executing Projections and Filters Using Extended Expressions <substrait>`
+        Projections and Filters using Substrait.
+
 Read Data from HDFS
 ===================
 

java/dataset/src/main/cpp/jni_wrapper.cc

@@ -29,6 +29,8 @@
 #include "arrow/filesystem/path_util.h"
 #include "arrow/filesystem/s3fs.h"
 #include "arrow/engine/substrait/util.h"
+#include "arrow/engine/substrait/serde.h"
+#include "arrow/engine/substrait/relation.h"
 #include "arrow/ipc/api.h"
 #include "arrow/util/iterator.h"
 #include "jni_util.h"
@@ -200,7 +202,6 @@ arrow::Result<std::shared_ptr<arrow::Schema>> SchemaFromColumnNames(
       return arrow::Status::Invalid("Partition column '", ref.ToString(), "' is not in dataset schema");
     }
   }
-
   return schema(std::move(columns))->WithMetadata(input->metadata());
 }
 }  // namespace
@@ -317,6 +318,14 @@ std::shared_ptr<arrow::Table> GetTableByName(const std::vector<std::string>& nam
   return it->second;
 }
 
+std::shared_ptr<arrow::Buffer> LoadArrowBufferFromByteBuffer(JNIEnv* env, jobject byte_buffer) {
+  const auto *buff = reinterpret_cast<jbyte*>(env->GetDirectBufferAddress(byte_buffer));
+  int length = env->GetDirectBufferCapacity(byte_buffer);
+  std::shared_ptr<arrow::Buffer> buffer = JniGetOrThrow(arrow::AllocateBuffer(length));
+  std::memcpy(buffer->mutable_data(), buff, length);
+  return buffer;
+}
+
 /*
  * Class:     org_apache_arrow_dataset_jni_NativeMemoryPool
  * Method:    getDefaultMemoryPool
@@ -455,11 +464,12 @@ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_closeDataset
 /*
  * Class:     org_apache_arrow_dataset_jni_JniWrapper
  * Method:    createScanner
- * Signature: (J[Ljava/lang/String;JJ)J
+ * Signature: (J[Ljava/lang/String;Ljava/nio/ByteBuffer;Ljava/nio/ByteBuffer;JJ)J
  */
 JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_createScanner(
-    JNIEnv* env, jobject, jlong dataset_id, jobjectArray columns, jlong batch_size,
-    jlong memory_pool_id) {
+    JNIEnv* env, jobject, jlong dataset_id, jobjectArray columns,
+    jobject substrait_projection, jobject substrait_filter,
+    jlong batch_size, jlong memory_pool_id) {
   JNI_METHOD_START
   arrow::MemoryPool* pool = reinterpret_cast<arrow::MemoryPool*>(memory_pool_id);
   if (pool == nullptr) {
@@ -474,6 +484,39 @@ JNIEXPORT jlong JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_createScann
     std::vector<std::string> column_vector = ToStringVector(env, columns);
     JniAssertOkOrThrow(scanner_builder->Project(column_vector));
   }
+  if (substrait_projection != nullptr) {
+    std::shared_ptr<arrow::Buffer> buffer = LoadArrowBufferFromByteBuffer(env,
+                                                            substrait_projection);
+    std::vector<arrow::compute::Expression> project_exprs;
+    std::vector<std::string> project_names;
+    arrow::engine::BoundExpressions bounded_expression =
+          JniGetOrThrow(arrow::engine::DeserializeExpressions(*buffer));
+    for(arrow::engine::NamedExpression& named_expression :
+                                        bounded_expression.named_expressions) {
+      if (!(named_expression.expression.type()->id() == arrow::Type::BOOL)) {
+        project_exprs.push_back(std::move(named_expression.expression));
+        project_names.push_back(std::move(named_expression.name));
+      }
+    }
+    JniAssertOkOrThrow(scanner_builder->Project(std::move(project_exprs), std::move(project_names)));
+  }
+  if (substrait_filter != nullptr) {
+    std::shared_ptr<arrow::Buffer> buffer = LoadArrowBufferFromByteBuffer(env,
+                                                                substrait_filter);
+    std::optional<arrow::compute::Expression> filter_expr;
+    arrow::engine::BoundExpressions bounded_expression =
+          JniGetOrThrow(arrow::engine::DeserializeExpressions(*buffer));
+    for(arrow::engine::NamedExpression& named_expression :
+                                        bounded_expression.named_expressions) {
+      if (named_expression.expression.type()->id() == arrow::Type::BOOL) {
+        if (filter_expr.has_value()) {
+          JniThrow("Only one filter expression may be provided");
+        }
+        filter_expr = named_expression.expression;
+      }
+    }
+    JniAssertOkOrThrow(scanner_builder->Filter(*filter_expr));
+  }
   JniAssertOkOrThrow(scanner_builder->BatchSize(batch_size));
 
   auto scanner = JniGetOrThrow(scanner_builder->Finish());
@@ -748,10 +791,7 @@ JNIEXPORT void JNICALL
   arrow::engine::ConversionOptions conversion_options;
   conversion_options.named_table_provider = std::move(table_provider);
   // mapping arrow::Buffer
-  auto *buff = reinterpret_cast<jbyte*>(env->GetDirectBufferAddress(plan));
-  int length = env->GetDirectBufferCapacity(plan);
-  std::shared_ptr<arrow::Buffer> buffer = JniGetOrThrow(arrow::AllocateBuffer(length));
-  std::memcpy(buffer->mutable_data(), buff, length);
+  std::shared_ptr<arrow::Buffer> buffer = LoadArrowBufferFromByteBuffer(env, plan);
   // execute plan
   std::shared_ptr<arrow::RecordBatchReader> reader_out =
     JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer, nullptr, nullptr, conversion_options));

java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java

@@ -17,6 +17,8 @@
 
 package org.apache.arrow.dataset.jni;
 
+import java.nio.ByteBuffer;
+
 /**
  * JNI wrapper for Dataset API's native implementation.
  */
@@ -66,15 +68,19 @@ private JniWrapper() {
 
   /**
    * Create Scanner from a Dataset and get the native pointer of the Dataset.
+   *
    * @param datasetId the native pointer of the arrow::dataset::Dataset instance.
    * @param columns desired column names.
    *                Columns not in this list will not be emitted when performing scan operation. Null equals
    *                to "all columns".
+   * @param substraitProjection substrait extended expression to evaluate for project new columns
+   * @param substraitFilter substrait extended expression to evaluate for apply filter
    * @param batchSize batch size of scanned record batches.
    * @param memoryPool identifier of memory pool used in the native scanner.
    * @return the native pointer of the arrow::dataset::Scanner instance.
    */
-  public native long createScanner(long datasetId, String[] columns, long batchSize, long memoryPool);
+  public native long createScanner(long datasetId, String[] columns, ByteBuffer substraitProjection,
+                                   ByteBuffer substraitFilter, long batchSize, long memoryPool);
 
   /**
    * Get a serialized schema from native instance of a Scanner.

java/dataset/src/main/java/org/apache/arrow/dataset/jni/NativeDataset.java

@@ -40,8 +40,12 @@ public synchronized NativeScanner newScan(ScanOptions options) {
     if (closed) {
       throw new NativeInstanceReleasedException();
     }
+
     long scannerId = JniWrapper.get().createScanner(datasetId, options.getColumns().orElse(null),
+        options.getSubstraitProjection().orElse(null),
+        options.getSubstraitFilter().orElse(null),
         options.getBatchSize(), context.getMemoryPool().getNativeInstanceId());
+
     return new NativeScanner(context, scannerId);
   }
 

java/dataset/src/main/java/org/apache/arrow/dataset/scanner/ScanOptions.java

@@ -17,6 +17,7 @@
 
 package org.apache.arrow.dataset.scanner;
 
+import java.nio.ByteBuffer;
 import java.util.Optional;
 
 import org.apache.arrow.util.Preconditions;
@@ -25,8 +26,10 @@
  * Options used during scanning.
  */
 public class ScanOptions {
-  private final Optional<String[]> columns;
   private final long batchSize;
+  private final Optional<String[]> columns;
+  private final Optional<ByteBuffer> substraitProjection;
+  private final Optional<ByteBuffer> substraitFilter;
 
   /**
    * Constructor.
@@ -56,6 +59,8 @@ public ScanOptions(long batchSize, Optional<String[]> columns) {
     Preconditions.checkNotNull(columns);
     this.batchSize = batchSize;
     this.columns = columns;
+    this.substraitProjection = Optional.empty();
+    this.substraitFilter = Optional.empty();
   }
 
   public ScanOptions(long batchSize) {
@@ -69,4 +74,77 @@ public Optional<String[]> getColumns() {
   public long getBatchSize() {
     return batchSize;
   }
+
+  public Optional<ByteBuffer> getSubstraitProjection() {
+    return substraitProjection;
+  }
+
+  public Optional<ByteBuffer> getSubstraitFilter() {
+    return substraitFilter;
+  }
+
+  /**
+   * Builder for Options used during scanning.
+   */
+  public static class Builder {
+    private final long batchSize;
+    private Optional<String[]> columns;
+    private ByteBuffer substraitProjection;
+    private ByteBuffer substraitFilter;
+
+    /**
+     * Constructor.
+     * @param batchSize Maximum row number of each returned {@link org.apache.arrow.vector.ipc.message.ArrowRecordBatch}
+     */
+    public Builder(long batchSize) {
+      this.batchSize = batchSize;
+    }
+
+    /**
+     * Set the Projected columns. Empty for scanning all columns.
+     *
+     * @param columns Projected columns. Empty for scanning all columns.
+     * @return the ScanOptions configured.
+     */
+    public Builder columns(Optional<String[]> columns) {
+      Preconditions.checkNotNull(columns);
+      this.columns = columns;
+      return this;
+    }
+
+    /**
+     * Set the Substrait extended expression for Projection new columns.
+     *
+     * @param substraitProjection Expressions to evaluate for project new columns.
+     * @return the ScanOptions configured.
+     */
+    public Builder substraitProjection(ByteBuffer substraitProjection) {
+      Preconditions.checkNotNull(substraitProjection);
+      this.substraitProjection = substraitProjection;
+      return this;
+    }
+
+    /**
+     * Set the Substrait extended expression for Filter.
+     *
+     * @param substraitFilter Expressions to evaluate for apply Filter.
+     * @return the ScanOptions configured.
+     */
+    public Builder substraitFilter(ByteBuffer substraitFilter) {
+      Preconditions.checkNotNull(substraitFilter);
+      this.substraitFilter = substraitFilter;
+      return this;
+    }
+
+    public ScanOptions build() {
+      return new ScanOptions(this);
+    }
+  }
+
+  private ScanOptions(Builder builder) {
+    batchSize = builder.batchSize;
+    columns = builder.columns;
+    substraitProjection = Optional.ofNullable(builder.substraitProjection);
+    substraitFilter = Optional.ofNullable(builder.substraitFilter);
+  }
 }