comphead commented on code in PR #1494:
URL: https://github.com/apache/datafusion-comet/pull/1494#discussion_r1990030870
##########
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:
should we also have a test with hdfs enabled?
--
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
