Lazily import arrow types and models

ccflow/__init__.py

@@ -5,7 +5,6 @@
 # which, in turn, import `ccflow`).
 from .exttypes import *  # noqa: I001
 
-from .arrow import *
 from .base import *
 from .compose import *
 from .callable import *
@@ -19,3 +18,29 @@
 from .publisher import *
 from .result import *
 from .utils import FileHandler, StreamHandler
+
+# Arrow models and types import pyarrow, so load them on first access.
+_LAZY_ARROW = {
+    "ArrowDateFilter": "arrow",
+    "ArrowDatetimeFilter": "arrow",
+    "ArrowFileSystem": "arrow",
+    "ArrowFilter": "arrow",
+    "ArrowLocalFileSystem": "arrow",
+    "ArrowPartitioning": "arrow",
+    "ArrowS3FileSystem": "arrow",
+    "ArrowSchemaModel": "arrow",
+    "ArrowTemplateFilter": "arrow",
+    "ArrowTimeFilter": "arrow",
+    "render_filters": "arrow",
+    "ArrowSchema": "exttypes.arrow",
+    "ArrowTable": "exttypes.arrow",
+    "PyArrowDatatype": "exttypes.arrow",
+}
+
+
+def __getattr__(name):
+    if name in _LAZY_ARROW:
+        from importlib import import_module
+
+        return getattr(import_module(f".{_LAZY_ARROW[name]}", __name__), name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

ccflow/examples/narwhals_pipelines.ipynb

@@ -192,7 +192,7 @@
    ],
    "source": [
     "def q1_monolith(lineitem):\n",
-    "    cutoff = datetime(1998, 9, 2)  # noqa: DTZ001\n",
+    "    cutoff = datetime(1998, 9, 2)\n",
     "    return (\n",
     "        lineitem.filter(nw.col(\"l_shipdate\") <= cutoff)\n",
     "        .with_columns(\n",
@@ -370,7 +370,7 @@
     "        return df.filter(nw.col(\"l_shipdate\") <= self.cutoff)\n",
     "\n",
     "\n",
-    "filtered = lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2)))  # noqa: DTZ001\n",
+    "filtered = lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2)))\n",
     "filtered.collect().shape"
    ]
   },
@@ -434,7 +434,7 @@
     "        )\n",
     "\n",
     "\n",
-    "lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2))).pipe(DeriveLineItemMetrics()).collect().columns  # noqa: DTZ001"
+    "lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2))).pipe(DeriveLineItemMetrics()).collect().columns"
    ]
   },
   {
@@ -541,7 +541,7 @@
     }
    ],
    "source": [
-    "result = lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2))).pipe(DeriveLineItemMetrics()).pipe(SummarizeByReturnStatus())  # noqa: DTZ001\n",
+    "result = lineitem.pipe(FilterByShipDate(cutoff=datetime(1998, 9, 2))).pipe(DeriveLineItemMetrics()).pipe(SummarizeByReturnStatus())\n",
     "result.collect().to_native()"
    ]
   },
@@ -710,7 +710,7 @@
    "source": [
     "q1_preproc = SequenceTransform(\n",
     "    transforms=[\n",
-    "        FilterByShipDate(cutoff=datetime(1998, 9, 2)),  # noqa: DTZ001\n",
+    "        FilterByShipDate(cutoff=datetime(1998, 9, 2)),\n",
     "        DeriveLineItemMetrics(),\n",
     "        SummarizeByReturnStatus(),\n",
     "    ],\n",
@@ -860,7 +860,7 @@
     "q1_pipeline = NarwhalsPipelineModel(\n",
     "    source=registry[\"/table/lineitem\"],\n",
     "    transforms=[\n",
-    "        FilterByShipDate(cutoff=datetime(1998, 9, 2)),  # noqa: DTZ001\n",
+    "        FilterByShipDate(cutoff=datetime(1998, 9, 2)),\n",
     "        DeriveLineItemMetrics(),\n",
     "        SummarizeByReturnStatus(),\n",
     "    ],\n",
@@ -1133,7 +1133,7 @@
     "        )\n",
     "\n",
     "\n",
-    "cutoff = datetime(1995, 3, 15)  # noqa: DTZ001\n",
+    "cutoff = datetime(1995, 3, 15)\n",
     "\n",
     "q3_pipeline = NarwhalsPipelineModel(\n",
     "    source=registry[\"/table/customer\"],\n",

ccflow/exttypes/__init__.py

@@ -1,8 +1,18 @@
-from .arrow import *
 from .exprtk import *
 from .frequency import *
 from .jinja import *
 
 # Do NOT import .polars. We don't want ccflow (without flow) to have a dependency on polars!
 from .pydantic_numpy import *
 from .pyobjectpath import *
+
+_LAZY_ARROW = ("ArrowSchema", "ArrowTable", "PyArrowDatatype")
+
+
+def __getattr__(name):
+    # Defer pyarrow until an arrow type is used.
+    if name in _LAZY_ARROW:
+        from . import arrow
+
+        return getattr(arrow, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

ccflow/tests/models/test_narwhals.py

@@ -243,7 +243,7 @@ def test_basic_join(self):
         lf = nw.from_native(pl.LazyFrame({"x": [1, 2, 3]}))
         right_src = FrameSource(data={"x": [1, 2, 3], "y": ["a", "b", "c"]})
         jt = JoinTransform(other=right_src, on="x", how="left")
-        out = lf.pipe(jt).collect().to_native()
+        out = lf.pipe(jt).collect().to_native().sort("x")
         assert out.columns == ["x", "y"]
         assert out["y"].to_list() == ["a", "b", "c"]
 
@@ -255,7 +255,7 @@ def test_json_roundtrip(self):
         jt = JoinTransform(other=right_src, on="x", how="left", suffix="_r")
         jt2 = JoinTransform.model_validate_json(jt.model_dump_json())
         lf = nw.from_native(pl.LazyFrame({"x": [1, 2, 3]}))
-        assert lf.pipe(jt).collect().to_native().equals(lf.pipe(jt2).collect().to_native())
+        assert lf.pipe(jt).collect().to_native().sort("x").equals(lf.pipe(jt2).collect().to_native().sort("x"))
 
     def test_rejects_non_narwhals_other(self):
         class BadSrc(CallableModel):
@@ -271,15 +271,15 @@ def test_in_pipeline(self):
         main = FrameSource(data={"x": [1, 2, 3], "v": [10, 20, 30]})
         side = FrameSource(data={"x": [1, 2, 3], "label": ["a", "b", "c"]})
         p = NarwhalsPipelineModel(source=main, transforms=[JoinTransform(other=side, on="x")])
-        out = p(NullContext()).df.collect().to_native()
+        out = p(NullContext()).df.collect().to_native().sort("x")
         assert out.columns == ["x", "v", "label"]
         assert out["label"].to_list() == ["a", "b", "c"]
 
     def test_left_on_right_on(self):
         lf = nw.from_native(pl.LazyFrame({"a": [1, 2, 3]}))
         right_src = FrameSource(data={"b": [1, 2, 3], "y": ["a", "b", "c"]})
         jt = JoinTransform(other=right_src, left_on="a", right_on="b", how="inner")
-        out = lf.pipe(jt).collect().to_native()
+        out = lf.pipe(jt).collect().to_native().sort("a")
         assert out["y"].to_list() == ["a", "b", "c"]
 
     def test_cross_join(self):
@@ -318,7 +318,7 @@ def __call__(self, context: TableCtx) -> NarwhalsFrameResult:
         src = TableSource(tables={"r": {"x": [1, 2], "y": ["a", "b"]}})
         lf = nw.from_native(pl.LazyFrame({"x": [1, 2]}))
         jt = JoinTransform(other=src, other_context=TableCtx(table="r"), on="x")
-        out = lf.pipe(jt).collect().to_native()
+        out = lf.pipe(jt).collect().to_native().sort("x")
         assert out["y"].to_list() == ["a", "b"]
 
     def test_rejects_other_context_type_mismatch(self):
@@ -346,7 +346,7 @@ def __call__(self, df: nw.LazyFrame) -> nw.LazyFrame:
                 return df.select(nw.col("x"), (nw.col("v") * 100).alias("hund"))
 
         jb = JoinBackTransform(inner=Project(), on="x")
-        out = lf.pipe(jb).collect().to_native()
+        out = lf.pipe(jb).collect().to_native().sort("x")
         assert out.columns == ["x", "v", "hund"]
         assert out["hund"].to_list() == [1000, 2000, 3000]
 

ccflow/tests/test_context.py

@@ -102,7 +102,7 @@ def test_datetime_validation(self):
         dt = datetime(2022, 1, 1, 12, 0, tzinfo=timezone.utc)
         c = DatetimeContext(dt=dt)
         self.assertEqual(DatetimeContext(dt=str(dt)), c)
-        self.assertEqual(DatetimeContext(dt=dt.date()), DatetimeContext(dt=datetime(2022, 1, 1)))  # noqa: DTZ001
+        self.assertEqual(DatetimeContext(dt=dt.date()), DatetimeContext(dt=datetime(2022, 1, 1)))
 
         # Test validation from other types
         self.assertEqual(TypeAdapter(DatetimeContext).validate_python({"dt": dt}), c)

ccflow/tests/test_import.py

@@ -22,6 +22,7 @@ def test(self):
         "matplotlib",
         "mlflow",
         "plotly",
+        "pyarrow",
         "pyarrow.dataset",
         # These aren't necessarily expensive, just things we don't want to import.
         "cexprtk",

ccflow/tests/test_validators.py

@@ -18,16 +18,16 @@ def test_normalize_date(self):
         c1 = datetime.now().astimezone().date() - timedelta(1)
         self.assertEqual(normalize_date("-1d"), c1)
 
-        self.assertEqual(normalize_date(datetime.now()), c)  # noqa: DTZ005
-        self.assertEqual(normalize_date(datetime.now().isoformat()), c)  # noqa: DTZ005
+        self.assertEqual(normalize_date(datetime.now()), c)
+        self.assertEqual(normalize_date(datetime.now().isoformat()), c)
 
         self.assertEqual(normalize_date("foo"), "foo")
         self.assertEqual(normalize_date(None), None)
 
     def test_normalize_datetime(self):
         today = datetime.today()  # noqa: DTZ002
-        now = datetime.now()  # noqa: DTZ005
-        c = datetime(today.year, today.month, today.day)  # noqa: DTZ001
+        now = datetime.now()
+        c = datetime(today.year, today.month, today.day)
 
         self.assertEqual(normalize_datetime(c), c)
         self.assertEqual(normalize_datetime("0d"), c)

ccflow/tests/utils/test_tokenize.py

@@ -636,7 +636,7 @@ def test_primitives(self, value, expected):
         "value,expected",
         [
             (date(2024, 1, 15), ("date", "2024-01-15")),
-            (datetime(2024, 1, 15, 10, 30, 0), ("datetime", "2024-01-15T10:30:00")),  # noqa: DTZ001
+            (datetime(2024, 1, 15, 10, 30, 0), ("datetime", "2024-01-15T10:30:00")),
             (time(10, 30, 0), ("time", "10:30:00")),
             (timedelta(hours=1, minutes=30), ("timedelta", 0, 5400, 0)),
         ],