Reuse ORC multi columns within single column write approach

trinodb · Jan 10, 2023 · 07152c0 · 07152c0
1 parent b133dca
commit 07152c0
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/lib/trino-orc/src/test/java/io/trino/orc/TestCachingOrcDataSource.java b/lib/trino-orc/src/test/java/io/trino/orc/TestCachingOrcDataSource.java
@@ -38,11 +38,11 @@
 
 import java.io.File;
 import java.io.IOException;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Random;
+import java.util.function.Function;
 import java.util.stream.Stream;
 
 import static io.airlift.testing.Assertions.assertGreaterThanOrEqual;
@@ -55,12 +55,11 @@
 import static io.trino.orc.OrcTester.Format.ORC_12;
 import static io.trino.orc.OrcTester.HIVE_STORAGE_TIME_ZONE;
 import static io.trino.orc.OrcTester.READER_OPTIONS;
-import static io.trino.orc.OrcTester.writeOrcColumnHive;
+import static io.trino.orc.OrcTester.writeOrcColumnsHiveFile;
 import static io.trino.orc.metadata.CompressionKind.NONE;
 import static io.trino.orc.metadata.CompressionKind.ZLIB;
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static java.lang.String.format;
-import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaStringObjectInspector;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.fail;
@@ -77,11 +76,14 @@ public void setUp()
     {
         tempFile = new TempFile();
         Random random = new Random();
-        Iterator<String> iterator = Stream.generate(() -> Long.toHexString(random.nextLong())).limit(POSITION_COUNT).iterator();
-        writeOrcColumnHive(
-                createOrcRecordWriter(tempFile.getFile(), ORC_12, ZLIB, javaStringObjectInspector),
-                VARCHAR,
-                iterator);
+        writeOrcColumnsHiveFile(
+                tempFile.getFile(),
+                ORC_12,
+                ZLIB,
+                ImmutableList.of("test"),
+                ImmutableList.of(VARCHAR),
+                Stream.generate(() -> (Function<Integer, Object>) (fieldIndex) -> Long.toHexString(random.nextLong()))
+                        .limit(POSITION_COUNT).iterator());
     }
 
     @AfterClass(alwaysRun = true)