Add readArrowReader method to allow loading a dataframe from an ArrowReader #528

Author: fb64

dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/arrowReading.kt

@@ -1,6 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.io
 
 import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.ipc.ArrowReader
 import org.apache.commons.compress.utils.SeekableInMemoryByteChannel
 import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.DataFrame
@@ -170,3 +171,11 @@ public fun DataFrame.Companion.readArrowFeather(
 } else {
     readArrowFeather(File(path), nullability)
 }
+
+/**
+ * Read [Arrow any format](https://arrow.apache.org/docs/java/ipc.html#reading-writing-ipc-formats) data from existing [reader]
+ */
+public fun DataFrame.Companion.readArrowReader(
+    reader: ArrowReader,
+    nullability: NullabilityOptions = NullabilityOptions.Infer
+): AnyFrame = readArrowReaderImpl(reader, nullability)

dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/arrowReadingImpl.kt

@@ -32,6 +32,7 @@ import org.apache.arrow.vector.VarCharVector
 import org.apache.arrow.vector.VectorSchemaRoot
 import org.apache.arrow.vector.complex.StructVector
 import org.apache.arrow.vector.ipc.ArrowFileReader
+import org.apache.arrow.vector.ipc.ArrowReader
 import org.apache.arrow.vector.ipc.ArrowStreamReader
 import org.apache.arrow.vector.types.pojo.Field
 import org.apache.arrow.vector.util.DateUtility
@@ -262,17 +263,7 @@ internal fun DataFrame.Companion.readArrowIPCImpl(
     allocator: RootAllocator = Allocator.ROOT,
     nullability: NullabilityOptions = NullabilityOptions.Infer,
 ): AnyFrame {
-    ArrowStreamReader(channel, allocator).use { reader ->
-        val flattened = buildList {
-            val root = reader.vectorSchemaRoot
-            val schema = root.schema
-            while (reader.loadNextBatch()) {
-                val df = schema.fields.map { f -> readField(root, f, nullability) }.toDataFrame()
-                add(df)
-            }
-        }
-        return flattened.concatKeepingSchema()
-    }
+    return readArrowReaderImpl(ArrowStreamReader(channel, allocator), nullability)
 }
 
 /**
@@ -283,14 +274,36 @@ internal fun DataFrame.Companion.readArrowFeatherImpl(
     allocator: RootAllocator = Allocator.ROOT,
     nullability: NullabilityOptions = NullabilityOptions.Infer,
 ): AnyFrame {
-    ArrowFileReader(channel, allocator).use { reader ->
+    return readArrowReaderImpl(ArrowFileReader(channel, allocator), nullability)
+}
+
+/**
+ * Read [Arrow any format](https://arrow.apache.org/docs/java/ipc.html#reading-writing-ipc-formats) data from existing [reader]
+ */
+internal fun DataFrame.Companion.readArrowReaderImpl(
+    reader: ArrowReader,
+    nullability: NullabilityOptions = NullabilityOptions.Infer
+): AnyFrame {
+    reader.use {
         val flattened = buildList {
-            reader.recordBlocks.forEach { block ->
-                reader.loadRecordBatch(block)
-                val root = reader.vectorSchemaRoot
-                val schema = root.schema
-                val df = schema.fields.map { f -> readField(root, f, nullability) }.toDataFrame()
-                add(df)
+            when (reader) {
+                is ArrowFileReader -> {
+                    reader.recordBlocks.forEach { block ->
+                        reader.loadRecordBatch(block)
+                        val root = reader.vectorSchemaRoot
+                        val schema = root.schema
+                        val df = schema.fields.map { f -> readField(root, f, nullability) }.toDataFrame()
+                        add(df)
+                    }
+                }
+                is ArrowStreamReader -> {
+                    val root = reader.vectorSchemaRoot
+                    val schema = root.schema
+                    while (reader.loadNextBatch()) {
+                        val df = schema.fields.map { f -> readField(root, f, nullability) }.toDataFrame()
+                        add(df)
+                    }
+                }
             }
         }
         return flattened.concatKeepingSchema()

docs/StardustDocs/topics/read.md

@@ -445,7 +445,7 @@ val df = DataFrame.readArrowFeather(file)
 
 [`DataFrame`](DataFrame.md) supports reading [Arrow interprocess streaming format](https://arrow.apache.org/docs/java/ipc.html#writing-and-reading-streaming-format)
 and [Arrow random access format](https://arrow.apache.org/docs/java/ipc.html#writing-and-reading-random-access-files)
-from raw Channel (ReadableByteChannel for streaming and SeekableByteChannel for random access), InputStream, File or ByteArray.
+from raw Channel (ReadableByteChannel for streaming and SeekableByteChannel for random access), ArrowReader, InputStream, File or ByteArray.
 
 > If you use Java 9+, follow the [Apache Arrow Java compatibility](https://arrow.apache.org/docs/java/install.html#java-compatibility) guide.
 >