Add Storage Layer In Python (#3224)

This PR adds a storage layer implementation on the Python side of
Texera's codebase, mirroring the implementation of our Java-based
storage layer.

## Motivation
- The primary motivation of having a storage layer in Python so that we
can let Python UDF operators' ports write directly to result tables
without needing to send the results back to Java.
- In the future we will also use the Python storage layer for UDF logs
and workflow runtime statistics.

## Storage APIs
- There are 3 abstract classes in Java's storage implementation:
  - `ReadOnlyVirtualDocument` for read-only tables
- `VirtualDocument` for tables supporting both read and write
operations.
  - `BufferedItemWriter` as a writer class of `VirtualDocument`
- We mirror the implementation in Python, but keep only the APIs
relevant to table storage (e.g., APIs related to dataset storage are not
kept in Python.)

## Iceberg Document
Following #3147, we add a table-storage implementation based on Apache
Iceberg (pyiceberg), including `IcebergDocument`, `IcebergTableWriter`,
`IcebergCatalogInstance`, and related util functions and tests.

### Limitations of / TODOs for python implementation 
pyiceberg is less mature than its java-based counterpart. As a result
there are a few functionalities not supported in our current Python
storage implementation.

#### Incremental Read
Incremental Read is not supported by pyiceberg. It will be supported [in
the future](apache/iceberg-python#533). Before
then we will not include incremental read in our Python codebase (it is
also not currently needed)
#### Concurrent writers
Iceberg uses optimistic concurrency control for concurrent writers. Java
Iceberg natively supports retry with configurable retry parameters,
using exponential backoff (without randomness). However pyiceberg does
not currently support retry. We implemented an ad-hoc custom retry
mechanism in `IcebergTableWriter`, using exponential random backoff
based on the [tenacity](https://tenacity.readthedocs.io/en/latest/)
library. It has a good speed (~0.6s for 10 concurrent writers writing
20K tuples) and is faster than Java’s iceberg-native retry (~6 seconds
for the same test). We may need to re-evaluate this custom
implementation if pyiceberg supports retry natively in the future.

## Iceberg Catalog
pyiceberg only supports SQL catalog (postgreSQL to be specific) and REST
catalog for production. We use postgresql based SQL catalog in this
implementation for the following reasons:
- It supports local storage.
- We tested that it is works with both Java and Python iceberg storage.
- It is easier to set up for developers (compared to REST services).

### PostgreSQL setup
Python storage layer requires a running postgreSQL service in the
environment, and an empty database for iceberg to work.
- **A script to set up a new postgres database for Texera's iceberg
storage has been added for CI tests.**
- The database will be used by pyiceberg to manage the catalog.
- The logic to setup the database is added in GitHub CI config.
- Java side can continue using Hadoop-based catalog for now until we add
storage on operator ports for both Java and Python.
- As the Python storage is not currently used by Python workers, no
action is required for developers for now.

### REST catalogs (feel free to skip this section)
I also explored 3 major REST catalog implementations
([lakekeeper](https://lakekeeper.io),
[polaris](https://polaris.apache.org), and
[gravitino](https://gravitino.apache.org)) and here are some
observations:
- REST catalogs are the trend primarily because different query engines
(Spark, Flink, Snowflake, etc.) relying on iceberg need a central place
to keep and manage the catalogs. Under the hood they all still use some
database as their storage layer.
- Most of them support / recommend cloud storage only in production and
do not support local storage.
- They are incubating projects and lack documentation. For example I
find it very hard to set up authentication (as pyiceberg requires
authentication to work with REST catalogs) using gravitino, and using
them will add a lot more burden to our developers.
- I have successfully made polaris work with our implementation after
setting up auth, but somehow it was very very slow.
- As postgres catalog is working, we will explore more about REST
catalog in the future if have migrated to cloud storage and have
scalability issues.

## Storage configurations

A static class `StorageConfigs` is added to manage storage-related
configurations. We do NOT read the configs from files. Instead we will
let Java pass the configs to Python worker, and the config will be
filled when initializing the worker. The storage config is hardcoded in
CI tests.

## Other items

`VFSURIFactory` and `DocumentFactory` are added in Python storage layer
mirroring the Java implementations.

## TODO for Java Storage
- Add SQL catalog as another type of iceberg catalog

---------

Co-authored-by: Jiadong Bai <[email protected]>

Author: 2 people

.github/workflows/github-action-build.yml

@@ -111,6 +111,13 @@ jobs:
           if [ -f core/amber/requirements.txt ]; then pip install -r core/amber/requirements.txt; fi
           if [ -f core/amber/r-requirements.txt ]; then pip install -r core/amber/r-requirements.txt; fi
           if [ -f core/amber/operator-requirements.txt ]; then pip install -r core/amber/operator-requirements.txt; fi
+      - name: Install PostgreSQL
+        run: sudo apt-get update && sudo apt-get install -y postgresql
+      - name: Start PostgreSQL Service
+        run: sudo systemctl start postgresql
+      - name: Create Database and User
+        run: |
+          cd core/scripts/sql && sudo -u postgres psql -f iceberg_postgres_catalog.sql
       - name: Lint with flake8 and black
         run: |
           cd core/amber/src/main/python && flake8 && black . --check

core/amber/requirements.txt

@@ -26,4 +26,9 @@ bidict==0.22.0
 cached_property==1.5.2
 psutil==5.9.0
 transformers==4.44.2
-tzlocal==2.1
+tzlocal==2.1
+pyiceberg==0.8.1
+readerwriterlock==1.0.9
+tenacity==8.5.0
+SQLAlchemy==2.0.37
+psycopg2==2.9.10

core/amber/src/main/python/core/models/schema/attribute_type.py

@@ -42,16 +42,18 @@ class AttributeType(Enum):
     AttributeType.DOUBLE: pa.float64(),
     AttributeType.BOOL: pa.bool_(),
     AttributeType.BINARY: pa.binary(),
-    AttributeType.TIMESTAMP: pa.timestamp("ms", tz="UTC"),
+    AttributeType.TIMESTAMP: pa.timestamp("us"),
 }
 
 FROM_ARROW_MAPPING = {
     lib.Type_INT32: AttributeType.INT,
     lib.Type_INT64: AttributeType.LONG,
     lib.Type_STRING: AttributeType.STRING,
+    lib.Type_LARGE_STRING: AttributeType.STRING,
     lib.Type_DOUBLE: AttributeType.DOUBLE,
     lib.Type_BOOL: AttributeType.BOOL,
     lib.Type_BINARY: AttributeType.BINARY,
+    lib.Type_LARGE_BINARY: AttributeType.BINARY,
     lib.Type_TIMESTAMP: AttributeType.TIMESTAMP,
 }
 

core/amber/src/main/python/core/models/schema/test_schema.py

@@ -27,7 +27,7 @@ def arrow_schema(self):
                 pa.field("field-3", pa.int64()),
                 pa.field("field-4", pa.float64()),
                 pa.field("field-5", pa.bool_()),
-                pa.field("field-6", pa.timestamp("ms", tz="UTC")),
+                pa.field("field-6", pa.timestamp("us")),
                 pa.field("field-7", pa.binary()),
             ]
         )

core/amber/src/main/python/core/storage/__init__.py

@@ -0,0 +1,103 @@
+from urllib.parse import urlparse
+
+from typing import Optional
+
+from core.models import Schema, Tuple
+from core.storage.iceberg.iceberg_catalog_instance import IcebergCatalogInstance
+from core.storage.iceberg.iceberg_document import IcebergDocument
+from core.storage.iceberg.iceberg_utils import (
+    create_table,
+    amber_tuples_to_arrow_table,
+    arrow_table_to_amber_tuples,
+    load_table_metadata,
+)
+from core.storage.model.virtual_document import VirtualDocument
+from core.storage.storage_config import StorageConfig
+from core.storage.vfs_uri_factory import VFSURIFactory, VFSResourceType
+
+
+class DocumentFactory:
+    """
+    Factory class to create and open documents.
+    Currently only iceberg documents are supported.
+    """
+
+    ICEBERG = "iceberg"
+
+    @staticmethod
+    def sanitize_uri_path(uri):
+        return uri.path.lstrip("/").replace("/", "_")
+
+    @staticmethod
+    def create_document(uri: str, schema: Schema) -> VirtualDocument:
+        parsed_uri = urlparse(uri)
+        if parsed_uri.scheme == VFSURIFactory.VFS_FILE_URI_SCHEME:
+            _, _, _, _, resource_type = VFSURIFactory.decode_uri(uri)
+
+            if resource_type in {
+                VFSResourceType.RESULT,
+                VFSResourceType.MATERIALIZED_RESULT,
+            }:
+                storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
+
+                iceberg_schema = Schema.as_arrow_schema(schema)
+
+                create_table(
+                    IcebergCatalogInstance.get_instance(),
+                    StorageConfig.ICEBERG_TABLE_NAMESPACE,
+                    storage_key,
+                    iceberg_schema,
+                    override_if_exists=True,
+                )
+
+                return IcebergDocument[Tuple](
+                    StorageConfig.ICEBERG_TABLE_NAMESPACE,
+                    storage_key,
+                    iceberg_schema,
+                    amber_tuples_to_arrow_table,
+                    arrow_table_to_amber_tuples,
+                )
+            else:
+                raise ValueError(f"Resource type {resource_type} is not supported")
+        else:
+            raise NotImplementedError(
+                f"Unsupported URI scheme: {parsed_uri.scheme} for creating the document"
+            )
+
+    @staticmethod
+    def open_document(uri: str) -> (VirtualDocument, Optional[Schema]):
+        parsed_uri = urlparse(uri)
+        if parsed_uri.scheme == "vfs":
+            _, _, _, _, resource_type = VFSURIFactory.decode_uri(uri)
+
+            if resource_type in {
+                VFSResourceType.RESULT,
+                VFSResourceType.MATERIALIZED_RESULT,
+            }:
+                storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
+
+                table = load_table_metadata(
+                    IcebergCatalogInstance.get_instance(),
+                    StorageConfig.ICEBERG_TABLE_NAMESPACE,
+                    storage_key,
+                )
+
+                if table is None:
+                    raise ValueError("No storage is found for the given URI")
+
+                amber_schema = Schema(table.schema().as_arrow())
+
+                document = IcebergDocument(
+                    StorageConfig.ICEBERG_TABLE_NAMESPACE,
+                    storage_key,
+                    table.schema(),
+                    amber_tuples_to_arrow_table,
+                    arrow_table_to_amber_tuples,
+                )
+                return document, amber_schema
+            else:
+                raise ValueError(f"Resource type {resource_type} is not supported")
+        else:
+            raise NotImplementedError(
+                f"Unsupported URI scheme: {parsed_uri.scheme} for opening the document"
+            )

core/amber/src/main/python/core/storage/document_factory.py

@@ -0,0 +1,43 @@
+from pyiceberg.catalog import Catalog
+from typing import Optional
+
+from core.storage.iceberg.iceberg_utils import create_postgres_catalog
+from core.storage.storage_config import StorageConfig
+
+
+class IcebergCatalogInstance:
+    """
+    IcebergCatalogInstance is a singleton that manages the Iceberg catalog instance.
+    Currently only postgres SQL catalog is supported.
+    - Provides a single shared catalog for all Iceberg table-related operations.
+    - Lazily initializes the catalog on first access.
+    - Supports replacing the catalog instance for testing or reconfiguration.
+    """
+
+    _instance: Optional[Catalog] = None
+
+    @classmethod
+    def get_instance(cls):
+        """
+        Retrieves the singleton Iceberg catalog instance.
+        - If the catalog is not initialized, it is lazily created using the configured
+        properties.
+        :return: the Iceberg catalog instance.
+        """
+        if cls._instance is None:
+            cls._instance = create_postgres_catalog(
+                "texera_iceberg",
+                StorageConfig.ICEBERG_FILE_STORAGE_DIRECTORY_PATH,
+                StorageConfig.ICEBERG_POSTGRES_CATALOG_USERNAME,
+                StorageConfig.ICEBERG_POSTGRES_CATALOG_PASSWORD,
+            )
+        return cls._instance
+
+    @classmethod
+    def replace_instance(cls, catalog: Catalog):
+        """
+        Replaces the existing Iceberg catalog instance.
+        - This method is useful for testing or dynamically updating the catalog.
+        :param catalog: the new Iceberg catalog instance to replace the current one.
+        """
+        cls._instance = catalog