deltacat 1.1.35__py3-none-any.whl → 2.0__py3-none-any.whl

deltacat/tests/utils/test_pyarrow.py

@@ -2,33 +2,20 @@ from unittest import TestCase
 from deltacat.utils.pyarrow import (
     s3_partial_parquet_file_to_table,
     pyarrow_read_csv,
-    ContentTypeValidationError,
     content_type_to_reader_kwargs,
     _add_column_kwargs,
-    logger,
     s3_file_to_table,
-    s3_file_to_parquet,
     ReadKwargsProviderPyArrowSchemaOverride,
     RAISE_ON_EMPTY_CSV_KWARG,
-    RAISE_ON_DECIMAL_OVERFLOW,
-    OVERRIDE_CONTENT_ENCODING_FOR_PARQUET_KWARG,
 )
-import decimal
 from deltacat.types.media import ContentEncoding, ContentType
 from deltacat.types.partial_download import PartialParquetParameters
 from pyarrow.parquet import ParquetFile
 import pyarrow as pa
 
 PARQUET_FILE_PATH = "deltacat/tests/utils/data/test_file.parquet"
-PARQUET_GZIP_COMPRESSED_FILE_PATH = "deltacat/tests/utils/data/test_file.parquet.gz"
 EMPTY_UTSV_PATH = "deltacat/tests/utils/data/empty.csv"
 NON_EMPTY_VALID_UTSV_PATH = "deltacat/tests/utils/data/non_empty_valid.csv"
-OVERFLOWING_DECIMAL_PRECISION_UTSV_PATH = (
-    "deltacat/tests/utils/data/overflowing_decimal_precision.csv"
-)
-OVERFLOWING_DECIMAL_SCALE_UTSV_PATH = (
-    "deltacat/tests/utils/data/overflowing_decimal_scale.csv"
-)
 GZIP_COMPRESSED_FILE_UTSV_PATH = "deltacat/tests/utils/data/non_empty_compressed.gz"
 BZ2_COMPRESSED_FILE_UTSV_PATH = "deltacat/tests/utils/data/non_empty_compressed.bz2"
 
@@ -420,253 +407,6 @@ class TestReadCSV(TestCase):
             ),
         )
 
-    def test_read_csv_when_decimal_precision_overflows_and_raise_kwarg_specified(self):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal128(4, 2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-        self.assertRaises(
-            pa.lib.ArrowInvalid,
-            lambda: pyarrow_read_csv(
-                OVERFLOWING_DECIMAL_PRECISION_UTSV_PATH,
-                **{**kwargs, RAISE_ON_DECIMAL_OVERFLOW: True},
-            ),
-        )
-
-    def test_read_csv_when_decimal_precision_overflows_sanity(self):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal128(4, 2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        self.assertRaises(
-            pa.lib.ArrowInvalid,
-            lambda: pyarrow_read_csv(OVERFLOWING_DECIMAL_PRECISION_UTSV_PATH, **kwargs),
-        )
-
-    def test_read_csv_when_decimal_scale_overflows_and_raise_kwarg_specified(self):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal128(20, 2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        self.assertRaises(
-            pa.lib.ArrowInvalid,
-            lambda: pyarrow_read_csv(
-                OVERFLOWING_DECIMAL_SCALE_UTSV_PATH,
-                **{**kwargs, RAISE_ON_DECIMAL_OVERFLOW: True},
-            ),
-        )
-
-    def test_read_csv_when_decimal_scale_overflows_sanity(self):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal128(20, 2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        result = pyarrow_read_csv(OVERFLOWING_DECIMAL_SCALE_UTSV_PATH, **kwargs)
-
-        self.assertEqual(len(result), 3)
-        self.assertEqual(
-            result[1][0].as_py(), decimal.Decimal("322236.66")
-        )  # rounding decimal
-        self.assertEqual(result[1][1].as_py(), decimal.Decimal("32.33"))  # not rounded
-        self.assertEqual(len(result.column_names), 2)
-        result_schema = result.schema
-        self.assertEqual(result_schema.field(0).type, "string")
-        self.assertEqual(result_schema.field(1).type, pa.decimal128(20, 2))
-
-    def test_read_csv_when_decimal_scale_overflows_and_negative_scale(self):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal128(20, -2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        result = pyarrow_read_csv(OVERFLOWING_DECIMAL_SCALE_UTSV_PATH, **kwargs)
-
-        self.assertEqual(len(result), 3)
-        self.assertEqual(
-            result[1][0].as_py(),
-            decimal.Decimal("322200"),  # consequence of negative scale
-        )  # rounding decimal
-        self.assertEqual(result[1][1].as_py(), decimal.Decimal("00"))
-        self.assertEqual(len(result.column_names), 2)
-        result_schema = result.schema
-        self.assertEqual(result_schema.field(0).type, "string")
-        self.assertEqual(result_schema.field(1).type, pa.decimal128(20, -2))
-
-    def test_read_csv_when_decimal_scale_overflows_with_decimal256(self):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal256(20, 2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        result = pyarrow_read_csv(OVERFLOWING_DECIMAL_SCALE_UTSV_PATH, **kwargs)
-
-        self.assertEqual(len(result), 3)
-        self.assertEqual(
-            result[1][0].as_py(), decimal.Decimal("322236.66")
-        )  # rounding decimal
-        self.assertEqual(result[1][1].as_py(), decimal.Decimal("32.33"))  # not rounded
-        self.assertEqual(len(result.column_names), 2)
-        result_schema = result.schema
-        self.assertEqual(result_schema.field(0).type, "string")
-        self.assertEqual(result_schema.field(1).type, pa.decimal256(20, 2))
-
-    def test_read_csv_when_decimal_scale_overflows_with_decimal256_and_raise_on_overflow(
-        self,
-    ):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal256(20, 2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        self.assertRaises(
-            pa.lib.ArrowNotImplementedError,
-            lambda: pyarrow_read_csv(
-                OVERFLOWING_DECIMAL_SCALE_UTSV_PATH,
-                **{**kwargs, RAISE_ON_DECIMAL_OVERFLOW: True},
-            ),
-        )
-
-    def test_read_csv_when_decimal_scale_overflows_without_any_schema_then_infers(self):
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=None)
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        result = pyarrow_read_csv(OVERFLOWING_DECIMAL_SCALE_UTSV_PATH, **kwargs)
-
-        # The default behavior of pyarrow is to invalid skip rows
-        self.assertEqual(len(result), 2)
-        self.assertEqual(result[1][0].as_py(), 32.33)  # rounding decimal
-        self.assertEqual(result[1][1].as_py(), 0.4)  # not rounded
-        self.assertEqual(len(result.column_names), 2)
-        result_schema = result.schema
-        self.assertEqual(result_schema.field(0).type, "string")
-        self.assertEqual(result_schema.field(1).type, pa.float64())
-
-    def test_read_csv_when_decimal_scale_and_precision_overflow_and_raise_on_overflow(
-        self,
-    ):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal128(5, 2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        self.assertRaises(
-            pa.lib.ArrowInvalid,
-            lambda: pyarrow_read_csv(OVERFLOWING_DECIMAL_SCALE_UTSV_PATH, **kwargs),
-        )
-
-    def test_read_csv_when_decimal_scale_overflow_and_file_like_obj_passed(self):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("decimal_value", pa.decimal128(15, 2))]
-        )
-        kwargs = content_type_to_reader_kwargs(ContentType.UNESCAPED_TSV.value)
-        _add_column_kwargs(
-            ContentType.UNESCAPED_TSV.value,
-            ["is_active", "decimal_value"],
-            ["is_active", "decimal_value"],
-            kwargs,
-        )
-
-        read_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(schema=schema)
-
-        kwargs = read_kwargs_provider(ContentType.UNESCAPED_TSV.value, kwargs)
-
-        with open(OVERFLOWING_DECIMAL_SCALE_UTSV_PATH, "rb") as file:
-            result = pyarrow_read_csv(file, **kwargs)
-
-            self.assertEqual(len(result), 3)
-            self.assertEqual(
-                result[1][0].as_py(), decimal.Decimal("322236.66")
-            )  # rounding decimal
-            self.assertEqual(
-                result[1][1].as_py(), decimal.Decimal("32.33")
-            )  # not rounded
-            self.assertEqual(len(result.column_names), 2)
-            result_schema = result.schema
-            self.assertEqual(result_schema.field(0).type, "string")
-            self.assertEqual(result_schema.field(1).type, pa.decimal128(15, 2))
-
 
 class TestS3FileToTable(TestCase):
     def test_s3_file_to_table_identity_sanity(self):
@@ -794,175 +534,3 @@ class TestS3FileToTable(TestCase):
             self.assertEqual(field.name, schema.field(index).name)
 
         self.assertEqual(result.schema.field(1).type, "string")
-
-    def test_s3_file_to_table_when_parquet_gzip(self):
-
-        pa_kwargs_provider = lambda content_type, kwargs: {
-            "reader_type": "pyarrow",
-            **kwargs,
-        }
-
-        result = s3_file_to_table(
-            PARQUET_GZIP_COMPRESSED_FILE_PATH,
-            ContentType.PARQUET.value,
-            ContentEncoding.GZIP.value,
-            ["n_legs", "animal"],
-            ["n_legs"],
-            pa_read_func_kwargs_provider=pa_kwargs_provider,
-        )
-
-        self.assertEqual(len(result), 6)
-        self.assertEqual(len(result.column_names), 1)
-        schema = result.schema
-        schema_index = schema.get_field_index("n_legs")
-        self.assertEqual(schema.field(schema_index).type, "int64")
-
-    def test_s3_file_to_table_when_utsv_gzip_and_content_type_overridden(self):
-        schema = pa.schema(
-            [("is_active", pa.string()), ("ship_datetime_utc", pa.timestamp("us"))]
-        )
-        # OVERRIDE_CONTENT_ENCODING_FOR_PARQUET_KWARG has no effect on uTSV files
-        pa_kwargs_provider = lambda content_type, kwargs: {
-            "reader_type": "pyarrow",
-            **kwargs,
-        }
-        pa_kwargs_provider = lambda content_type, kwargs: {
-            "reader_type": "pyarrow",
-            OVERRIDE_CONTENT_ENCODING_FOR_PARQUET_KWARG: ContentEncoding.IDENTITY.value,
-            **kwargs,
-        }
-
-        result = s3_file_to_table(
-            GZIP_COMPRESSED_FILE_UTSV_PATH,
-            ContentType.UNESCAPED_TSV.value,
-            ContentEncoding.GZIP.value,
-            ["is_active", "ship_datetime_utc"],
-            None,
-            pa_read_func_kwargs_provider=pa_kwargs_provider,
-        )
-
-        self.assertEqual(len(result), 3)
-        self.assertEqual(len(result.column_names), 2)
-        result_schema = result.schema
-        for index, field in enumerate(result_schema):
-            self.assertEqual(field.name, schema.field(index).name)
-
-        self.assertEqual(result.schema.field(0).type, "string")
-
-    def test_s3_file_to_table_when_parquet_gzip_and_encoding_overridden(self):
-        pa_kwargs_provider = lambda content_type, kwargs: {
-            "reader_type": "pyarrow",
-            OVERRIDE_CONTENT_ENCODING_FOR_PARQUET_KWARG: ContentEncoding.IDENTITY.value,
-            **kwargs,
-        }
-
-        result = s3_file_to_table(
-            PARQUET_FILE_PATH,
-            ContentType.PARQUET.value,
-            ContentEncoding.GZIP.value,
-            ["n_legs", "animal"],
-            ["n_legs"],
-            pa_read_func_kwargs_provider=pa_kwargs_provider,
-        )
-
-        self.assertEqual(len(result), 6)
-        self.assertEqual(len(result.column_names), 1)
-        schema = result.schema
-        schema_index = schema.get_field_index("n_legs")
-        self.assertEqual(schema.field(schema_index).type, "int64")
-
-
-class TestS3FileToParquet(TestCase):
-    def test_s3_file_to_parquet_sanity(self):
-        test_s3_url = PARQUET_FILE_PATH
-        test_content_type = ContentType.PARQUET.value
-        test_content_encoding = ContentEncoding.IDENTITY.value
-        pa_kwargs_provider = lambda content_type, kwargs: {
-            "reader_type": "pyarrow",
-            **kwargs,
-        }
-        with self.assertLogs(logger=logger.name, level="DEBUG") as cm:
-            result_parquet_file: ParquetFile = s3_file_to_parquet(
-                test_s3_url,
-                test_content_type,
-                test_content_encoding,
-                ["n_legs", "animal"],
-                ["n_legs"],
-                pa_read_func_kwargs_provider=pa_kwargs_provider,
-            )
-        log_message_log_args = cm.records[0].getMessage()
-        log_message_presanitize_kwargs = cm.records[1].getMessage()
-        self.assertIn(
-            f"Reading {test_s3_url} to PyArrow ParquetFile. Content type: {test_content_type}. Encoding: {test_content_encoding}",
-            log_message_log_args,
-        )
-        self.assertIn("{'reader_type': 'pyarrow'}", log_message_presanitize_kwargs)
-        for index, field in enumerate(result_parquet_file.schema_arrow):
-            self.assertEqual(
-                field.name, result_parquet_file.schema_arrow.field(index).name
-            )
-        self.assertEqual(result_parquet_file.schema_arrow.field(0).type, "int64")
-
-    def test_s3_file_to_parquet_when_parquet_gzip_encoding_and_overridden_returns_success(
-        self,
-    ):
-        test_s3_url = PARQUET_FILE_PATH
-        test_content_type = ContentType.PARQUET.value
-        test_content_encoding = ContentEncoding.GZIP.value
-        pa_kwargs_provider = lambda content_type, kwargs: {
-            "reader_type": "pyarrow",
-            OVERRIDE_CONTENT_ENCODING_FOR_PARQUET_KWARG: ContentEncoding.IDENTITY.value,
-            **kwargs,
-        }
-        with self.assertLogs(logger=logger.name, level="DEBUG") as cm:
-            result_parquet_file: ParquetFile = s3_file_to_parquet(
-                test_s3_url,
-                test_content_type,
-                test_content_encoding,
-                ["n_legs", "animal"],
-                ["n_legs"],
-                pa_read_func_kwargs_provider=pa_kwargs_provider,
-            )
-        log_message_log_args = cm.records[0].getMessage()
-        log_message_log_new_content_encoding = cm.records[1].getMessage()
-        log_message_presanitize_kwargs = cm.records[2].getMessage()
-        self.assertIn(
-            f"Reading {test_s3_url} to PyArrow ParquetFile. Content type: {test_content_type}. Encoding: {test_content_encoding}",
-            log_message_log_args,
-        )
-        self.assertIn(
-            f"Overriding {test_s3_url} content encoding from {ContentEncoding.GZIP.value} to {ContentEncoding.IDENTITY.value}",
-            log_message_log_new_content_encoding,
-        )
-        self.assertIn("{'reader_type': 'pyarrow'}", log_message_presanitize_kwargs)
-        for index, field in enumerate(result_parquet_file.schema_arrow):
-            self.assertEqual(
-                field.name, result_parquet_file.schema_arrow.field(index).name
-            )
-        self.assertEqual(result_parquet_file.schema_arrow.field(0).type, "int64")
-
-    def test_s3_file_to_parquet_when_parquet_gzip_encoding_not_overridden_throws_error(
-        self,
-    ):
-        test_s3_url = PARQUET_FILE_PATH
-        test_content_type = ContentType.PARQUET.value
-        test_content_encoding = ContentEncoding.GZIP.value
-        pa_kwargs_provider = lambda content_type, kwargs: {
-            "reader_type": "pyarrow",
-            **kwargs,
-        }
-        with self.assertRaises(ContentTypeValidationError):
-            with self.assertLogs(logger=logger.name, level="DEBUG") as cm:
-                s3_file_to_parquet(
-                    test_s3_url,
-                    test_content_type,
-                    test_content_encoding,
-                    ["n_legs", "animal"],
-                    ["n_legs"],
-                    pa_read_func_kwargs_provider=pa_kwargs_provider,
-                )
-        log_message_log_args = cm.records[0].getMessage()
-        self.assertIn(
-            f"Reading {test_s3_url} to PyArrow ParquetFile. Content type: {test_content_type}. Encoding: {test_content_encoding}",
-            log_message_log_args,
-        )

deltacat/types/partial_download.py

@@ -63,7 +63,7 @@ class PartialParquetParameters(PartialFileDownloadParams):
         return self["row_groups_to_download"]
 
     @property
-    def num_row_groups(self) -> List[int]:
+    def num_row_groups(self) -> int:
         return self["num_row_groups"]
 
     @property

deltacat/types/tables.py

@@ -89,7 +89,7 @@ class TableWriteMode(str, Enum):
     Enum controlling how a given dataset will be written to a table.
 
     AUTO: CREATE if the table doesn't exist, APPEND if the table exists
-    without primary keys, and MERGE if the table exists with primary keys.
+    without merge keys, and MERGE if the table exists with merge keys.
     CREATE: Create the table if it doesn't exist, throw an error if it does.
     APPEND: Append to the table if it exists, throw an error if it doesn't.
     REPLACE: Replace existing table contents with the data to write.

deltacat/utils/export.py

@@ -0,0 +1,59 @@
+import logging
+import json
+import pyarrow as pa
+import pyarrow.parquet
+import pyarrow.feather
+from typing import Callable, Dict
+
+from deltacat.storage.rivulet.reader.query_expression import QueryExpression
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def export_parquet(dataset, file_uri: str, query: QueryExpression = QueryExpression()):
+    records = dataset.scan(query).to_arrow()
+    table = pa.Table.from_batches(records)
+    pyarrow.parquet.write_table(table, file_uri)
+
+
+def export_feather(dataset, file_uri: str, query: QueryExpression = QueryExpression()):
+    records = dataset.scan(query).to_arrow()
+    table = pa.Table.from_batches(records)
+    pyarrow.feather.write_feather(table, file_uri)
+
+
+def export_json(dataset, file_uri: str, query: QueryExpression = QueryExpression()):
+    with open(file_uri, "w") as f:
+        for batch in dataset.scan(query).to_pydict():
+            json.dump(batch, f, indent=2)
+            f.write("\n")
+
+
+def export_dataset(dataset, file_uri: str, format: str = "parquet", query=None):
+    """
+    Export the dataset to a file.
+
+    TODO: Make this pluggable for custom formats.
+
+    Args:
+        dataset: The dataset to export.
+        file_uri: The URI to write the dataset to.
+        format: The format to write the dataset in. Options are [parquet, feather, json].
+        query: QueryExpression to filter the dataset before exporting.
+    """
+    # Supported format handlers
+    export_handlers: Dict[str, Callable] = {
+        "parquet": export_parquet,
+        "feather": export_feather,
+        "json": export_json,
+    }
+
+    if format not in export_handlers:
+        raise ValueError(
+            f"Unsupported format: {format}. Supported formats are {list(export_handlers.keys())}"
+        )
+
+    export_handlers[format](dataset, file_uri, query or QueryExpression())
+
+    logger.info(f"Dataset exported to {file_uri} in {format} format.")