BSE-4358: BodoSQL support for S3 tables

BodoSQL/bodosql/__init__.py

@@ -19,6 +19,11 @@
     GlueCatalogType,
     get_glue_connection,
 )
+from bodosql.bodosql_types.s3_tables_catalog import (
+    S3TablesCatalog,
+    S3TablesCatalogType,
+    get_s3_tables_connection,
+)
 
 import bodosql.context_ext
 import bodosql.ddl_ext

BodoSQL/bodosql/bodosql_types/s3_tables_catalog.py

@@ -0,0 +1,182 @@
+"""Python and JIT class for describing a S3 Tables Iceberg catalog. A S3 Tables
+catalog contains all information needed to connect to S3 Tables for organizing and modifying tables.
+"""
+
+import numba
+from numba.core import cgutils, types
+from numba.extending import (
+    NativeValue,
+    box,
+    intrinsic,
+    make_attribute_wrapper,
+    models,
+    overload,
+    register_model,
+    typeof_impl,
+    unbox,
+)
+
+from bodo.io.iceberg import IcebergConnectionType
+from bodo.utils.typing import get_literal_value, raise_bodo_error
+from bodosql import DatabaseCatalog, DatabaseCatalogType
+from bodosql.imported_java_classes import JavaEntryPoint
+
+
+def _create_java_s3_tables_catalog(warehouse: str):
+    """
+    Create a Java BodoS3Tables object.
+    Args:
+        warehouse (str): The warehouse to connect to.
+    Returns:
+        JavaObject: A Java BodoS3Tables object.
+    """
+    return JavaEntryPoint.buildBodoS3TablesCatalog(warehouse)
+
+
+class S3TablesCatalog(DatabaseCatalog):
+    """
+    Python class for storing the information
+        needed to connect to a S3 Tables Iceberg catalog.
+    """
+
+    def __init__(self, warehouse: str):
+        """
+        Create a S3 Tables catalog from a connection string to a S3 Tables catalog.
+        Args:
+            warehouse (str): The warehouse to connect to.
+        """
+        self.warehouse = warehouse
+
+    def get_java_object(self):
+        return _create_java_s3_tables_catalog(self.warehouse)
+
+    def __eq__(self, other):
+        if not isinstance(other, S3TablesCatalog):
+            return False
+        return self.warehouse == other.warehouse
+
+
+@overload(S3TablesCatalog, no_unliteral=True)
+def overload_s3_tables_catalog_constructor(warehouse: str):
+    raise_bodo_error("S3TablesCatalog: Cannot be created in JIT mode.")
+
+
+class S3TablesCatalogType(DatabaseCatalogType):
+    def __init__(self, warehouse: str):
+        self.warehouse = warehouse
+        super().__init__(name=f"S3TablesCatalog({self.warehouse=})")
+
+    def get_java_object(self):
+        return _create_java_s3_tables_catalog(self.warehouse)
+
+    @property
+    def key(self):
+        return self.warehouse
+
+
+@typeof_impl.register(S3TablesCatalog)
+def typeof_s3_tables_catalog(val, c):
+    return S3TablesCatalogType(warehouse=val.warehouse)
+
+
+register_model(S3TablesCatalogType)(models.OpaqueModel)
+
+
+@box(S3TablesCatalogType)
+def box_s3_tables_catalog_type(typ, val, c):
+    """
+    Box a S3 Tables Catalog native representation into a Python object. We populate
+    the contents based on typing information.
+    """
+    warehouse_obj = c.pyapi.from_native_value(
+        types.unicode_type,
+        c.context.get_constant_generic(c.builder, types.unicode_type, typ.warehouse),
+        c.env_manager,
+    )
+    s3_tables_catalog_obj = c.pyapi.unserialize(
+        c.pyapi.serialize_object(S3TablesCatalog)
+    )
+    res = c.pyapi.call_function_objargs(s3_tables_catalog_obj, (warehouse_obj,))
+    c.pyapi.decref(warehouse_obj)
+    c.pyapi.decref(s3_tables_catalog_obj)
+    return res
+
+
+@unbox(S3TablesCatalogType)
+def unbox_s3_tables_catalog_type(typ, val, c):
+    """
+    Unbox a S3 Tables Catalog Python object into its native representation.
+    Since the actual model is opaque we can just generate a dummy.
+    """
+    return NativeValue(c.context.get_dummy_value())
+
+
+@numba.jit(types.unicode_type(types.unicode_type))
+def get_conn_str(warehouse):
+    """Get the connection string for a S3 Tables Iceberg catalog."""
+    return warehouse
+
+
+class S3TablesConnectionType(IcebergConnectionType):
+    """
+    Python class for storing the information needed to connect to a S3 Tables Iceberg catalog.
+    The compiler can get a connection string using the get_conn_str function.
+    The runtime can get a connection string using the conn_str attribute.
+    """
+
+    def __init__(self, warehouse):
+        self.warehouse = warehouse
+        self.conn_str = get_conn_str(warehouse)
+
+        super().__init__(name=f"S3TablesConnectionType({warehouse=})")
+
+    def get_conn_str(self) -> str:
+        return "iceberg+" + self.conn_str
+
+
+@intrinsic
+def _get_s3_tables_connection(typingctx, warehouse, conn_str):
+    """Create a struct model for a  S3TablesonnectionType from a warehouse and connection string."""
+    literal_warehouse = get_literal_value(warehouse)
+    s3_tables_connection_type = S3TablesConnectionType(literal_warehouse)
+
+    def codegen(context, builder, sig, args):
+        """lowering code to initialize a S3TablesConnectionType"""
+        s3_tables_connection_type = sig.return_type
+        s3_tables_connection_struct = cgutils.create_struct_proxy(
+            s3_tables_connection_type
+        )(context, builder)
+        context.nrt.incref(builder, sig.args[1], args[1])
+        s3_tables_connection_struct.conn_str = args[1]
+        return s3_tables_connection_struct._getvalue()
+
+    return s3_tables_connection_type(warehouse, conn_str), codegen
+
+
+def get_s3_tables_connection(warehouse: str):
+    pass
+
+
+@overload(get_s3_tables_connection, no_unliteral=True)
+def overload_get_s3_tables_connection(warehouse: str):
+    """Overload for get_s3_tables_connection that creates a S3TablesConnectionType."""
+
+    def impl(warehouse: str):  # pragma: no cover
+        conn_str = get_conn_str(warehouse)
+        return _get_s3_tables_connection(warehouse, conn_str)
+
+    return impl
+
+
+@register_model(S3TablesConnectionType)
+class S3TablesConnectionModel(models.StructModel):
+    """Model for S3TablesConnectionType that has one member, conn_str."""
+
+    def __init__(self, dmm, fe_type):
+        members = [
+            ("conn_str", types.unicode_type),
+        ]
+        super().__init__(dmm, fe_type, members)
+
+
+make_attribute_wrapper(S3TablesConnectionType, "conn_str", "conn_str")

BodoSQL/bodosql/context.py

@@ -1249,7 +1249,9 @@ def _create_generator(self, hide_credentials: bool):
             self.default_tz,
         )
 
-    def add_or_replace_view(self, name: str, table: pd.DataFrame | TablePath):
+    def add_or_replace_view(
+        self, name: str, table: pd.DataFrame | TablePath
+    ) -> BodoSQLContext:
         """Create a new BodoSQLContext that contains all of the old DataFrames and the
         new table being provided. If there is a DataFrame in the old BodoSQLContext with
         the same name, it is replaced by the new table in the new BodoSQLContext. Otherwise

BodoSQL/bodosql/tests/conftest.py

@@ -27,7 +27,7 @@
     tabular_connection,
 )
 from bodo.tests.iceberg_database_helpers.utils import get_spark
-from bodo.tests.utils import gen_nonascii_list
+from bodo.tests.utils import gen_nonascii_list, temp_env_override
 
 # Patch to avoid PySpark's Py4j exception handler in testing.
 # See:
@@ -1973,6 +1973,17 @@ def glue_catalog():
     return bodosql.GlueCatalog(warehouse=warehouse)
 
 
+@pytest.fixture
+def s3_tables_catalog():
+    """
+    Returns a s3 tables catalog object
+    """
+
+    warehouse = "arn:aws:s3tables:us-east-2:427443013497:bucket/unittest-bucket"
+    with temp_env_override({"AWS_REGION": "us-east-2"}):
+        yield bodosql.S3TablesCatalog(warehouse=warehouse)
+
+
 @pytest.fixture(
     params=[
         pytest.param(

BodoSQL/bodosql/tests/test_types/test_s3_tables_iceberg_catalog.py

@@ -0,0 +1,111 @@
+import pandas as pd
+
+import bodo
+import bodosql
+from bodo.tests.utils import (
+    assert_tables_equal,
+    check_func,
+    gen_unique_table_id,
+    pytest_s3_tables,
+    run_rank0,
+)
+from bodosql.bodosql_types.s3_tables_catalog import S3TablesConnectionType
+
+pytestmark = pytest_s3_tables
+
+
+def test_basic_read(memory_leak_check, s3_tables_catalog):
+    """
+    Test reading an entire Iceberg table from S3 Tables in SQL
+    """
+    bc = bodosql.BodoSQLContext(catalog=s3_tables_catalog)
+
+    def impl(bc, query):
+        return bc.sql(query)
+
+    py_out = pd.DataFrame(
+        {
+            "A": ["ally", "bob", "cassie", "david", None],
+            "B": [10.5, -124.0, 11.11, 456.2, -8e2],
+            "C": [True, None, False, None, None],
+        }
+    )
+
+    query = 'SELECT * FROM "read_namespace"."bodo_iceberg_read_test"'
+    check_func(
+        impl,
+        (bc, query),
+        py_output=py_out,
+        sort_output=True,
+        reset_index=True,
+    )
+
+
+def test_s3_tables_catalog_iceberg_write(s3_tables_catalog, memory_leak_check):
+    """tests that writing tables works"""
+    import bodo_iceberg_connector as bic
+
+    in_df = pd.DataFrame(
+        {
+            "ints": list(range(100)),
+            "floats": [float(x) for x in range(100)],
+            "str": [str(x) for x in range(100)],
+            "dict_str": ["abc", "df"] * 50,
+        }
+    )
+    bc = bodosql.BodoSQLContext(catalog=s3_tables_catalog)
+    bc = bc.add_or_replace_view("TABLE1", in_df)
+    con_str = S3TablesConnectionType(s3_tables_catalog.warehouse).get_conn_str()
+    table_name = run_rank0(
+        lambda: gen_unique_table_id("bodosql_catalog_write_iceberg_table").upper()
+    )().lower()
+
+    def impl(bc, query):
+        bc.sql(query)
+        # Return an arbitrary value. This is just to enable a py_output
+        # so we can reuse the check_func infrastructure.
+        return 5
+
+    ctas_query = f'CREATE OR REPLACE TABLE "write_namespace"."{table_name}" AS SELECT * from __bodolocal__."TABLE1"'
+    exception_occurred_in_test_body = False
+    try:
+        # Only test with only_1D=True so we only insert into the table once.
+        check_func(
+            impl,
+            (bc, ctas_query),
+            only_1D=True,
+            py_output=5,
+            use_table_format=True,
+        )
+
+        @bodo.jit
+        def read_results(con_str, schema, table_name):
+            output_df = pd.read_sql_table(table_name, con=con_str, schema=schema)
+            return bodo.allgatherv(output_df)
+
+        output_df = read_results(con_str, "write_namespace", table_name)
+        assert_tables_equal(output_df, in_df, check_dtype=False)
+
+    except Exception as e:
+        # In the case that another exception ocurred within the body of the try,
+        # We may not have created a table to drop.
+        # because of this, we call delete_table in a try/except, to avoid
+        # masking the original exception
+        exception_occurred_in_test_body = True
+        raise e
+    finally:
+        if exception_occurred_in_test_body:
+            try:
+                run_rank0(bic.delete_table)(
+                    bodo.io.iceberg.format_iceberg_conn(con_str),
+                    "write_namespace",
+                    table_name,
+                )
+            except Exception:
+                pass
+        else:
+            run_rank0(bic.delete_table)(
+                bodo.io.iceberg.format_iceberg_conn(con_str),
+                "write_namespace",
+                table_name,
+            )

BodoSQL/calcite_sql/bodosql-calcite-application/pom.xml

@@ -33,7 +33,7 @@
         <!-- Start Iceberg Dependencies. We may remove these as we cleanup the relationship with the Iceberg Connector. -->
         <iceberg.version>1.5.2</iceberg.version>
         <hadoop.version>3.3.3</hadoop.version>
-        <aws.version>2.20.26</aws.version>
+        <aws.version>2.29.26</aws.version>
         <janino.version>3.1.10</janino.version>
         <!-- End Iceberg Dependencies -->
     </properties>
@@ -220,6 +220,18 @@
             <artifactId>s3</artifactId>
             <version>${aws.version}</version>
         </dependency>
+        <!-- S3Tables Dependencies -->
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>s3tables</artifactId>
+            <version>${aws.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>software.amazon.s3tables</groupId>
+            <artifactId>s3-tables-catalog-for-iceberg</artifactId>
+            <version>0.1.3</version>
+        </dependency>
+
         <dependency>
             <groupId>org.testcontainers</groupId>
             <artifactId>testcontainers</artifactId>

BodoSQL/calcite_sql/bodosql-calcite-application/src/main/java/com/bodosql/calcite/application/BodoSQLTypeSystems/BodoSQLRelDataTypeSystem.java

@@ -51,17 +51,17 @@ public class BodoSQLRelDataTypeSystem extends RelDataTypeSystemImpl {
   public boolean enableStreamingSortLimitOffset;
 
   public static class CatalogContext {
-    public String currentDatabase;
+    public @Nullable String currentDatabase;
     public String currentAccount;
 
-    public CatalogContext(String database, String account) {
+    public CatalogContext(@Nullable String database, String account) {
       this.currentDatabase = database;
       this.currentAccount = account;
     }
   }
 
   /*
-  CURRENT_DATABSE: Name of the database in use for the current session.
+  CURRENT_DATABASE: Name of the database in use for the current session.
 
   */
   private CatalogContext currentDatabase;