address comments

apache · May 20, 2024 · 0f35c70 · 0f35c70
1 parent 1223f36
commit 0f35c70
Show file tree

Hide file tree

Showing 7 changed files with 40 additions and 42 deletions.
diff --git a/cpp/src/arrow/dataset/file_csv.cc b/cpp/src/arrow/dataset/file_csv.cc
@@ -527,7 +527,7 @@ Result<std::shared_ptr<FragmentScanOptions>> CsvFragmentScanOptions::from(
         column_types[field->name()] = field->type();
       }
     } else {
-      return Status::Invalid("Not support this config " + it.first);
+      return Status::Invalid("Config " + it.first + "is not supported.");
     }
   }
   return options;

diff --git a/cpp/src/arrow/engine/substrait/extension_internal.cc b/cpp/src/arrow/engine/substrait/extension_internal.cc
@@ -26,19 +26,16 @@ namespace engine {
 
 Status FromProto(const substrait::extensions::AdvancedExtension& extension,
                  std::unordered_map<std::string, std::string>& out) {
-  if (!extension.has_enhancement()) {
-    return Status::Invalid("AdvancedExtension does not have enhancement");
-  }
+  ARROW_RETURN_IF(!extension.has_enhancement(),
+                  Status::Invalid("AdvancedExtension does not have enhancement."));
   const auto& enhancement = extension.enhancement();
-  substrait::Expression_Literal literal;
-
-  if (!enhancement.UnpackTo(&literal)) {
-    return Status::Invalid("Unpack the literal failed");
-  }
-
-  if (!literal.has_map()) {
-    return Status::Invalid("Literal does not have map");
-  }
+  substrait::Expression expression;
+  ARROW_RETURN_IF(!enhancement.UnpackTo(&expression),
+                  Status::Invalid("Unpack the Expression failed."));
+  ARROW_RETURN_IF(!expression.has_literal(),
+                  Status::Invalid("Expression does not have a literal."));
+  auto literal = expression.literal();
+  ARROW_RETURN_IF(!literal.has_map(), Status::Invalid("Literal does not have a map."));
   auto literalMap = literal.map();
   auto size = literalMap.key_values_size();
   for (auto i = 0; i < size; i++) {

diff --git a/java/dataset/src/main/cpp/jni_wrapper.cc b/java/dataset/src/main/cpp/jni_wrapper.cc
@@ -164,9 +164,7 @@ GetFragmentScanOptions(jint file_format_id,
       return arrow::dataset::CsvFragmentScanOptions::from(configs);
 #endif
     default:
-      std::string error_message =
-          "illegal file format id: " + std::to_string(file_format_id);
-      return arrow::Status::Invalid(error_message);
+      return arrow::Status::Invalid("Illegal file format id: " ,file_format_id);
   }
 }
 

diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/scanner/FragmentScanOptions.java b/java/dataset/src/main/java/org/apache/arrow/dataset/scanner/FragmentScanOptions.java
@@ -20,9 +20,7 @@
 import java.nio.ByteBuffer;
 import java.util.Map;
 
-import org.apache.arrow.dataset.substrait.StringMapNode;
-
-import com.google.protobuf.Any;
+import org.apache.arrow.dataset.substrait.util.ConvertUtil;
 
 import io.substrait.proto.AdvancedExtension;
 
@@ -43,11 +41,8 @@ default ByteBuffer serializeMap(Map<String, String> config) {
     if (config.isEmpty()) {
       return null;
     }
-    StringMapNode stringMapNode = new StringMapNode(config);
-    AdvancedExtension.Builder extensionBuilder = AdvancedExtension.newBuilder();
-    Any.Builder builder = extensionBuilder.getEnhancementBuilder();
-    builder.setValue(stringMapNode.toProtobuf().toByteString());
-    AdvancedExtension extension = extensionBuilder.build();
+
+    AdvancedExtension extension = ConvertUtil.expressionToExtension(ConvertUtil.mapToExpression(config));
     ByteBuffer buf = ByteBuffer.allocateDirect(extension.getSerializedSize());
     buf.put(extension.toByteArray());
     return buf;

diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/scanner/csv/CsvConvertOptions.java b/java/dataset/src/main/java/org/apache/arrow/dataset/scanner/csv/CsvConvertOptions.java
@@ -36,10 +36,6 @@ public Optional<ArrowSchema> getArrowSchema() {
     return cSchema;
   }
 
-  public long getArrowSchemaAddress() {
-    return cSchema.isPresent() ? cSchema.get().memoryAddress() : -1;
-  }
-
   public Map<String, String> getConfigs() {
     return configs;
   }

diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/scanner/csv/CsvFragmentScanOptions.java b/java/dataset/src/main/java/org/apache/arrow/dataset/scanner/csv/CsvFragmentScanOptions.java
@@ -71,7 +71,10 @@ public ByteBuffer serialize() {
             parseOptions.entrySet().stream()),
         convertOptions.getConfigs().entrySet().stream()).collect(
         Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
-    options.put("column_type", Long.toString(convertOptions.getArrowSchemaAddress()));
+
+    if (convertOptions.getArrowSchema().isPresent()) {
+      options.put("column_type", Long.toString(convertOptions.getArrowSchema().get().memoryAddress()));
+    }
     return serializeMap(options);
   }
 

diff --git a/...rrow/dataset/substrait/StringMapNode.java → ...w/dataset/substrait/util/ConvertUtil.java b/...rrow/dataset/substrait/StringMapNode.java → ...w/dataset/substrait/util/ConvertUtil.java
@@ -15,27 +15,23 @@
  * limitations under the License.
  */
 
-package org.apache.arrow.dataset.substrait;
+package org.apache.arrow.dataset.substrait.util;
 
-import java.io.Serializable;
-import java.util.HashMap;
 import java.util.Map;
 
-import io.substrait.proto.Expression;
+import com.google.protobuf.Any;
 
-public class StringMapNode implements Serializable {
-  private final Map<String, String> values = new HashMap<>();
+import io.substrait.proto.AdvancedExtension;
+import io.substrait.proto.Expression;
 
-  public StringMapNode(Map<String, String> values) {
-    this.values.putAll(values);
-  }
+public class ConvertUtil {
 
   /**
-   * Serialize String map.
+   * Convert map to substrait Expression.
    *
-   * @return Substrait Literal
+   * @return Substrait Expression
    */
-  public Expression.Literal toProtobuf() {
+  public static Expression mapToExpression(Map<String, String> values) {
     Expression.Literal.Builder literalBuilder = Expression.Literal.newBuilder();
     Expression.Literal.Map.KeyValue.Builder keyValueBuilder =
         Expression.Literal.Map.KeyValue.newBuilder();
@@ -48,6 +44,19 @@ public Expression.Literal toProtobuf() {
       mapBuilder.addKeyValues(keyValueBuilder.build());
     }
     literalBuilder.setMap(mapBuilder.build());
-    return literalBuilder.build();
+    return Expression.newBuilder().setLiteral(literalBuilder.build()).build();
+  }
+
+  /**
+   * Add substrait expression to AdvancedExtension.
+   *
+   * @param expr Substrait Expression.
+   * @return Substrait AdvancedExtension
+   */
+  public static AdvancedExtension expressionToExtension(Expression expr) {
+    AdvancedExtension.Builder extensionBuilder = AdvancedExtension.newBuilder();
+    Any.Builder builder = extensionBuilder.getEnhancementBuilder();
+    builder.setValue(expr.toByteString());
+    return extensionBuilder.build();
   }
 }