Add import csv arrow test (#20082)

3 files changed, 25 insertions, 4 deletions

diff --git a/ydb/tests/functional/tpc/medium/test_import_csv.py b/ydb/tests/functional/tpc/medium/test_import_csv.py
index 905b81448a8..e0a662b03f7 100644
--- a/ydb/tests/functional/tpc/medium/test_import_csv.py
+++ b/ydb/tests/functional/tpc/medium/test_import_csv.py
@@ -2,10 +2,17 @@ from ydb.tests.olap.load.lib.import_csv import ImportFileCsvBase # noqa
 from ydb.tests.functional.tpc.lib.conftest import FunctionalTestBase
 
 
-class TestExternalImportCsv(ImportFileCsvBase, FunctionalTestBase):
-    external_folder: str = 'e1'
-
+class FunctionalImportFileCsvBase(ImportFileCsvBase, FunctionalTestBase):
     @classmethod
     def setup_class(cls) -> None:
         cls.setup_cluster()
         super().setup_class()
+
+
+class TestExternalImportCsv(FunctionalImportFileCsvBase):
+    external_folder: str = 'e1'
+
+
+class TestExternalImportCsvArrow(FunctionalImportFileCsvBase):
+    external_folder: str = 'e1'
+    send_format: str = 'arrow'
diff --git a/ydb/tests/olap/load/lib/import_csv.py b/ydb/tests/olap/load/lib/import_csv.py
index 28dafaab0c8..588ccb46ad1 100644
--- a/ydb/tests/olap/load/lib/import_csv.py
+++ b/ydb/tests/olap/load/lib/import_csv.py
@@ -15,6 +15,7 @@ class ImportFileCsvBase(UploadSuiteBase):
     iterations: int = 1
     cpu_cores: float = 0.0
     cpu_time: float = 0.0
+    send_format: str = ''  # Optional parameter for send format
 
     def init(self):
         # Create tables
@@ -35,7 +36,14 @@ class ImportFileCsvBase(UploadSuiteBase):
         if not csv_files:
             raise RuntimeError(f'No .csv files found in {import_dir}')
         import_path = os.path.join(import_dir, csv_files[0])
-        result = yatest.common.execute(['/usr/bin/time'] + YdbCliHelper.get_cli_command() + ['import', 'file', 'csv', '-p', self.table_path, import_path, '--header'])
+
+        cmd = ['/usr/bin/time'] + YdbCliHelper.get_cli_command() + ['import', 'file', 'csv', '-p', self.table_path, import_path, '--header']
+        if self.send_format:
+            cmd.extend(['--send-format', self.send_format])
+
+        result = yatest.common.execute(cmd)
+
+        assert result.returncode == 0, f'Import failed with return code {result.returncode} and stderr: {result.stderr.decode("utf-8")}'
 
         stderr_output = result.stderr.decode('utf-8')
         for line in stderr_output.split('\n'):
@@ -89,3 +97,8 @@ class ImportFileCsvBase(UploadSuiteBase):
 
 class TestImportFileCsv(ImportFileCsvBase):
     external_folder: str = 'ecommerce'
+
+
+class TestImportFileCsvArrow(ImportFileCsvBase):
+    external_folder: str = 'ecommerce'
+    send_format: str = 'arrow'
diff --git a/ydb/tests/olap/load/lib/upload.py b/ydb/tests/olap/load/lib/upload.py
index 81934b4804c..7d0ba1b995c 100644
--- a/ydb/tests/olap/load/lib/upload.py
+++ b/ydb/tests/olap/load/lib/upload.py
@@ -42,6 +42,7 @@ class UploadSuiteBase(LoadSuiteBase):
             logging.error(f'Error: {e}')
             result.add_error(str(e))
             result.traceback = e.__traceback__
+            raise e
         result.iterations[0].time = time() - start_time
         self.validate(result)
         self.save_result_additional_info(result)