From e89f5627e4f1824c36e2e60cfaee3e9f3e977e87 Mon Sep 17 00:00:00 2001
From: xi chaomin <36392121+xicm@users.noreply.github.com>
Date: Mon, 13 Jun 2022 22:22:12 +0800
Subject: [PATCH] [HUDI-3682] testReaderFilterRowKeys fails in
 TestHoodieOrcReaderWriter (#5790)

TestReaderFilterRowKeys needs to get the key from RECORD_KEY_METADATA_FIELD, but the writer in current UT does not populate the meta field and the schema does not contains meta fields.

This fix writes data with schema which contains meta fields and calls writeAvroWithMetadata for writing.

Co-authored-by: xicm <xicm@asiainfo.com>
---
 .../io/storage/TestHoodieOrcReaderWriter.java  | 10 +++++-----
 .../io/storage/TestHoodieReaderWriterBase.java | 18 +++++++++++++++++-
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieOrcReaderWriter.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieOrcReaderWriter.java
index 282f10269..817fc25a5 100644
--- a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieOrcReaderWriter.java
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieOrcReaderWriter.java
@@ -34,6 +34,7 @@ import org.apache.orc.Reader;
 import org.mockito.Mockito;
 
 import java.io.IOException;
+import java.util.function.Supplier;
 
 import static org.apache.hudi.avro.HoodieAvroWriteSupport.HOODIE_AVRO_BLOOM_FILTER_METADATA_KEY;
 import static org.apache.hudi.avro.HoodieAvroWriteSupport.HOODIE_MAX_RECORD_KEY_FOOTER;
@@ -41,6 +42,7 @@ import static org.apache.hudi.avro.HoodieAvroWriteSupport.HOODIE_MIN_RECORD_KEY_
 import static org.apache.hudi.io.storage.HoodieOrcConfig.AVRO_SCHEMA_METADATA_KEY;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.when;
 
 public class TestHoodieOrcReaderWriter extends TestHoodieReaderWriterBase {
 
@@ -59,6 +61,9 @@ public class TestHoodieOrcReaderWriter extends TestHoodieReaderWriterBase {
     int maxFileSize = Integer.parseInt(HoodieStorageConfig.ORC_FILE_MAX_SIZE.defaultValue());
     HoodieOrcConfig config = new HoodieOrcConfig(conf, CompressionKind.ZLIB, orcStripSize, orcBlockSize, maxFileSize, filter);
     TaskContextSupplier mockTaskContextSupplier = Mockito.mock(TaskContextSupplier.class);
+    Supplier<Integer> partitionSupplier = Mockito.mock(Supplier.class);
+    when(mockTaskContextSupplier.getPartitionIdSupplier()).thenReturn(partitionSupplier);
+    when(partitionSupplier.get()).thenReturn(10);
     String instantTime = "000";
     return new HoodieOrcWriter<>(instantTime, getFilePath(), config, avroSchema, mockTaskContextSupplier);
   }
@@ -92,9 +97,4 @@ public class TestHoodieOrcReaderWriter extends TestHoodieReaderWriterBase {
           orcReader.getSchema().toString());
     }
   }
-
-  @Override
-  public void testReaderFilterRowKeys() {
-    // TODO(HUDI-3682): fix filterRowKeys test for ORC due to a bug in ORC logic
-  }
 }
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieReaderWriterBase.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieReaderWriterBase.java
index 4617eb93a..902f42e38 100644
--- a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieReaderWriterBase.java
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieReaderWriterBase.java
@@ -20,6 +20,7 @@
 package org.apache.hudi.io.storage;
 
 import org.apache.hudi.common.bloom.BloomFilter;
+import org.apache.hudi.common.model.HoodieKey;
 
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericData;
@@ -157,7 +158,7 @@ public abstract class TestHoodieReaderWriterBase {
 
   @Test
   public void testReaderFilterRowKeys() throws Exception {
-    writeFileWithSimpleSchema();
+    writeFileWithSchemaWithMeta();
     Configuration conf = new Configuration();
     verifyMetadata(conf);
     verifyFilterRowKeys(createReader(conf));
@@ -177,6 +178,21 @@ public abstract class TestHoodieReaderWriterBase {
     writer.close();
   }
 
+  protected void writeFileWithSchemaWithMeta() throws Exception {
+    Schema avroSchema = getSchemaFromResource(TestHoodieReaderWriterBase.class, "/exampleSchemaWithMetaFields.avsc");
+    HoodieFileWriter<GenericRecord> writer = createWriter(avroSchema, true);
+    for (int i = 0; i < NUM_RECORDS; i++) {
+      GenericRecord record = new GenericData.Record(avroSchema);
+      String key = "key" + String.format("%02d", i);
+      record.put("_row_key", key);
+      record.put("time", Integer.toString(i));
+      record.put("number", i);
+      writer.writeAvroWithMetadata(new HoodieKey((String) record.get("_row_key"),
+          Integer.toString((Integer) record.get("number"))), record);
+    }
+    writer.close();
+  }
+
   protected void verifySimpleRecords(Iterator<GenericRecord> iterator) {
     int index = 0;
     while (iterator.hasNext()) {