diff --git a/vortex-duckdb/src/copy.rs b/vortex-duckdb/src/copy.rs index ec1d3b30775..f158b2e45c5 100644 --- a/vortex-duckdb/src/copy.rs +++ b/vortex-duckdb/src/copy.rs @@ -1,11 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use async_fs::OpenOptions; use futures::SinkExt; use futures::TryStreamExt; use futures::channel::mpsc; use futures::channel::mpsc::Sender; +use object_store::ObjectStore; +use object_store::registry::ObjectStoreRegistry; use parking_lot::Mutex; use static_assertions::assert_impl_all; use vortex::array::ArrayRef; @@ -20,11 +24,16 @@ use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteSummary; +use vortex::file::multi::parse_uri_or_path; +use vortex::io::VortexWrite; +use vortex::io::compat::Compat; +use vortex::io::object_store::ObjectStoreWrite; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Task; use vortex::io::runtime::current::CurrentThreadWorkerPool; use vortex::io::session::RuntimeSessionExt; +use crate::REGISTRY; use crate::RUNTIME; use crate::SESSION; use crate::convert::FromLogicalType; @@ -119,15 +128,35 @@ pub fn copy_to_initialize_global( let handle = SESSION.handle(); - let write_task = handle.spawn(async move { - let writer = OpenOptions::new() - .write(true) - .truncate(true) - .create(true) - .open(file_path) - .await?; - SESSION.write_options().write(writer, array_stream).await - }); + let url = parse_uri_or_path(&file_path)?; + let write_task = if url.scheme() == "file" { + handle.spawn(async move { + let mut writer = OpenOptions::new() + .write(true) + .truncate(true) + .create(true) + .open(file_path) + .await?; + let summary = SESSION + .write_options() + .write(&mut writer, array_stream) + .await?; + writer.shutdown().await?; + Ok(summary) + }) + } else { + let (object_store, path) = REGISTRY.resolve(&url)?; + let object_store = Arc::new(Compat::new(object_store)) as Arc; + handle.spawn(async move { + let mut writer = ObjectStoreWrite::new(object_store, &path).await?; + let summary = SESSION + .write_options() + .write(&mut writer, array_stream) + .await?; + writer.shutdown().await?; + Ok(summary) + }) + }; let worker_pool = RUNTIME.new_pool(); worker_pool.set_workers_to_available_parallelism(); diff --git a/vortex-duckdb/src/lib.rs b/vortex-duckdb/src/lib.rs index fe14c86f8f6..b16d41e9e75 100644 --- a/vortex-duckdb/src/lib.rs +++ b/vortex-duckdb/src/lib.rs @@ -11,6 +11,7 @@ use std::sync::LazyLock; use std::sync::OnceLock; use vortex::VortexSessionDefault; +use vortex::cloud::Registry; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::io::runtime::BlockingRuntime; @@ -43,6 +44,8 @@ mod e2e_test; // A global runtime for Vortex operations within DuckDB. static RUNTIME: LazyLock = LazyLock::new(CurrentThreadRuntime::new); +/// Process-wide registry, so repeated scans against the same bucket share one client. +static REGISTRY: LazyLock = LazyLock::new(Registry::new); static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::default().with_handle(RUNTIME.handle()); vortex_spatial::initialize(&session); diff --git a/vortex-duckdb/src/multi_file.rs b/vortex-duckdb/src/multi_file.rs index b8ddc076b8e..f5bb3323e08 100644 --- a/vortex-duckdb/src/multi_file.rs +++ b/vortex-duckdb/src/multi_file.rs @@ -2,12 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; -use std::sync::LazyLock; use itertools::Itertools; use object_store::registry::ObjectStoreRegistry; use url::Url; -use vortex::cloud::Registry; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; @@ -19,14 +17,12 @@ use vortex::io::object_store::ObjectStoreFileSystem; use vortex::io::runtime::BlockingRuntime; use vortex::layout::scan::multi::MultiLayoutDataSource; +use crate::REGISTRY; use crate::RUNTIME; use crate::SESSION; use crate::duckdb::BindInputRef; use crate::duckdb::ExtractedValue; -/// Process-wide registry, so repeated scans against the same bucket share one client. -static REGISTRY: LazyLock = LazyLock::new(Registry::new); - fn resolve_filesystem(glob_url: &Url) -> VortexResult<(FileSystemRef, String)> { // Compat makes us use tokio which is very bad for local reads on // high-core machines because reads go into blocking pool