Add Set type test for On Demand Feature Views

Copilot · franciscojavierarceo · Copilot · commit ba4c68edea6d · 2026-01-22T18:19:35.000Z
Co-authored-by: franciscojavierarceo &lt;4163062+franciscojavierarceo@users.noreply.github.com&gt;
diff --git a/sdk/python/tests/unit/test_on_demand_python_transformation.py b/sdk/python/tests/unit/test_on_demand_python_transformation.py
@@ -33,6 +33,7 @@
     Float64,
     Int64,
     PdfBytes,
+    Set,
     String,
     UnixTimestamp,
     ValueType,
@@ -1731,3 +1732,170 @@ def docling_transform_docs(inputs: dict[str, Any]):
                     "Let's have fun with Natural Language Processing on PDFs."
                 ],
             }
+
+
+def test_python_transformation_with_set_types():
+    """Test that Set types work correctly in on-demand feature views."""
+    with tempfile.TemporaryDirectory() as data_dir:
+        store = FeatureStore(
+            config=RepoConfig(
+                project="test_set_types",
+                registry=os.path.join(data_dir, "registry.db"),
+                provider="local",
+                entity_key_serialization_version=3,
+                online_store=SqliteOnlineStoreConfig(
+                    path=os.path.join(data_dir, "online.db")
+                ),
+            )
+        )
+
+        # Create a simple driver entity
+        driver = Entity(
+            name="driver", join_keys=["driver_id"], value_type=ValueType.INT64
+        )
+
+        # Generate test data
+        end_date = datetime.now().replace(microsecond=0, second=0, minute=0)
+        start_date = end_date - timedelta(days=15)
+        driver_entities = [1001, 1002, 1003]
+        driver_df = create_driver_hourly_stats_df(driver_entities, start_date, end_date)
+        driver_stats_path = os.path.join(data_dir, "driver_stats.parquet")
+        driver_df.to_parquet(path=driver_stats_path, allow_truncated_timestamps=True)
+
+        driver_stats_source = FileSource(
+            name="driver_hourly_stats_source",
+            path=driver_stats_path,
+            timestamp_field="event_timestamp",
+            created_timestamp_column="created",
+        )
+
+        driver_stats_fv = FeatureView(
+            name="driver_hourly_stats",
+            entities=[driver],
+            ttl=timedelta(days=0),
+            schema=[
+                Field(name="conv_rate", dtype=Float32),
+                Field(name="acc_rate", dtype=Float32),
+                Field(name="avg_daily_trips", dtype=Int64),
+            ],
+            online=True,
+            source=driver_stats_source,
+        )
+
+        # Request source with Set types
+        request_source = RequestSource(
+            name="request_source",
+            schema=[
+                Field(name="visited_locations", dtype=Set(String)),
+                Field(name="favorite_numbers", dtype=Set(Int64)),
+            ],
+        )
+
+        # On-demand feature view that processes sets
+        @on_demand_feature_view(
+            sources=[request_source, driver_stats_fv],
+            schema=[
+                Field(name="unique_locations", dtype=Set(String)),
+                Field(name="location_count", dtype=Int64),
+                Field(name="unique_numbers", dtype=Set(Int64)),
+                Field(name="number_count", dtype=Int64),
+                Field(name="has_favorite_location", dtype=Bool),
+            ],
+            mode="python",
+        )
+        def set_processor_view(inputs: dict[str, Any]) -> dict[str, Any]:
+            output = {}
+            # Sets automatically deduplicate
+            output["unique_locations"] = inputs["visited_locations"]
+            output["location_count"] = [
+                len(locs) for locs in inputs["visited_locations"]
+            ]
+            output["unique_numbers"] = inputs["favorite_numbers"]
+            output["number_count"] = [len(nums) for nums in inputs["favorite_numbers"]]
+            output["has_favorite_location"] = [
+                "NYC" in locs for locs in inputs["visited_locations"]
+            ]
+            return output
+
+        # Apply the feature store objects
+        store.apply([driver, driver_stats_source, driver_stats_fv, set_processor_view])
+
+        # Write to online store
+        store.write_to_online_store(feature_view_name="driver_hourly_stats", df=driver_df)
+
+        # Test online feature retrieval with Set types
+        entity_rows = [
+            {
+                "driver_id": 1001,
+                "visited_locations": {"NYC", "LA", "SF", "NYC"},  # Duplicate NYC
+                "favorite_numbers": {1, 2, 3, 2, 1},  # Duplicates
+            }
+        ]
+
+        online_response = store.get_online_features(
+            entity_rows=entity_rows,
+            features=[
+                "driver_hourly_stats:conv_rate",
+                "driver_hourly_stats:avg_daily_trips",
+                "set_processor_view:unique_locations",
+                "set_processor_view:location_count",
+                "set_processor_view:unique_numbers",
+                "set_processor_view:number_count",
+                "set_processor_view:has_favorite_location",
+            ],
+        ).to_dict()
+
+        result = {name: value[0] for name, value in online_response.items()}
+
+        # Type assertions - verify Set types are returned as sets
+        assert isinstance(result["unique_locations"], set)
+        assert all(isinstance(loc, str) for loc in result["unique_locations"])
+
+        assert isinstance(result["unique_numbers"], set)
+        assert all(isinstance(num, int) for num in result["unique_numbers"])
+
+        assert isinstance(result["location_count"], int)
+        assert isinstance(result["number_count"], int)
+        assert isinstance(result["has_favorite_location"], bool)
+
+        # Value assertions - verify deduplication worked
+        assert result["unique_locations"] == {"NYC", "LA", "SF"}
+        assert result["location_count"] == 3  # Duplicate "NYC" was removed
+
+        assert result["unique_numbers"] == {1, 2, 3}
+        assert result["number_count"] == 3  # Duplicates were removed
+
+        assert result["has_favorite_location"] is True  # NYC is in the set
+
+        # Test with list input that gets converted to set
+        entity_rows_with_list = [
+            {
+                "driver_id": 1002,
+                "visited_locations": ["Boston", "Boston", "Seattle", "Portland"],  # List with duplicates
+                "favorite_numbers": [7, 8, 9, 7],  # List with duplicates
+            }
+        ]
+
+        online_response_list = store.get_online_features(
+            entity_rows=entity_rows_with_list,
+            features=[
+                "set_processor_view:unique_locations",
+                "set_processor_view:location_count",
+                "set_processor_view:unique_numbers",
+                "set_processor_view:number_count",
+                "set_processor_view:has_favorite_location",
+            ],
+        ).to_dict()
+
+        result_list = {name: value[0] for name, value in online_response_list.items()}
+
+        # Verify list input was converted to set and deduplicated
+        assert isinstance(result_list["unique_locations"], set)
+        assert result_list["unique_locations"] == {"Boston", "Seattle", "Portland"}
+        assert result_list["location_count"] == 3
+
+        assert isinstance(result_list["unique_numbers"], set)
+        assert result_list["unique_numbers"] == {7, 8, 9}
+        assert result_list["number_count"] == 3
+
+        assert result_list["has_favorite_location"] is False  # NYC not in the set
