Add full path metadata test and last_modified formatting

Author: BohuTANG

databend_aiserver/udfs/docparse.py

@@ -239,6 +239,11 @@ def ai_parse_document(stage_location: StageLocation, path: str) -> Dict[str, Any
         # Output shape:
         # { "chunks": [...], "metadata": {...}, "error_information": [...] }
         resolved_path = resolve_stage_subpath(stage_location, path)
+        root_prefix = ""
+        storage_root = stage_location.storage.get("root") if stage_location.storage else ""
+        if storage_root:
+            root_prefix = storage_root.rstrip("/")
+        full_path = f"{root_prefix}/{resolved_path}" if root_prefix else resolved_path
 
         # Keep metadata first for predictable JSON ordering.
         payload: Dict[str, Any] = {
@@ -248,7 +253,7 @@ def ai_parse_document(stage_location: StageLocation, path: str) -> Dict[str, Any
                 "duration_ms": duration_ms,
                 "file_size": file_size if file_size is not None else 0,
                 "filename": Path(path).name,
-                "path": resolved_path or path,
+                "path": full_path or path,
                 "timings_ms": {
                     "convert": (t_convert_end_ns - t_convert_start_ns) / 1_000_000.0,
                     "chunk": (t_chunk_end_ns - t_convert_end_ns) / 1_000_000.0,

databend_aiserver/udfs/stage.py

@@ -30,6 +30,15 @@
 )
 
 
+def _format_last_modified(value: Any) -> Optional[str]:
+    if value is None:
+        return None
+    # OpenDAL returns datetime objects; fall back to string otherwise.
+    if hasattr(value, "isoformat"):
+        return value.isoformat()
+    return str(value)
+
+
 def _collect_stage_files(
     stage_location: StageLocation, max_files: Optional[int]
 ) -> tuple[List[Dict[str, Any]], bool]:
@@ -85,9 +94,7 @@ def _collect_stage_files(
             if metadata.etag:
                 file_info["etag"] = metadata.etag
             if hasattr(metadata, "last_modified"):
-                lm = metadata.last_modified
-                if lm:
-                    file_info["last_modified"] = lm
+                file_info["last_modified"] = _format_last_modified(metadata.last_modified)
 
         entries.append(file_info)
         if max_entries is not None and len(entries) >= max_entries:
@@ -167,7 +174,7 @@ def ai_list_files(
 
             # Convert mode to string if it exists, otherwise None
             mode_str = str(metadata.mode) if metadata.mode is not None else None
-            last_modified = getattr(metadata, "last_modified", None)
+            last_modified = _format_last_modified(getattr(metadata, "last_modified", None))
 
             yield {
                 "stage_name": stage_location.stage_name,

tests/unit/conftest.py

@@ -52,3 +52,20 @@ def memory_stage() -> StageLocation:
     operator.write(f"{RELATIVE_PATH}/lorem_ipsum.docx", DOCX_SRC.read_bytes())
     operator.write(f"{RELATIVE_PATH}/subdir/note.txt", b"hello from memory")
     return stage
+
+
+@pytest.fixture
+def memory_stage_with_root() -> StageLocation:
+    stage = StageLocation(
+        name="stage",
+        stage_name="memory_stage_root",
+        stage_type="External",
+        storage={"type": "memory", "root": "s3://wizardbend/dataset"},
+        relative_path=RELATIVE_PATH,
+        raw_info={},
+    )
+
+    operator = get_operator(stage)
+    operator.create_dir(f"{RELATIVE_PATH}/")
+    operator.write(f"{RELATIVE_PATH}/2206.01062.pdf", PDF_SRC.read_bytes())
+    return stage

tests/unit/test_docparse_path.py

@@ -0,0 +1,22 @@
+# Copyright 2025 Databend Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from databend_aiserver.udfs.docparse import ai_parse_document
+
+
+def test_docparse_metadata_path_uses_root(memory_stage_with_root):
+    payload = ai_parse_document(memory_stage_with_root, "2206.01062.pdf")
+    meta = payload.get("metadata", {})
+    assert meta["path"] == "s3://wizardbend/dataset/data/2206.01062.pdf"
+    assert meta["filename"] == "2206.01062.pdf"