[FSTORE-1805][APPEND] Engine.add_file() got an unexpected keyword argument 'distribute' (#751)

Author: bubriks

python/hsfs/engine/python.py

@@ -1553,7 +1553,7 @@ def _get_app_options(
             spark_job_configuration=spark_job_configuration,
         )
 
-    def add_file(self, file: Optional[str]) -> Optional[str]:
+    def add_file(self, file: Optional[str], distribute=True) -> Optional[str]:
         if not file:
             return file
 

python/tests/engine/test_python.py

@@ -3253,14 +3253,27 @@ def test_get_app_options(self, mocker):
         assert mock_ingestion_job_conf.call_count == 1
         assert mock_ingestion_job_conf.call_args[1]["write_options"] == {"test": 2}
 
-    def test_add_file(self):
+    @pytest.mark.parametrize(
+        "distribute_arg",
+        [
+            None,  # Test without providing distribute argument (uses default)
+            True,  # Test with distribute=True
+            False,  # Test with distribute=False
+        ],
+    )
+    def test_add_file(self, distribute_arg):
         # Arrange
         python_engine = python.Engine()
 
         file = None
 
         # Act
-        result = python_engine.add_file(file=file)
+        if distribute_arg is None:
+            # Call without distribute argument
+            result = python_engine.add_file(file=file)
+        else:
+            # Call with distribute argument
+            result = python_engine.add_file(file=file, distribute=distribute_arg)
 
         # Assert
         assert result == file

python/tests/engine/test_spark.py

@@ -4379,8 +4379,17 @@ def test_read_stream_kafka_message_format_avro_include_metadata(self, mocker):
         assert "name" in result.columns
         assert result.schema["name"].dataType == StringType()
 
-    def test_add_file(self, mocker):
+    @pytest.mark.parametrize(
+        "distribute_arg",
+        [
+            None,  # Test without providing distribute argument (uses default)
+            True,  # Test with distribute=True
+            False,  # Test with distribute=False
+        ],
+    )
+    def test_add_file(self, mocker, distribute_arg):
         # Arrange
+        mock_dataset_api = mocker.patch("hsfs.core.dataset_api.DatasetApi")
         mock_pyspark_files_get = mocker.patch("pyspark.files.SparkFiles.get")
         mocker.patch("hopsworks_common.client._is_external", return_value=False)
         mocker.patch("shutil.copy")
@@ -4389,14 +4398,31 @@ def test_add_file(self, mocker):
 
         spark_engine = spark.Engine()
 
+        # Mock dataset API and file I/O for distribute=False case
+        if distribute_arg is False:
+            mock_dataset_api.return_value.read_content.return_value.content = bytes()
+            mocker.patch("builtins.open", mocker.mock_open())
+
         # Act
-        spark_engine.add_file(
-            file="test_file",
-        )
+        if distribute_arg is None:
+            # Call without distribute argument
+            spark_engine.add_file(file="test_file")
+        else:
+            # Call with distribute argument
+            spark_engine.add_file(file="test_file", distribute=distribute_arg)
 
         # Assert
-        mock_add_file.assert_called_once_with("hdfs://test_file")
-        mock_pyspark_files_get.assert_called_once_with("test_file")
+        if distribute_arg is False:
+            # When distribute=False, read_content should be called once
+            mock_dataset_api.return_value.read_content.assert_called_once()
+            # addFile and SparkFiles.get should NOT be called
+            mock_add_file.assert_not_called()
+            mock_pyspark_files_get.assert_not_called()
+        else:
+            # When distribute is True or None (default), addFile should be called
+            mock_dataset_api.return_value.read_content.assert_not_called()
+            mock_add_file.assert_called_once_with("hdfs://test_file")
+            mock_pyspark_files_get.assert_called_once_with("test_file")
 
     def test_add_file_if_present_in_job_configuration(self, mocker):
         # Arrange