test coverage

Author: sfc-gh-mayliu

src/snowflake/snowpark/_internal/analyzer/analyzer.py

@@ -916,13 +916,10 @@ def _process_partition_by_in_iceberg_config(
         Process partition_by expressions from iceberg_config, converting Column objects to SQL strings.
         Returns a new iceberg_config dict with partition_by as a list of SQL strings, or the original config if no processing needed.
         """
-        if not iceberg_config:
+        if iceberg_config is None or iceberg_config.get("partition_by") is None:
             return iceberg_config
 
         iceberg_config = {k.lower(): v for k, v in iceberg_config.items()}
-
-        if not iceberg_config.get("partition_by"):
-            return iceberg_config
         pb = iceberg_config["partition_by"]
 
         # Convert to list and filter out empty expressions

tests/integ/scala/test_dataframe_writer_suite.py

@@ -295,6 +295,7 @@ def test_iceberg_partition_by(session, local_testing_mode):
         ),
     )
 
+    # Test 1: Single string value
     table_name_1 = Utils.random_table_name()
     df.write.save_as_table(
         table_name_1,
@@ -316,6 +317,7 @@ def test_iceberg_partition_by(session, local_testing_mode):
     finally:
         session.table(table_name_1).drop_table()
 
+    # Test 2: Empty list
     table_name_2 = Utils.random_table_name()
     df.write.save_as_table(
         table_name_2,
@@ -331,6 +333,53 @@ def test_iceberg_partition_by(session, local_testing_mode):
     finally:
         session.table(table_name_2).drop_table()
 
+    # Test 3: Single Column object
+    table_name_3 = Utils.random_table_name()
+    df.write.save_as_table(
+        table_name_3,
+        iceberg_config={
+            "external_volume": "PYTHON_CONNECTOR_ICEBERG_EXVOL",
+            "catalog": "SNOWFLAKE",
+            "partition_by": col("b"),
+        },
+    )
+    try:
+        ddl = session._run_query(f"select get_ddl('table', '{table_name_3}')")
+        assert "PARTITION BY (B)" in ddl[0][0]
+    finally:
+        session.table(table_name_3).drop_table()
+
+    # Test 4: Mix of strings and Column objects with empty strings
+    table_name_4 = Utils.random_table_name()
+    df.write.save_as_table(
+        table_name_4,
+        iceberg_config={
+            "external_volume": "PYTHON_CONNECTOR_ICEBERG_EXVOL",
+            "catalog": "SNOWFLAKE",
+            "partition_by": ["a", "", col("b")],
+        },
+    )
+    try:
+        ddl = session._run_query(f"select get_ddl('table', '{table_name_4}')")
+        assert "PARTITION BY (A, B)" in ddl[0][0]
+    finally:
+        session.table(table_name_4).drop_table()
+
+    # Test 5: No partition_by
+    table_name_5 = Utils.random_table_name()
+    df.write.save_as_table(
+        table_name_5,
+        iceberg_config={
+            "external_volume": "PYTHON_CONNECTOR_ICEBERG_EXVOL",
+            "catalog": "SNOWFLAKE",
+        },
+    )
+    try:
+        ddl = session._run_query(f"select get_ddl('table', '{table_name_5}')")
+        assert "PARTITION BY" not in ddl[0][0]
+    finally:
+        session.table(table_name_5).drop_table()
+
 
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",