Skip to content

Commit

Permalink
refactor(duckdb): simplify to_pyarrow_batches implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
jcrist committed Mar 23, 2023
1 parent 343fb37 commit d6235ee
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 36 deletions.
11 changes: 7 additions & 4 deletions ibis/backends/duckdb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import itertools
import os
import warnings
import weakref
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, MutableMapping
Expand Down Expand Up @@ -555,9 +556,6 @@ def to_pyarrow_batches(
!!! warning "DuckDB returns 1024 size batches regardless of what argument is passed."
"""
self._import_pyarrow()

from ibis.backends.duckdb.pyarrow import IbisRecordBatchReader

self._register_in_memory_tables(expr)
query_ast = self.compiler.to_ast_ensure_limit(expr, limit, params=params)
sql = query_ast.compile()
Expand All @@ -574,7 +572,12 @@ def to_pyarrow_batches(
cursor = con.execute(sql)

reader = cursor.cursor.fetch_record_batch(chunk_size=chunk_size)
return IbisRecordBatchReader(reader, cursor)
# Use a weakref finalizer to keep the cursor alive until the record
# batch reader is garbage collected. It would be nicer if we could make
# the cursor cleanup happen when the reader is closed, but that's not
# currently possible with pyarrow.
weakref.finalize(reader, lambda: cursor.close())
return reader

def to_pyarrow(
self,
Expand Down
32 changes: 0 additions & 32 deletions ibis/backends/duckdb/pyarrow.py

This file was deleted.

0 comments on commit d6235ee

Please sign in to comment.