Fix map() when no function used at all.

ArneBinder · Sep 30, 2024 · 0f6ed10 · 0f6ed10
1 parent c5ca512
commit 0f6ed10
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/pie_datasets/core/dataset.py b/src/pie_datasets/core/dataset.py
@@ -376,7 +376,9 @@ def map(
         result_document_type: Optional[Type[Document]] = None,
     ) -> "Dataset":
         dataset = super().map(
-            function=decorate_convert_to_dict_of_lists(function) if as_documents else function,
+            function=decorate_convert_to_dict_of_lists(function)
+            if as_documents and function is not None
+            else function,
             with_indices=with_indices,
             with_rank=with_rank,
             input_columns=input_columns,
@@ -582,7 +584,7 @@ def map(  # type: ignore
             function=decorate_convert_to_document_and_back(
                 function, document_type=self.document_type, batched=batched
             )
-            if as_documents
+            if as_documents and function is not None
             else function,
             batched=batched,
             **kwargs,