Fix merge_index dropping nested directories from shard paths

streaming/base/util.py

@@ -327,20 +327,24 @@ def _merge_index_from_list(index_file_urls: Sequence[Union[str, tuple[str, str]]
             if not os.path.exists(dest):
                 raise FileNotFoundError(f'Index file {dest} does not exist or not accessible.')
 
-            partitions.append(dest)
+            partitions.append((dest, _get_partition_dirname(url, out)))
 
         # merge shards from all index files
         shards = []
-        for partition_index in partitions:
+        for partition_index, partition_dirname in partitions:
             p = Path(partition_index)
+            if partition_dirname is None:
+                # The partition does not reside under ``out``, so its path relative to ``out``
+                # cannot be determined. Fall back to the partition's parent directory name.
+                partition_dirname = os.path.basename(p.parent)
             obj = json.load(open(partition_index))
             for i in range(len(obj['shards'])):
                 shard = obj['shards'][i]
                 for key in ('raw_data', 'zip_data', 'raw_meta', 'zip_meta'):
                     if shard.get(key):
                         basename = shard[key]['basename']
                         obj['shards'][i][key]['basename'] = os.path.join(
-                            os.path.basename(p.parent), basename)
+                            partition_dirname, basename)
             shards += obj['shards']
 
         # Save merged index locally
@@ -363,6 +367,46 @@ def _merge_index_from_list(index_file_urls: Sequence[Union[str, tuple[str, str]]
             shutil.rmtree(cu.local, ignore_errors=True)
 
 
+def _get_partition_dirname(index_file_url: Union[str, tuple[str, str]],
+                           out: Union[str, tuple[str, str]]) -> Optional[str]:
+    """Get the directory of a partition index file relative to the merge root ``out``.
+
+    Args:
+        index_file_url (Union[str, Tuple[str,str]]): a partition index file url, either a single
+            path string or a (local, remote) tuple.
+        out (Union[str, Tuple[str,str]]): the merge root, either a single path string or a
+            (local, remote) tuple.
+
+    Returns:
+        Optional[str]: the partition directory relative to ``out``, e.g. ``group1/subdir2`` for
+            an index file at ``<out>/group1/subdir2/index.json``. ``None`` if the partition does
+            not reside under ``out``.
+    """
+    urls = index_file_url if isinstance(index_file_url, tuple) else (index_file_url,)
+    roots = out if isinstance(out, tuple) else (out,)
+    for url in urls:
+        url_obj = urllib.parse.urlparse(url)
+        for root in roots:
+            root_obj = urllib.parse.urlparse(root)
+            # Only compare urls that live in the same place, e.g. the same bucket or the local
+            # filesystem.
+            if (url_obj.scheme, url_obj.netloc) != (root_obj.scheme, root_obj.netloc):
+                continue
+            root_path = root_obj.path
+            if root_obj.scheme and not root_path:
+                root_path = '/'
+            try:
+                rel = os.path.relpath(os.path.dirname(url_obj.path), root_path)
+            except ValueError:
+                continue
+            if rel == os.curdir:
+                # The partition index sits directly at the root, so do not prefix basenames.
+                return ''
+            if rel != os.pardir and not rel.startswith(os.pardir + os.sep):
+                return rel
+    return None
+
+
 def _not_merged_index(index_file_path: str, out: str):
     """Check if index_file_path is the merged index at folder out.
 

tests/test_util.py

@@ -199,6 +199,64 @@ def test_format_remote_index_files(scheme: str):
         assert obj.scheme == scheme
 
 
+@pytest.mark.parametrize(('index_file_url', 'out', 'expected'), [
+    ('/foo/group1/subdir2/index.json', '/foo', os.path.join('group1', 'subdir2')),
+    ('/foo/subdir1/index.json', '/foo', 'subdir1'),
+    ('/foo/index.json', '/foo', ''),
+    ('/elsewhere/subdir1/index.json', '/foo', None),
+    ('s3://bucket/foo/group1/subdir2/index.json', 's3://bucket/foo',
+     os.path.join('group1', 'subdir2')),
+    ('s3://bucket/foo/subdir1/index.json', 's3://bucket', os.path.join('foo', 'subdir1')),
+    ('s3://bucket/foo/subdir1/index.json', 'gs://bucket/foo', None),
+    ('s3://other/foo/subdir1/index.json', 's3://bucket/foo', None),
+    (('/foo/subdir1/index.json', 's3://bucket/foo/subdir1/index.json'),
+     ('/foo', 's3://bucket/foo'), 'subdir1'),
+    (('/elsewhere/subdir1/index.json', 's3://bucket/foo/group1/subdir1/index.json'),
+     's3://bucket/foo', os.path.join('group1', 'subdir1')),
+])
+def test_get_partition_dirname(index_file_url: Union[str, tuple[str, str]],
+                               out: Union[str, tuple[str, str]], expected: Optional[str]):
+    """Validate partition directories are resolved relative to the merge root."""
+    from streaming.base.util import _get_partition_dirname
+
+    assert _get_partition_dirname(index_file_url, out) == expected
+
+
+@pytest.mark.parametrize('keep_local', [True, False])
+def test_merge_index_from_root_local_nested(local_remote_dir: tuple[str, str], keep_local: bool):
+    """Validate the merged index keeps full relative paths for nested partitions."""
+    from streaming import MDSWriter, StreamingDataset
+
+    out, _ = local_remote_dir
+    n_samples = 0
+    for group in ('group1', 'group2'):
+        for subdir in ('subdir1', 'subdir2'):
+            with MDSWriter(out=os.path.join(out, group, subdir),
+                           columns={'id': 'int'},
+                           keep_local=True) as writer:
+                for _ in range(3):
+                    writer.write({'id': n_samples})
+                    n_samples += 1
+
+    merge_index(out, keep_local=keep_local)
+    integrity_check(out, keep_local=keep_local, expected_n_shard_files=4)
+
+    if not keep_local:
+        return
+
+    merged_index = json.load(open(os.path.join(out, 'index.json')))
+    basenames = sorted(shard['raw_data']['basename'] for shard in merged_index['shards'])
+    assert basenames == [
+        os.path.join('group1', 'subdir1', 'shard.00000.mds'),
+        os.path.join('group1', 'subdir2', 'shard.00000.mds'),
+        os.path.join('group2', 'subdir1', 'shard.00000.mds'),
+        os.path.join('group2', 'subdir2', 'shard.00000.mds'),
+    ]
+
+    dataset = StreamingDataset(local=out, batch_size=1)
+    assert sorted(sample['id'] for sample in dataset) == list(range(n_samples))
+
+
 @pytest.mark.parametrize('index_file_urls_pattern', [1, 2, 3])
 @pytest.mark.parametrize('keep_local', [True, False])
 @pytest.mark.parametrize('scheme', ['gs://', 's3://', 'oci://'])