From 998a3350ed399a3173a4a6c5219a18541b5cbca4 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Fri, 16 Oct 2026 23:31:12 +0000 Subject: [PATCH 1/2] Reuse column lookup preparation for repeated table schemas --- sdm/tensor/table.py | 26 +++++++++++++++++--------- test/tensor/test_table.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/sdm/tensor/table.py b/sdm/tensor/table.py index ab99c27e92..a7a90a4134 100644 --- a/sdm/tensor/table.py +++ b/sdm/tensor/table.py @@ -125,6 +125,20 @@ def __hash__(self) -> int: return hash(frozenset(self.columns.items())) +@functools.lru_cache(maxsize=128) +def _column_locations( + columns: tuple[tuple[Stype, tuple[str, ...]], ...], +) -> dict[str, tuple[Stype, int]]: + locations = { + name: (stype, index) + for stype, names in columns + for index, name in enumerate(names) + } + if sum(len(names) for _, names in columns) != len(locations): + raise ValueError("Expected column names to be unique") + return locations + + class TableTensor(Tensor): r"""A :class:`torch.Tensor` for tensorized, lossless table data. @@ -328,18 +342,12 @@ def __new__( f"{len(columns[stype])} {_columns} (got {block.size(-1)})" ) - column_names = list(chain.from_iterable(columns.values())) - column_to_loc: dict[str, tuple[Stype, int]] = {} - for stype_like, names in columns.items(): - stype = Stype(stype_like) - for i, name in enumerate(names): - column_to_loc[name] = (stype, i) - if len(column_names) != len(column_to_loc): - raise ValueError("Expected column names to be unique") + # Only immutable schema keys are retained; each table owns its lookup. + column_to_loc = _column_locations(tuple(columns.items())).copy() out = Tensor._make_wrapper_subclass( cls, - size=(*size, len(column_names)), + size=(*size, len(column_to_loc)), dtype=numerical.dtype, device=numerical.device, requires_grad=False, diff --git a/test/tensor/test_table.py b/test/tensor/test_table.py index 44adcc5b5f..9cf8261837 100644 --- a/test/tensor/test_table.py +++ b/test/tensor/test_table.py @@ -103,6 +103,36 @@ def test_column_names() -> None: ) +def test_repeated_schema_preserves_column_ownership_and_validation() -> None: + names = ["left", "right"] + values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + first = TableTensor(columns={"numerical": names}, numerical=values) + columns = {str(stype): names for stype, names in first.columns.items()} + second = TableTensor(columns=columns, numerical=values + 1) + names.reverse() + changed = columns.copy() + changed["numerical"] = tuple(names) + reversed_table = TableTensor(columns=changed, numerical=values) + + assert first.columns[Stype.numerical] == ("left", "right") + assert second.columns == first.columns + assert reversed_table.columns[Stype.numerical] == ("right", "left") + torch.testing.assert_close(first[["left"]].numerical, values[:, :1]) + torch.testing.assert_close(second[["left"]].numerical, values[:, :1] + 1) + torch.testing.assert_close( + reversed_table[["left"]].numerical, values[:, 1:] + ) + + with pytest.raises(ValueError, match="hold 2 columns"): + TableTensor(columns=columns, numerical=torch.ones(2, 3)) + with pytest.raises(ValueError, match="to be unique"): + TableTensor( + columns={"numerical": ["left"], "datetime": ["left"]}, + numerical=torch.ones(2, 1), + datetime=torch.ones(2, 1, dtype=torch.int64), + ) + + def test_equal() -> None: numerical = torch.randn(2, 2) categorical = CategoricalTensor( From 317dc1049cb12fc224c1f604a6b41c04705e0894 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Fri, 16 Oct 2026 23:52:50 +0000 Subject: [PATCH 2/2] Batch category metadata preparation with bounded temporary storage --- sdm/processing/categorical/align.py | 121 ++++++++++++++++------ test/processing/categorical/test_align.py | 110 +++++++++++++++++++- 2 files changed, 200 insertions(+), 31 deletions(-) diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 6d146bca49..e89bd4c372 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -15,6 +15,7 @@ Stype, TableTensor, ) +from sdm._memory import split_size from sdm.nn._buffer import BufferList from sdm.processing import EnsembleProcessor from sdm.relational.join import join_index @@ -396,9 +397,8 @@ def _category_lookup( fitted_categories: Tensor, codes: Tensor, ) -> Tensor: - lookup = codes.new_full((input_categories.numel(),), -1) if fitted_categories.numel() == 0: - return lookup + return codes.new_full(input_categories.size(), -1) comparable_categories = input_categories sorted_categories = fitted_categories @@ -414,12 +414,9 @@ def _category_lookup( sorted_categories, comparable_categories, ) - position = position.clamp(max=sorted_categories.numel() - 1) - match = sorted_categories[position] == comparable_categories - left_index = match.nonzero().view(-1) - right_index = perm[position[left_index]] - lookup[left_index] = right_index.to(codes.dtype) - return lookup + position = position.clamp(max=sorted_categories.size(-1) - 1) + match = sorted_categories.gather(-1, position) == comparable_categories + return torch.where(match, perm.gather(-1, position), -1).to(codes) def _align_to_categories( self, @@ -435,12 +432,22 @@ def _align_to_categories( {} for _ in table.categorical.categories ] string_requests: list[tuple[int, int, StringTensor, Tensor]] = [] + numeric_requests: dict[ + tuple[ + int, int, torch.dtype, torch.dtype, torch.device, torch.device + ], + list[tuple[int, int, Tensor, Tensor]], + ] = {} + columns_by_width: dict[int, list[int]] = {} for column_index, input_categories in enumerate( table.categorical.categories ): if input_categories.numel() == 0: continue + columns_by_width.setdefault(input_categories.numel(), []).append( + column_index + ) lookup_by_categories = lookups_by_column[column_index] pending_categories: set[int] = set() for batch_categories in fitted_categories: @@ -473,11 +480,50 @@ def _align_to_categories( ) ) continue - lookup_by_categories[identity] = self._category_lookup( - input_categories, - column_categories, + pending_categories.add(identity) + key = ( + input_categories.numel(), + column_categories.numel(), + input_categories.dtype, + column_categories.dtype, + input_categories.device, + column_categories.device, + ) + numeric_requests.setdefault(key, []).append( + ( + column_index, + identity, + input_categories, + column_categories, + ) + ) + + # Equal vocabulary widths need no padding and preserve row-wise sorts. + for requests in numeric_requests.values(): + input_categories, column_categories = requests[0][2:] + # Bound stacked vocabularies, sort/search buffers and lookups. + size = split_size( + num_items=len(requests), + item_bytes=8 + * (input_categories.numel() + column_categories.numel()) + * max( + input_categories.element_size(), + column_categories.element_size(), + 8, + ), + device=input_categories.device, + ) + for start in range(0, len(requests), size): + chunk = requests[start : start + size] + lookups = self._category_lookup( + torch.stack([request[2] for request in chunk]), + torch.stack([request[3] for request in chunk]), codes, ) + for (column_index, identity, _, _), lookup in zip( + chunk, lookups.unbind(), strict=True + ): + lookups_by_column[column_index][identity] = lookup string_lookups = self._string_category_lookups( input_categories=tuple( @@ -497,25 +543,42 @@ def _align_to_categories( ): lookups_by_column[column_index][identity] = lookup - for column_index, input_categories in enumerate( - table.categorical.categories - ): - if input_categories.numel() == 0: - continue - lookup_by_categories = lookups_by_column[column_index] - lookups = [ - lookup_by_categories[id(batch_categories[column_index])] - for batch_categories in fitted_categories - ] - column_codes = codes[..., column_index] - mask = column_codes >= 0 - lookup = torch.stack(lookups) - indices = column_codes.clamp_min(0).long() - aligned_codes[..., column_index] = torch.where( - mask, - lookup.gather(1, indices), - -1, + for width, columns in columns_by_width.items(): + # Bound gathered codes, integer indices, masks and mapped outputs. + size = split_size( + num_items=len(columns), + item_bytes=8 + * (codes[..., :1].numel() + len(fitted_categories) * width) + * max(codes.element_size(), 8), + device=codes.device, ) + for start in range(0, len(columns), size): + column_indices = columns[start : start + size] + # Map [column, batch, row] codes with vocabulary lookups. + lookup = torch.stack( + [ + lookups_by_column[column_index][ + id(batch_categories[column_index]) + ] + for column_index in column_indices + for batch_categories in fitted_categories + ] + ).view(len(column_indices), len(fitted_categories), width) + # Upload known column metadata without waiting for device work. + column_index = torch.tensor(column_indices, device="cpu").to( + device=codes.device, non_blocking=True + ) + column_codes = codes.index_select(-1, column_index).movedim( + -1, 0 + ) + mask = column_codes >= 0 + indices = column_codes.clamp_min(0).long() + mapped_codes = torch.where( + mask, + lookup.gather(-1, indices), + -1, + ).movedim(0, -1) + aligned_codes.index_copy_(-1, column_index, mapped_codes) return tuple( (table if single_table else table[batch_index]).replace_blocks( diff --git a/test/processing/categorical/test_align.py b/test/processing/categorical/test_align.py index b40bf306e4..b30fe915cd 100644 --- a/test/processing/categorical/test_align.py +++ b/test/processing/categorical/test_align.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Literal +from typing import Literal, cast import pandas as pd import pytest @@ -9,7 +9,7 @@ from sdm import CategoricalTensor, EnsembleTable, StringTensor, TableTensor from sdm.processing import AlignCategories -from sdm.testing import withCUDA +from sdm.testing import onlyCUDA, withCUDA def _table( @@ -110,6 +110,39 @@ def test_align_categories_removes_query_only_joint_vocabulary() -> None: assert query.categorical.code.squeeze(-1).tolist() == [-1, 1] +@onlyCUDA +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_align_numeric_categories_without_host_synchronization( + dtype: torch.dtype, +) -> None: + context = TableTensor( + categorical=CategoricalTensor( + code=torch.tensor([[0], [1], [0]], device="cuda", dtype=dtype), + categories=(torch.tensor([9, 3, 7], device="cuda"),), + ), + ) + query = TableTensor( + categorical=CategoricalTensor( + code=torch.tensor( + [[0], [1], [2], [-1]], device="cuda", dtype=dtype + ), + categories=(torch.tensor([3, 7, 9], device="cuda"),), + ), + ) + processor = AlignCategories(sort_by="value").fit(context) + previous = torch.cuda.get_sync_debug_mode() + try: + torch.cuda.set_sync_debug_mode("error") + with torch.device("cuda"): + output = processor.transform(query) + finally: + torch.cuda.set_sync_debug_mode(previous) + + assert output.categorical.code.dtype == dtype + assert output.categorical.code.tolist() == [[0], [-1], [1], [-1]] + assert output.categorical.categories[0].tolist() == [3, 9] + + @pytest.mark.parametrize("sort_by", ["code", "frequency", "value"]) def test_align_categories_orders_joint_vocabulary( sort_by: Literal["code", "frequency", "value"], @@ -504,3 +537,76 @@ def test_align_categories_ensemble_matches_member_fits( expected_query = reference.transform(query) assert query_output[member_id].equal(expected_query) assert fitted_query_output[member_id].equal(expected_query) + + +@withCUDA +@pytest.mark.parametrize("sort_by", ["code", "frequency", "value"]) +@pytest.mark.parametrize("min_frequency", [1, 2]) +@pytest.mark.parametrize("memory_fraction", ["0.05", "1e-10"]) +def test_align_numeric_columns_preserves_independent_vocabularies( + device: torch.device, + sort_by: Literal["code", "frequency", "value"], + min_frequency: int, + memory_fraction: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SDM_CHUNK_MEMORY_FRACTION", memory_fraction) + context_columns: list[torch.Tensor] = [] + query_columns: list[torch.Tensor] = [] + expected_columns: list[torch.Tensor] = [] + for index in range(7): + dtype = torch.float64 if index % 2 else torch.int64 + context_values = [5 + 20 * index, 1 + 20 * index, 9 + 20 * index] + query_values = [1 + 20 * index, 7 + 20 * index, 5 + 20 * index] + if index == 6: + context_values.append(13 + 20 * index) + query_values.append(13 + 20 * index) + context = TableTensor( + columns={"categorical": [str(index)]}, + categorical=CategoricalTensor( + code=torch.tensor( + [[0], [1], [2], [1], [2], [2], [-1]], + dtype=torch.int32, + device=device, + ), + categories=( + torch.tensor(context_values, dtype=dtype, device=device), + ), + ), + ) + query = TableTensor( + columns={"categorical": [str(index)]}, + categorical=CategoricalTensor( + code=torch.tensor( + [[0], [1], [2], [-1]], dtype=torch.int32, device=device + ), + categories=( + torch.tensor(query_values, dtype=dtype, device=device), + ), + ), + ) + expected_columns.append( + AlignCategories(sort_by, min_frequency=min_frequency) + .fit(context) + .transform(query) + ) + context_columns.append(context) + query_columns.append(query) + + context = cast(TableTensor, torch.cat(context_columns, dim=-1)) + query = cast(TableTensor, torch.cat(query_columns, dim=-1)) + processor = AlignCategories(sort_by, min_frequency=min_frequency).fit( + context + ) + with torch.no_grad(): + actual = processor.transform(query) + expected = cast(TableTensor, torch.cat(expected_columns, dim=-1)) + assert torch.equal(actual.categorical.code, expected.categorical.code) + assert all( + torch.equal(actual_values, expected_values) + for actual_values, expected_values in zip( + actual.categorical.categories, + expected.categorical.categories, + strict=True, + ) + )