Merge pull request #27 from cody-scott/polars-null-fix

replace nulls for polars
cody-scott · Dec 17, 2024 · 901c72e · 901c72e
2 parents 837565e + 9ccd077
commit 901c72e
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 11 deletions.
diff --git a/dagster_mssql_bcp_tests/bcp_pandas/test_io_manager.py b/dagster_mssql_bcp_tests/bcp_pandas/test_io_manager.py
@@ -387,8 +387,8 @@ def my_asset(context):
 
         data = pd.DataFrame(
             {
-                "a": [1, 1],
-                "b": ["a\t\nb", "a\t\tb"],
+                "a": [1, 1, 1],
+                "b": ["a\t\nb", "a\t\tb", ""],
             }
         )
 
@@ -400,7 +400,7 @@ def my_asset(context):
         with self.connect_mssql() as connection:
             data = connection.exec_driver_sql(
                 f'SELECT * FROM {schema}.{table}')
-            assert data.fetchall() == [(1, 'a\t\nb'), (1, 'a\t\tb')]
+            assert data.fetchall() == [(1, 'a\t\nb'), (1, 'a\t\tb'), (1, None)]
 
     def test_absent_identity(self):
         schema = "test_pandas_bcp_schema"

diff --git a/dagster_mssql_bcp_tests/bcp_polars/test_bcp.py b/dagster_mssql_bcp_tests/bcp_polars/test_bcp.py
@@ -276,13 +276,13 @@ def test_replace_values(self, polars_io):
         expected = pl.DataFrame(
             {
                 "c": [
-                    '',
-                    '',
+                    None,
+                    None,
                     "c",
                     "abc__TAB__def",
                     "abc__TAB____NEWLINE__def",
                     "abc__NEWLINE__def",
-                    '',
+                    None,
                     "somenanthing",
                 ],
                 "should_process_replacements": [0, 0, 0, 1, 1, 1, 0, 0],
@@ -337,7 +337,7 @@ def test_replace_values(self, polars_io):
             [{"name": "a", "type": "NVARCHAR", "length": 20}]
         )
         df = pl.DataFrame({"a": ['""', "a", '', '"adsf"']})
-        expected = pl.DataFrame({"a": ['', "a", '', '"adsf"'], "should_process_replacements": [0,0,0,0]})
+        expected = pl.DataFrame({"a": [None, "a", None, '"adsf"'], "should_process_replacements": [0,0,0,0]})
         df = polars_io._replace_values(df, schema)
         pl_testing.assert_frame_equal(df, expected)
 

diff --git a/dagster_mssql_bcp_tests/bcp_polars/test_io_manager.py b/dagster_mssql_bcp_tests/bcp_polars/test_io_manager.py
@@ -379,8 +379,8 @@ def my_asset(context):
 
         data = pl.DataFrame(
             {
-                "a": [1, 1],
-                "b": ["a\t\nb", "a\t\tb"],
+                "a": [1, 1, 1],
+                "b": ["a\t\nb", "a\t\tb", ""],
             }
         )
 
@@ -392,7 +392,7 @@ def my_asset(context):
         )
         with self.connect_mssql() as connection:
             data = connection.exec_driver_sql(f'SELECT * FROM {schema}.{table}')
-            assert data.fetchall() == [(1, 'a\t\nb'), (1, 'a\t\tb')]
+            assert data.fetchall() == [(1, 'a\t\nb'), (1, 'a\t\tb'), (1, None)]
 
     def test_absent_identity(self):
         schema = "test_polars_bcp_schema"

diff --git a/src/dagster_mssql_bcp/bcp_polars/polars_mssql_bcp.py b/src/dagster_mssql_bcp/bcp_polars/polars_mssql_bcp.py
@@ -97,7 +97,7 @@ def _replace_values(self, data: pl.LazyFrame, asset_schema: AssetSchema):
             ]
         )
 
-        # data = data.with_columns(pl.col(pl.String).replace("", None))
+        data = data.with_columns(pl.col(pl.String).replace("", None))
 
         return data
-Original file line number
+Diff line change
@@ Expand Up @@
                 ]
             )
-            # data = data.with_columns(pl.col(pl.String).replace("", None))
+            data = data.with_columns(pl.col(pl.String).replace("", None))
             return data
@@ Expand Down @@