Skip to content

feat(storage,driver,adapters): add streaming Arrow batches, binary COPY, staged upserts, and index lifecycle - #868

Merged
cofin merged 14 commits into
mainfrom
research_databakes_data_flow
Oct 31, 2026
Merged

cofin merged 14 commits into
mainfrom
research_databakes_data_flow

Conversation

@cofin

@cofin cofin commented Oct 30, 2026 •

Copy link
Copy Markdown
Member

Summary

Implements the complete high-performance-data-flow specification across SQLSpec's Arrow utilities, storage pipeline, driver base classes, database adapters, and cache subsystem:

  1. Streaming Arrow RecordBatch Coercion & Stage-Split Telemetry (t1):
    • Adds coerce_arrow_batches, arrow_batch_to_rows, and ingest_telemetry_from_batches in sqlspec/utils/arrow_helpers.py (directly imported and re-exported with canonical names in sqlspec/driver/_storage_helpers.py).
    • Extends StorageTelemetry (sqlspec/storage/pipeline.py) with stage-level metrics (read_duration_s, encode_duration_s, write_duration_s, commit_duration_s, shards_processed, shards_skipped, wal_bytes).
  2. Columnar PostgreSQL COPY (FORMAT BINARY) & Streaming pgvector Framing (t2):
    • Implements encode_arrow_batch_pg_binary, iter_pgcopy_binary_chunks, and iter_pgcopy_binary_chunks_async in sqlspec/utils/arrow_helpers.py with a zero-Python-loop structured NumPy dtype fast path (releasing the GIL inside NumPy C kernels for fixed-width and pgvector FixedSizeList<float32>/FixedSizeList<float16> columns) and direct Arrow buffer slicing (memoryview + np.diff(offsets)) for variable-width (utf8, binary, jsonb) columns.
    • Wires streaming binary COPY across psycopg (sync & async), asyncpg, psqlpy, cockroach_psycopg, and cockroach_asyncpg inside each adapter's own core.py and driver.py.
  3. Bounded-Memory Storage Streaming, Vector-Aware Parquet & Schema-Pinned Readers (t3):
    • Adds read_arrow_batches / read_arrow_batches_async and write_arrow_batches / write_arrow_batches_async on SyncStoragePipeline and AsyncStoragePipeline.
    • Adds resolve_parquet_use_dictionary to automatically disable Parquet dictionary encoding on floating-point list/fixed-size-list vector columns while preserving dictionary encoding on scalar/string columns.
    • Adds iter_streamed_arrow_batches supporting arrow_schema-pinned streaming jsonl/ndjson (pa.json.ParseOptions(explicit_schema=arrow_schema)), csv (pa.csv.ConvertOptions(column_types=arrow_schema)), parquet, and arrow-ipc, plus ObStoreBackend range-read read_buffer_bytes and _ObstoreSink bytearray copy-on-write safety.
  4. Parallel Multi-Shard Ingestion, Transactional Checkpointing & Dynamic Cloud Auth (t4):
    • Adds list_sources / list_sources_async glob/prefix discovery and execute_load_from_storage_sync / execute_load_from_storage_async with database-backed checkpoint_table / checkpoint_phase resumability, max_workers, and executor="thread" | "process" (multiprocessing.get_context("spawn") + compute_worker_thread_budget / init_spawn_worker_thread_budget).
    • Adds dynamic auth_token_provider support across PsycopgSyncConfig, PsycopgAsyncConfig (with async AlloyDB connector parity), AsyncpgConfig, PsqlpyConfig, and AdbcConfig, plus DuckDBSecretConfig value_provider / refresh_interval_seconds and native DuckDB GCS bearer_token + glob pushdown.
  5. Cross-Dialect Staged Bulk Load & Upsert (t5, t6):
    • Adds stage_table, conflict_keys, on_conflict="update" | "nothing", exclude_update_columns, drop_stage_table, and type_casts across load_from_arrow, load_from_storage, and load_from_records on all 19 database adapters (UNLOGGED on PostgreSQL/CockroachDB, TEMP on DuckDB/SQLite/MySQL, NOLOGGING on Oracle, SELECT TOP 0 * INTO on MSSQL, WITH NO DATA on Db2, LIKE on BigQuery, native insert_or_update / insert_or_ignore on Spanner, and UNLOGGED array-to-UDT cast staging on ADBC).
  6. Cross-Driver Bulk-Load Index Lifecycle (t7):
    • Adds driver.without_indexes(table, *, schema=None, concurrent=False) (sync & async context manager) and drop_indexes / concurrent_indexes on load_from_arrow and load_from_storage backed by driver.data_dictionary.get_indexes and get_ddl.
  7. DuckDB Out-of-Core Resource Controls (t8):
    • Adds detect_container_memory_limit_bytes / resolve_duckdb_memory_limit (Linux cgroup v1/v2 memory clamping), validate_duckdb_temp_directory (/proc/mounts tmpfs/ramfs detection), and post-extension parquet_metadata_cache application in DuckDBConnectionPool.
  8. L1/L2 Cache Consolidation (t9):
    • Consolidates CachedQuery, QueryCache, StatementCacheHolder, and resolve_statement_cache_size from sqlspec/driver/_query_cache.py into sqlspec/core/cache.py so L1 per-config QueryCache instances and L2 process-global NamespacedCache / StatementPipelineRegistry caches are managed together by clear_all_caches(), update_cache_config(), reset_stats_only(), and get_cache_statistics()["query"].

@cofin
cofin force-pushed the research_databakes_data_flow branch from f8dec2c to edbf197 Compare October 30, 2026 18:58
@codecov-commenter

codecov-commenter commented Oct 30, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 70.09932% with 572 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.40%. Comparing base (e7f6a78) to head (c93ce8d).

Files with missing lines Patch % Lines
sqlspec/driver/_storage_helpers.py 70.52% 106 Missing and 67 partials ⚠️
sqlspec/utils/arrow_helpers.py 69.82% 79 Missing and 23 partials ⚠️
sqlspec/storage/_arrow_stream.py 31.25% 39 Missing and 5 partials ⚠️
sqlspec/storage/pipeline.py 55.29% 38 Missing ⚠️
sqlspec/adapters/duckdb/core.py 76.92% 13 Missing and 14 partials ⚠️
sqlspec/adapters/duckdb/pool.py 75.32% 10 Missing and 9 partials ⚠️
sqlspec/core/cache.py 90.00% 12 Missing and 7 partials ⚠️
sqlspec/adapters/spanner/driver.py 23.80% 11 Missing and 5 partials ⚠️
sqlspec/adapters/oracledb/driver.py 25.00% 8 Missing and 7 partials ⚠️
sqlspec/adapters/bigquery/driver.py 14.28% 7 Missing and 5 partials ⚠️
... and 25 more
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #868      +/-   ##
==========================================
- Coverage   80.58%   80.40%   -0.18%     
==========================================
  Files         563      562       -1     
  Lines       99315   100697    +1382     
  Branches    19155    19414     +259     
==========================================
+ Hits        80030    80970     +940     
- Misses      14287    14594     +307     
- Partials     4998     5133     +135     
Flag Coverage Δ
cloud-adbc-storage 30.97% <19.28%> (-0.25%) ⬇️
integration 55.73% <28.85%> (-0.60%) ⬇️
mssql-db2 32.37% <17.30%> (-0.29%) ⬇️
mysql-mariadb-embedded 38.03% <24.20%> (-0.39%) ⬇️
oracle 31.51% <17.35%> (-0.28%) ⬇️
postgres-cockroach 40.36% <25.30%> (-0.34%) ⬇️
py3.11 78.99% <69.10%> (-0.16%) ⬇️
py3.12 79.00% <69.10%> (-0.16%) ⬇️
py3.13 79.00% <69.10%> (-0.16%) ⬇️
py3.14 79.79% <69.30%> (-0.19%) ⬇️
unit 73.71% <68.79%> (-0.03%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
sqlspec/adapters/adbc/config.py 97.13% <100.00%> (+0.03%) ⬆️
sqlspec/adapters/aiosqlite/driver.py 87.85% <100.00%> (+12.92%) ⬆️
sqlspec/adapters/cockroach_asyncpg/driver.py 87.87% <ø> (ø)
sqlspec/adapters/cockroach_psycopg/core.py 89.37% <100.00%> (ø)
sqlspec/adapters/cockroach_psycopg/driver.py 82.59% <ø> (ø)
sqlspec/adapters/duckdb/config.py 96.36% <100.00%> (+0.13%) ⬆️
sqlspec/adapters/psycopg/core.py 95.23% <100.00%> (+0.03%) ⬆️
sqlspec/adapters/sqlite/driver.py 93.59% <100.00%> (+11.92%) ⬆️
sqlspec/config.py 92.51% <100.00%> (-0.01%) ⬇️
sqlspec/core/__init__.py 100.00% <ø> (ø)
... and 39 more

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@cofin
cofin marked this pull request as ready for review October 30, 2026 21:09
@cofin
cofin changed the base branch from main to research_sybase_adbc_driver October 31, 2026 01:55
@cofin
cofin force-pushed the research_sybase_adbc_driver branch from 3f9c0f7 to 74fc6be Compare October 31, 2026 02:19
@cofin
cofin force-pushed the research_databakes_data_flow branch 2 times, most recently from 5063633 to f1e1a47 Compare October 31, 2026 02:31
@cofin
cofin added this pull request to stack #873 October 31, 2026 04:44
@cofin
cofin removed this pull request from stack #873 October 31, 2026 04:44
Base automatically changed from research_sybase_adbc_driver to main October 31, 2026 04:46
cofin added 14 commits October 31, 2026 04:49
@cofin
cofin force-pushed the research_databakes_data_flow branch from 93db25e to c93ce8d Compare October 31, 2026 15:32
@cofin
cofin merged commit 07adca4 into main Oct 31, 2026
80 of 82 checks passed
@cofin
cofin deleted the research_databakes_data_flow branch October 31, 2026 16:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants