fix and test more of Polars adapter

Author: JohnMount

data_algebra/data_model.py

@@ -72,6 +72,18 @@ def table_is_keyed_by_columns(self, table, *, column_names: Iterable[str]) -> bo
         :return: True if rows are uniquely keyed by values in named columns
         """
 
+    @abc.abstractmethod
+    def concat_rows(self, frame_list: List):
+        """
+        Concatenate rows from frame_list
+        """
+
+    @abc.abstractmethod
+    def concat_columns(self, frame_list):
+        """
+        Concatenate columns from frame_list
+        """
+
     # evaluate
 
     @abc.abstractmethod

data_algebra/pandas_base.py

@@ -456,14 +456,24 @@ def bad_column_positions(self, x):
                 self.pd.isnull(x), numpy.logical_or(numpy.isnan(x), numpy.isinf(x))
             )
         return self.pd.isnull(x)
+    
+    def concat_rows(self, frame_list: List):
+        """
+        Concatenate rows from frame_list
+        """
+        frame_list = list(frame_list)
+        assert len(frame_list) > 0
+        if len(frame_list) == 1:
+            return self.clean_copy(frame_list[0])
+        res = self.pd.concat(frame_list, axis=0)
+        return res
 
-    def concat_columns(self, frame_list):
+    def concat_columns(self, frame_list: List):
         """
-        Concatinate columns from frame_list
+        Concatenate columns from frame_list
         """
         frame_list = list(frame_list)
-        if len(frame_list) <= 0:
-            return None
+        assert len(frame_list) > 0
         if len(frame_list) == 1:
             return self.clean_copy(frame_list[0])
         res = self.pd.concat(frame_list, axis=1)

data_algebra/polars_model.py

@@ -289,7 +289,6 @@ def _populate_expr_impl_map() -> Dict[int, Dict[str, Callable]]:
         "%/%": lambda a, b: a / b,
         "around": lambda a, b: a.round(b),
         "coalesce": lambda a, b: pl.when(a.is_null()).then(b).otherwise(a),
-        "concat": lambda a, b: a.concat(b),
         "date_diff": lambda a, b: a.date_diff(b),
         "is_in": lambda a, b: a.is_in(b),
         "mod": lambda a, b: a % b,
@@ -367,6 +366,7 @@ def __init__(self, *, use_lazy_eval: bool = True):
         }
         self._expr_impl_map = _populate_expr_impl_map()
         self._impl_map_arbitrary_arity = {
+            "concat": lambda *args: pl.concat_str(args),
             "fmax": lambda *args: pl.max(args),
             "fmin": lambda *args: pl.min(args),
             "maximum": lambda *args: pl.max(args),
@@ -440,9 +440,19 @@ def bad_column_positions(self, x):
         """
         return x.is_null()
 
+    def concat_rows(self, frame_list: List):
+        """
+        Concatenate rows from frame_list
+        """
+        frame_list = list(frame_list)
+        assert len(frame_list) > 0
+        if len(frame_list) == 1:
+            return frame_list[0]
+        pl.concat(frame_list, how="vertical")
+
     def concat_columns(self, frame_list):
         """
-        Concatinate columns from frame_list
+        Concatenate columns from frame_list
         """
         frame_list = list(frame_list)
         if len(frame_list) <= 0:

data_algebra/solutions.py

@@ -224,7 +224,7 @@ def xicor_score_variables_plan(
     assert isinstance(n_rep, int)
     record_map = RecordMap(
         blocks_out=RecordSpecification(
-            control_table=data_algebra.data_model.default_data_model().pd.DataFrame(
+            control_table=data_algebra.data_model.default_data_model().data_frame(
                 {
                     "variable_name": x_vars,
                     "x": x_vars,
@@ -237,7 +237,7 @@ def xicor_score_variables_plan(
         ),
         strict=False,
     )
-    rep_frame = data_algebra.data_model.default_data_model().pd.DataFrame({"rep": range(n_rep)})
+    rep_frame = data_algebra.data_model.default_data_model().data_frame({"rep": range(n_rep)})
     grouped_calc = (
         xicor_query(
             d
@@ -529,13 +529,12 @@ def replicate_rows_query(
     assert power_key_colname not in d.column_names
     # get a pandas namespace
     local_data_model = data_algebra.data_model.default_data_model()
-    pd = local_data_model.pd
     # build powers of 2 until max_count is met or exceeded
     powers = list(range(int(numpy.ceil(numpy.log(max_count) / numpy.log(2))) + 1))
     # replicate each power the number of times it specifies
-    count_frame = pd.concat(
+    count_frame = local_data_model.concat_rows(
         [
-            pd.DataFrame(
+            local_data_model.data_frame(
                 {
                     power_key_colname: f"p{p}",
                     seq_column_name: range(int(2**p)),

tests/test_braid.py

@@ -12,11 +12,11 @@
 def test_braid():
     pd = data_algebra.data_model.default_data_model().pd
     d_state = pd.DataFrame({
-        't': [1, 3, 5],
+        't': [1., 3., 5.],
         'state': ['a', 'b', 'c'],
     })
     d_event = pd.DataFrame({
-        't': [1, 4],
+        't': [1., 4.],
         'value': [10, 20],
     })
     ops = data_algebra.solutions.braid_data(
@@ -31,7 +31,7 @@ def test_braid():
     res = ops.eval({'d_state': d_state, 'd_event': d_event})
     # print(data_algebra.util.pandas_to_example_str(res))
     expect = pd.DataFrame({
-        't': [1, 1, 3, 4, 5],
+        't': [1., 1., 3., 4., 5.],
         'state': ['a', 'a', 'b', 'b', 'c'],
         'value': [None, 10, None, 20, None],
         'record_type': ['state_row', 'event_row', 'state_row', 'event_row', 'state_row'],

tests/test_expand_rows.py

@@ -31,5 +31,4 @@ def test_replicate_rows_query():
         ops=ops,
         data={'d': d, 'rt': rt},
         expect=expect,
-        try_on_Polars=False,  # TODO: turn this on
     )

tests/test_idioms.py

@@ -389,7 +389,6 @@ def test_idiom_concat_op():
     )
 
     data_algebra.test_util.check_transform(ops=ops, data=d, expect=expect,
-        try_on_Polars=False,  # TODO: turn this on
     )