parthchandra commented on code in PR #1494: URL: https://github.com/apache/datafusion-comet/pull/1494#discussion_r1990123376
########## native/core/src/parquet/mod.rs: ########## @@ -649,11 +650,11 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat let runtime = tokio::runtime::Builder::new_multi_thread() .enable_all() .build()?; + let session_ctx = SessionContext::new(); + let (object_store_url, object_store_path) = + prepare_object_store(session_ctx.runtime_env(), path.clone()).unwrap(); Review Comment: Done ########## native/core/src/execution/planner.rs: ########## @@ -1208,7 +1208,18 @@ impl PhysicalPlanner { // By default, local FS object store registered // if `hdfs` feature enabled then HDFS file object store registered - let object_store_url = register_object_store(Arc::clone(&self.session_ctx))?; + // Get one file from the list of files + let one_file = scan + .file_partitions + .first() + .unwrap() + .partitioned_file + .first() + .unwrap() + .file_path + .clone(); Review Comment: Done ########## native/core/src/parquet/parquet_support.rs: ########## @@ -199,38 +202,121 @@ fn cast_struct_to_struct( } } -// Default object store which is local filesystem -#[cfg(not(feature = "hdfs"))] -pub(crate) fn register_object_store( - session_context: Arc<SessionContext>, -) -> Result<ObjectStoreUrl, ExecutionError> { - let object_store = object_store::local::LocalFileSystem::new(); - let url = ObjectStoreUrl::parse("file://")?; - session_context - .runtime_env() - .register_object_store(url.as_ref(), Arc::new(object_store)); - Ok(url) -} - -// HDFS object store +// Mirrors object_store::parse::parse_url for the hdfs object store #[cfg(feature = "hdfs")] -pub(crate) fn register_object_store( - session_context: Arc<SessionContext>, -) -> Result<ObjectStoreUrl, ExecutionError> { - // TODO: read the namenode configuration from file schema or from spark.defaultFS - let url = ObjectStoreUrl::parse("hdfs://namenode:9000")?; - if let Some(object_store) = - datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref()) +fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> { + match datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref()) { - session_context - .runtime_env() - .register_object_store(url.as_ref(), Arc::new(object_store)); + Some(object_store) => { + let path = object_store.get_path(url.as_str()); + Ok((Box::new(object_store), path)) + } + _ => { + return Err(object_store::Error::Generic { + store: "HadoopFileSystem", + source: Box::new(datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::HdfsErr::Generic( + "Could not create hdfs object store".to_string(), + )), + }); + } + } +} - return Ok(url); +#[cfg(not(feature = "hdfs"))] +fn parse_hdfs_url(_url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> { + Err(object_store::Error::Generic { + store: "HadoopFileSystem", + source: "Hdfs support is not enabled in this build".into(), + }) +} + +/// Parses the url, registers the object store, and returns a tuple of the object store url and object store path +pub(crate) fn prepare_object_store( + runtime_env: Arc<RuntimeEnv>, + url: String, +) -> Result<(ObjectStoreUrl, Path), ExecutionError> { + let mut url = Url::parse(url.as_str()).unwrap(); + let mut scheme = url.scheme(); + if scheme == "s3a" { + scheme = "s3"; + url.set_scheme("s3") + .expect("Could not convert scheme from s3a to s3"); Review Comment: Done ########## native/core/src/parquet/parquet_support.rs: ########## @@ -199,38 +202,121 @@ fn cast_struct_to_struct( } } -// Default object store which is local filesystem -#[cfg(not(feature = "hdfs"))] -pub(crate) fn register_object_store( - session_context: Arc<SessionContext>, -) -> Result<ObjectStoreUrl, ExecutionError> { - let object_store = object_store::local::LocalFileSystem::new(); - let url = ObjectStoreUrl::parse("file://")?; - session_context - .runtime_env() - .register_object_store(url.as_ref(), Arc::new(object_store)); - Ok(url) -} - -// HDFS object store +// Mirrors object_store::parse::parse_url for the hdfs object store #[cfg(feature = "hdfs")] -pub(crate) fn register_object_store( - session_context: Arc<SessionContext>, -) -> Result<ObjectStoreUrl, ExecutionError> { - // TODO: read the namenode configuration from file schema or from spark.defaultFS - let url = ObjectStoreUrl::parse("hdfs://namenode:9000")?; - if let Some(object_store) = - datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref()) +fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> { + match datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref()) { - session_context - .runtime_env() - .register_object_store(url.as_ref(), Arc::new(object_store)); + Some(object_store) => { + let path = object_store.get_path(url.as_str()); + Ok((Box::new(object_store), path)) + } + _ => { + return Err(object_store::Error::Generic { + store: "HadoopFileSystem", + source: Box::new(datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::HdfsErr::Generic( + "Could not create hdfs object store".to_string(), + )), + }); + } + } +} - return Ok(url); +#[cfg(not(feature = "hdfs"))] +fn parse_hdfs_url(_url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> { + Err(object_store::Error::Generic { + store: "HadoopFileSystem", + source: "Hdfs support is not enabled in this build".into(), + }) +} + +/// Parses the url, registers the object store, and returns a tuple of the object store url and object store path +pub(crate) fn prepare_object_store( + runtime_env: Arc<RuntimeEnv>, + url: String, +) -> Result<(ObjectStoreUrl, Path), ExecutionError> { + let mut url = Url::parse(url.as_str()).unwrap(); Review Comment: Done ########## native/core/src/parquet/parquet_support.rs: ########## @@ -199,38 +202,121 @@ fn cast_struct_to_struct( } } -// Default object store which is local filesystem -#[cfg(not(feature = "hdfs"))] -pub(crate) fn register_object_store( - session_context: Arc<SessionContext>, -) -> Result<ObjectStoreUrl, ExecutionError> { - let object_store = object_store::local::LocalFileSystem::new(); - let url = ObjectStoreUrl::parse("file://")?; - session_context - .runtime_env() - .register_object_store(url.as_ref(), Arc::new(object_store)); - Ok(url) -} - -// HDFS object store +// Mirrors object_store::parse::parse_url for the hdfs object store #[cfg(feature = "hdfs")] -pub(crate) fn register_object_store( - session_context: Arc<SessionContext>, -) -> Result<ObjectStoreUrl, ExecutionError> { - // TODO: read the namenode configuration from file schema or from spark.defaultFS - let url = ObjectStoreUrl::parse("hdfs://namenode:9000")?; - if let Some(object_store) = - datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref()) +fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> { + match datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref()) { - session_context - .runtime_env() - .register_object_store(url.as_ref(), Arc::new(object_store)); + Some(object_store) => { + let path = object_store.get_path(url.as_str()); + Ok((Box::new(object_store), path)) + } + _ => { + return Err(object_store::Error::Generic { + store: "HadoopFileSystem", + source: Box::new(datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::HdfsErr::Generic( + "Could not create hdfs object store".to_string(), + )), + }); + } + } +} - return Ok(url); +#[cfg(not(feature = "hdfs"))] +fn parse_hdfs_url(_url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> { + Err(object_store::Error::Generic { + store: "HadoopFileSystem", + source: "Hdfs support is not enabled in this build".into(), + }) +} + +/// Parses the url, registers the object store, and returns a tuple of the object store url and object store path +pub(crate) fn prepare_object_store( + runtime_env: Arc<RuntimeEnv>, + url: String, +) -> Result<(ObjectStoreUrl, Path), ExecutionError> { + let mut url = Url::parse(url.as_str()).unwrap(); + let mut scheme = url.scheme(); + if scheme == "s3a" { + scheme = "s3"; + url.set_scheme("s3") + .expect("Could not convert scheme from s3a to s3"); } + let url_key = format!( + "{}://{}", + scheme, + &url[url::Position::BeforeHost..url::Position::AfterPort], + ); + + let (object_store, object_store_path): (Box<dyn ObjectStore>, Path) = if scheme == "hdfs" { + match parse_hdfs_url(&url) { + Ok(r) => r, + Err(e) => return Err(ExecutionError::GeneralError(e.to_string())), + } + } else { + match parse_url(&url) { + Ok(r) => r, + Err(e) => return Err(ExecutionError::GeneralError(e.to_string())), + } + }; + let object_store_url = ObjectStoreUrl::parse(url_key.clone())?; + runtime_env.register_object_store(&url, Arc::from(object_store)); + Ok((object_store_url, object_store_path)) +} + +#[cfg(test)] +mod tests { + use crate::execution::operators::ExecutionError; + use crate::parquet::parquet_support::prepare_object_store; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_execution::runtime_env::RuntimeEnv; + use object_store::path::Path; + use std::sync::Arc; + use url::Url; + + #[test] + #[cfg(not(feature = "hdfs"))] Review Comment: Added the test. I have trouble running it from the command line (it runs fine in Rust Rover) so it would be nice if someone can double check it for me. ########## native/core/src/parquet/parquet_support.rs: ########## @@ -199,38 +202,121 @@ fn cast_struct_to_struct( } } -// Default object store which is local filesystem -#[cfg(not(feature = "hdfs"))] -pub(crate) fn register_object_store( - session_context: Arc<SessionContext>, -) -> Result<ObjectStoreUrl, ExecutionError> { - let object_store = object_store::local::LocalFileSystem::new(); - let url = ObjectStoreUrl::parse("file://")?; - session_context - .runtime_env() - .register_object_store(url.as_ref(), Arc::new(object_store)); - Ok(url) -} - -// HDFS object store +// Mirrors object_store::parse::parse_url for the hdfs object store #[cfg(feature = "hdfs")] -pub(crate) fn register_object_store( - session_context: Arc<SessionContext>, -) -> Result<ObjectStoreUrl, ExecutionError> { - // TODO: read the namenode configuration from file schema or from spark.defaultFS - let url = ObjectStoreUrl::parse("hdfs://namenode:9000")?; - if let Some(object_store) = - datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref()) +fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> { + match datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref()) { - session_context - .runtime_env() - .register_object_store(url.as_ref(), Arc::new(object_store)); + Some(object_store) => { + let path = object_store.get_path(url.as_str()); + Ok((Box::new(object_store), path)) + } + _ => { + return Err(object_store::Error::Generic { + store: "HadoopFileSystem", + source: Box::new(datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::HdfsErr::Generic( + "Could not create hdfs object store".to_string(), + )), + }); + } + } +} - return Ok(url); +#[cfg(not(feature = "hdfs"))] +fn parse_hdfs_url(_url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> { + Err(object_store::Error::Generic { + store: "HadoopFileSystem", + source: "Hdfs support is not enabled in this build".into(), + }) +} + +/// Parses the url, registers the object store, and returns a tuple of the object store url and object store path +pub(crate) fn prepare_object_store( + runtime_env: Arc<RuntimeEnv>, + url: String, +) -> Result<(ObjectStoreUrl, Path), ExecutionError> { + let mut url = Url::parse(url.as_str()).unwrap(); + let mut scheme = url.scheme(); + if scheme == "s3a" { + scheme = "s3"; + url.set_scheme("s3") + .expect("Could not convert scheme from s3a to s3"); } + let url_key = format!( + "{}://{}", + scheme, + &url[url::Position::BeforeHost..url::Position::AfterPort], + ); + + let (object_store, object_store_path): (Box<dyn ObjectStore>, Path) = if scheme == "hdfs" { + match parse_hdfs_url(&url) { + Ok(r) => r, + Err(e) => return Err(ExecutionError::GeneralError(e.to_string())), + } + } else { + match parse_url(&url) { + Ok(r) => r, + Err(e) => return Err(ExecutionError::GeneralError(e.to_string())), + } + }; Review Comment: Done -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For additional commands, e-mail: github-h...@datafusion.apache.org