diff --git a/mdio/dataset.h b/mdio/dataset.h index af5cab7..f9b5710 100644 --- a/mdio/dataset.h +++ b/mdio/dataset.h @@ -101,29 +101,8 @@ inline Future write_zmetadata( inline Future dataset_kvs_store( const std::string& dataset_path, tensorstore::Context context = tensorstore::Context::Default()) { - ::nlohmann::json kvstore; - - // Use shared utility to infer driver from path prefix - std::string driver = zarr::InferDriverFromPath(dataset_path); - kvstore["driver"] = driver; - - if (driver == "file") { - // Local file system - just normalize with trailing slash - kvstore["path"] = zarr::NormalizePathWithSlash(dataset_path); - return tensorstore::kvstore::Open(kvstore, context); - } - - // Cloud storage (GCS or S3) - extract bucket and path - auto [bucket, path] = zarr::ExtractCloudPath(dataset_path); - if (bucket.empty()) { - return absl::InvalidArgumentError( - "gcs/s3 drivers requires [s3/gs]://[bucket]/[path_to_file]"); - } - - kvstore["bucket"] = bucket; - kvstore["path"] = zarr::NormalizePathWithSlash(path); - - return tensorstore::kvstore::Open(kvstore, context); + MDIO_ASSIGN_OR_RETURN(auto loc, zarr::ResolveKvStoreLocation(dataset_path)); + return tensorstore::kvstore::Open(zarr::BuildKvStoreSpec(loc), context); } /** diff --git a/mdio/dataset_factory.h b/mdio/dataset_factory.h index cfa99eb..3d67050 100644 --- a/mdio/dataset_factory.h +++ b/mdio/dataset_factory.h @@ -422,28 +422,9 @@ inline void transform_shape( */ inline absl::Status transform_metadata(const std::string& path, nlohmann::json& variable /*NOLINT*/) { - // Use shared utilities for driver inference and path handling - std::string driver = mdio::zarr::InferDriverFromPath(path); std::string var_name = variable["kvstore"]["path"].get(); - - variable["kvstore"]["driver"] = driver; - - if (driver == "file") { - // Local filesystem - normalize path with trailing slash - variable["kvstore"]["path"] = - mdio::zarr::NormalizePathWithSlash(path) + var_name; - } else { - // Cloud storage (GCS or S3) - extract bucket and path - auto [bucket, cloud_path] = mdio::zarr::ExtractCloudPath(path); - if (bucket.empty()) { - return absl::InvalidArgumentError( - "Cloud path requires [gs/s3]://[bucket]/[path to file] name"); - } - variable["kvstore"]["bucket"] = bucket; - variable["kvstore"]["path"] = - mdio::zarr::NormalizePathWithSlash(cloud_path) + var_name; - } - + MDIO_ASSIGN_OR_RETURN(auto loc, mdio::zarr::ResolveKvStoreLocation(path)); + variable["kvstore"] = mdio::zarr::BuildKvStoreSpec(loc, var_name); return absl::OkStatus(); } diff --git a/mdio/zarr/zarr_driver.h b/mdio/zarr/zarr_driver.h index 830b121..b63af13 100644 --- a/mdio/zarr/zarr_driver.h +++ b/mdio/zarr/zarr_driver.h @@ -223,10 +223,8 @@ inline std::string GetConsolidatedMetadataFileName() { return ".zmetadata"; } * @return The driver name ("file", "gcs", or "s3"). */ inline std::string InferDriverFromPath(const std::string& path) { - if (path.length() > 5) { - if (path.substr(0, 5) == "gs://") return "gcs"; - if (path.substr(0, 5) == "s3://") return "s3"; - } + if (path.compare(0, 5, "gs://") == 0) return "gcs"; + if (path.compare(0, 5, "s3://") == 0) return "s3"; return "file"; } @@ -244,15 +242,21 @@ inline std::string NormalizePath(const std::string& path) { } /** - * @brief Normalizes a path by ensuring it has a trailing slash. + * @brief Normalizes a path so it ends in exactly one trailing slash. + * + * Repeated trailing slashes are collapsed. Cloud object keys are not path + * normalized by the store, so "a//var" and "a/var" are distinct objects. + * * @param path The input path. - * @return The normalized path with a trailing slash. + * @return The normalized path with a single trailing slash. */ inline std::string NormalizePathWithSlash(const std::string& path) { - std::string result = path; - if (!result.empty() && result.back() != '/') { - result.push_back('/'); + std::string result = NormalizePath(path); + if (result.empty()) { + // Preserve an empty path, and collapse a root-only path back to "/". + return path.empty() ? path : "/"; } + result.push_back('/'); return result; } @@ -273,6 +277,76 @@ inline std::pair ExtractCloudPath( without_scheme.substr(bucket_end + 1)}; } +/** + * @brief Resolved TensorStore kvstore location for a dataset path. + * + * Cloud URLs (`gs://`, `s3://`) are split into `bucket` plus a bucket-relative + * `path`. File paths keep the original path. `path` has a trailing slash when + * non-empty so a child name can be appended. + */ +struct KvStoreLocation { + std::string driver; + std::string bucket; + std::string path; +}; + +/** + * @brief Splits a dataset path into TensorStore kvstore fields. + * + * This is the single place that maps a user path onto `driver` / `bucket` / + * `path`. + * + * @param dataset_path Local path or `gs://` / `s3://` URI. + * @return The resolved location, or InvalidArgument for a missing cloud bucket. + */ +inline Result ResolveKvStoreLocation( + const std::string& dataset_path) { + std::string driver = InferDriverFromPath(dataset_path); + if (driver == "file") { + return KvStoreLocation{driver, "", NormalizePathWithSlash(dataset_path)}; + } + + auto [bucket, path] = ExtractCloudPath(dataset_path); + if (bucket.empty()) { + return absl::InvalidArgumentError( + "Cloud path requires [gs/s3]://[bucket]/[path to file]"); + } + return KvStoreLocation{driver, bucket, NormalizePathWithSlash(path)}; +} + +/** + * @brief Builds a TensorStore kvstore JSON spec from a resolved location. + * + * For non-file drivers, `bucket` is always set. `child` is appended to `path` + * (typically a variable / array name). + * + * @param loc The resolved dataset location. + * @param child Optional child key appended to `path`. + * @return The kvstore JSON object. + */ +inline nlohmann::json BuildKvStoreSpec(const KvStoreLocation& loc, + const std::string& child = "") { + nlohmann::json kvstore = {{"driver", loc.driver}, {"path", loc.path + child}}; + if (loc.driver != "file") { + kvstore["bucket"] = loc.bucket; + } + return kvstore; +} + +/** + * @brief Builds a TensorStore Variable open spec (zarr/zarr3 + kvstore). + * @param zarr_driver TensorStore array driver (`zarr` or `zarr3`). + * @param loc The resolved dataset location. + * @param var_name Variable / array name appended to the kvstore path. + * @return The Variable spec JSON object. + */ +inline nlohmann::json BuildVariableSpec(const std::string& zarr_driver, + const KvStoreLocation& loc, + const std::string& var_name) { + return {{"driver", zarr_driver}, + {"kvstore", BuildKvStoreSpec(loc, var_name)}}; +} + // ============================================================================ // Shared JSON Utilities // ============================================================================ diff --git a/mdio/zarr/zarr_test.cc b/mdio/zarr/zarr_test.cc index b709ca0..1c1ae85 100644 --- a/mdio/zarr/zarr_test.cc +++ b/mdio/zarr/zarr_test.cc @@ -266,10 +266,13 @@ TEST(ZarrDriver, NormalizePathWithSlash_AddsTrailingSlash) { TEST(ZarrDriver, NormalizePathWithSlash_AlreadyHasSlash) { EXPECT_EQ(mdio::zarr::NormalizePathWithSlash("/path/to/data/"), "/path/to/data/"); + EXPECT_EQ(mdio::zarr::NormalizePathWithSlash("/path/to/data///"), + "/path/to/data/"); } TEST(ZarrDriver, NormalizePathWithSlash_EmptyPath) { EXPECT_EQ(mdio::zarr::NormalizePathWithSlash(""), ""); + EXPECT_EQ(mdio::zarr::NormalizePathWithSlash("///"), "/"); } TEST(ZarrDriver, ExtractCloudPath_GCS) { @@ -298,6 +301,69 @@ TEST(ZarrDriver, ExtractCloudPath_ShortUrl) { EXPECT_EQ(path, ""); } +TEST(ZarrDriver, ResolveKvStoreLocation_LocalFile) { + auto loc = mdio::zarr::ResolveKvStoreLocation("/path/to/data"); + ASSERT_TRUE(loc.ok()) << loc.status(); + EXPECT_EQ(loc->driver, "file"); + EXPECT_EQ(loc->bucket, ""); + EXPECT_EQ(loc->path, "/path/to/data/"); +} + +TEST(ZarrDriver, ResolveKvStoreLocation_GCS) { + auto loc = mdio::zarr::ResolveKvStoreLocation("gs://bucket/path"); + ASSERT_TRUE(loc.ok()) << loc.status(); + EXPECT_EQ(loc->driver, "gcs"); + EXPECT_EQ(loc->bucket, "bucket"); + EXPECT_EQ(loc->path, "path/"); +} + +TEST(ZarrDriver, ResolveKvStoreLocation_S3) { + auto loc = + mdio::zarr::ResolveKvStoreLocation("s3://my-bucket/nested/path/volume"); + ASSERT_TRUE(loc.ok()) << loc.status(); + EXPECT_EQ(loc->driver, "s3"); + EXPECT_EQ(loc->bucket, "my-bucket"); + EXPECT_EQ(loc->path, "nested/path/volume/"); +} + +TEST(ZarrDriver, ResolveKvStoreLocation_CollapsesTrailingSlashes) { + auto loc = mdio::zarr::ResolveKvStoreLocation("s3://my-bucket/nested/path//"); + ASSERT_TRUE(loc.ok()) << loc.status(); + EXPECT_EQ(loc->path, "nested/path/"); +} + +TEST(ZarrDriver, ResolveKvStoreLocation_RejectsMissingBucket) { + EXPECT_FALSE(mdio::zarr::ResolveKvStoreLocation("gs://").ok()); + EXPECT_FALSE(mdio::zarr::ResolveKvStoreLocation("s3:///path").ok()); +} + +TEST(ZarrDriver, BuildKvStoreSpec_LocalFile) { + auto loc = mdio::zarr::ResolveKvStoreLocation("/path/to/dataset"); + ASSERT_TRUE(loc.ok()) << loc.status(); + auto spec = mdio::zarr::BuildKvStoreSpec(loc.value(), "myvar"); + EXPECT_EQ(spec["driver"], "file"); + EXPECT_EQ(spec["path"], "/path/to/dataset/myvar"); + EXPECT_FALSE(spec.contains("bucket")); +} + +TEST(ZarrDriver, BuildKvStoreSpec_GCS) { + auto loc = mdio::zarr::ResolveKvStoreLocation("gs://bucket/path"); + ASSERT_TRUE(loc.ok()) << loc.status(); + auto spec = mdio::zarr::BuildKvStoreSpec(loc.value(), "myvar"); + EXPECT_EQ(spec["driver"], "gcs"); + EXPECT_EQ(spec["bucket"], "bucket"); + EXPECT_EQ(spec["path"], "path/myvar"); +} + +TEST(ZarrDriver, BuildKvStoreSpec_S3) { + auto loc = mdio::zarr::ResolveKvStoreLocation("s3://my-bucket/nested/path"); + ASSERT_TRUE(loc.ok()) << loc.status(); + auto spec = mdio::zarr::BuildKvStoreSpec(loc.value(), "amplitude"); + EXPECT_EQ(spec["driver"], "s3"); + EXPECT_EQ(spec["bucket"], "my-bucket"); + EXPECT_EQ(spec["path"], "nested/path/amplitude"); +} + // ============================================================================= // JSON Utility Tests // ============================================================================= @@ -895,20 +961,36 @@ TEST(ZarrV3, ExtractChildArrayCandidates_Empty) { } TEST(ZarrV3, BuildVariableSpec_LocalFile) { - auto spec = - mdio::zarr::v3::BuildVariableSpec("file", "/path/to/dataset", "myvar"); + auto loc = mdio::zarr::ResolveKvStoreLocation("/path/to/dataset"); + ASSERT_TRUE(loc.ok()) << loc.status(); + auto spec = mdio::zarr::BuildVariableSpec("zarr3", loc.value(), "myvar"); EXPECT_EQ(spec["driver"], "zarr3"); EXPECT_EQ(spec["kvstore"]["driver"], "file"); EXPECT_EQ(spec["kvstore"]["path"], "/path/to/dataset/myvar"); + EXPECT_FALSE(spec["kvstore"].contains("bucket")); } TEST(ZarrV3, BuildVariableSpec_GCS) { - auto spec = mdio::zarr::v3::BuildVariableSpec("gcs", "bucket/path", "myvar"); + auto loc = mdio::zarr::ResolveKvStoreLocation("gs://bucket/path"); + ASSERT_TRUE(loc.ok()) << loc.status(); + auto spec = mdio::zarr::BuildVariableSpec("zarr3", loc.value(), "myvar"); EXPECT_EQ(spec["driver"], "zarr3"); EXPECT_EQ(spec["kvstore"]["driver"], "gcs"); - EXPECT_EQ(spec["kvstore"]["path"], "bucket/path/myvar"); + EXPECT_EQ(spec["kvstore"]["bucket"], "bucket"); + EXPECT_EQ(spec["kvstore"]["path"], "path/myvar"); +} + +TEST(ZarrV3, BuildVariableSpec_S3) { + auto loc = mdio::zarr::ResolveKvStoreLocation("s3://my-bucket/nested/path"); + ASSERT_TRUE(loc.ok()) << loc.status(); + auto spec = mdio::zarr::BuildVariableSpec("zarr3", loc.value(), "amplitude"); + + EXPECT_EQ(spec["driver"], "zarr3"); + EXPECT_EQ(spec["kvstore"]["driver"], "s3"); + EXPECT_EQ(spec["kvstore"]["bucket"], "my-bucket"); + EXPECT_EQ(spec["kvstore"]["path"], "nested/path/amplitude"); } } // namespace ZarrV3Tests diff --git a/mdio/zarr/zarr_v2.h b/mdio/zarr/zarr_v2.h index 39e53be..194c616 100644 --- a/mdio/zarr/zarr_v2.h +++ b/mdio/zarr/zarr_v2.h @@ -326,22 +326,12 @@ struct V2MetadataState { PromiseType promise; tensorstore::KvStore kvs; std::string dataset_path; - std::string normalized_path; - std::string driver; - std::string bucket; - std::string cloud_path; + KvStoreLocation location; - explicit V2MetadataState(PromiseType p, const std::string& path) + V2MetadataState(PromiseType p, std::string path, KvStoreLocation loc) : promise(std::move(p)), - dataset_path(path), - normalized_path(NormalizePathWithSlash(path)), - driver(InferDriverFromPath(path)) { - if (driver != "file") { - auto cloud_parts = ExtractCloudPath(path); - bucket = cloud_parts.first; - cloud_path = NormalizePathWithSlash(cloud_parts.second); - } - } + dataset_path(std::move(path)), + location(std::move(loc)) {} /// Completes with an error status. void Fail(absl::Status status) { promise.SetResult(std::move(status)); } @@ -352,16 +342,8 @@ struct V2MetadataState { } /// Builds a variable spec for the given variable name. - nlohmann::json BuildVariableSpec(const std::string& var_name) const { - nlohmann::json spec = {{"driver", std::string(kDriverName)}, - {"kvstore", {{"driver", driver}}}}; - if (driver != "file") { - spec["kvstore"]["bucket"] = bucket; - spec["kvstore"]["path"] = cloud_path + var_name; - } else { - spec["kvstore"]["path"] = normalized_path + var_name; - } - return spec; + nlohmann::json MakeVariableSpec(const std::string& var_name) const { + return BuildVariableSpec(std::string(kDriverName), location, var_name); } }; @@ -413,7 +395,7 @@ inline void OnZmetadataRead( if (element.value().contains("dtype") && IsMetadataOnlyDataType(element.value()["dtype"])) { std::string var_name = ExtractVariableName(key); - auto spec = state->BuildVariableSpec(var_name); + auto spec = state->MakeVariableSpec(var_name); spec["_mdio_header_only"] = true; spec["_mdio_zarray"] = element.value(); const std::string zattrs_key = var_name + "/.zattrs"; @@ -426,7 +408,7 @@ inline void OnZmetadataRead( continue; } std::string var_name = ExtractVariableName(key); - json_vars.push_back(state->BuildVariableSpec(var_name)); + json_vars.push_back(state->MakeVariableSpec(var_name)); } if (json_vars.empty()) { @@ -467,11 +449,13 @@ inline void OnV2KvStoreReady( inline Future>> ReadConsolidatedMetadata(const std::string& dataset_path, tensorstore::Future kvs_future) { + MDIO_ASSIGN_OR_RETURN(auto location, ResolveKvStoreLocation(dataset_path)); + auto pair = tensorstore::PromiseFuturePair< std::tuple<::nlohmann::json, std::vector<::nlohmann::json>>>::Make(); auto state = std::make_shared( - std::move(pair.promise), dataset_path); + std::move(pair.promise), dataset_path, std::move(location)); kvs_future.ExecuteWhenReady( [state](tensorstore::ReadyFuture ready) { diff --git a/mdio/zarr/zarr_v3.h b/mdio/zarr/zarr_v3.h index 1b27fa2..6508048 100644 --- a/mdio/zarr/zarr_v3.h +++ b/mdio/zarr/zarr_v3.h @@ -200,21 +200,6 @@ inline std::vector ExtractChildArrayCandidates( return candidates; } -/** - * @brief Builds a variable spec for a Zarr V3 array. - * @param driver The kvstore driver name. - * @param base_path The base dataset path. - * @param var_name The variable/array name. - * @return The JSON spec for opening this variable. - */ -inline nlohmann::json BuildVariableSpec(const std::string& driver, - const std::string& base_path, - const std::string& var_name) { - return { - {"driver", std::string(kDriverName)}, - {"kvstore", {{"driver", driver}, {"path", base_path + "/" + var_name}}}}; -} - // ============================================================================ // Dtype Conversion // ============================================================================ @@ -487,17 +472,12 @@ struct V3MetadataState { PromiseType promise; tensorstore::KvStore kvs; nlohmann::json dataset_metadata; - std::string dataset_path; - std::string driver; - std::string normalized_path; + KvStoreLocation location; std::vector candidates; std::shared_ptr> read_futures; - explicit V3MetadataState(PromiseType p, const std::string& path) - : promise(std::move(p)), - dataset_path(path), - driver(InferDriverFromPath(path)), - normalized_path(NormalizePath(path)) {} + V3MetadataState(PromiseType p, KvStoreLocation loc) + : promise(std::move(p)), location(std::move(loc)) {} /// Completes with an error status. void Fail(absl::Status status) { promise.SetResult(std::move(status)); } @@ -509,7 +489,7 @@ struct V3MetadataState { /// Builds a variable spec for the given variable name. nlohmann::json MakeVariableSpec(const std::string& var_name) const { - return BuildVariableSpec(driver, normalized_path, var_name); + return BuildVariableSpec(std::string(kDriverName), location, var_name); } /// Filters read results to build variable specs for arrays only. @@ -638,11 +618,13 @@ inline void OnV3KvStoreReady( inline Future>> ReadMetadata(const std::string& dataset_path, tensorstore::Future kvs_future) { + MDIO_ASSIGN_OR_RETURN(auto location, ResolveKvStoreLocation(dataset_path)); + auto pair = tensorstore::PromiseFuturePair< std::tuple<::nlohmann::json, std::vector<::nlohmann::json>>>::Make(); auto state = std::make_shared( - std::move(pair.promise), dataset_path); + std::move(pair.promise), std::move(location)); kvs_future.ExecuteWhenReady( [state](tensorstore::ReadyFuture ready) {