Skip to content

Commit 203bc9d

Browse files
timsaucerclaude
andcommitted
feat: ADBC-backed Spark DataSource for DataFusion table providers
Add a Spark DataSourceV2 connector that reads from a DataFusion TableProvider through a standard ADBC driver. Spark talks to the upstream arrow-adbc Java driver manager (adbc-core + adbc-driver-jni), which loads a native DataFusion ADBC cdylib and returns arrow-java ArrowReaders consumed zero-copy as ArrowColumnVectors on the cluster-provided Arrow. Spark connector (spark/): - AdbcDatafusionTableProvider: registers the `adbc-datafusion` format; schema probed once on the driver via AdbcConnection.getTableSchema. - Scan with projection/filter/limit pushed into the scan as a Substrait plan, with a SQL fallback path. - Multi-partition reads via executePartitioned/readPartition; the target_partitions option tunes scan parallelism. - Per-executor AdbcConnectionPool caches the AdbcDatabase per driver + options key, with a per-task connection (level 1) as the safe default; connection sharing (level 2) is gated and off by default. Example driver (examples/adbc-datafusion-driver/): a DataFusion ADBC cdylib exercising the full stack, with a PySpark end-to-end script and partitioning tests. Docs: ADBC-to-Spark partitioning guidance and connector notes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent ab57cb9 commit 203bc9d

28 files changed

Lines changed: 3525 additions & 0 deletions

‎.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ rust-target/
66
*.iml
77
.DS_Store
88
tpch-data/
9+
# Lockfile for the example ADBC driver crate (a demo cdylib, not a published lib).
10+
examples/adbc-datafusion-driver/Cargo.lock
911
.claude
1012
docs/superpowers
1113
docs/build/

‎docs/adbc-spark-partitioning.md‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
# Partitioning: DataFusion table providers → Spark tasks
2+
3+
How DataFusion partitions become Spark tasks through the `adbc-datafusion`
4+
connector, how to size them, and what a `TableProvider` author should do. Pairs
5+
with the connector docs under `examples/adbc-datafusion-driver/`
6+
(`PARTITIONS.md`).
7+
8+
## The mapping is 1:1, and the provider owns it
9+
10+
```
11+
provider.scan() output partitions
12+
→ DataFusion physical plan partition count N
13+
→ executePartitioned() = N descriptors
14+
→ N Spark InputPartitions = N Spark tasks (scan stage)
15+
```
16+
17+
For a pushed scan (projection / filter / limit only, no DataFusion-side join or
18+
aggregation), the physical plan's output partitioning **is** the provider's
19+
`scan()` partitioning. So **the provider's `scan()` partition count is the Spark
20+
task count.** Spark never splits or merges a partition afterward; the cluster's
21+
core count only bounds how many of the `N` run concurrently (the rest run in
22+
waves). `df.rdd().getNumPartitions() == N`.
23+
24+
## The Spark-task sweet spot
25+
26+
- **Floor:** `N` ≥ total executor cores, or cores sit idle. ~2–4× cores gives
27+
slack for skew and stragglers.
28+
- **Per-partition size:** large enough to amortize per-task overhead. Canonical
29+
target ≈ **128–256 MB** of data (Spark's `spark.sql.files.maxPartitionBytes`
30+
default is 128 MB), or ≥ ~100 ms–seconds of work.
31+
- **Ceiling — and ADBC makes it lower.** Beyond Spark's generic per-task
32+
scheduling cost, this connector adds two per-partition costs: each ADBC
33+
partition descriptor carries the serialized query bytes (duplicated ×N on the
34+
driver), and **each task re-plans the whole query** in `read_partition`. So
35+
thousands of partitions is far costlier here than for a native Spark file
36+
source. Keep `N` in the hundreds, not tens of thousands; prefer bigger
37+
partitions over a huge count.
38+
39+
Rule of thumb: `N ≈ clamp(total_bytes / target_bytes, floor, ceiling)`, balanced
40+
by **bytes**, not by split count.
41+
42+
## The model mismatch you must account for
43+
44+
DataFusion and Spark size scans differently:
45+
46+
- **Spark native file source — byte-bound.** `N = Σ ceil(file_size / 128MB)`.
47+
Large data ⇒ many partitions, `N ≫ cores`, run in waves; each task ≈ 128 MB.
48+
- **DataFusion scan — parallelism-bound.** `ListingTable` packs files into
49+
**`target_partitions` groups** (≈ machine cores), with
50+
`repartition_file_min_size` (~10 MB) as a floor and optional single-file
51+
splitting to *reach* `target_partitions`. So DataFusion's default
52+
`N ≈ target_partitions`, **not** `bytes / 128MB`.
53+
54+
Consequence: if a provider mimics DataFusion's default (`N = target_partitions`)
55+
and the connector sets `target_partitions = defaultParallelism` (≈ cores), then
56+
**large data → ~cores giant tasks** (`total/cores` each) → memory pressure and
57+
poor sizing. For feeding Spark you usually want the **byte-bound** model — more
58+
partitions than cores for big data.
59+
60+
**So a custom provider feeding Spark should size by bytes, not blindly use
61+
`target_partitions` as the literal partition count.**
62+
63+
## Provider recipe
64+
65+
`scan(state, …)` receives the session. Read the connector-supplied hint, then
66+
decide by what the provider knows:
67+
68+
```text
69+
T = state.config().target_partitions() // connector hint ≈ desired parallelism
70+
// (set from SparkContext.defaultParallelism)
71+
target_bytes = ~128–256 MB // bias larger for ADBC to keep N down
72+
min_bytes = floor so partitions never get tiny // ~ tens of MB
73+
74+
if total_bytes known AND splittable (files / row groups / key ranges):
75+
N = clamp(ceil(total_bytes / target_bytes), 1, CEILING) // byte-bound
76+
bin-pack splits into N groups BALANCED BY BYTES // not by count
77+
do not split below min_bytes; merge small splits
78+
79+
elif split_count known but not bytes (shards / external partitions):
80+
N = min(split_count, CEILING) // natural partitioning
81+
if split_count >> T: coalesce shards into ~T balanced groups
82+
83+
else (size and splits unknown; opaque / streaming):
84+
N = T // fall back to the hint
85+
lean slightly high — stragglers hurt more than overhead
86+
```
87+
88+
- **Balance by bytes, not split count.** Spark can't rebalance a fat partition
89+
(each task is atomic). One 10 GB partition among 100 small ones stalls the
90+
whole stage. Bin-pack splits so partition bytes are even.
91+
- `T` (the connector's `target_partitions`) is the **hint / cap / fallback** —
92+
not the data-bound count. Use it when bytes are unknown, or as a parallelism
93+
ceiling. For known data, **bytes drive `N`**.
94+
95+
## Determinism caveat (specific to the ADBC partition path)
96+
97+
`execute_partitions` counts `N` on the **driver**; `read_partition` re-plans on
98+
the **executor**. Partition `index i` must mean the same slice both times, so
99+
**`scan()` must produce the identical `N` in both places.**
100+
101+
- If the provider keys `N` off `target_partitions`, the driver **pins it in the
102+
descriptor** and `read_partition` re-plans with it held fixed → determinism for
103+
free. (That is why the pin exists.)
104+
- If the provider sizes by **bytes** (ignoring `target_partitions`), the pin does
105+
not help — compute `N` from **stable** inputs (a cached file listing / a fixed
106+
stats snapshot). Otherwise stat drift between driver-plan and executor-read
107+
changes `N` and shifts partitions.
108+
109+
## Connector knobs
110+
111+
- **`target_partitions`** (implemented): defaults to
112+
`SparkContext.defaultParallelism` (total executor cores). The connector issues
113+
`SET datafusion.execution.target_partitions = N` on the planning session; the
114+
driver pins it into each descriptor, so executors re-plan identically. Reaches
115+
the provider via `state.config()`. Optionally pass `k×cores`.
116+
- **`batch_size`** (follow-up, connector-only): would `SET
117+
datafusion.execution.batch_size` on the executor connection in the reader;
118+
`read_partition` inherits it via `copied_config()`. Affects Arrow batch
119+
granularity, not partition count, so no pinning needed.
120+
- **`maxPartitionBytes`**: not a DataFusion session config and the connector
121+
cannot derive it (it does not know the data size). Byte-sizing lives in the
122+
provider's `scan()`; the connector can only forward it as a provider-specific
123+
option.
124+
- **Reported partitioning / storage-partitioned joins
125+
(`SupportsReportPartitioning`)**: not viable over ADBC today — the descriptor
126+
carries no distribution metadata, and DataFusion's hash partitioning does not
127+
match Spark's key-grouped/`bucket` model.
128+
129+
## Recommendation (summary)
130+
131+
- **Connector:** keep `target_partitions` defaulted from `defaultParallelism`; it
132+
is the parallelism hint that reaches `scan()` and is pinned for determinism.
133+
- **Provider (the real lever):**
134+
- bytes known + splittable → **byte-bound**: `N = ceil(total_bytes/target_bytes)`,
135+
bin-packed by bytes, merged above `min_bytes`, capped well below "thousands"
136+
(ADBC re-plans per partition); compute `N` from a stable stat snapshot.
137+
- only splits known → natural partitioning, coalesce toward `T` if huge.
138+
- nothing known → `N = T`, lean slightly high.
139+
- default `target_bytes ≈ 128–256 MB`; raise it (fewer, bigger partitions)
140+
before letting `N` reach the thousands.
141+
142+
One line: **the provider's `scan()` partition count is the Spark task count; size
143+
it by bytes toward ~128–256 MB balanced partitions, floored at executor cores and
144+
capped well below thousands because ADBC re-plans per partition — and make that
145+
count a deterministic function of a stat snapshot so the driver plan and executor
146+
reads agree.**
Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
[package]
19+
name = "adbc-datafusion-example-driver"
20+
version = "0.1.0"
21+
edition = "2021"
22+
publish = false
23+
description = "End-to-end example: an ADBC driver that exposes a custom DataFusion TableProvider"
24+
25+
# Standalone crate that stands in for a real, separate driver repo. The empty
26+
# [workspace] table makes it its own workspace root so it is NOT pulled into the
27+
# parent datafusion-java workspace (whose members are listed explicitly).
28+
[workspace]
29+
30+
[dependencies]
31+
# The ready-made ADBC-over-DataFusion driver. We depend on it as a LIBRARY and
32+
# only customize the SessionContext via its new_with_context_init hook -- no
33+
# fork, no datafusion-ffi module loading. Pulled from its own repo (git) since
34+
# that is where it lives.
35+
adbc-driver-datafusion = { git = "https://github.com/adbc-drivers/datafusion" }
36+
adbc_core = "0.23"
37+
adbc_ffi = "0.23"
38+
datafusion = "53"
39+
async-trait = "0.1"
40+
41+
# Build both an rlib (so the provider/driver can be unit-tested or embedded) and
42+
# a cdylib (the artifact the arrow-adbc Java driver manager loads at runtime).
43+
[lib]
44+
crate-type = ["cdylib", "lib"]
45+
46+
# Temporary: redirect adbc-driver-datafusion to the PR branch that implements
47+
# execute_partitions / read_partition (adbc-drivers/datafusion#32), so we can run
48+
# a full multi-partition system test before it merges. The crate is not on
49+
# crates.io, so this patches the git source (not [patch.crates-io]). Remove once
50+
# the PR is merged and the canonical dependency carries the change.
51+
[patch."https://github.com/adbc-drivers/datafusion"]
52+
adbc-driver-datafusion = { git = "https://github.com/timsaucer/adbc-driver-datafusion", branch = "feat/execute-partitions" }

0 commit comments

Comments
 (0)