Merge pull request #185 from MITLibraries/USE-449-handle-duckdb-v1.5

USE 449 - Update for DuckDB 1.5 compatibility

Author: ghukill

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "timdex_dataset_api"
-version = "4.0.0"
+version = "4.1.0"
 description = "Python library for interacting with a TIMDEX parquet dataset"
 readme = "README.md"
 requires-python = ">=3.12"

timdex_dataset_api/dataset.py

@@ -622,9 +622,7 @@ def _iter_data_chunks(self, data_query: str) -> Iterator[pa.RecordBatch]:
             self.conn.execute("""set threads=16;""")
         try:
             cursor = self.conn.execute(data_query)
-            yield from cursor.fetch_record_batch(
-                rows_per_batch=self.config.read_batch_size
-            )
+            yield from cursor.to_arrow_reader(batch_size=self.config.read_batch_size)
         finally:
             if self.location_scheme == "s3":
                 self.conn.execute(f"""set threads={self.conn_factory.threads};""")

timdex_dataset_api/embeddings.py

@@ -372,8 +372,8 @@ def read_batches_iter(
             **filters,
         )
         cursor = self.conn.execute(data_query)
-        yield from cursor.fetch_record_batch(
-            rows_per_batch=self.timdex_dataset.config.read_batch_size
+        yield from cursor.to_arrow_reader(
+            batch_size=self.timdex_dataset.config.read_batch_size
         )
 
         logger.debug(f"read() elapsed: {round(time.perf_counter() - start_time, 2)}s")

timdex_dataset_api/utils.py

@@ -161,6 +161,7 @@ def _install_extensions(self, conn: DuckDBPyConnection) -> None:
             secrets_dir.mkdir(parents=True, exist_ok=True)
             extensions_dir.mkdir(parents=True, exist_ok=True)
 
+            conn.execute(f"set home_directory='{duckdb_home.as_posix()}';")
             conn.execute(f"set secret_directory='{secrets_dir.as_posix()}';")
             conn.execute(f"set extension_directory='{extensions_dir.as_posix()}';")