[BSE-3976] Disable SF Iceberg Prefetching Behind a Feature Flag (#8355)

bodo-ai · Oct 7, 2024 · d7d30ba · d7d30ba
1 parent cedf05c
commit d7d30ba
Show file tree

Hide file tree

Showing 10 changed files with 99 additions and 27 deletions.
diff --git a/BodoSQL/bodosql/context.py b/BodoSQL/bodosql/context.py
@@ -1198,6 +1198,7 @@ def _create_generator(self, hide_credentials: bool):
                 bodo.enable_streaming_sort_limit_offset,
                 bodo.bodo_sql_style,
                 bodo.bodosql_full_caching,
+                bodo.prefetch_sf_iceberg,
             )
         extra_args = () if self.default_tz is None else (self.default_tz,)
         generator = RelationalAlgebraGeneratorClass(
@@ -1214,6 +1215,7 @@ def _create_generator(self, hide_credentials: bool):
             bodo.enable_streaming_sort_limit_offset,
             bodo.bodo_sql_style,
             bodo.bodosql_full_caching,
+            bodo.prefetch_sf_iceberg,
             *extra_args,
         )
         return generator

diff --git a/BodoSQL/bodosql/context_ext.py b/BodoSQL/bodosql/context_ext.py
@@ -454,6 +454,7 @@ def _gen_sql_plan_pd_func_text_and_lowered_globals(
                     bodo.enable_streaming_sort_limit_offset,
                     bodo.bodo_sql_style,
                     bodo.bodosql_full_caching,
+                    bodo.prefetch_sf_iceberg,
                 )
             else:
                 extra_args = (
@@ -476,6 +477,7 @@ def _gen_sql_plan_pd_func_text_and_lowered_globals(
                     bodo.enable_streaming_sort_limit_offset,
                     bodo.bodo_sql_style,
                     bodo.bodosql_full_caching,
+                    bodo.prefetch_sf_iceberg,
                     *extra_args,
                 )
         except Exception as e:

diff --git a/BodoSQL/bodosql/tests/test_types/test_snowflake_iceberg_catalog.py b/BodoSQL/bodosql/tests/test_types/test_snowflake_iceberg_catalog.py
@@ -57,20 +57,13 @@ def impl(bc, query):
     )
 
     query = "SELECT A, B, C FROM BODOSQL_ICEBERG_READ_TEST"
-    stream = StringIO()
-    logger = create_string_io_logger(stream)
-    with set_logging_stream(logger, 1):
-        check_func(
-            impl,
-            (bc, query),
-            py_output=py_out,
-            sort_output=True,
-            reset_index=True,
-        )
-        check_logger_msg(
-            stream,
-            'Execution time for prefetching SF-managed Iceberg metadata "TEST_DB"."PUBLIC"."BODOSQL_ICEBERG_READ_TEST"',
-        )
+    check_func(
+        impl,
+        (bc, query),
+        py_output=py_out,
+        sort_output=True,
+        reset_index=True,
+    )
 
 
 @temp_env_override({"AWS_REGION": "us-east-1"})
@@ -667,3 +660,59 @@ def check_table_type():
         ), f"Table type is not as expected. Expected MANAGED but found {output_table_type}"
     finally:
         drop_snowflake_table(table_name, db, schema, user=3)
+
+
+@temp_env_override({"AWS_REGION": "us-east-1"})
+def test_prefetch_flag(memory_leak_check):
+    """
+    Test that if the prefetch flag is set, a prefetch occurs
+    TODO: [BSE-3977] There is a bug that causes the prefetch
+    to not be used when the query is compiled in the same process
+    """
+
+    old_prefetch_flag = bodo.prefetch_sf_iceberg
+
+    try:
+        bodo.prefetch_sf_iceberg = True
+
+        catalog = bodosql.SnowflakeCatalog(
+            os.environ["SF_USERNAME"],
+            os.environ["SF_PASSWORD"],
+            "bodopartner.us-east-1",
+            "DEMO_WH",
+            "TEST_DB",
+            connection_params={"schema": "PUBLIC", "role": "ACCOUNTADMIN"},
+            iceberg_volume="exvol",
+        )
+        bc = bodosql.BodoSQLContext(catalog=catalog)
+
+        def impl(bc, query):
+            return bc.sql(query)
+
+        py_out = pd.DataFrame(
+            {
+                "A": ["ally", "bob", "cassie", "david", pd.NA],
+                "B": [10.5, -124.0, 11.11, 456.2, -8e2],
+                "C": [True, pd.NA, False, pd.NA, pd.NA],
+            }
+        )
+
+        query = "SELECT A, B, C FROM BODOSQL_ICEBERG_READ_TEST"
+        stream = StringIO()
+        logger = create_string_io_logger(stream)
+        with set_logging_stream(logger, 1):
+            check_func(
+                impl,
+                (bc, query),
+                py_output=py_out,
+                sort_output=True,
+                reset_index=True,
+            )
+
+            check_logger_msg(
+                stream,
+                'Execution time for prefetching SF-managed Iceberg metadata "TEST_DB"."PUBLIC"."BODOSQL_ICEBERG_READ_TEST"',
+            )
+
+    finally:
+        bodo.prefetch_sf_iceberg = old_prefetch_flag
diff --git a/...ql-calcite-application/src/main/java/com/bodosql/calcite/adapter/bodo/StreamingOptions.kt b/...ql-calcite-application/src/main/java/com/bodosql/calcite/adapter/bodo/StreamingOptions.kt
@@ -1,3 +1,3 @@
 package com.bodosql.calcite.adapter.bodo
 
-data class StreamingOptions(val chunkSize: Int)
+data class StreamingOptions(val chunkSize: Int, val prefetchSFIceberg: Boolean)
diff --git a/...ation/src/main/java/com/bodosql/calcite/adapter/iceberg/IcebergToBodoPhysicalConverter.kt b/...ation/src/main/java/com/bodosql/calcite/adapter/iceberg/IcebergToBodoPhysicalConverter.kt
@@ -165,11 +165,13 @@ class IcebergToBodoPhysicalConverter(cluster: RelOptCluster, traits: RelTraitSet
 
         // Store table name and base connection string in builder for prefetch gen
         // Note: Only when reading from a Snowflake-managed Iceberg table (using Snowflake catalog)
-        if (tableScanNode.getCatalogTable().getCatalog() is SnowflakeCatalog) {
-            var staticConExpr = relInput.generatePythonConnStr(ImmutableList.of("", ""))
-            staticConExpr =
-                if (staticConExpr is StringLiteral) StringLiteral("iceberg+${staticConExpr.arg}") else staticConExpr
-            ctx.builder().addSfIcebergTablePath(staticConExpr, tableScanNode.getCatalogTable().getQualifiedName())
+        if (ctx.streamingOptions().prefetchSFIceberg) {
+            if (tableScanNode.getCatalogTable().getCatalog() is SnowflakeCatalog) {
+                var staticConExpr = relInput.generatePythonConnStr(ImmutableList.of("", ""))
+                staticConExpr =
+                    if (staticConExpr is StringLiteral) StringLiteral("iceberg+${staticConExpr.arg}") else staticConExpr
+                ctx.builder().addSfIcebergTablePath(staticConExpr, tableScanNode.getCatalogTable().getQualifiedName())
+            }
         }
 
         val columnsArg = Expr.List(cols.map { v -> StringLiteral(v) })

diff --git a/...calcite-application/src/main/java/com/bodosql/calcite/application/BodoCodeGenVisitor.java b/...calcite-application/src/main/java/com/bodosql/calcite/application/BodoCodeGenVisitor.java
@@ -197,7 +197,8 @@ public BodoCodeGenVisitor(
       List<RelDataType> dynamicParamTypes,
       Map<String, RelDataType> namedParamTypeMap,
       Map<Integer, Integer> idMapping,
-      boolean hideOperatorIDs) {
+      boolean hideOperatorIDs,
+      boolean prefetchSFIceberg) {
     super();
     this.loweredGlobals = loweredGlobalVariablesMap;
     this.originalSQLQuery = originalSQLQuery;
@@ -209,7 +210,7 @@ public BodoCodeGenVisitor(
     this.tracingLevel = tracingLevel;
     this.generatedCode = new Module.Builder(verboseLevel);
     this.generatedCode.setHideOperatorIDs(hideOperatorIDs);
-    this.streamingOptions = new StreamingOptions(batchSize);
+    this.streamingOptions = new StreamingOptions(batchSize, prefetchSFIceberg);
     this.dynamicParamTypes = dynamicParamTypes;
     this.namedParamTypeMap = namedParamTypeMap;
     this.generatedCode.setIDMapping(idMapping);

diff --git a/...application/src/main/java/com/bodosql/calcite/application/RelationalAlgebraGenerator.java b/...application/src/main/java/com/bodosql/calcite/application/RelationalAlgebraGenerator.java
@@ -143,6 +143,9 @@ public class RelationalAlgebraGenerator {
   /** Should we use the covering expression approach to cache or only exact matches. */
   public static boolean coveringExpressionCaching = false;
 
+  /** Should we prefetch metadata location for Snowflake-managed Iceberg tables. */
+  public static boolean prefetchSFIceberg = false;
+
   /**
    * Helper method for RelationalAlgebraGenerator constructors to create a SchemaPlus object from a
    * list of BodoSqlSchemas.
@@ -195,7 +198,8 @@ public RelationalAlgebraGenerator(
       boolean enableStreamingSort,
       boolean enableStreamingSortLimitOffset,
       String sqlStyle,
-      boolean coveringExpressionCaching) {
+      boolean coveringExpressionCaching,
+      boolean prefetchSFIceberg) {
     this.catalog = null;
     this.plannerType = choosePlannerType(plannerType);
     this.verboseLevel = verboseLevel;
@@ -217,6 +221,7 @@ public RelationalAlgebraGenerator(
     this.enableTimestampTz = enableTimestampTz;
     this.enableRuntimeJoinFilters = enableRuntimeJoinFilters;
     this.coveringExpressionCaching = coveringExpressionCaching;
+    this.prefetchSFIceberg = prefetchSFIceberg;
   }
 
   /** Constructor for the relational algebra generator class that takes in the default timezone. */
@@ -234,6 +239,7 @@ public RelationalAlgebraGenerator(
       boolean enableStreamingSortLimitOffset,
       String sqlStyle,
       boolean coveringExpressionCaching,
+      boolean prefetchSFIceberg,
       String defaultTz) {
     this.catalog = null;
     this.plannerType = choosePlannerType(plannerType);
@@ -258,6 +264,7 @@ public RelationalAlgebraGenerator(
     this.enableTimestampTz = enableTimestampTz;
     this.enableRuntimeJoinFilters = enableRuntimeJoinFilters;
     this.coveringExpressionCaching = coveringExpressionCaching;
+    this.prefetchSFIceberg = prefetchSFIceberg;
   }
 
   /**
@@ -283,7 +290,8 @@ public RelationalAlgebraGenerator(
       boolean enableStreamingSort,
       boolean enableStreamingSortLimitOffset,
       String sqlStyle,
-      boolean coveringExpressionCaching) {
+      boolean coveringExpressionCaching,
+      boolean prefetchSFIceberg) {
     this.catalog = catalog;
     this.plannerType = choosePlannerType(plannerType);
     this.verboseLevel = verboseLevel;
@@ -295,6 +303,7 @@ public RelationalAlgebraGenerator(
     this.enableRuntimeJoinFilters = enableRuntimeJoinFilters;
     this.sqlStyle = sqlStyle;
     this.coveringExpressionCaching = coveringExpressionCaching;
+    this.prefetchSFIceberg = prefetchSFIceberg;
     System.setProperty("calcite.default.charset", "UTF-8");
     List<String> catalogDefaultSchema = catalog.getDefaultSchema(0);
     final @Nullable String currentDatabase;
@@ -557,7 +566,8 @@ private String getDDLPandasString(
             dynamicTypes,
             namedParamTypes,
             Map.of(),
-            this.hideOperatorIDs);
+            this.hideOperatorIDs,
+            this.prefetchSFIceberg);
     codegen.generateDDLCode(ddlNode, new GenerateDDLTypes(this.planner.getTypeFactory()));
     return codegen.getGeneratedCode();
   }
@@ -601,7 +611,8 @@ private String getPandasStringFromPlan(
             dynamicTypes,
             namedParamTypes,
             v.getIDMapping(),
-            this.hideOperatorIDs);
+            this.hideOperatorIDs,
+            this.prefetchSFIceberg);
     codegen.go(rel);
     return codegen.getGeneratedCode();
   }

diff --git a/...lcite-application/src/main/java/com/bodosql/calcite/application/Testing/GenTestFixture.kt b/...lcite-application/src/main/java/com/bodosql/calcite/application/Testing/GenTestFixture.kt
@@ -46,6 +46,7 @@ abstract class GenTestFixture {
                 // Maintain case sensitivity in the Snowflake style by default
                 "SNOWFLAKE",
                 false,
+                false,
             )
 
         println("SQL query:")

diff --git a/...cite-application/src/main/java/com/bodosql/calcite/application/Testing/PandasGenTest.java b/...cite-application/src/main/java/com/bodosql/calcite/application/Testing/PandasGenTest.java
@@ -102,7 +102,8 @@ public static void main(String[] args) throws Exception {
             true, // Disable Streaming Sort for Testing
             false, // Disable Streaming Sort Limit Offset for Testing
             "SNOWFLAKE", // Maintain case sensitivity in the Snowflake style by default
-            false // Only cache identical nodes
+            false, // Only cache identical nodes
+            false // Generate a prefetch call at the beginning of SQL queries
             );
     List<ColumnDataTypeInfo> paramTypes =
         List.of(new ColumnDataTypeInfo(BodoSQLColumnDataType.INT64, false));

diff --git a/bodo/__init__.py b/bodo/__init__.py
@@ -145,6 +145,9 @@ def _global_except_hook(exctype, value, traceback):
 bodo_disable_streaming_window_sort = (
     os.environ.get("BODO_DISABLE_STREAMING_WINDOW_SORT", "0") != "0"
 )
+# If enabled, generate a prefetch function call to load metadata paths for
+# Snowflake-managed Iceberg tables in the BodoSQL plan.
+prefetch_sf_iceberg = os.environ.get("BODO_PREFETCH_SF_ICEBERG", "0") != "0"
 
 
 def get_sql_config_str() -> str:
-Original file line number
+Diff line change
@@ Expand Up / @@ -46,6 +46,7 @@ abstract class GenTestFixture { @@
                     // Maintain case sensitivity in the Snowflake style by default
                     "SNOWFLAKE",
                     false,
+                    false,
                 )
             println("SQL query:")
@@ Expand Down @@