Missing dimension_names metadata in opening zarr (xr.open_zarr)

[DagonArises]

Dependencies:

[tool.poetry.dependencies]
python = "^3.11"
xarray = "^2025.4.0"
scipy = "^1.15.3"
netcdf4 = "^1.7.2"
dask = "^2025.5.1"
gcsfs = "^2025.5.1"
google-cloud-storage = "^3.1.0"
psutil = "^7.0.0"
zarr = "^3.0.8"

Error observed.
'Zarr object is missing the dimension_names metadata which is required for xarray to determine variable dimensions.'

Script.

import os
import xarray as xr
import zarr
import time
import psutil
import logging
import tempfile
import glob
import sys
import re
import numpy as np
from google.cloud import storage
import dask
from numcodecs.zarr3 import Blosc

# Get configuration from environment variables
INPUT_BUCKET_PATH = os.environ.get("INPUT_BUCKET_PATH") 
OUTPUT_BUCKET_NAME = os.environ.get("OUTPUT_BUCKET_NAME")
OUTPUT_BLOB_PATH = os.environ.get("OUTPUT_BLOB_PATH")  
MOUNT_PATH = os.environ.get("MOUNT_PATH", "/mnt/gcs-1")

def extract_indices(filename):
    basename = os.path.basename(filename)
    # Match pattern ending with _XX_YY.nc where XX and YY are numbers
    match = re.search(r'_(\d+)_(\d+)\.nc$', basename)
    if match:
        return int(match.group(1)), int(match.group(2))
    else:
        raise ValueError(f"Could not extract indices from filename: {filename}")


def main():
    start_time = time.perf_counter()
    process = psutil.Process(os.getpid())

    try:
        # Step 1. Access mounted files directly
        input_dir = os.path.join(MOUNT_PATH, INPUT_BUCKET_PATH)
        nc_files = glob.glob(os.path.join(input_dir, FILE_PATTERN))
        logger.info(f"Found {len(nc_files)} files to merge.")

        # Step 2. Find the dimensions from filenames
        max_x_idx = -1
        max_y_idx = -1
        for nc_file in nc_files:
            x_idx, y_idx = extract_indices(nc_file)
            max_x_idx = max(max_x_idx, x_idx)
            max_y_idx = max(max_y_idx, y_idx)
        
        nx = max_x_idx + 1  # Number of files in x direction
        ny = max_y_idx + 1  # Number of files in y direction
        logger.info(f"Grid dimensions: nx={nx}, ny={ny}")

        # Step 3. Compute output size
        sample_ds = xr.open_dataset(nc_files[0])
        t_i = len(sample_ds.time)
        x_i = len(sample_ds.x)
        y_i = len(sample_ds.y)
        CHUNK_SIZE_X = x_i//2
        CHUNK_SIZE_Y = y_i//2
        
        # Get all variable names
        data_vars = list(sample_ds.data_vars)
        logger.info(f"Variables to process: {data_vars}")
        sample_ds.close()
        
        # Compute overall dimensions
        total_x = nx * x_i
        total_y = ny * y_i
        logger.info(f"Individual file dimensions: time={t_i}, x={x_i}, y={y_i}")
        logger.info(f"Total merged dimensions: time={t_i}, x={total_x}, y={total_y}")

        # Step 4. Create Zarr store with compression in temporary directory
        with tempfile.TemporaryDirectory() as temp_dir:
            zarr_path = os.path.join(temp_dir, "merged.zarr")
            
            # Create zarr store with appropriate chunking and compression
            store = zarr.storage.LocalStore(zarr_path)
            root = zarr.group(store=store, overwrite=True)

            # Create arrays for all variables with compression
            zarr_arrays = {}
            for var_name in data_vars:
                zarr_arrays[var_name] = root.create_array(
                    var_name,
                    shape=(t_i, total_x, total_y),
                    chunks=(10, CHUNK_SIZE_X, CHUNK_SIZE_Y),
                    dtype='float32',
                    fill_value=-9999.0,
                    compressors=Blosc(cname='zstd', clevel=3, shuffle=2)  # Fixed: compressor not compressors
                )
                # Set dimension_names for each variable array
                zarr_arrays[var_name].attrs['dimension_names'] = ['time', 'x', 'y']
            
            # Create coordinate arrays
            time_array = root.create_array('time', shape=(t_i,), dtype='float32')
            x_array = root.create_array('x', shape=(total_x,), dtype='float32')
            y_array = root.create_array('y', shape=(total_y,), dtype='float32')
            
            # Set dimension_names for coordinate arrays
            time_array.attrs['dimension_names'] = ['time']
            x_array.attrs['dimension_names'] = ['x']
            y_array.attrs['dimension_names'] = ['y']
            
            # Set group-level dimension_names for all variables
            dimension_names_dict = {
                'time': ['time'],
                'x': ['x'],
                'y': ['y']
            }
            # Add all data variables
            for var_name in data_vars:
                dimension_names_dict[var_name] = ['time', 'x', 'y']
            
            root.attrs['dimension_names'] = dimension_names_dict
            
            # Step 5. Write NetCDF files into Zarr one by one
            logger.info("Writing NetCDF files to Zarr...")
            files_processed = 0
            
            for nc_file in nc_files:
                x_idx, y_idx = extract_indices(nc_file)
                
                # Open dataset
                ds = xr.open_dataset(nc_file)
                
                # Write data to zarr at relative position
                x_start = x_idx * x_i
                x_end = x_start + x_i
                y_start = y_idx * y_i
                y_end = y_start + y_i
                
                # Write all variables
                for var_name in data_vars:
                    zarr_arrays[var_name][:, x_start:x_end, y_start:y_end] = ds[var_name].values
                
                # Write coordinates (only need to do this once)
                if files_processed == 0:
                    time_array[:] = ds.time.values
                    # Create spatial coordinates
                    x_coords = np.arange(total_x) * (ds['x'].values[1] - ds['x'].values[0])
                    y_coords = np.arange(total_y) * (ds['y'].values[1] - ds['y'].values[0])
                    x_array[:] = x_coords
                    y_array[:] = y_coords
                
                ds.close()
                
                files_processed += 1
                if files_processed % 10 == 0:
                    logger.info(f"Processed {files_processed}/{len(nc_files)} files")
                
                # Remove the debug break
                if files_processed == 5:
                    break
            
            logger.info(f"All {files_processed} files written to Zarr")
            
            # Consolidate metadata before converting
            zarr.consolidate_metadata(zarr_path)

            # Step 6. Convert Zarr to NetCDF with chunking for low memory usage
            with tempfile.NamedTemporaryFile(suffix='.nc', delete=False) as tmp_file:
                temp_nc_path = tmp_file.name
            
            # Open zarr as xarray dataset with chunking
            logger.info('Opening zarr dataset...')
            ds_zarr = xr.open_zarr(
                zarr_path, 
                consolidated=True, 
                chunks={'time': 10, 'x': CHUNK_SIZE_X, 'y': CHUNK_SIZE_Y},
                zarr_format=3  
            )
            logger.info('Zarr dataset opened successfully')

            # Apply encoding to all variables
            encoding = {}
            for var in ds_zarr.data_vars:  # Fixed: use ds_zarr not merged_ds
                encoding[var] = {
                    "dtype": 'float32', 
                    "zlib": True, 
                    "complevel": 4, 
                    "_FillValue": -9999.0, 
                    'chunksizes': (10, CHUNK_SIZE_X, CHUNK_SIZE_Y)
                }

            logger.info(f"Encoding defined for variables: {list(encoding.keys())}")
            logger.info("Converting to NetCDF...")
            
            # Use dask single-threaded scheduler to minimize memory usage
            with dask.config.set(scheduler='synchronous'):
                ds_zarr.to_netcdf(temp_nc_path, format="NETCDF4", encoding=encoding)

            ds_zarr.close()
            
            # Step 7. Upload to destination bucket
            logger.info(f"Uploading to gs://{OUTPUT_BUCKET_NAME}/{OUTPUT_BLOB_PATH}")
            
            with open(temp_nc_path, 'rb') as f:
                merged_netcdf_bytes = f.read()
            storage_client = storage.Client()
            bucket = storage_client.bucket(OUTPUT_BUCKET_NAME)
            blob = bucket.blob(OUTPUT_BLOB_PATH)
            blob.upload_from_string(merged_netcdf_bytes, content_type='application/x-netcdf')
            
            # Clean up temporary file
            os.unlink(temp_nc_path)
            
            # Memory and time reporting
            peak_memory = process.memory_info().rss / (1024 * 1024)  # MB
            elapsed_time = time.perf_counter() - start_time
            output_size = len(merged_netcdf_bytes) / (1024 * 1024)  # MB
            
            logger.info(f'Processing completed:')
            logger.info(f'  Elapsed time: {elapsed_time:.2f} seconds')
            logger.info(f'  Peak memory usage: {peak_memory:.2f} MB')
            logger.info(f'  Output file size: {output_size:.2f} MB')
            logger.info(f'  Variables processed: {data_vars}')
            logger.info(f'  Successfully uploaded to gs://{OUTPUT_BUCKET_NAME}/{OUTPUT_BLOB_PATH}')

    except Exception as e:
        logger.error(f"Error during processing: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    main()

I set dimension_names for each attribute but the step of opening zarr still reported missing dimension_names:

ds_zarr = xr.open_zarr(
                zarr_path, 
                consolidated=True, 
                chunks={'time': 10, 'x': CHUNK_SIZE_X, 'y': CHUNK_SIZE_Y},
                zarr_format = 3
            )

Would like some help in correctly preparing zarray for to-netcdf conversion. Thanks in advance!