From aced3225de1c7fd50f014f2aa623b1472620520c Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli Date: Tue, 5 Nov 2024 15:19:55 +0000 Subject: [PATCH] perf: use Narwhals expressions (and stable API) (#528) * use Narwhals stable api, use expressions * just use .dt.convert_time_zone --- vegafusion-python/vegafusion/runtime.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/vegafusion-python/vegafusion/runtime.py b/vegafusion-python/vegafusion/runtime.py index 4cfe45eb..b915f65e 100644 --- a/vegafusion-python/vegafusion/runtime.py +++ b/vegafusion-python/vegafusion/runtime.py @@ -4,7 +4,7 @@ from types import ModuleType from typing import TYPE_CHECKING, Any, Literal, TypedDict, Union, cast -import narwhals as nw +import narwhals.stable.v1 as nw from arro3.core import Table from vegafusion._vegafusion import get_cpu_count, get_virtual_memory @@ -566,18 +566,15 @@ def pre_transform_datasets( "transformed data" ) - # Localize datetime columns to UTC, then extract the native DataFrame - # to return + # Convert to `local_tz` (or, set to UTC and then convert if starting + # from time-zone-naive data), then extract the native DataFrame to return. processed_datasets = [] for df in nw_dataframes: - for name in df.columns: - dtype = df[name].dtype - if dtype == nw.Datetime: - df = df.with_columns( - df[name] - .dt.replace_time_zone("UTC") - .dt.convert_time_zone(local_tz) - ) + df = df.with_columns( + nw.col(col).dt.convert_time_zone(local_tz) + for col, dtype in df.schema.items() + if dtype == nw.Datetime + ) processed_datasets.append(df.to_native()) return processed_datasets, warnings