Skip to content

Commit 78351ce

Browse files
igor-suhorukovigor-suhorukov
authored
ARROW-17303: [Java][Dataset] Read Arrow IPC files by NativeDatasetFactory (apache#13760) (apache#13811)
This PR allow developers to create Dataset from ARROW IPC files in JVM code like: `FileSystemDatasetFactory factory = new FileSystemDatasetFactory(rootAllocator(), NativeMemoryPool.getDefault(), FileFormat.ARROW_IPC, arrowDatasetURL);` It is foundation for Apache Spark arrow data source to process huge existing partitioned datasets in ARROW file format without additional data format conversion Lead-authored-by: Igor Suhorukov <igor.suhorukov@gmail.com> Co-authored-by: igor.suhorukov <igor.suhorukov@gmail.com> Signed-off-by: David Li <li.davidm96@gmail.com>
1 parent a2f3666 commit 78351ce

3 files changed

Lines changed: 44 additions & 0 deletions

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,8 @@ arrow::Result<std::shared_ptr<arrow::dataset::FileFormat>> GetFileFormat(
8989
switch (file_format_id) {
9090
case 0:
9191
return std::make_shared<arrow::dataset::ParquetFileFormat>();
92+
case 1:
93+
return std::make_shared<arrow::dataset::IpcFileFormat>();
9294
default:
9395
std::string error_message =
9496
"illegal file format id: " + std::to_string(file_format_id);

‎java/dataset/src/main/java/org/apache/arrow/dataset/file/FileFormat.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
*/
2323
public enum FileFormat {
2424
PARQUET(0),
25+
ARROW_IPC(1),
2526
NONE(-1);
2627

2728
private final int id;

‎java/dataset/src/test/java/org/apache/arrow/dataset/file/TestFileSystemDataset.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import static org.junit.jupiter.api.Assertions.assertThrows;
2424

2525
import java.io.File;
26+
import java.io.FileOutputStream;
2627
import java.io.IOException;
2728
import java.util.ArrayList;
2829
import java.util.Collections;
@@ -46,11 +47,15 @@
4647
import org.apache.arrow.dataset.scanner.ScanOptions;
4748
import org.apache.arrow.util.AutoCloseables;
4849
import org.apache.arrow.vector.FieldVector;
50+
import org.apache.arrow.vector.IntVector;
4951
import org.apache.arrow.vector.VectorLoader;
5052
import org.apache.arrow.vector.VectorSchemaRoot;
53+
import org.apache.arrow.vector.ipc.ArrowFileWriter;
5154
import org.apache.arrow.vector.ipc.ArrowReader;
5255
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
5356
import org.apache.arrow.vector.types.Types;
57+
import org.apache.arrow.vector.types.pojo.ArrowType;
58+
import org.apache.arrow.vector.types.pojo.Field;
5459
import org.apache.arrow.vector.types.pojo.Schema;
5560
import org.apache.avro.generic.GenericRecord;
5661
import org.apache.avro.generic.GenericRecordBuilder;
@@ -316,6 +321,42 @@ public void testErrorThrownWhenIterateOnIteratorAfterTaskClose() throws Exceptio
316321
AutoCloseables.close(factory);
317322
}
318323

324+
@Test
325+
public void testBaseArrowIpcRead() throws Exception {
326+
File dataFile = TMP.newFile();
327+
Schema sourceSchema = new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true))));
328+
try (VectorSchemaRoot root = VectorSchemaRoot.create(sourceSchema, rootAllocator());
329+
FileOutputStream sink = new FileOutputStream(dataFile);
330+
ArrowFileWriter writer = new ArrowFileWriter(root, /*dictionaryProvider=*/null, sink.getChannel())) {
331+
IntVector ints = (IntVector) root.getVector(0);
332+
ints.setSafe(0, 0);
333+
ints.setSafe(1, 1024);
334+
ints.setSafe(2, Integer.MAX_VALUE);
335+
root.setRowCount(3);
336+
writer.start();
337+
writer.writeBatch();
338+
writer.end();
339+
}
340+
341+
String arrowDataURI = dataFile.toURI().toString();
342+
FileSystemDatasetFactory factory = new FileSystemDatasetFactory(rootAllocator(), NativeMemoryPool.getDefault(),
343+
FileFormat.ARROW_IPC, arrowDataURI);
344+
ScanOptions options = new ScanOptions(100);
345+
Schema schema = inferResultSchemaFromFactory(factory, options);
346+
List<ArrowRecordBatch> datum = collectResultFromFactory(factory, options);
347+
348+
assertSingleTaskProduced(factory, options);
349+
assertEquals(1, datum.size());
350+
assertEquals(1, schema.getFields().size());
351+
assertEquals("ints", schema.getFields().get(0).getName());
352+
353+
String expectedJsonUnordered = String.format("[[0],[1024],[%d]]", Integer.MAX_VALUE);
354+
checkParquetReadResult(schema, expectedJsonUnordered, datum);
355+
356+
AutoCloseables.close(datum);
357+
AutoCloseables.close(factory);
358+
}
359+
319360
private void checkParquetReadResult(Schema schema, String expectedJson, List<ArrowRecordBatch> actual)
320361
throws IOException {
321362
final ObjectMapper json = new ObjectMapper();

0 commit comments

Comments
 (0)