chalkpy 2.89.22__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/df/ast_parser.py

@@ -93,9 +93,7 @@ def parse_dataframe_getitem():
     )
     assert isinstance(func_node, ast.Subscript)
     slc = func_node.slice
-    if isinstance(slc, ast.Index):
-        slc = slc.value  # type: ignore
-        assert isinstance(slc, ast.expr)
+    assert isinstance(slc, ast.expr)
     converted_slice = convert_slice(slc)
     return eval_converted_expr(converted_slice, glbs=func_frame.f_globals, lcls=func_frame.f_locals)
 
@@ -227,13 +225,7 @@ def _convert_maybe_tuple(slc: ast.expr):
         return _convert_ops(slc)
 
 
-def convert_slice(slc: Union[ast.expr, ast.Index]):
-    if isinstance(slc, ast.Index):
-        # Index is deprecated in Python 3.9+
-        slc = slc.value  # type: ignore
-        assert isinstance(slc, ast.expr)
-        slc = _convert_maybe_tuple(slc)
-        return ast.Index(value=slc)  # pyright: ignore[reportCallIssue]
+def convert_slice(slc: ast.expr):
     return _convert_maybe_tuple(slc)
 
 

chalk/features/_class_property.py

@@ -5,6 +5,8 @@ import functools
 from typing import Any, Callable, List, Type, TypeVar, cast
 
 from chalk._lsp.error_builder import FeatureClassErrorBuilder
+from chalk.features.feature_wrapper import UnresolvedFeature
+from chalk.utils.notebook import is_notebook
 
 T = TypeVar("T")
 V = TypeVar("V")
@@ -54,6 +56,11 @@ def classproperty_support(cls: Type[T]) -> Type[T]:
             if (res := self.__chalk_notebook_feature_expressions__.get(item)) is not None:
                 return res
 
+            # If in notebook, fallback to constructing FQN string instead of raising error
+            if is_notebook():
+                fqn = f"{self.namespace}.{item}"
+                return UnresolvedFeature(fqn)
+
             builder: FeatureClassErrorBuilder = self.__chalk_error_builder__
             builder.invalid_attribute(
                 root_feature_str=self.namespace,

chalk/features/_embedding/embedding.py

@@ -25,6 +25,7 @@ from chalk.utils.collections import ensure_tuple
 SUPPORTED_LOCAL_MODELS = {
     "all-MiniLM-L6-v2",  # https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
     "sample-bert",  # For internal Chalk use
+    "sample-linear-nn",  # For internal Chalk use
 }
 
 # This will eventually be included in SUPPORTED_LOCAL_MODELS

chalk/features/_embedding/sentence_transformer.py

@@ -111,7 +111,7 @@ class SentenceTransformerProvider(EmbeddingProvider):
                     raise ValueError(
                         f"Expected to find an embedding for input at position {idx}, but the response data was exhausted."
                     )
-            yield create_fixedsize_with_nulls(response, self.dimensions)
+            yield create_fixedsize_with_nulls(values_with_nulls, self.dimensions)
 
     def get_vector_class(self) -> Type[Vector]:
         return Vector[self.dimensions]

chalk/features/_encoding/converter.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import io
 import json
 import types
 import typing
@@ -50,6 +51,7 @@ from chalk.features._encoding.pyarrow import (
     rich_to_pyarrow,
 )
 from chalk.features._encoding.rich import structure_primitive_to_rich, unstructure_rich_to_primitive
+from chalk.features.feature_wrapper import UnresolvedFeature
 from chalk.utils.collections import unwrap_annotated_if_needed, unwrap_optional_and_annotated_if_needed
 from chalk.utils.df_utils import pa_array_to_pl_series
 from chalk.utils.json import JSON, TJSON, is_pyarrow_json_type, pyarrow_json_type
@@ -870,6 +872,79 @@ class PrimitiveFeatureConverter(Generic[_TPrim]):
         else:
             raise TypeError(f"Could not convert the pyarrow dtype {dtype} to a protobuf message")
 
+    @classmethod
+    def convert_pa_field_to_proto_field(cls, field: pa.Field) -> pb.Field:
+        """Convert a PyArrow Field to proto Field."""
+        field_proto = pb.Field(
+            name=field.name, arrow_type=cls.convert_pa_dtype_to_proto_dtype(field.type), nullable=field.nullable
+        )
+
+        if field.metadata:
+            # field.metadata is of types dict[bytes, bytes]
+            for k, v in field.metadata.items():
+                field_proto.metadata[k.decode("utf-8")] = v.decode("utf-8")
+
+        return field_proto
+
+    @classmethod
+    def convert_proto_field_to_pa_field(cls, proto_field: pb.Field) -> pa.Field:
+        """Convert a proto Field to PyArrow Field."""
+        arrow_type = cls.convert_proto_dtype_to_pa_dtype(proto_field.arrow_type)
+
+        # don't have to convert back to dict[bytes, bytes] as can initialize with dict[str, str]
+        metadata = dict(proto_field.metadata) if proto_field.metadata else None
+
+        return pa.field(
+            name=proto_field.name,
+            type=arrow_type,
+            nullable=proto_field.nullable,
+            metadata=metadata,
+        )
+
+    @classmethod
+    def convert_pa_schema_to_proto_schema(cls, schema: pa.Schema) -> pb.Schema:
+        schema_proto = pb.Schema(
+            columns=[cls.convert_pa_field_to_proto_field(field) for field in schema],
+        )
+
+        if schema.metadata:
+            # schema.metadata is of types dict[bytes, bytes]
+            for k, v in schema.metadata.items():
+                schema_proto.metadata[k.decode("utf-8")] = v.decode("utf-8")
+
+        return schema_proto
+
+    @classmethod
+    def convert_proto_schema_to_pa_schema(cls, proto_schema: pb.Schema) -> pa.Schema:
+        fields = [cls.convert_proto_field_to_pa_field(proto_field) for proto_field in proto_schema.columns]
+
+        # don't have to convert back to dict[bytes, bytes] as can initialize with dict[str, str]
+        metadata = dict(proto_schema.metadata) if proto_schema.metadata else None
+
+        return pa.schema(fields, metadata=metadata)
+
+    @staticmethod
+    def convert_arrow_table_to_proto(table: pa.Table | pa.RecordBatch) -> pb.TableParquetBytes:
+        if isinstance(table, pa.RecordBatch):
+            table = pa.Table.from_batches([table])
+        elif isinstance(table, pa.Table):
+            pass
+        else:
+            raise TypeError(f"expected pa.Table or pa.RecordBatch, got {type(table)!r}")
+
+        sink = io.BytesIO()
+        import pyarrow.parquet
+
+        pyarrow.parquet.write_table(table, sink)
+        return pb.TableParquetBytes(encoded_parquet_bytes=sink.getvalue())
+
+    @staticmethod
+    def convert_arrow_table_from_proto(proto: pb.TableParquetBytes) -> pa.Table:
+        import pyarrow.parquet
+
+        pf = pyarrow.parquet.ParquetFile(io.BytesIO(proto.encoded_parquet_bytes))
+        return pyarrow.parquet.read_table(pf)
+
     @staticmethod
     def _serialize_pa_decimal_to_pb(value: Union[pa.Decimal128Scalar, pa.Decimal256Scalar]) -> pb.ScalarValue:
         dec_val = value.as_py()
@@ -1183,8 +1258,14 @@ class FeatureConverter(PrimitiveFeatureConverter[_TPrim], Generic[_TPrim, _TRich
         # because it is also used for error handling inside of `from_rich_to_primitive`.
         self._name = name
         if rich_default != ...:
-            # The missing value strategy doesn't really matter because rich_default is not missing
-            primitive_default = self.from_rich_to_primitive(rich_default, missing_value_strategy="allow")
+            # In notebook environments, UnresolvedFeature may be used as a placeholder
+            # for features that can't be resolved due to a stale registry.
+            # Treat these as missing defaults since they're not concrete values.
+            if isinstance(rich_default, UnresolvedFeature):
+                rich_default = ...
+            else:
+                # The missing value strategy doesn't really matter because rich_default is not missing
+                primitive_default = self.from_rich_to_primitive(rich_default, missing_value_strategy="allow")
         super().__init__(
             name, is_nullable=is_nullable, pyarrow_dtype=pyarrow_dtype, primitive_default=primitive_default
         )

chalk/features/_encoding/pyarrow.py

@@ -8,12 +8,28 @@ import ipaddress
 import typing
 import uuid
 from datetime import date, datetime, time, timedelta
-from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Mapping, Set, Tuple, Type, cast
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Dict,
+    FrozenSet,
+    List,
+    Literal,
+    Mapping,
+    Set,
+    Tuple,
+    Type,
+    TypeGuard,
+    cast,
+    get_args,
+    get_origin,
+    is_typeddict,
+)
 
 import attrs
 import google.protobuf.message
 import pyarrow as pa
-from typing_extensions import Annotated, Literal, TypeGuard, get_args, get_origin, is_typeddict
 
 from chalk.features._encoding.http import HttpResponse, get_http_response_as_pyarrow
 from chalk.features._encoding.primitive import ChalkStructType, TPrimitive
@@ -24,7 +40,7 @@ from chalk.utils.collections import is_namedtuple, is_optional, unwrap_optional_
 from chalk.utils.enum import get_enum_value_type
 from chalk.utils.json import JSON, is_pyarrow_json_type
 from chalk.utils.missing_dependency import missing_dependency_exception
-from chalk.utils.pl_helpers import is_new_polars
+from chalk.utils.pl_helpers import is_new_polars, pl_array
 from chalk.utils.pydanticutil.pydantic_compat import is_pydantic_basemodel
 
 if TYPE_CHECKING:
@@ -418,7 +434,7 @@ def pyarrow_to_polars(
         underlying = pa_type.value_type
         if is_new_polars and use_fixed_size_list:
             # pl.Array is only available in polars >=0.18
-            return pl.Array(inner=pyarrow_to_polars(underlying, name=f"{name}[]"), width=pa_type.list_size)
+            return pl_array(inner=pyarrow_to_polars(underlying, name=f"{name}[]"), size=pa_type.list_size)
         else:
             return pl.List(pyarrow_to_polars(underlying, name=f"{name}[]"))
     if pa.types.is_struct(pa_type):

chalk/features/_encoding/rich.py

@@ -8,7 +8,7 @@ import enum
 import ipaddress
 import uuid
 from datetime import date, datetime, time, timedelta
-from typing import Any, FrozenSet, List, Set, Tuple, Type, TypeVar, Union, cast
+from typing import Any, FrozenSet, List, Set, Tuple, Type, TypeVar, Union, cast, get_args, get_origin, is_typeddict
 
 import attrs
 import cattrs
@@ -23,8 +23,6 @@ try:
 except ImportError:
     V1BaseModel = None
 
-from typing_extensions import get_args, get_origin, is_typeddict
-
 from chalk.features._encoding.primitive import ChalkStructType, TPrimitive
 from chalk.utils.cached_type_hints import cached_get_type_hints
 from chalk.utils.collections import is_namedtuple, unwrap_optional_and_annotated_if_needed

chalk/features/_tensor.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Any, Tuple, Type, Union, overload
+from typing import Any, Tuple, Type, TypeGuard, Union, overload
 
 import numpy as np
 import pyarrow as pa
-from typing_extensions import TypeGuard
 
 TensorDimension = Union[int, str]
 

chalk/features/dataframe/_filters.py

@@ -4,16 +4,17 @@ import collections.abc
 import datetime
 import enum
 import functools
-from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Protocol, Sequence, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Protocol, Sequence, TypeGuard, TypeVar, Union, cast
 
 import pyarrow as pa
-from typing_extensions import Self, TypeGuard
+from typing_extensions import Self
 
 from chalk.features._encoding.converter import pyarrow_to_polars
 from chalk.features.feature_field import Feature
 from chalk.features.feature_wrapper import FeatureWrapper, unwrap_feature
 from chalk.features.filter import Filter, TimeDelta, get_filter_now
 from chalk.utils.collections import ensure_tuple
+from chalk.utils.pl_helpers import polars_lazy_frame_collect_schema
 
 if TYPE_CHECKING:
     import polars as pl
@@ -442,7 +443,7 @@ class _PolarsStructAdapter(StructAdapter["pl.Expr"]):
 
 def filter_data_frame(
     item: Any,
-    underlying: Union[pl.DataFrame, pl.LazyFrame],
+    underlying: pl.LazyFrame,
     namespace: Optional[str],
 ) -> Union[pl.DataFrame, pl.LazyFrame]:
 
@@ -463,7 +464,10 @@ def filter_data_frame(
             )
     now = get_filter_now()
     if len(projections) > 0:
-        key_error_or_none = dataframe_missing_key_error(projections, underlying.columns)
+        key_error_or_none = dataframe_missing_key_error(
+            projections,
+            (underlying.collect_schema().names() if polars_lazy_frame_collect_schema else underlying.columns),
+        )
         if key_error_or_none is not None:
             raise key_error_or_none
     # now = datetime.datetime.now(tz=datetime.timezone.utc)
@@ -472,7 +476,12 @@ def filter_data_frame(
     timestamp_feature = (
         None if namespace is None else CURRENT_FEATURE_REGISTRY.get().get_feature_sets()[namespace].__chalk_ts__
     )
-    pl_expr = convert_filters_to_pl_expr(filters, underlying.schema, timestamp_feature, now)
+    pl_expr = convert_filters_to_pl_expr(
+        filters,
+        (underlying.collect_schema() if polars_lazy_frame_collect_schema else underlying.schema),
+        timestamp_feature,
+        now,
+    )
     df = underlying
     if pl_expr is not None:
         df = df.filter(pl_expr)

chalk/features/dataframe/_impl.py

@@ -32,7 +32,6 @@ from typing import (
     overload,
 )
 
-import packaging.version
 import pyarrow as pa
 
 from chalk.features._chalkop import Aggregation
@@ -56,6 +55,13 @@ from chalk.utils.df_utils import (
 )
 from chalk.utils.duration import Duration, parse_chalk_duration
 from chalk.utils.missing_dependency import missing_dependency_exception
+from chalk.utils.pl_helpers import (
+    polars_group_by_instead_of_groupby,
+    polars_lazy_frame_collect_schema,
+    polars_name_dot_suffix_instead_of_suffix,
+    polars_uses_schema_overrides,
+    schema_compat,
+)
 from chalk.utils.pydanticutil.pydantic_compat import is_pydantic_basemodel
 
 if TYPE_CHECKING:
@@ -473,7 +479,12 @@ class DataFrame(metaclass=DataFrameMeta):
             raise ValueError(f"Unable to convert data of type {type(data).__name__} into a DataFrame")
         # Rename / validate that all column names are root fqns
         if self._pydantic_model is None:
-            self.columns = tuple(Feature.from_root_fqn(str(c)) for c in underlying.columns)
+            self.columns = tuple(
+                Feature.from_root_fqn(str(c))
+                for c in (
+                    underlying.collect_schema().names() if polars_lazy_frame_collect_schema else underlying.columns
+                )
+            )
         else:
             self.columns = ()
 
@@ -505,7 +516,13 @@ class DataFrame(metaclass=DataFrameMeta):
         """
         import polars as pl
 
-        rename_map = {x: Distance.fqn for x in underlying.columns if x in self._distance_feature_fqns}
+        rename_map = {
+            x: Distance.fqn
+            for x in (
+                underlying.collect_schema().names() if polars_lazy_frame_collect_schema else underlying.columns
+            )  # pyright: ignore
+            if x in self._distance_feature_fqns
+        }
 
         underlying = underlying.rename(rename_map)
         if len(rename_map) > 0:
@@ -620,7 +637,14 @@ class DataFrame(metaclass=DataFrameMeta):
             )
         elif all(isinstance(col, str) for col in ensure_tuple(item)):
             # Select the columns with `.select()` since they're by name.
-            key_error_or_none = dataframe_missing_key_error(ensure_tuple(item), self._underlying.columns)
+            key_error_or_none = dataframe_missing_key_error(
+                ensure_tuple(item),
+                (
+                    self._underlying.collect_schema().names()
+                    if polars_lazy_frame_collect_schema
+                    else self._underlying.columns
+                ),
+            )
             if key_error_or_none is not None:
                 raise key_error_or_none
             materialized = self._materialize()
@@ -700,7 +724,7 @@ class DataFrame(metaclass=DataFrameMeta):
             if len(operation.filters) > 0:
                 f = convert_filters_to_pl_expr(
                     operation.filters,
-                    self._underlying.schema,
+                    schema_compat(self._underlying),
                     timestamp_feature,
                     now,
                 )
@@ -711,10 +735,10 @@ class DataFrame(metaclass=DataFrameMeta):
 
         data = self._underlying.lazy()
 
-        if packaging.version.parse(pl.__version__) <= packaging.version.parse("0.19.0"):
-            data = data.groupby(groupby)
-        else:
+        if polars_group_by_instead_of_groupby:
             data = data.group_by(groupby)
+        else:
+            data = data.groupby(groupby)  # pyright: ignore
 
         data = data.agg(cols).collect()
 
@@ -778,7 +802,7 @@ class DataFrame(metaclass=DataFrameMeta):
 
         col_str = str(column)
 
-        col_dtype = self._underlying.schema[col_str]
+        col_dtype = schema_compat(self._underlying)[col_str]
         underlying = self._underlying
         if col_dtype != pl.Float64() and col_dtype != pl.Float32():
             underlying = underlying.select(pl.col(col_str).cast(pl.Float32))
@@ -965,7 +989,7 @@ class DataFrame(metaclass=DataFrameMeta):
             if len(operation.filters) > 0:
                 f = convert_filters_to_pl_expr(
                     operation.filters,
-                    self._underlying.schema,
+                    schema_compat(self._underlying),
                     timestamp_feature,
                     now,
                 )
@@ -974,22 +998,40 @@ class DataFrame(metaclass=DataFrameMeta):
 
             cols.append(operation.fn(c).alias(str(alias)))
 
-        return DataFrame(
-            self._underlying.lazy()
-            .sort(str(index), descending=False)
-            .groupby_dynamic(
-                index_column=str(index),
-                by=groupby,
-                offset=offset,
-                every=every,
-                period=period,
-                start_by=start_by,
+        if polars_group_by_instead_of_groupby:
+            return DataFrame(
+                self._underlying.lazy()
+                .sort(str(index), descending=False)
+                .group_by_dynamic(
+                    index_column=str(index),
+                    group_by=groupby,
+                    offset=offset,
+                    every=every,
+                    period=period,
+                    start_by=start_by,
+                )
+                .agg(cols)
+                .collect(),
+                convert_dtypes=self._convert_dtypes,
+                pydantic_model=self._pydantic_model,
+            )
+        else:
+            return DataFrame(
+                self._underlying.lazy()
+                .sort(str(index), descending=False)
+                .groupby_dynamic(  # pyright: ignore
+                    index_column=str(index),
+                    by=groupby,
+                    offset=offset,
+                    every=every,
+                    period=period,
+                    start_by=start_by,
+                )
+                .agg(cols)
+                .collect(),
+                convert_dtypes=self._convert_dtypes,
+                pydantic_model=self._pydantic_model,
             )
-            .agg(cols)
-            .collect(),
-            convert_dtypes=self._convert_dtypes,
-            pydantic_model=self._pydantic_model,
-        )
 
     def join(
         self,
@@ -1200,7 +1242,7 @@ class DataFrame(metaclass=DataFrameMeta):
                 if len(operation.filters) > 0:
                     f = convert_filters_to_pl_expr(
                         operation.filters,
-                        self._underlying.schema,
+                        schema_compat(self._underlying),
                         timestamp_feature,
                         now,
                     )
@@ -1424,15 +1466,25 @@ class DataFrame(metaclass=DataFrameMeta):
         else:
             cols_to_select, dtypes, new_columns = cls._parse_columns(columns)
 
-        # 'dtypes' deprecated for 'schema_overrides' in polars 0.20+, but parameter renamed without breaking
-        data = pl.read_csv(
-            source=path,
-            has_header=has_header,
-            columns=cols_to_select,
-            dtypes=dtypes,  # pyright: ignore[reportCallIssue]
-            new_columns=new_columns,
-            storage_options=DataFrame._get_storage_options(),
-        )
+        # 'dtypes' deprecated for 'schema_overrides' in polars 0.20.31+
+        if polars_uses_schema_overrides:
+            data = pl.read_csv(
+                source=path,
+                has_header=has_header,
+                columns=cols_to_select,
+                schema_overrides=dtypes,  # pyright: ignore[reportCallIssue]
+                new_columns=new_columns,
+                storage_options=DataFrame._get_storage_options(),
+            )
+        else:
+            data = pl.read_csv(
+                source=path,
+                has_header=has_header,
+                columns=cols_to_select,
+                dtypes=dtypes,  # pyright: ignore[reportCallIssue]
+                new_columns=new_columns,
+                storage_options=DataFrame._get_storage_options(),
+            )
         return cls(data)
 
     @classmethod
@@ -1922,7 +1974,10 @@ class DataFrame(metaclass=DataFrameMeta):
             raise ValueError("DataFrame dimensions do not match")
 
         suffix = "__POLARS_CMP_OTHER"
-        other_renamed = other.select(pl.all().suffix(suffix))
+        if polars_name_dot_suffix_instead_of_suffix:
+            other_renamed = other.select(pl.all().name.suffix(suffix))
+        else:
+            other_renamed = other.select(pl.all().suffix(suffix))  # pyright: ignore
         combined = pl.concat([materialized, other_renamed], how="horizontal")
 
         if op == "eq":

chalk/features/dataframe/_validation.py

@@ -8,6 +8,7 @@ import isodate
 from chalk.features._encoding.missing_value import MissingValueStrategy
 from chalk.features.feature_field import Feature, FeatureNotFoundException
 from chalk.utils.collections import get_unique_item
+from chalk.utils.pl_helpers import apply_compat, schema_compat, str_json_decode_compat
 
 if TYPE_CHECKING:
     import polars as pl
@@ -67,7 +68,7 @@ def validate_df_schema(underlying: Union[pl.DataFrame, pl.LazyFrame]):
     # This is called from within DataFrame.__init__, which validates that polars is installed
     import polars as pl
 
-    for root_fqn, actual_dtype in underlying.schema.items():
+    for root_fqn, actual_dtype in schema_compat(underlying).items():
         feature = Feature.from_root_fqn(root_fqn)
         if feature.is_has_one or feature.is_has_many:
             continue
@@ -87,7 +88,7 @@ def validate_df_schema(underlying: Union[pl.DataFrame, pl.LazyFrame]):
             isinstance(expected_dtype, pl.List)
             and actual_dtype == pl.Utf8  # pyright: ignore[reportUnnecessaryComparison]
         ):
-            col = pl.col(root_fqn).str.json_extract(expected_dtype)
+            col = str_json_decode_compat(pl.col(root_fqn), expected_dtype)
             try:
                 underlying = underlying.with_columns(col.cast(expected_dtype))
             except (Exception, pl.PolarsPanicError) as e:
@@ -123,21 +124,24 @@ def validate_df_schema(underlying: Union[pl.DataFrame, pl.LazyFrame]):
                     if isinstance(expected_dtype, pl.Datetime):
                         # tzinfo = None if expected_dtype.time_zone is None else zoneinfo.ZoneInfo(expected_dtype.time_zone)
                         underlying = underlying.with_columns(pl.col(root_fqn).str.strptime(pl.Datetime).alias(root_fqn))
-                        if cast(pl.Datetime, underlying.schema[root_fqn]).time_zone is not None:
+                        if cast(pl.Datetime, schema_compat(underlying)[root_fqn]).time_zone is not None:
                             assert expected_dtype.time_zone is not None
                             cast_expr = pl.col(root_fqn).dt.convert_time_zone(expected_dtype.time_zone)
                         else:
                             cast_expr = pl.col(root_fqn).dt.replace_time_zone(expected_dtype.time_zone)
                     elif expected_dtype == pl.Date:
-                        cast_expr = pl.col(root_fqn).apply(
+                        cast_expr = apply_compat(
+                            pl.col(root_fqn),
                             lambda x: None if x is None else isodate.parse_date(x),
                         )
                     elif expected_dtype == pl.Time:
-                        cast_expr = pl.col(root_fqn).apply(
+                        cast_expr = apply_compat(
+                            pl.col(root_fqn),
                             lambda x: None if x is None else isodate.parse_time(x),
                         )
                     elif expected_dtype == pl.Duration:
-                        cast_expr = pl.col(root_fqn).apply(
+                        cast_expr = apply_compat(
+                            pl.col(root_fqn),
                             lambda x: None if x is None else isodate.parse_duration(x),
                         )
                     else:
@@ -168,7 +172,7 @@ def validate_nulls(
 
     if isinstance(underlying, pl.LazyFrame):
         underlying = underlying.collect()
-    schema = underlying.schema
+    schema = schema_compat(underlying)
     null_count_rows = underlying.null_count().to_dicts()
     if len(null_count_rows) == 0:
         return underlying  # Empty dataframe