Support TIMESTAMP WITH TIMEZONE and DECIMAL types with DuckDB connect…

…ion (#283) * Handle duckdb's TIMESTAMP WITH TIME ZONE type * Handle duckdb's DECIMAL type * Bump version to 1.1.1
vega · Mar 25, 2023 · d2d13d9 · d2d13d9
1 parent 2c88627
commit d2d13d9
Show file tree

Hide file tree

Showing 23 changed files with 123 additions and 55 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/automation/bump_version.py b/automation/bump_version.py
@@ -53,7 +53,6 @@ def bump_version(version):
     package_json_dirs = [
         root / "vegafusion-wasm",
         root / "javascript" / "vegafusion-embed",
-        root / "javascript" / "vegafusion-chart-editor",
         root / "python" / "vegafusion-jupyter"
     ]
     for package_json_dir in package_json_dirs:

diff --git a/javascript/vegafusion-embed/package-lock.json b/javascript/vegafusion-embed/package-lock.json
diff --git a/javascript/vegafusion-embed/package.json b/javascript/vegafusion-embed/package.json
@@ -1,6 +1,6 @@
 {
   "name": "vegafusion-embed",
-  "version": "1.1.0",
+  "version": "1.1.1",
   "description": "Library to embed vegafusion visualizations",
   "keywords": [
     "vega",

diff --git a/python/vegafusion-jupyter/package-lock.json b/python/vegafusion-jupyter/package-lock.json
diff --git a/python/vegafusion-jupyter/package.json b/python/vegafusion-jupyter/package.json
@@ -1,6 +1,6 @@
 {
   "name": "vegafusion-jupyter",
-  "version": "1.1.0",
+  "version": "1.1.1",
   "description": "Altair Jupyter Widget library that relies on VegaFusion for serverside calculations",
   "keywords": [
     "jupyter",

diff --git a/python/vegafusion-jupyter/vegafusion_jupyter/_frontend.py b/python/vegafusion-jupyter/vegafusion_jupyter/_frontend.py
@@ -2,4 +2,4 @@
 Information about the frontend package of the widgets.
 """    
 module_name = "vegafusion-jupyter"
-module_version = "^1.1.0"
+module_version = "^1.1.1"
diff --git a/python/vegafusion-jupyter/vegafusion_jupyter/_version.py b/python/vegafusion-jupyter/vegafusion_jupyter/_version.py
@@ -1 +1 @@
-__version__ = '1.1.0'
+__version__ = '1.1.1'
diff --git a/python/vegafusion/setup.cfg b/python/vegafusion/setup.cfg
@@ -4,7 +4,7 @@ universal = 0
 [metadata]
 name = vegafusion
 description = Core tools for using VegaFusion from Python
-version = 1.1.0
+version = 1.1.1
 long_description = file: README.md
 long_description_content_type = text/markdown
 keywords = vega, altair, vegafusion, arrow
@@ -34,6 +34,6 @@ install_requires =
 
 [options.extras_require]
 embed = 
-	vegafusion-python-embed==1.1.0
+	vegafusion-python-embed==1.1.1
 	vl-convert-python>=0.7.0
 
diff --git a/python/vegafusion/tests/test_pretransform.py b/python/vegafusion/tests/test_pretransform.py
@@ -5,6 +5,7 @@
 import json
 import polars as pl
 from datetime import date
+import decimal
 
 
 def order_items_spec():
@@ -1045,6 +1046,68 @@ def test_pre_transform_dataset_duckdb_conn():
         vf.runtime.set_connection("datafusion")
 
 
+def test_pre_transform_dataset_duckdb_with_decimal_conn():
+    import duckdb
+
+    n = 4050
+    # Input a polars DataFrame (which follows the DataFrame Interface Protocol)
+    order_items = pd.DataFrame({
+        "menu_item_int": [0] * n + [1] * (2 * n) + [2] * (3 * n)
+    })
+
+    try:
+        # Create duckdb connection and register order_items with duckdb
+        conn = duckdb.connect()
+        conn.register("order_items_int", order_items)
+        conn.query(
+            "SELECT menu_item_int::DECIMAL(12,2) as menu_item from order_items_int"
+        ).to_view("order_items")
+
+        # Set this as the active connection
+        vf.runtime.set_connection(conn)
+
+        # order_items includes a table://order_items data url
+        vega_spec = order_items_spec()
+        datasets, warnings = vf.runtime.pre_transform_datasets(
+            vega_spec,
+            ["data_0"],
+            "UTC",
+        )
+        assert len(warnings) == 0
+        assert len(datasets) == 1
+
+        result = datasets[0]
+        expected = pd.DataFrame({
+            "menu_item": [decimal.Decimal(0), decimal.Decimal(1), decimal.Decimal(2)],
+            "__count": [n, 2 * n, 3 * n]
+        })
+        pd.testing.assert_frame_equal(result, expected)
+    finally:
+        vf.runtime.set_connection("datafusion")
+
+
+def test_duckdb_timestamp_with_timezone():
+    try:
+        vf.runtime.set_connection("duckdb")
+        dates_df = pd.DataFrame({
+            "date_col": [date(2022, 1, 1), date(2022, 1, 2), date(2022, 1, 3)],
+        })
+        dates_df["date_col"] = pd.to_datetime(dates_df.date_col).dt.tz_localize("UTC")
+        spec = date_column_spec()
+
+        (output_ds,), warnings = vf.runtime.pre_transform_datasets(
+            spec, ["data_0"], "America/New_York", default_input_tz="UTC", inline_datasets=dict(dates=dates_df)
+        )
+
+        # Timestamps are in the local timezone, so they should be midnight local time
+        assert list(output_ds.date_col) == [
+            pd.Timestamp('2022-01-01 00:00:00', tz='UTC'),
+            pd.Timestamp('2022-01-02 00:00:00', tz='UTC'),
+            pd.Timestamp('2022-01-03 00:00:00', tz='UTC')
+        ]
+    finally:
+        vf.runtime.set_connection("datafusion")
+
 def test_gh_268_hang():
     """
     Tests for hang reported in https://github.com/hex-inc/vegafusion/issues/268

diff --git a/python/vegafusion/vegafusion/_version.py b/python/vegafusion/vegafusion/_version.py
@@ -1 +1 @@
-__version__ = '1.1.0'
+__version__ = '1.1.1'
diff --git a/python/vegafusion/vegafusion/connection/duckdb.py b/python/vegafusion/vegafusion/connection/duckdb.py
@@ -1,3 +1,4 @@
+import re
 import warnings
 
 from . import SqlConnection, CsvReadOptions
@@ -26,6 +27,11 @@ def duckdb_type_name_to_pyarrow_type(duckdb_type: str) -> pa.DataType:
         return pa.int32()
     elif duckdb_type in ("BIGINT", "INT8", "LONG"):
         return pa.int64()
+    elif duckdb_type.startswith("DECIMAL"):
+        matches = re.findall(r"\d+", duckdb_type)
+        precision = int(matches[0])
+        scale = int(matches[1])
+        return pa.decimal128(precision, scale)
     elif duckdb_type == "UTINYINT":
         return pa.uint8()
     elif duckdb_type == "USMALLINT":
@@ -34,14 +40,14 @@ def duckdb_type_name_to_pyarrow_type(duckdb_type: str) -> pa.DataType:
         return pa.uint32()
     elif duckdb_type == "UBIGINT":
         return pa.uint64()
-    elif duckdb_type == "DOUBLE":
-        return pa.float64()
     elif duckdb_type == "BOOLEAN":
         return pa.bool_()
     elif duckdb_type == "DATE":
         return pa.date32()
     elif duckdb_type == "TIMESTAMP":
         return pa.timestamp("ms")
+    elif duckdb_type == "TIMESTAMP WITH TIME ZONE":
+        return pa.timestamp("ms", tz="UTC")
     else:
         raise ValueError(f"Unexpected DuckDB type: {duckdb_type}")
 

diff --git a/vegafusion-common/Cargo.toml b/vegafusion-common/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "vegafusion-common"
-version = "1.1.0"
+version = "1.1.1"
 edition = "2021"
 description = "Common components required by multiple VegaFusion crates"
 license = "BSD-3-Clause"

diff --git a/vegafusion-core/Cargo.toml b/vegafusion-core/Cargo.toml
@@ -2,7 +2,7 @@
 name = "vegafusion-core"
 license = "BSD-3-Clause"
 edition = "2021"
-version = "1.1.0"
+version = "1.1.1"
 description = "Core components required by multiple VegaFusion crates, with WASM compatibility"
 
 [features]
@@ -35,7 +35,7 @@ features = [ "preserve_order",]
 [dependencies.vegafusion-common]
 path = "../vegafusion-common"
 features = [ "json", "sqlparser",]
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.datafusion-common]
 version = "18.0.0"

diff --git a/vegafusion-dataframe/Cargo.toml b/vegafusion-dataframe/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "vegafusion-dataframe"
 license = "BSD-3-Clause"
-version = "1.1.0"
+version = "1.1.1"
 edition = "2021"
 description = "VegaFusion's DataFrame and Connection traits"
 
@@ -10,7 +10,7 @@ async-trait = "0.1.53"
 
 [dependencies.vegafusion-common]
 path = "../vegafusion-common"
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.datafusion-common]
 version = "18.0.0"

diff --git a/vegafusion-datafusion-udfs/Cargo.toml b/vegafusion-datafusion-udfs/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "vegafusion-datafusion-udfs"
 license = "BSD-3-Clause"
-version = "1.1.0"
+version = "1.1.1"
 edition = "2021"
 description = "Custom DataFusion UDFs used by VegaFusion"
 
@@ -14,7 +14,7 @@ regex = "^1.5.5"
 
 [dependencies.vegafusion-common]
 path = "../vegafusion-common"
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.datafusion-physical-expr]
 version = "18.0.0"
diff --git a/vegafusion-python-embed/Cargo.toml b/vegafusion-python-embed/Cargo.toml
@@ -2,7 +2,7 @@
 name = "vegafusion-python-embed"
 license = "BSD-3-Clause"
 edition = "2021"
-version = "1.1.0"
+version = "1.1.1"
 description = "vegafusion-python-embed PyO3 Python Package"
 
 [lib]
@@ -36,21 +36,21 @@ features = [ "pyarrow",]
 [dependencies.vegafusion-common]
 path = "../vegafusion-common"
 features = [ "pyo3",]
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.vegafusion-core]
 path = "../vegafusion-core"
 features = [ "pyarrow",]
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.vegafusion-runtime]
 path = "../vegafusion-runtime"
 features = [ "pyarrow",]
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.vegafusion-sql]
 path = "../vegafusion-sql"
-version = "1.1.0"
+version = "1.1.1"
 features = [ "datafusion-conn",]
 
 [dependencies.tokio]

diff --git a/vegafusion-runtime/Cargo.toml b/vegafusion-runtime/Cargo.toml
@@ -6,7 +6,7 @@ harness = false
 name = "vegafusion-runtime"
 license = "BSD-3-Clause"
 edition = "2021"
-version = "1.1.0"
+version = "1.1.1"
 description = "VegaFusion Runtime"
 
 [features]
@@ -59,20 +59,20 @@ version = "0.30.0"
 [dependencies.vegafusion-common]
 path = "../vegafusion-common"
 features = [ "json", "sqlparser", "prettyprint",]
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.vegafusion-core]
 path = "../vegafusion-core"
 features = [ "sqlparser",]
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.vegafusion-datafusion-udfs]
 path = "../vegafusion-datafusion-udfs"
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.vegafusion-dataframe]
 path = "../vegafusion-dataframe"
-version = "1.1.0"
+version = "1.1.1"
 
 [dependencies.serde]
 version = "1.0.137"
@@ -110,5 +110,5 @@ features = [ "async_tokio",]
 
 [dev-dependencies.vegafusion-sql]
 path = "../vegafusion-sql"
-version = "1.1.0"
+version = "1.1.1"
 features = [ "datafusion-conn",]