Can we stop copying the Arrow schema over FFI for every batch?

[andygrove]

What is the problem the feature request solves?

In CometBatchIterator we export the schema with each batch via Arrow FFI:

          val arrowSchema = ArrowSchema.wrap(schemaAddrs(index))
          val arrowArray = ArrowArray.wrap(arrayAddrs(index))
          Data.exportVector(
            allocator,
            getFieldVector(valueVector, "export"),
            provider,
            arrowArray,
            arrowSchema)

Exporting the schema seems quite expensive since it involves string copies and memory allocation. It gets more expensive for complex schemas, especially when nested types are involved.

Internally in Data.exportVector, the schema is exported with:

exportField(allocator, vector.getField(), provider, outSchema);

I wonder if we could refactor CometBatchIterator to just export the schema once, with the first batch, and then have the native side re-use that schema for subsequent batches.

Describe the potential solution

No response

Additional context

No response

[andygrove]

I refactored the code to make separate FFI calls for exporting schema and array and measured the time of each:

          val arrowSchema = ArrowSchema.wrap(schemaAddrs(index))
          val arrowArray = ArrowArray.wrap(arrayAddrs(index))
          val export = getFieldVector(valueVector, "export")
          // export schema
          val t1 = System.nanoTime()
          Data.exportField(allocator, export.getField, provider, arrowSchema)
          val t2 = System.nanoTime()
          // export array
          Data.exportVector(allocator, export, provider, arrowArray)
          val t3 = System.nanoTime()
          // scalastyle:off println
          println(s"Exported schema in ${t2 - t1} ns and array in ${t3 - t2} ns")

Exporting the schema seems to be more expensive than exporting the data. It seems like this would be worth optimizing.

Exported schema in 1773 ns and array in 421 ns
Exported schema in 1402 ns and array in 401 ns
Exported schema in 1704 ns and array in 891 ns
Exported schema in 1884 ns and array in 351 ns
Exported schema in 1232 ns and array in 1072 ns
Exported schema in 1553 ns and array in 450 ns
Exported schema in 1923 ns and array in 1032 ns
Exported schema in 1392 ns and array in 1663 ns
Exported schema in 481 ns and array in 561 ns
Exported schema in 1022 ns and array in 551 ns
Exported schema in 1082 ns and array in 942 ns

[andygrove]

Schema can vary between batches due to dictionary encoding, so maybe we cannot avoid serializing it each time.

[andygrove]

Another approach could be to use the Arrow C streaming interface

[andygrove]

Spark has an ArrowBatchStreamWriter, but I don't think it supports dictionaries.

[andygrove]

After more metrics improvements in #1133 it is clear that improving FFI performance is not a high priority for now, although moving to the Arrow C stream interface could make it more efficient