deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/tests/utils/ray_utils/test_dataset.py

@@ -6,6 +6,8 @@ from fsspec import AbstractFileSystem
 from ray.data.datasource import FilenameProvider
 from deltacat.types.media import ContentType
 import ray
+import gzip
+import json
 
 
 class TestDatasetToFile:
@@ -20,7 +22,13 @@ class TestDatasetToFile:
 
     @pytest.fixture(scope="module")
     def mock_dataset(self):
-        return from_items([{"col1": i, "col2": i * 2} for i in range(1000)])
+        # Include data that would need escaping to test quoting behavior
+        return from_items([{"col1": "a,b\tc|d", "col2": 0} for _ in range(5)])
+
+    @pytest.fixture(scope="module")
+    def mock_unescaped_dataset(self):
+        # Use data without delimiters for unescaped TSV test
+        return from_items([{"col1": "abc", "col2": 0} for _ in range(5)])
 
     @pytest.fixture(scope="module")
     def mock_filename_provider(self):
@@ -35,12 +43,12 @@ class TestDatasetToFile:
     def test_parquet_sanity(self, mock_dataset, mock_filename_provider):
         from deltacat.utils.ray_utils.dataset import dataset_to_file
 
-        fs: AbstractFileSystem = fsspec.filesystem("local")
+        fs: AbstractFileSystem = fsspec.filesystem("file")
 
         dataset_to_file(
             mock_dataset,
             self.BASE_PATH,
-            file_system=fs,
+            filesystem=fs,
             block_path_provider=mock_filename_provider,
         )
 
@@ -51,16 +59,126 @@ class TestDatasetToFile:
     def test_csv_sanity(self, mock_dataset, mock_filename_provider):
         from deltacat.utils.ray_utils.dataset import dataset_to_file
 
-        fs: AbstractFileSystem = fsspec.filesystem("local")
+        fs: AbstractFileSystem = fsspec.filesystem("file")
 
         dataset_to_file(
             mock_dataset,
             self.BASE_PATH,
-            file_system=fs,
+            filesystem=fs,
             block_path_provider=mock_filename_provider,
             content_type=ContentType.CSV.value,
         )
 
         file_expected_at = f"{self.BASE_PATH}/{self.SUB_PATH}"
         assert fs.exists(file_expected_at), "file was not written"
+
+        # Verify CSV format and content
+        with fs.open(file_expected_at, "rb") as f:
+            with gzip.GzipFile(fileobj=f) as gz:
+                content = gz.read().decode("utf-8")
+                # Should be quoted due to commas in data
+                assert '"a,b\tc|d",0' in content
+
+        fs.delete(file_expected_at)
+
+    def test_tsv_sanity(self, mock_dataset, mock_filename_provider):
+        from deltacat.utils.ray_utils.dataset import dataset_to_file
+
+        fs: AbstractFileSystem = fsspec.filesystem("file")
+
+        dataset_to_file(
+            mock_dataset,
+            self.BASE_PATH,
+            filesystem=fs,
+            block_path_provider=mock_filename_provider,
+            content_type=ContentType.TSV.value,
+        )
+
+        file_expected_at = f"{self.BASE_PATH}/{self.SUB_PATH}"
+        assert fs.exists(file_expected_at), "file was not written"
+
+        # Verify TSV format and content
+        with fs.open(file_expected_at, "rb") as f:
+            with gzip.GzipFile(fileobj=f) as gz:
+                content = gz.read().decode("utf-8")
+                # Should be quoted due to tabs in data
+                assert '"a,b\tc|d"\t0' in content
+
+        fs.delete(file_expected_at)
+
+    def test_psv_sanity(self, mock_dataset, mock_filename_provider):
+        from deltacat.utils.ray_utils.dataset import dataset_to_file
+
+        fs: AbstractFileSystem = fsspec.filesystem("file")
+
+        dataset_to_file(
+            mock_dataset,
+            self.BASE_PATH,
+            filesystem=fs,
+            block_path_provider=mock_filename_provider,
+            content_type=ContentType.PSV.value,
+        )
+
+        file_expected_at = f"{self.BASE_PATH}/{self.SUB_PATH}"
+        assert fs.exists(file_expected_at), "file was not written"
+
+        # Verify PSV format and content
+        with fs.open(file_expected_at, "rb") as f:
+            with gzip.GzipFile(fileobj=f) as gz:
+                content = gz.read().decode("utf-8")
+                # Should be quoted due to pipes in data
+                assert '"a,b\tc|d"|0' in content
+
+        fs.delete(file_expected_at)
+
+    def test_unescaped_tsv_sanity(self, mock_unescaped_dataset, mock_filename_provider):
+        from deltacat.utils.ray_utils.dataset import dataset_to_file
+
+        fs: AbstractFileSystem = fsspec.filesystem("file")
+
+        dataset_to_file(
+            mock_unescaped_dataset,
+            self.BASE_PATH,
+            filesystem=fs,
+            block_path_provider=mock_filename_provider,
+            content_type=ContentType.UNESCAPED_TSV.value,
+        )
+
+        file_expected_at = f"{self.BASE_PATH}/{self.SUB_PATH}"
+        assert fs.exists(file_expected_at), "file was not written"
+
+        # Verify UNESCAPED_TSV format and content
+        with fs.open(file_expected_at, "rb") as f:
+            with gzip.GzipFile(fileobj=f) as gz:
+                content = gz.read().decode("utf-8")
+                # Should NOT be quoted since data has no delimiters
+                assert "abc\t0" in content
+
+        fs.delete(file_expected_at)
+
+    def test_json_sanity(self, mock_dataset, mock_filename_provider):
+        from deltacat.utils.ray_utils.dataset import dataset_to_file
+
+        fs: AbstractFileSystem = fsspec.filesystem("file")
+
+        dataset_to_file(
+            mock_dataset,
+            self.BASE_PATH,
+            filesystem=fs,
+            block_path_provider=mock_filename_provider,
+            content_type=ContentType.JSON.value,
+        )
+
+        file_expected_at = f"{self.BASE_PATH}/{self.SUB_PATH}"
+        assert fs.exists(file_expected_at), "file was not written"
+
+        # Verify JSON format and content
+        with fs.open(file_expected_at, "rb") as f:
+            with gzip.GzipFile(fileobj=f) as gz:
+                content = gz.read().decode("utf-8")
+                # Each line should be a valid JSON object
+                first_line = content.split("\n")[0]
+                record = json.loads(first_line)
+                assert record == {"col1": "a,b\tc|d", "col2": 0}
+
         fs.delete(file_expected_at)

deltacat/tests/utils/test_daft.py

@@ -1,6 +1,9 @@
 import unittest
 from deltacat.types.media import ContentEncoding, ContentType
-from deltacat.utils.daft import daft_s3_file_to_table, s3_files_to_dataframe
+from deltacat.utils.daft import (
+    daft_file_to_pyarrow_table,
+    files_to_dataframe,
+)
 from deltacat.utils.pyarrow import ReadKwargsProviderPyArrowSchemaOverride
 from deltacat.types.partial_download import PartialParquetParameters
 import pyarrow as pa
@@ -8,11 +11,11 @@ import pyarrow as pa
 from pyarrow import parquet as pq
 
 
-class TestDaftS3FileToTable(unittest.TestCase):
+class TestDaftFileToPyarrowTable(unittest.TestCase):
     MVP_PATH = "deltacat/tests/utils/data/mvp.parquet"
 
-    def test_read_from_s3_all_columns(self):
-        table = daft_s3_file_to_table(
+    def test_read_from_local_all_columns(self):
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -20,8 +23,8 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.names, ["a", "b"])
         self.assertEqual(table.num_rows, 100)
 
-    def test_read_from_s3_single_column_via_include_columns(self):
-        table = daft_s3_file_to_table(
+    def test_read_from_local_single_column_via_include_columns(self):
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -30,8 +33,8 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.names, ["b"])
         self.assertEqual(table.num_rows, 100)
 
-    def test_read_from_s3_single_column_via_column_names(self):
-        table = daft_s3_file_to_table(
+    def test_read_from_local_single_column_via_column_names(self):
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -40,12 +43,12 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.names, ["b"])
         self.assertEqual(table.num_rows, 100)
 
-    def test_read_from_s3_single_column_with_schema(self):
+    def test_read_from_local_single_column_with_schema(self):
         schema = pa.schema([("a", pa.int8()), ("b", pa.string())])
         pa_read_func_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(
             schema=schema
         )
-        table = daft_s3_file_to_table(
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -56,12 +59,12 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.field("a").type, pa.int8())
         self.assertEqual(table.num_rows, 100)
 
-    def test_read_from_s3_single_column_with_schema_reverse_order(self):
+    def test_read_from_local_single_column_with_schema_reverse_order(self):
         schema = pa.schema([("b", pa.string()), ("a", pa.int8())])
         pa_read_func_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(
             schema=schema
         )
-        table = daft_s3_file_to_table(
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -71,12 +74,12 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.field("a").type, pa.int8())
         self.assertEqual(table.num_rows, 100)
 
-    def test_read_from_s3_single_column_with_schema_subset_cols(self):
+    def test_read_from_local_single_column_with_schema_subset_cols(self):
         schema = pa.schema([("a", pa.int8())])
         pa_read_func_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(
             schema=schema
         )
-        table = daft_s3_file_to_table(
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -86,12 +89,12 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.field("a").type, pa.int8())
         self.assertEqual(table.num_rows, 100)
 
-    def test_read_from_s3_single_column_with_schema_extra_cols(self):
+    def test_read_from_local_single_column_with_schema_extra_cols(self):
         schema = pa.schema([("a", pa.int8()), ("MISSING", pa.string())])
         pa_read_func_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(
             schema=schema
         )
-        table = daft_s3_file_to_table(
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -104,12 +107,12 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.field("MISSING").type, pa.string())
         self.assertEqual(table.num_rows, 100)
 
-    def test_read_from_s3_single_column_with_schema_extra_cols_column_names(self):
+    def test_read_from_local_single_column_with_schema_extra_cols_column_names(self):
         schema = pa.schema([("a", pa.int8()), ("MISSING", pa.string())])
         pa_read_func_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(
             schema=schema
         )
-        table = daft_s3_file_to_table(
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -123,12 +126,12 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.field("MISSING").type, pa.string())
         self.assertEqual(table.num_rows, 100)
 
-    def test_read_from_s3_single_column_with_schema_only_missing_col(self):
+    def test_read_from_local_single_column_with_schema_only_missing_col(self):
         schema = pa.schema([("a", pa.int8()), ("MISSING", pa.string())])
         pa_read_func_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(
             schema=schema
         )
-        table = daft_s3_file_to_table(
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -142,12 +145,12 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.schema.field("MISSING").type, pa.string())
         self.assertEqual(table.num_rows, 0)
 
-    def test_read_from_s3_single_column_with_row_groups(self):
+    def test_read_from_local_single_column_with_row_groups(self):
 
         metadata = pq.read_metadata(self.MVP_PATH)
         ppp = PartialParquetParameters.of(pq_metadata=metadata)
         ppp["row_groups_to_download"] = ppp.row_groups_to_download[1:2]
-        table = daft_s3_file_to_table(
+        table = daft_file_to_pyarrow_table(
             self.MVP_PATH,
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
@@ -158,45 +161,132 @@ class TestDaftS3FileToTable(unittest.TestCase):
         self.assertEqual(table.num_rows, 10)
 
 
-class TestDaftS3FilesToDataFrame(unittest.TestCase):
+class TestFilesToDataFrame(unittest.TestCase):
     MVP_PATH = "deltacat/tests/utils/data/mvp.parquet"
 
-    def test_read_from_s3_all_columns(self):
-        df = s3_files_to_dataframe(
+    def test_read_local_files_all_columns(self):
+        df = files_to_dataframe(
             uris=[self.MVP_PATH],
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
-            ray_init_options={"local_mode": True},
+            ray_init_options={"local_mode": True, "ignore_reinit_error": True},
         )
 
         table = df.to_arrow()
         self.assertEqual(table.schema.names, ["a", "b"])
         self.assertEqual(table.num_rows, 100)
 
-    def test_does_not_read_from_s3_if_not_materialized(self):
-        df = s3_files_to_dataframe(
+    def test_read_local_files_with_column_selection(self):
+        df = files_to_dataframe(
             uris=[self.MVP_PATH],
             content_encoding=ContentEncoding.IDENTITY.value,
             content_type=ContentType.PARQUET.value,
-            ray_init_options={"local_mode": True},
+            include_columns=["b"],
+            ray_init_options={"local_mode": True, "ignore_reinit_error": True},
+        )
+
+        table = df.to_arrow()
+        self.assertEqual(table.schema.names, ["b"])
+        self.assertEqual(table.num_rows, 100)
+
+    def test_read_local_files_does_not_materialize_by_default(self):
+        df = files_to_dataframe(
+            uris=[self.MVP_PATH],
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            ray_init_options={"local_mode": True, "ignore_reinit_error": True},
         )
 
+        # Should raise RuntimeError because df is not materialized yet
         self.assertRaises(RuntimeError, lambda: len(df))
+
+        # After collecting, it should work
         df.collect()
         self.assertEqual(len(df), 100)
 
-    def test_raises_error_if_not_supported_content_type(self):
+    def test_supports_unescaped_tsv_content_type(self):
+        # Test that UNESCAPED_TSV is now supported (was previously unsupported)
+        # Use a CSV file since we're testing TSV reader functionality
+        csv_path = "deltacat/tests/utils/data/non_empty_valid.csv"
+        df = files_to_dataframe(
+            uris=[csv_path],
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.UNESCAPED_TSV.value,
+            ray_init_options={"local_mode": True, "ignore_reinit_error": True},
+        )
+        # Should succeed without raising an exception - this tests that UNESCAPED_TSV is supported
+        table = df.to_arrow()
+        # Just verify we got some data back, don't assert specific schema since we're reading CSV as TSV
+        self.assertGreater(table.num_rows, 0)
+        self.assertGreater(len(table.schema.names), 0)
 
+    def test_supports_gzip_content_encoding(self):
+        # Test that GZIP encoding is now supported (was previously unsupported)
+        df = files_to_dataframe(
+            uris=[self.MVP_PATH],
+            content_encoding=ContentEncoding.GZIP.value,
+            content_type=ContentType.PARQUET.value,
+            ray_init_options={"local_mode": True, "ignore_reinit_error": True},
+        )
+        # Should succeed without raising an exception
+        table = df.to_arrow()
+        self.assertEqual(table.schema.names, ["a", "b"])
+        self.assertEqual(table.num_rows, 100)
+
+    def test_raises_error_if_not_supported_content_type(self):
+        # Test that truly unsupported content types raise NotImplementedError
         self.assertRaises(
-            AssertionError,
-            lambda: s3_files_to_dataframe(
+            NotImplementedError,
+            lambda: files_to_dataframe(
                 uris=[self.MVP_PATH],
                 content_encoding=ContentEncoding.IDENTITY.value,
-                content_type=ContentType.UNESCAPED_TSV.value,
-                ray_init_options={"local_mode": True},
+                content_type=ContentType.AVRO.value,  # AVRO is actually unsupported
+                ray_init_options={"local_mode": True, "ignore_reinit_error": True},
             ),
         )
 
+    def test_raises_error_if_not_supported_content_encoding(self):
+        # Test that truly unsupported content encodings raise NotImplementedError
+        self.assertRaises(
+            NotImplementedError,
+            lambda: files_to_dataframe(
+                uris=[self.MVP_PATH],
+                content_encoding=ContentEncoding.ZSTD.value,  # ZSTD is actually unsupported
+                content_type=ContentType.PARQUET.value,
+                ray_init_options={"local_mode": True, "ignore_reinit_error": True},
+            ),
+        )
+
+    def test_accepts_custom_kwargs(self):
+        # Test that custom kwargs are passed through to daft.read_parquet
+        df = files_to_dataframe(
+            uris=[self.MVP_PATH],
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            ray_init_options={"local_mode": True, "ignore_reinit_error": True},
+            # Custom kwarg that should be passed to daft.read_parquet
+            coerce_int96_timestamp_unit="ns",
+        )
+
+        table = df.to_arrow()
+        self.assertEqual(table.schema.names, ["a", "b"])
+        self.assertEqual(table.num_rows, 100)
+
+    def test_accepts_io_config(self):
+        # Test that io_config parameter is accepted and passed correctly
+        df = files_to_dataframe(
+            uris=[self.MVP_PATH],
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            ray_init_options={"local_mode": True, "ignore_reinit_error": True},
+            # io_config=None should work fine for local files
+            io_config=None,
+        )
+
+        table = df.to_arrow()
+        self.assertEqual(table.schema.names, ["a", "b"])
+        self.assertEqual(table.num_rows, 100)
+
 
 if __name__ == "__main__":
     unittest.main()