Join problems with custom TableProviders

[robo-todd]

I'm making an assessment if we can use DataFusion for our application which would fall into the "Custom Database" category use case. We will be creating lots of custom TableProviders for this work. (Thanks for the brilliantly flexible ways we can extend DataFusion!).

My problem is that when I start to create joins with our custom TableProviders I get crippling performance problems that basically eliminate DataFusion for our use case. What is happening (from printing out what I get from running joins, etc.) is that the Inner (and all others I have tested, RightSemi, LeftSemi, etc.) returns a large number (18 or more) tiny record batches (like 1 row each) of results from joins pulling from our custom TableProviders. All the extra overhead and allocation is really killing performance.

I have written essentially the same test using built-in MemTable providers and these tests perform a minimum of 10x faster for the same data sizes (in our case small batches < 64 at the moment), basically same data layouts, etc. and I've compared the constraints (each of our tables has a primary index in column 0 in all cases). When I run the same join logical plan structure in the MemTable test it returns single RecordBatches of a reasonable number of rows (for the small number of rows in this test). Whereas, the same join (as far as I can tell) on our custom TableProviders returns a giant pile of ~ 20 record batches of row size = 1 or maybe sometimes 2.

When I compare just running logical scans pulling columns from our TableProviders they are comparable to MemTable performance.

Any ideas on how to figure out what I'm doing wrong here?

[robo-todd]

Adding some details: This is with DataFusion 48 and 49 versions.

Also for a join of this size (64 rows approximately?) I would not expect to see more than one scan of each incoming table.

The custom TableProviders fully obey the column projection for scan requests, but are otherwise pretty simple internally. They are returning complete single record batches of roughly 64 rows, so I'm mystified how join is producing streams of ultra-tiny 1 and 2 row record batches as output.

[robo-todd]

Adding further information:
Here are the logical plans of the two different joins. My custom TableProviders are based on the table provider example and I'm still trying to understand why the join with custom table providers is returning a large number of tiny 1 or 2 row record batches.

Plan From MemTable test: (Results in 1 batch with 10's of rows depending on data).
Join(Join { left: TableScan(TableScan { table_name: Partial { schema: "public", table: "values" }, source: "...", projection: None, projected_schema: DFSchema { inner: Schema { fields: [Field { name: "id", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "timestamp", data_type: Timestamp(Nanosecond, None), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: FixedSizeList(Field { name: "item", data_type: Float32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, 64), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }, field_qualifiers: [Some(Partial { schema: "public", table: "values" }), Some(Partial { schema: "public", table: "values" }), Some(Partial { schema: "public", table: "values" })], functional_dependencies: FunctionalDependencies { deps: [FunctionalDependence { source_indices: [0], target_indices: [0, 1, 2], nullable: false, mode: Single }] } }, filters: [], fetch: None, .. }), right: Filter(Filter { predicate: IsNotNull(Column(Column { relation: Some(Partial { schema: "public", table: "entities" }), name: "super_entity" })), input: TableScan(TableScan { table_name: Partial { schema: "public", table: "entities" }, source: "...", projection: Some([0, 2]), projected_schema: DFSchema { inner: Schema { fields: [Field { name: "id", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "super_entity", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }, field_qualifiers: [Some(Partial { schema: "public", table: "entities" }), Some(Partial { schema: "public", table: "entities" })], functional_dependencies: FunctionalDependencies { deps: [FunctionalDependence { source_indices: [0], target_indices: [0, 1], nullable: false, mode: Single }] } }, filters: [], fetch: None, .. }) }), on: [(Column(Column { relation: Some(Partial { schema: "public", table: "values" }), name: "id" }),Column(Column { relation: Some(Partial { schema: "public", table: "entities" }), name: "id" }))], filter: None, join_type: Inner, join_constraint: On, schema: DFSchema { inner: Schema { fields: [Field { name: "id", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "timestamp", data_type: Timestamp(Nanosecond, None), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: FixedSizeList(Field { name: "item", data_type: Float32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, 64), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "id", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "super_entity", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }, field_qualifiers: [Some(Partial { schema: "public", table: "values" }), Some(Partial { schema: "public", table: "values" }), Some(Partial { schema: "public", table: "values" }), Some(Partial { schema: "public", table: "entities" }), Some(Partial { schema: "public", table: "entities" })], functional_dependencies: FunctionalDependencies { deps: [FunctionalDependence { source_indices: [0], target_indices: [0, 1, 2], nullable: false, mode: Multi }, FunctionalDependence { source_indices: [3], target_indices: [3, 4], nullable: false, mode: Multi }] } }, null_equality: NullEqualsNothing }) Results in 1 batches:

Plan From Custom TableProvider Test: (Results in 18 batches of 1-2 row data)

Join(Join { left: TableScan(TableScan { table_name: Partial { schema: "testcase", table: "values" }, source: "...", projection: None, projected_schema: DFSchema { inner: Schema { fields: [Field { name: "id", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "timestamp", data_type: Timestamp(Microsecond, None), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: List(Field { name: "item", data_type: Float32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }, field_qualifiers: [Some(Partial { schema: "testcase", table: "values" }), Some(Partial { schema: "testcase", table: "values" }), Some(Partial { schema: "testcase", table: "values" })], functional_dependencies: FunctionalDependencies { deps: [FunctionalDependence { source_indices: [0], target_indices: [0, 1, 2], nullable: false, mode: Single }] } }, filters: [], fetch: None, .. }), right: Filter(Filter { predicate: IsNotNull(Column(Column { relation: Some(Partial { schema: "testcase", table: "entities" }), name: "super_entity" })), input: TableScan(TableScan { table_name: Partial { schema: "testcase", table: "entities" }, source: "...", projection: Some([0, 2]), projected_schema: DFSchema { inner: Schema { fields: [Field { name: "id", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "super_entity", data_type: UInt64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }, field_qualifiers: [Some(Partial { schema: "testcase", table: "entities" }), Some(Partial { schema: "testcase", table: "entities" })], functional_dependencies: FunctionalDependencies { deps: [FunctionalDependence { source_indices: [0], target_indices: [0, 1], nullable: false, mode: Single }] } }, filters: [], fetch: None, .. }) }), on: [(Column(Column { relation: Some(Partial { schema: "testcase", table: "values" }), name: "id" }), Column(Column { relation: Some(Partial { schema: "testcase", table: "entities" }), name: "id" }))], filter: None, join_type: Inner, join_constraint: On, schema: DFSchema { inner: Schema { fields: [Field { name: "id", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "timestamp", data_type: Timestamp(Microsecond, None), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: List(Field { name: "item", data_type: Float32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "id", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "super_entity", data_type: UInt64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }, field_qualifiers: [Some(Partial { schema: "testcase", table: "values" }), Some(Partial { schema: "testcase", table: "values" }), Some(Partial { schema: "testcase", table: "values" }), Some(Partial { schema: "testcase", table: "entities" }), Some(Partial { schema: "testcase", table: "entities" })], functional_dependencies: FunctionalDependencies { deps: [FunctionalDependence { source_indices: [0], target_indices: [0, 1, 2], nullable: false, mode: Multi }, FunctionalDependence { source_indices: [3], target_indices: [3, 4], nullable: false, mode: Multi }] } }, null_equality: NullEqualsNothing }) Join in 18 batches: