feat(go/adbc/driver/flightsql): support bulk ingest #3808
Conversation
…ht SQL driver. Addresses apache#1107
prmoore77
changed the title
There was a problem hiding this comment.
We currently use docker-compose to manage this and I'd rather keep it consistent...
There was a problem hiding this comment.
Hi, @lidavidm - I've changed the code to use docker-compose, per your feedback.
|
FWIW - I tested the locally built driver/wheel from Python against a remote (Azure) GizmoSQL server - and it seems to work pretty well: import os
import time
import duckdb
from adbc_driver_flightsql import dbapi as gizmosql
from codetiming import Timer
from dotenv import load_dotenv
from config import get_logger
# Timer logging setup
TIMER_TEXT = "{name}: Elapsed time: {:.4f} seconds"
def main():
load_dotenv()
logger = get_logger()
timer_logger = logger.info
with Timer(name=f"Overall program",
text=TIMER_TEXT,
initial_text=True,
logger=timer_logger
):
with Timer(name=f" Generate TPCH data and load into DuckDB (1GB)",
text=TIMER_TEXT,
initial_text=True,
logger=timer_logger
):
# Connect to DuckDB (memory only)
duckdb_conn = duckdb.connect()
duckdb_conn.install_extension("tpch")
duckdb_conn.load_extension("tpch")
duckdb_conn.execute(query="CALL dbgen(sf=1.0)")
with Timer(name=f" Get RecordBatch reader for the DuckDB lineitem table",
text=TIMER_TEXT,
initial_text=True,
logger=timer_logger
):
lineitem_arrow_reader = duckdb_conn.table("lineitem").fetch_arrow_reader(batch_size=10_000)
with Timer(name=f" Bulk ingest the data into GizmoSQL",
text=TIMER_TEXT,
initial_text=True,
logger=timer_logger
):
with gizmosql.connect(
uri="grpc+tls://try-gizmosql-adbc.gizmodata.com:31337",
db_kwargs={"username": os.environ["GIZMOSQL_USERNAME"],
"password": os.environ["GIZMOSQL_PASSWORD"]
},
autocommit=True
).cursor() as cursor:
ingest_start = time.perf_counter()
rows_loaded = cursor.adbc_ingest(
table_name="bulk_ingest_lineitem",
data=lineitem_arrow_reader,
mode="replace"
)
ingest_seconds = time.perf_counter() - ingest_start
rows_per_sec = (rows_loaded / ingest_seconds) if ingest_seconds > 0 else float("inf")
logger.info(msg=f"Loaded rows: {rows_loaded:,}")
logger.info(msg=f"Ingest time: {ingest_seconds:.4f} s")
logger.info(msg=f"Rows/sec: {rows_per_sec:,.2f}")
if __name__ == "__main__":
main()Result: 
|
|
Well, the integration tests are failing in the pipeline. They "worked on my laptop" - but I'll investigate to see what is happening. |
I believe the integration tests related to Bulk Ingestion are working now... |
|
hi @zeroshade - I think the integration tests are working well now. I do see some CI failures, but I think they are unrelated to my changes. |
lidavidm
changed the title
| // executeIngestWithReader performs the bulk ingest operation with the given record reader. | ||
| func executeIngestWithReader( | ||
| ctx context.Context, | ||
| client *flightsql.Client, | ||
| rdr array.RecordReader, | ||
| opts *flightsql.ExecuteIngestOpts, | ||
| callOpts ...grpc.CallOption, | ||
| ) (int64, error) { | ||
| return client.ExecuteIngest(ctx, rdr, opts, callOpts...) | ||
| } |
There was a problem hiding this comment.
This is a function with one line that is only ever used once; just inline it?
| rdr, err := array.NewRecordReader(batch.Schema(), []arrow.RecordBatch{batch}) | ||
| if err != nil { | ||
| return nil, adbc.Error{ | ||
| Msg: fmt.Sprintf("[Flight SQL Statement] failed to create record reader: %s", err.Error()), |
There was a problem hiding this comment.
nit: I think there's no need to explicitly call err.Error()
| case adbc.OptionKeyIngestTargetTable: | ||
| s.query.sqlQuery = "" | ||
| s.query.substraitPlan = nil | ||
| s.targetTable = val |
There was a problem hiding this comment.
We need to clear s.prepared too?
There was a problem hiding this comment.
setSqlQuery needs to clear targetTable
| class TestBulkIngest: | ||
| """Test bulk ingest functionality.""" |
There was a problem hiding this comment.
nit: is there need for the class? It's not required in pytest
| result = reader.read_all() | ||
| assert result.column("cnt")[0].as_py() == 5 | ||
|
|
||
| def test_ingest_various_types(self, gizmosql): |
There was a problem hiding this comment.
Is this super necessary? It's more a test of gizmosql than the driver
| class TestBulkIngestDBAPI: | ||
| """Test bulk ingest using the DBAPI interface.""" |
There was a problem hiding this comment.
Same question here. The test names all encode the difference anyways, so why add the extra layer?
|
hi @lidavidm - I've made edits per your feedback. Thanks for your help. |
| } | ||
| } | ||
| case adbc.OptionKeyIngestTargetTable: | ||
| s.prepared = nil |
| func (srv *BulkIngestTestServer) GetIngestedData() []arrow.RecordBatch { | ||
| srv.mu.Lock() | ||
| defer srv.mu.Unlock() | ||
| return srv.ingestedData | ||
| } |
There was a problem hiding this comment.
The lock doesn't really do anything here unless you also copy the slice
| func (srv *BulkIngestTestServer) GetIngestRequests() []flightsql.StatementIngest { | ||
| srv.mu.Lock() | ||
| defer srv.mu.Unlock() | ||
| return srv.ingestRequests | ||
| } |
2 participants
Addresses #1107