[BSE-4458] Refactor Iceberg IO to Use PyIceberg Under the Hood by srilman · Pull Request #235 · bodo-ai/Bodo

srilman · 2025-02-12T19:34:15Z

Changes included in this PR

Refactor our Iceberg IO to support PyIceberg as a backend. Replaces a lot of metadata-related operations to use PyIceberg instead, delete any helper code that PyIceberg already does, and limit the use of the Iceberg connector.

Testing strategy

PR Ci and Nightly both

User facing changes

Users no longer need to install the Iceberg connector for Iceberg support. Instead, they need to either install bodo[iceberg] or bodo and pyiceberg.

Checklist

Pipelines passed before requesting review. To run CI you must include [run CI] in your commit message.
I am familiar with the Contributing Guide
I have installed + ran pre-commit hooks.

srilman · 2025-02-13T16:44:20Z

Nightly: https://dev.azure.com/bodo-inc/Bodo/_build?definitionId=5&_a=summary

commit 8805b95 Author: Srinivas Lade <[email protected]> Date: Fri Jan 31 11:43:13 2025 -0500 [BSE-4471] Move SF Prefetch Support to Python (#189) commit 9821639 Author: Srinivas Lade <[email protected]> Date: Wed Jan 29 12:33:52 2025 -0500 [BSE-4470][BSE-4474] Add PyIceberg Read Support (#160)

srilman · 2025-02-18T16:33:33Z

Nightly Run that worked: https://dev.azure.com/bodo-inc/Bodo/_build/results?buildId=25556&view=results

srilman · 2025-02-18T16:36:05Z

-    // Benefits: Potentially some speed up in collecting repeated metadata like
-    // schema
-    // Downsides: Does not refresh if another program modifies the Catalog
-    return CachingCatalog.wrap(catalog);


Removed the CachingCatalog change because it was causing issues when jumping between PyIceberg and Iceberg-Java. This shouldn't be as much as a problem anymore because we use some caching on the PyIceberg side, which is still the predominant caller for getting metadata. Iceberg Java is only used for more heavy handed stuff like commits.

srilman · 2025-02-18T16:36:22Z

-      catalog = S3TablesBuilder.create(connStr);
-      return CachingCatalog.wrap(catalog);
+    if (connStr.startsWith("iceberg+arn:aws:s3tables") && catalogType.equals("s3tables")) {
+      catalog = S3TablesBuilder.create(connStr.replaceFirst("^iceberg\\+", ""));


Was a bug, I removed the part where we strip the prefix during compilation.

srilman · 2025-02-18T16:37:01Z

+
+    def _table_path(self, identifier: Identifier) -> str:
+        wh_path = self.warehouse_path.removesuffix("/")
+        return f"{wh_path}/{'/'.join(identifier)}"


Bug here to remove cases where {warehouse_path}//...

You mean this fixes a bug? Could you add comment in the code instead.

To clarify it was a bug in my own code in the base read PR that was discovered when running nightly.

srilman · 2025-02-18T16:38:31Z

@@ -0,0 +1,83 @@
+import os


Added a FileIO wrapper to call into PyArrow's AzureFileSystem

srilman · 2025-02-18T16:39:21Z

-        conn_str = (
-            f"snowflake://{account}.snowflakecomputing.com/?{urlencode(conn_contents)}"
-        )
+def _fs_from_file_path(file_path: str, io: FileIO) -> FileSystem:


Mostly a duplicate of a internal function from PyIceberg, but slightly modified to take use of our monkey-patch

Could you add a link to the original function that was copied?

srilman · 2025-02-18T16:40:08Z

@@ -0,0 +1,79 @@
+import pyiceberg.io


Sadly I needed to monkey patch a few things from PyIceberg to fix various issues. See below

We can make the monkey patch more robust by following the numba_compat.py format (check hash for code changes and add comments with "BODO change" notes).

Makes sense, I'm going to leave it like this for the time being because I will have to make some large changes when PyIceberg 0.9 comes out anyways

srilman · 2025-02-18T16:40:43Z

+    Monkey-patched version of pyiceberg.schema.build_position_accessors to
+    use the new _BuildPositionAccessors class.
+    """
+    return visit(schema_or_type, _BuildPositionAccessors())


Fix to apply filters to top-level struct fields. Right now, PyIceberg only works on inner struct fields

Let's add as comment in the code.

srilman · 2025-02-18T16:41:23Z

+pyiceberg.schema.build_position_accessors = build_position_accessors
+
+# Moneky-patch inferred IO to support Bodo's
+pyiceberg.io.pyarrow.PyArrowFileIO = BodoPyArrowFileIO


Hook to apply our AzureFileSystem changes to PyIceberg. I need to do this here because there are places where PyIceberg just calls PyArrowFileIO directly instead of using their own hooks for custom FileIO.

Let's add the comment in the code.

srilman · 2025-02-18T16:42:50Z

+    pa.Schema,
+    # Runtime Schema
+    list[str],
+    list[types.ArrayCompatible],


I messed up this part initially which caused some schema evolution tests to fail. Basically, this function should return the original schema info (names, types, PyArrow schema) and schema after apply BodoSQL column pruning.

srilman · 2025-02-18T16:43:26Z

+        return DateLiteral((val - datetime.date(1970, 1, 1)).days)
+    if isinstance(val, datetime.datetime):
+        return TimestampLiteral((val - datetime.datetime(1970, 1, 1)).microseconds)
+    if isinstance(val, (list, pd.core.arrays.ExtensionArray)):


Had to add pd.core.arrays.ExtensionArray

Why? This is for nested types?

A comment in the code would be good.

I'm not sure exactly why, just that sometimes Pandas arrays are used instead of lists especially when coming from BodoSQL

The list/array difference is fine, but my point is that we may need other nested types (struct/map) here too. A TODO may be enough for now.

IsaacWarren

Just reviewed the places commented on, assuming others are unchanged. Thanks Srinivas!

IsaacWarren · 2025-02-18T17:09:46Z

+                    "AZURE_STORAGE_ACCOUNT_NAME"
+                )
+
+            account_key = self.properties.get(ADLS_ACCOUNT_KEY) or os.environ.get(


Add a todo to pass sas token for arrow 19? Also we should support ADLS_CLIENT_ID, ADLS_TENANT_ID, and ADLS_CLIENT_SECRET

codecov · 2025-02-18T17:30:06Z

Codecov Report

Attention: Patch coverage is 64.54616% with 457 lines in your changes missing coverage. Please review.

Project coverage is 66.98%. Comparing base (6c7b124) to head (80e211e).
Report is 36 commits behind head on main.

❌ Your project status has failed because the head coverage (66.98%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #235      +/-   ##
==========================================
+ Coverage   61.48%   66.98%   +5.50%     
==========================================
  Files         169      176       +7     
  Lines       61920    62627     +707     
  Branches     8663     8729      +66     
==========================================
+ Hits        38070    41950    +3880     
+ Misses      21430    18077    -3353     
- Partials     2420     2600     +180

ehsantn

Thanks for the huge effort @srilman. I just skimmed through to see if anything stands out since most of the code is already reviewed.

ehsantn · 2025-02-18T19:56:44Z

+    pytest.mark.skip(
+        reason="[BSE-4569] MERGE INTO with PyIceberg is not supported yet"
+    ),


Are we disabling our current MERGE INTO support?

Yes for the time being

ehsantn · 2025-02-18T20:02:36Z

-
-    return impl
+this_module = sys.modules[__name__]
+PyArrowFSType, pyarrow_fs_type = install_opaque_class(


Did we decide that this doesn't need memory management?

Yes I believe so, because it is something that we don't explicitly create a reference to, and the transaction that we use later on in the function hold a reference to it

ehsantn · 2025-02-18T20:06:43Z

+
+    def _table_path(self, identifier: Identifier) -> str:
+        wh_path = self.warehouse_path.removesuffix("/")
+        return f"{wh_path}/{'/'.join(identifier)}"


You mean this fixes a bug? Could you add comment in the code instead.

ehsantn · 2025-02-18T20:07:24Z

+        # TODO
+        pass


This is supposed to be TODO? Raise an error instead?

ehsantn · 2025-02-18T20:07:49Z

+    def drop_table(self, identifier: str | Identifier) -> None:
+        # TODO
+        pass
+
+    def purge_table(self, identifier: str | Identifier) -> None:
+        # TODO
+        pass


ehsantn · 2025-02-18T20:33:50Z

        # that have been pushed down to I/O.
        rtjf_terms: list[tuple[ir.Var, tuple[int], tuple[int, int, str]]] | None = None,
    ):
+        # Info requireds to connect to the catalog and table


Suggested change

# Info requireds to connect to the catalog and table

# Info required to connect to the catalog and table

ehsantn · 2025-02-18T20:34:48Z

+        return DateLiteral((val - datetime.date(1970, 1, 1)).days)
+    if isinstance(val, datetime.datetime):
+        return TimestampLiteral((val - datetime.datetime(1970, 1, 1)).microseconds)
+    if isinstance(val, (list, pd.core.arrays.ExtensionArray)):


Why? This is for nested types?

ehsantn · 2025-02-18T20:34:55Z

+        return DateLiteral((val - datetime.date(1970, 1, 1)).days)
+    if isinstance(val, datetime.datetime):
+        return TimestampLiteral((val - datetime.datetime(1970, 1, 1)).microseconds)
+    if isinstance(val, (list, pd.core.arrays.ExtensionArray)):


A comment in the code would be good.

ehsantn · 2025-02-18T20:35:33Z

+                    rtjf_filters = pie.And(
+                        rtjf_filters, pie.In(col, literal(unique_vals))
                    )
+                    print(rtjf_filters, unique_vals)


Remove print?

Suggested change

print(rtjf_filters, unique_vals)

ehsantn · 2025-02-18T20:38:40Z

 from bodo.tests.utils import pytest_mark_one_rank

-pytestmark = pytest.mark.iceberg
+pytestmark = [pytest.mark.iceberg, pytest.mark.skip]


Why are we skipping this file. Add a comment in the code?

Oh just missed that, undid it

srilman force-pushed the slade/pyiceberg-refactor branch 2 times, most recently from 7362849 to ef8e93f Compare February 13, 2025 16:42

srilman mentioned this pull request Feb 17, 2025

0.8 Patch Release for Arrow 19 Support apache/iceberg-python#1672

Closed

srilman added 8 commits February 18, 2025 11:33

[BSE-4473] PyIceberg Write Support (#202)

03b6716

[BSE-4472] Add Puffin Support Back Through Iceberg Connector (#204)

60e3daa

Fix Struct Read Issues

33b3da7

Fix the PyIceberg marking

7e2e6e4

Theta Related Issues

bf88d35

Ready

5a51ff6

[run ci] fix path append issue

85fb58b

srilman force-pushed the slade/pyiceberg-refactor branch from 2c4a674 to 85fb58b Compare February 18, 2025 16:33

srilman commented Feb 18, 2025

View reviewed changes

[run ci] fix import problem

f9469f2

srilman requested review from IsaacWarren and ehsantn February 18, 2025 16:58

srilman marked this pull request as ready for review February 18, 2025 16:58

IsaacWarren approved these changes Feb 18, 2025

View reviewed changes

ehsantn approved these changes Feb 18, 2025

View reviewed changes

srilman added 3 commits February 19, 2025 09:37

Merge branch 'main' into slade/pyiceberg-refactor

165b766

[run ci] ok all

d1fe13f

[run ci] oops

80e211e

srilman merged commit 151eb1e into main Feb 19, 2025

srilman deleted the slade/pyiceberg-refactor branch February 19, 2025 17:43

	# Info requireds to connect to the catalog and table
	# Info required to connect to the catalog and table

Conversation

srilman commented Feb 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes included in this PR

Testing strategy

User facing changes

Checklist

Uh oh!

srilman commented Feb 13, 2025

Uh oh!

srilman commented Feb 18, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

IsaacWarren left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Feb 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ehsantn left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

srilman commented Feb 12, 2025 •

edited

Loading

codecov Bot commented Feb 18, 2025 •

edited

Loading