Skip to content

API docs: reference documentation for classes and functions of datarepo project.

Databases catalog module

Catalog

A catalog that manages multiple databases and provides access to their tables.

Source code in src/datarepo/core/catalog/catalog.py
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
class Catalog:
    """A catalog that manages multiple databases and provides access to their tables."""

    def __init__(
        self,
        dbs: dict[str, Database],
        package_name: Optional[str] = None,
        metadata: Optional[CatalogMetadata] = None,
    ):
        """Initialize the Catalog.

        Example usage:
            ``` py
            from datarepo.core import Catalog, CatalogMetadata
            catalog = Catalog(
                dbs={
                    "my_db": ModuleDatabase(my_database_module),
                },
                package_name="my_package",
                metadata=CatalogMetadata(jupyterhub_url="https://jupyterhub.example.com")
            )
            ```

        Args:
            dbs (dict[str, Database]): A dictionary of database names and their corresponding Database objects.
            package_name (str | None, optional): The name of the package that contains the catalog. Defaults to None.
            metadata (Optional[CatalogMetadata], optional): Metadata for the catalog. Defaults to None.
        """
        self._dbs = dbs
        self._package_name = package_name
        self._metadata = metadata or CatalogMetadata()
        self._global_args: dict[str, Any] | None = None

    def set_global_args(self, global_args: dict[str, Any]) -> None:
        """Set global arguments for all database queries.

        Args:
            global_args (dict[str, Any]): A dictionary of global arguments to apply to all database queries.
        """
        self._global_args = global_args

    def db(self, db_name: str) -> Database:
        """Get a database from the catalog.

        Args:
            db_name (str): The name of the database.

        Raises:
            KeyError: If the database is not found.

        Returns:
            Database: The requested database.
        """
        db = self._dbs.get(db_name)

        if db is None:
            raise KeyError(
                f"Database '{db_name}' not found. Available databases: {self.dbs()}"
            )

        if self._global_args is None:
            return db

        return DatabaseWithGlobalArgs(db, self._global_args)

    def dbs(self) -> list[str]:
        """Get a list of database names in the catalog.

        Returns:
            list[str]: A list of database names.
        """
        return list(self._dbs.keys())

__init__(dbs, package_name=None, metadata=None)

Initialize the Catalog.

Example usage
from datarepo.core import Catalog, CatalogMetadata
catalog = Catalog(
    dbs={
        "my_db": ModuleDatabase(my_database_module),
    },
    package_name="my_package",
    metadata=CatalogMetadata(jupyterhub_url="https://jupyterhub.example.com")
)

Parameters:

Name Type Description Default
dbs dict[str, Database]

A dictionary of database names and their corresponding Database objects.

required
package_name str | None

The name of the package that contains the catalog. Defaults to None.

None
metadata Optional[CatalogMetadata]

Metadata for the catalog. Defaults to None.

None
Source code in src/datarepo/core/catalog/catalog.py
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def __init__(
    self,
    dbs: dict[str, Database],
    package_name: Optional[str] = None,
    metadata: Optional[CatalogMetadata] = None,
):
    """Initialize the Catalog.

    Example usage:
        ``` py
        from datarepo.core import Catalog, CatalogMetadata
        catalog = Catalog(
            dbs={
                "my_db": ModuleDatabase(my_database_module),
            },
            package_name="my_package",
            metadata=CatalogMetadata(jupyterhub_url="https://jupyterhub.example.com")
        )
        ```

    Args:
        dbs (dict[str, Database]): A dictionary of database names and their corresponding Database objects.
        package_name (str | None, optional): The name of the package that contains the catalog. Defaults to None.
        metadata (Optional[CatalogMetadata], optional): Metadata for the catalog. Defaults to None.
    """
    self._dbs = dbs
    self._package_name = package_name
    self._metadata = metadata or CatalogMetadata()
    self._global_args: dict[str, Any] | None = None

db(db_name)

Get a database from the catalog.

Parameters:

Name Type Description Default
db_name str

The name of the database.

required

Raises:

Type Description
KeyError

If the database is not found.

Returns:

Name Type Description
Database Database

The requested database.

Source code in src/datarepo/core/catalog/catalog.py
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
def db(self, db_name: str) -> Database:
    """Get a database from the catalog.

    Args:
        db_name (str): The name of the database.

    Raises:
        KeyError: If the database is not found.

    Returns:
        Database: The requested database.
    """
    db = self._dbs.get(db_name)

    if db is None:
        raise KeyError(
            f"Database '{db_name}' not found. Available databases: {self.dbs()}"
        )

    if self._global_args is None:
        return db

    return DatabaseWithGlobalArgs(db, self._global_args)

dbs()

Get a list of database names in the catalog.

Returns:

Type Description
list[str]

list[str]: A list of database names.

Source code in src/datarepo/core/catalog/catalog.py
258
259
260
261
262
263
264
def dbs(self) -> list[str]:
    """Get a list of database names in the catalog.

    Returns:
        list[str]: A list of database names.
    """
    return list(self._dbs.keys())

set_global_args(global_args)

Set global arguments for all database queries.

Parameters:

Name Type Description Default
global_args dict[str, Any]

A dictionary of global arguments to apply to all database queries.

required
Source code in src/datarepo/core/catalog/catalog.py
226
227
228
229
230
231
232
def set_global_args(self, global_args: dict[str, Any]) -> None:
    """Set global arguments for all database queries.

    Args:
        global_args (dict[str, Any]): A dictionary of global arguments to apply to all database queries.
    """
    self._global_args = global_args

CatalogMetadata dataclass

Metadata for catalog.

Source code in src/datarepo/core/catalog/catalog.py
186
187
188
189
190
@dataclass
class CatalogMetadata:
    """Metadata for catalog."""

    jupyterhub_url: str | None = None

Database

Bases: Protocol

A protocol for a database that provides access to tables.

Source code in src/datarepo/core/catalog/catalog.py
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Database(Protocol):
    """A protocol for a database that provides access to tables."""

    def get_tables(self, show_deprecated: bool = False) -> dict[str, TableProtocol]:
        """Get a dictionary of tables in the database.

        Args:
            show_deprecated (bool, optional): Whether to include deprecated tables. Defaults to False.
        """
        ...

    def tables(self, show_deprecated: bool = False) -> list[str]:
        """Get a list of table names in the database.

        Args:
            show_deprecated (bool, optional): Whether to include deprecated tables. Defaults to False.

        Returns:
            list[str]: A list of table names.
        """
        return list(self.get_tables(show_deprecated).keys())

    def table(self, name: str, *args: Any, **kwargs: Any) -> NlkDataFrame:
        """Get a table from the database.

        Args:
            name (str): The name of the table.

        Returns:
            NlkDataFrame: The requested table.
        """
        ...

get_tables(show_deprecated=False)

Get a dictionary of tables in the database.

Parameters:

Name Type Description Default
show_deprecated bool

Whether to include deprecated tables. Defaults to False.

False
Source code in src/datarepo/core/catalog/catalog.py
13
14
15
16
17
18
19
def get_tables(self, show_deprecated: bool = False) -> dict[str, TableProtocol]:
    """Get a dictionary of tables in the database.

    Args:
        show_deprecated (bool, optional): Whether to include deprecated tables. Defaults to False.
    """
    ...

table(name, *args, **kwargs)

Get a table from the database.

Parameters:

Name Type Description Default
name str

The name of the table.

required

Returns:

Name Type Description
NlkDataFrame NlkDataFrame

The requested table.

Source code in src/datarepo/core/catalog/catalog.py
32
33
34
35
36
37
38
39
40
41
def table(self, name: str, *args: Any, **kwargs: Any) -> NlkDataFrame:
    """Get a table from the database.

    Args:
        name (str): The name of the table.

    Returns:
        NlkDataFrame: The requested table.
    """
    ...

tables(show_deprecated=False)

Get a list of table names in the database.

Parameters:

Name Type Description Default
show_deprecated bool

Whether to include deprecated tables. Defaults to False.

False

Returns:

Type Description
list[str]

list[str]: A list of table names.

Source code in src/datarepo/core/catalog/catalog.py
21
22
23
24
25
26
27
28
29
30
def tables(self, show_deprecated: bool = False) -> list[str]:
    """Get a list of table names in the database.

    Args:
        show_deprecated (bool, optional): Whether to include deprecated tables. Defaults to False.

    Returns:
        list[str]: A list of table names.
    """
    return list(self.get_tables(show_deprecated).keys())

ModuleDatabase

Bases: Database

A database that is implemented as a Python module.

Source code in src/datarepo/core/catalog/catalog.py
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
class ModuleDatabase(Database):
    """A database that is implemented as a Python module."""

    def __init__(self, db: ModuleType) -> None:
        """Initialize the ModuleDatabase.

        Example usage:
            ``` py
            import my_database_module
            db = ModuleDatabase(my_database_module)
            ```

        Args:
            db (ModuleType): The database module.
        """
        self.db = db

    def __getattr__(self, name: str):
        # HACK: to maintain backwards compatibility when accessing module attributes
        return self._get_table(name)

    def get_tables(self, show_deprecated: bool = False) -> dict[str, TableProtocol]:
        """Get a dictionary of tables in the database.

        Example usage:
            ``` py
            db = ModuleDatabase(my_database_module)
            tables = db.get_tables(show_deprecated=True)
            ```

        Args:
            show_deprecated (bool, optional): Whether to include deprecated tables. Defaults to False.

        Returns:
            dict[str, TableProtocol]: A dictionary of table names and their corresponding TableProtocol objects.
        """
        methods = dir(self.db)

        tables = {}
        for name in methods:
            table = self._get_table(name)
            if table is None:
                continue

            if table.table_metadata.is_deprecated and not show_deprecated:
                continue

            tables[name] = table

        return tables

    def table(self, name: str, *args: Any, **kwargs: Any) -> NlkDataFrame:
        """Get a table from the database.

        Example usage:
            ``` py
            db = ModuleDatabase(my_database_module)
            table = db.table("my_table")
            ```

        Args:
            name (str): The name of the table.

        Raises:
            KeyError: If the table is not found.

        Returns:
            NlkDataFrame: The requested table.
        """
        tbl = self._get_table(name)
        if tbl is None:
            raise KeyError(f"Table '{name}' not found in database")

        if tbl.table_metadata.is_deprecated:
            warnings.warn(f"The table '{name}' is deprecated", DeprecationWarning)

        return tbl(*args, **kwargs)

    def _get_table(self, name: str) -> TableProtocol | None:
        """Get a table from the database.

        Args:
            name (str): The name of the table.

        Returns:
            TableProtocol | None: The requested table or None if not found.
        """
        table = getattr(self.db, name)
        if not hasattr(table, "table_metadata"):
            return None

        return cast(TableProtocol, table)

__init__(db)

Initialize the ModuleDatabase.

Example usage
import my_database_module
db = ModuleDatabase(my_database_module)

Parameters:

Name Type Description Default
db ModuleType

The database module.

required
Source code in src/datarepo/core/catalog/catalog.py
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(self, db: ModuleType) -> None:
    """Initialize the ModuleDatabase.

    Example usage:
        ``` py
        import my_database_module
        db = ModuleDatabase(my_database_module)
        ```

    Args:
        db (ModuleType): The database module.
    """
    self.db = db

get_tables(show_deprecated=False)

Get a dictionary of tables in the database.

Example usage
db = ModuleDatabase(my_database_module)
tables = db.get_tables(show_deprecated=True)

Parameters:

Name Type Description Default
show_deprecated bool

Whether to include deprecated tables. Defaults to False.

False

Returns:

Type Description
dict[str, TableProtocol]

dict[str, TableProtocol]: A dictionary of table names and their corresponding TableProtocol objects.

Source code in src/datarepo/core/catalog/catalog.py
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def get_tables(self, show_deprecated: bool = False) -> dict[str, TableProtocol]:
    """Get a dictionary of tables in the database.

    Example usage:
        ``` py
        db = ModuleDatabase(my_database_module)
        tables = db.get_tables(show_deprecated=True)
        ```

    Args:
        show_deprecated (bool, optional): Whether to include deprecated tables. Defaults to False.

    Returns:
        dict[str, TableProtocol]: A dictionary of table names and their corresponding TableProtocol objects.
    """
    methods = dir(self.db)

    tables = {}
    for name in methods:
        table = self._get_table(name)
        if table is None:
            continue

        if table.table_metadata.is_deprecated and not show_deprecated:
            continue

        tables[name] = table

    return tables

table(name, *args, **kwargs)

Get a table from the database.

Example usage
db = ModuleDatabase(my_database_module)
table = db.table("my_table")

Parameters:

Name Type Description Default
name str

The name of the table.

required

Raises:

Type Description
KeyError

If the table is not found.

Returns:

Name Type Description
NlkDataFrame NlkDataFrame

The requested table.

Source code in src/datarepo/core/catalog/catalog.py
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def table(self, name: str, *args: Any, **kwargs: Any) -> NlkDataFrame:
    """Get a table from the database.

    Example usage:
        ``` py
        db = ModuleDatabase(my_database_module)
        table = db.table("my_table")
        ```

    Args:
        name (str): The name of the table.

    Raises:
        KeyError: If the table is not found.

    Returns:
        NlkDataFrame: The requested table.
    """
    tbl = self._get_table(name)
    if tbl is None:
        raise KeyError(f"Table '{name}' not found in database")

    if tbl.table_metadata.is_deprecated:
        warnings.warn(f"The table '{name}' is deprecated", DeprecationWarning)

    return tbl(*args, **kwargs)

Dataframes module

Tables module

ClickHouseTable

Bases: TableProtocol

A table implementation that reads data from ClickHouse.

Source code in src/datarepo/core/tables/clickhouse_table.py
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
class ClickHouseTable(TableProtocol):
    """A table implementation that reads data from ClickHouse."""

    def __init__(
        self,
        name: str,
        schema: pa.Schema,
        description: str = "",
        config: ClickHouseTableConfig | None = None,
        docs_filters: List[Filter] | None = None,
        docs_columns: Optional[List[str]] = None,
        roapi_opts: RoapiOptions | None = None,
        unique_columns: Optional[List[str]] = None,
        table_metadata_args: Optional[Dict[str, Any]] = None,
        stats_cols: Optional[List[str]] = None,
    ):
        """Initialize the ClickHouseTable.

        Example usage:
            ```python
            from datarepo.core.tables import ClickHouseTable, ClickHouseTableConfig

            config = ClickHouseTableConfig(
                host="localhost",
                port=8443,
                username="user",
                password="password",
                database="default",
                table_name="delta_my_table",  # actual table in ClickHouse
                secure=True,
                verify=True,
                settings={"max_result_rows": 1000}
            )

            table = ClickHouseTable(
                name="my_table",  # logical/API name
                schema=pa.schema(
                    [
                        ("implant_id", pa.int64()),
                        ("date", pa.string()),
                        ("uniq", pa.string()),
                        ("value", pa.int64()),
                    ]
                ),
                config=config,
                description="My ClickHouse table",
                docs_filters=[...],
                docs_columns=[...],
                unique_columns=["uniq"],
                table_metadata_args={"answer": "42"},
                stats_cols=["implant_id"]
            )
            ```

        Args:
            name: Logical name of the table (used in Python API/code)
            schema: Schema of the table
            config: Configuration for connecting to ClickHouse. The config.table_name
                specifies the actual ClickHouse table name (useful when migrating
                from delta-backed to native tables). Falls back to `name` if not set.
            description: Description of the table for documentation
            docs_filters: Filters to show in documentation
            docs_columns: Columns to show in documentation
            roapi_opts: Options for ROAPI integration
            unique_columns: Columns to use for deduplication
            table_metadata_args: Additional metadata arguments
            stats_cols: Statistics columns, used to define columns that have statistics
        """
        self.name = name
        self.schema = schema
        self.config = config or make_clickhouse_config()
        self.unique_columns = unique_columns or []
        self.docs_filters = docs_filters or []
        self.docs_columns = docs_columns
        self.stats_cols = stats_cols or []
        # for roapi
        self.uri = self.config.get_uri()

        self.table_metadata = TableMetadata(
            table_type="CLICKHOUSE",
            description=description,
            docs_args={"filters": self.docs_filters, "columns": self.docs_columns},
            roapi_opts=roapi_opts or RoapiOptions(),
            **(table_metadata_args or {}),
        )

    def get_schema(self) -> TableSchema:
        """Generate and return the schema of the table, including columns.

        Returns:
            TableSchema: table schema containing column information.
        """
        schema = self.schema

        columns = [
            TableColumn(
                name=name,
                type=str(schema.field(name).type),
                readonly=False,
                filter_only=False,
                has_stats=name in self.stats_cols,
            )
            for name in schema.names
        ]

        return TableSchema(
            partitions=[],  # Clickhouse does not have partitions exposed in schema.
            columns=columns,
        )

    def _build_query(
        self,
        filters: InputFilters | None = None,
        columns: Optional[List[str]] = None,
        config: ClickHouseTableConfig | None = None,
    ) -> str:
        """Build a SQL query for the ClickHouse table.

        Args:
            filters (InputFilters, optional): Filters to apply to the query. Defaults to None.
            columns (Optional[List[str]], optional): Columns to select in the query. Defaults to None.

        Returns:
            str: SQL query string to select data from the ClickHouse table.

        Raises:
            ValueError: If a filter uses an unsupported operator. Unary null
                predicates ignore Filter.value.
        """
        config = config or self.config
        table_name = config.table_name or self.name

        column_expr = "*"
        if columns:
            valid_columns = [c for c in columns if c in self.schema.names]
            if valid_columns:
                column_expr = ", ".join(f"`{c}`" for c in valid_columns)
            else:
                logger.warning(
                    f"No valid columns provided for table {self.name}. Using '*' to select all columns."
                )

        # Build WHERE clause from filters if provided
        where_clause = ""
        if filters:
            normalized_filters = normalize_filters(filters)
            filter_expressions = []

            for filter_set in normalized_filters:
                set_expressions = []
                for f in filter_set:
                    if f.operator == "=":
                        set_expressions.append(
                            f"`{f.column}` = {format_value_for_sql(f.value)}"
                        )
                    elif f.operator == "!=":
                        set_expressions.append(
                            f"`{f.column}` != {format_value_for_sql(f.value)}"
                        )
                    elif f.operator == ">":
                        set_expressions.append(
                            f"`{f.column}` > {format_value_for_sql(f.value)}"
                        )
                    elif f.operator == "<":
                        set_expressions.append(
                            f"`{f.column}` < {format_value_for_sql(f.value)}"
                        )
                    elif f.operator == ">=":
                        set_expressions.append(
                            f"`{f.column}` >= {format_value_for_sql(f.value)}"
                        )
                    elif f.operator == "<=":
                        set_expressions.append(
                            f"`{f.column}` <= {format_value_for_sql(f.value)}"
                        )
                    elif f.operator == "in":
                        values = ", ".join(
                            format_value_for_sql(v) for v in cast(list, f.value)
                        )
                        set_expressions.append(f"`{f.column}` IN ({values})")
                    elif f.operator == "not in":
                        values = ", ".join(
                            format_value_for_sql(v) for v in cast(list, f.value)
                        )
                        set_expressions.append(f"`{f.column}` NOT IN ({values})")
                    elif f.operator == "is null":
                        set_expressions.append(f"`{f.column}` IS NULL")
                    elif f.operator == "is not null":
                        set_expressions.append(f"`{f.column}` IS NOT NULL")
                    elif f.operator in [
                        "contains",
                        "includes",
                        "includes any",
                        "includes all",
                    ]:
                        set_expressions.append(
                            f"`{f.column}` LIKE {format_value_for_sql(f.value)}"
                        )
                    else:
                        raise ValueError(f"Unsupported filter operator: {f.operator!r}")

                if set_expressions:
                    filter_expressions.append("(" + " AND ".join(set_expressions) + ")")

            if filter_expressions:
                where_clause = "WHERE " + " OR ".join(filter_expressions)

        return f"SELECT {column_expr} FROM `{config.database}`.`{table_name}` {where_clause}"

    def __call__(  # type: ignore[override]
        self,
        filters: InputFilters | None = None,
        columns: List[str] | None = None,
        config: ClickHouseTableConfig | None = None,
        **kwargs: Dict[str, Any],
    ) -> NlkDataFrame:
        """Query data from the ClickHouse table.

        Example usage:
            ``` py
            from datarepo.core.tables import ClickHouseTable, ClickHouseTableConfig
            config = ClickHouseTableConfig(...)
            table = ClickHouseTable(...)
            df = table(filters=[Filter("implant_id", "=", 123)], columns=["date", "value"])
            ```

        Args:
            filters: Filters to apply to the data.
            columns: Columns to select from the table.
            config: Optional configuration for the ClickHouse table.
            **kwargs: Additional arguments.

        Returns:
            NlkDataFrame: A lazy Polars DataFrame with the requested data.

        Raises:
            ValueError: If a filter uses an unsupported operator, before creating
                a backend client or executing a query, or if a configured unique
                column is missing from the declared schema. Key validation is
                deferred until a read; constructing a table does not connect.
        """
        config = config or self.config

        query_columns = columns
        output_columns = None
        if self.unique_columns:
            missing_keys = [
                c for c in self.unique_columns if c not in self.schema.names
            ]
            if missing_keys:
                raise ValueError(
                    f"Unique columns missing from table schema: {missing_keys}"
                )
            if columns:
                output_columns = [c for c in columns if c in self.schema.names]
                if output_columns:
                    # Transport the complete identity, even for a partial public
                    # projection. Preserve the all-invalid '*' fallback below.
                    query_columns = output_columns + [
                        c for c in self.unique_columns if c not in output_columns
                    ]

        query = self._build_query(filters, query_columns, config)

        client = config.get_client()
        arrow_result = client.query_arrow(query)
        df = pl.from_arrow(arrow_result)
        if isinstance(df, pl.Series):
            df = df.to_frame()

        if self.unique_columns:
            df = df.filter(pl.struct(self.unique_columns).is_last_distinct())
            if output_columns:
                df = df.select(output_columns)

        return df.lazy()

__call__(filters=None, columns=None, config=None, **kwargs)

Query data from the ClickHouse table.

Example usage
from datarepo.core.tables import ClickHouseTable, ClickHouseTableConfig
config = ClickHouseTableConfig(...)
table = ClickHouseTable(...)
df = table(filters=[Filter("implant_id", "=", 123)], columns=["date", "value"])

Parameters:

Name Type Description Default
filters InputFilters | None

Filters to apply to the data.

None
columns List[str] | None

Columns to select from the table.

None
config ClickHouseTableConfig | None

Optional configuration for the ClickHouse table.

None
**kwargs Dict[str, Any]

Additional arguments.

{}

Returns:

Name Type Description
NlkDataFrame NlkDataFrame

A lazy Polars DataFrame with the requested data.

Raises:

Type Description
ValueError

If a filter uses an unsupported operator, before creating a backend client or executing a query, or if a configured unique column is missing from the declared schema. Key validation is deferred until a read; constructing a table does not connect.

Source code in src/datarepo/core/tables/clickhouse_table.py
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
def __call__(  # type: ignore[override]
    self,
    filters: InputFilters | None = None,
    columns: List[str] | None = None,
    config: ClickHouseTableConfig | None = None,
    **kwargs: Dict[str, Any],
) -> NlkDataFrame:
    """Query data from the ClickHouse table.

    Example usage:
        ``` py
        from datarepo.core.tables import ClickHouseTable, ClickHouseTableConfig
        config = ClickHouseTableConfig(...)
        table = ClickHouseTable(...)
        df = table(filters=[Filter("implant_id", "=", 123)], columns=["date", "value"])
        ```

    Args:
        filters: Filters to apply to the data.
        columns: Columns to select from the table.
        config: Optional configuration for the ClickHouse table.
        **kwargs: Additional arguments.

    Returns:
        NlkDataFrame: A lazy Polars DataFrame with the requested data.

    Raises:
        ValueError: If a filter uses an unsupported operator, before creating
            a backend client or executing a query, or if a configured unique
            column is missing from the declared schema. Key validation is
            deferred until a read; constructing a table does not connect.
    """
    config = config or self.config

    query_columns = columns
    output_columns = None
    if self.unique_columns:
        missing_keys = [
            c for c in self.unique_columns if c not in self.schema.names
        ]
        if missing_keys:
            raise ValueError(
                f"Unique columns missing from table schema: {missing_keys}"
            )
        if columns:
            output_columns = [c for c in columns if c in self.schema.names]
            if output_columns:
                # Transport the complete identity, even for a partial public
                # projection. Preserve the all-invalid '*' fallback below.
                query_columns = output_columns + [
                    c for c in self.unique_columns if c not in output_columns
                ]

    query = self._build_query(filters, query_columns, config)

    client = config.get_client()
    arrow_result = client.query_arrow(query)
    df = pl.from_arrow(arrow_result)
    if isinstance(df, pl.Series):
        df = df.to_frame()

    if self.unique_columns:
        df = df.filter(pl.struct(self.unique_columns).is_last_distinct())
        if output_columns:
            df = df.select(output_columns)

    return df.lazy()

__init__(name, schema, description='', config=None, docs_filters=None, docs_columns=None, roapi_opts=None, unique_columns=None, table_metadata_args=None, stats_cols=None)

Initialize the ClickHouseTable.

Example usage
from datarepo.core.tables import ClickHouseTable, ClickHouseTableConfig

config = ClickHouseTableConfig(
    host="localhost",
    port=8443,
    username="user",
    password="password",
    database="default",
    table_name="delta_my_table",  # actual table in ClickHouse
    secure=True,
    verify=True,
    settings={"max_result_rows": 1000}
)

table = ClickHouseTable(
    name="my_table",  # logical/API name
    schema=pa.schema(
        [
            ("implant_id", pa.int64()),
            ("date", pa.string()),
            ("uniq", pa.string()),
            ("value", pa.int64()),
        ]
    ),
    config=config,
    description="My ClickHouse table",
    docs_filters=[...],
    docs_columns=[...],
    unique_columns=["uniq"],
    table_metadata_args={"answer": "42"},
    stats_cols=["implant_id"]
)

Parameters:

Name Type Description Default
name str

Logical name of the table (used in Python API/code)

required
schema Schema

Schema of the table

required
config ClickHouseTableConfig | None

Configuration for connecting to ClickHouse. The config.table_name specifies the actual ClickHouse table name (useful when migrating from delta-backed to native tables). Falls back to name if not set.

None
description str

Description of the table for documentation

''
docs_filters List[Filter] | None

Filters to show in documentation

None
docs_columns Optional[List[str]]

Columns to show in documentation

None
roapi_opts RoapiOptions | None

Options for ROAPI integration

None
unique_columns Optional[List[str]]

Columns to use for deduplication

None
table_metadata_args Optional[Dict[str, Any]]

Additional metadata arguments

None
stats_cols Optional[List[str]]

Statistics columns, used to define columns that have statistics

None
Source code in src/datarepo/core/tables/clickhouse_table.py
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def __init__(
    self,
    name: str,
    schema: pa.Schema,
    description: str = "",
    config: ClickHouseTableConfig | None = None,
    docs_filters: List[Filter] | None = None,
    docs_columns: Optional[List[str]] = None,
    roapi_opts: RoapiOptions | None = None,
    unique_columns: Optional[List[str]] = None,
    table_metadata_args: Optional[Dict[str, Any]] = None,
    stats_cols: Optional[List[str]] = None,
):
    """Initialize the ClickHouseTable.

    Example usage:
        ```python
        from datarepo.core.tables import ClickHouseTable, ClickHouseTableConfig

        config = ClickHouseTableConfig(
            host="localhost",
            port=8443,
            username="user",
            password="password",
            database="default",
            table_name="delta_my_table",  # actual table in ClickHouse
            secure=True,
            verify=True,
            settings={"max_result_rows": 1000}
        )

        table = ClickHouseTable(
            name="my_table",  # logical/API name
            schema=pa.schema(
                [
                    ("implant_id", pa.int64()),
                    ("date", pa.string()),
                    ("uniq", pa.string()),
                    ("value", pa.int64()),
                ]
            ),
            config=config,
            description="My ClickHouse table",
            docs_filters=[...],
            docs_columns=[...],
            unique_columns=["uniq"],
            table_metadata_args={"answer": "42"},
            stats_cols=["implant_id"]
        )
        ```

    Args:
        name: Logical name of the table (used in Python API/code)
        schema: Schema of the table
        config: Configuration for connecting to ClickHouse. The config.table_name
            specifies the actual ClickHouse table name (useful when migrating
            from delta-backed to native tables). Falls back to `name` if not set.
        description: Description of the table for documentation
        docs_filters: Filters to show in documentation
        docs_columns: Columns to show in documentation
        roapi_opts: Options for ROAPI integration
        unique_columns: Columns to use for deduplication
        table_metadata_args: Additional metadata arguments
        stats_cols: Statistics columns, used to define columns that have statistics
    """
    self.name = name
    self.schema = schema
    self.config = config or make_clickhouse_config()
    self.unique_columns = unique_columns or []
    self.docs_filters = docs_filters or []
    self.docs_columns = docs_columns
    self.stats_cols = stats_cols or []
    # for roapi
    self.uri = self.config.get_uri()

    self.table_metadata = TableMetadata(
        table_type="CLICKHOUSE",
        description=description,
        docs_args={"filters": self.docs_filters, "columns": self.docs_columns},
        roapi_opts=roapi_opts or RoapiOptions(),
        **(table_metadata_args or {}),
    )

get_schema()

Generate and return the schema of the table, including columns.

Returns:

Name Type Description
TableSchema TableSchema

table schema containing column information.

Source code in src/datarepo/core/tables/clickhouse_table.py
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
def get_schema(self) -> TableSchema:
    """Generate and return the schema of the table, including columns.

    Returns:
        TableSchema: table schema containing column information.
    """
    schema = self.schema

    columns = [
        TableColumn(
            name=name,
            type=str(schema.field(name).type),
            readonly=False,
            filter_only=False,
            has_stats=name in self.stats_cols,
        )
        for name in schema.names
    ]

    return TableSchema(
        partitions=[],  # Clickhouse does not have partitions exposed in schema.
        columns=columns,
    )

ClickHouseTableConfig dataclass

Configuration for connecting to ClickHouse.

Source code in src/datarepo/core/tables/clickhouse_table.py
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
@dataclass
class ClickHouseTableConfig:
    """Configuration for connecting to ClickHouse."""

    host: str
    port: int = 8123
    username: Optional[str] = None
    password: Optional[str] = None
    database: str = "default"
    table_name: Optional[str] = None
    secure: bool = False
    verify: bool = False
    settings: Dict[str, Any] = field(default_factory=dict)

    def get_uri(self) -> str:
        """Construct the URI for the ClickHouse table.

        Returns:
            str: URI for the ClickHouse table.
        """
        # check if username and password are provided
        if not self.username or not self.password:
            return f"clickhouse://{self.host}:{self.port}/{self.database}"
        return f"clickhouse://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"

    def get_client(self) -> clickhouse_connect.driver.Client:
        """Get a ClickHouse client for this configuration.

        Returns:
            clickhouse_connect.driver.Client: ClickHouse client instance.
        """
        return clickhouse_connect.get_client(
            host=self.host,
            port=self.port,
            username=self.username or "default",
            password=self.password or "",
            database=self.database,
            secure=self.secure,
            verify=self.verify,
            settings=self.settings,
        )

get_client()

Get a ClickHouse client for this configuration.

Returns:

Type Description
Client

clickhouse_connect.driver.Client: ClickHouse client instance.

Source code in src/datarepo/core/tables/clickhouse_table.py
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def get_client(self) -> clickhouse_connect.driver.Client:
    """Get a ClickHouse client for this configuration.

    Returns:
        clickhouse_connect.driver.Client: ClickHouse client instance.
    """
    return clickhouse_connect.get_client(
        host=self.host,
        port=self.port,
        username=self.username or "default",
        password=self.password or "",
        database=self.database,
        secure=self.secure,
        verify=self.verify,
        settings=self.settings,
    )

get_uri()

Construct the URI for the ClickHouse table.

Returns:

Name Type Description
str str

URI for the ClickHouse table.

Source code in src/datarepo/core/tables/clickhouse_table.py
41
42
43
44
45
46
47
48
49
50
def get_uri(self) -> str:
    """Construct the URI for the ClickHouse table.

    Returns:
        str: URI for the ClickHouse table.
    """
    # check if username and password are provided
    if not self.username or not self.password:
        return f"clickhouse://{self.host}:{self.port}/{self.database}"
    return f"clickhouse://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"

DeltaCacheOptions dataclass

Source code in src/datarepo/core/tables/deltalake_table.py
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@dataclass
class DeltaCacheOptions:
    # Path to the directory where files are cached. This can be the same for all tables
    # since the S3 table prefix is included in the cached paths.
    file_cache_path: str
    # Duration for which the _last_checkpoint file is cached. A longer duration means
    # we can reuse cached checkpoints for longer, which could improve loading performance
    # from not downloading new checkpoint parquets. However, we also need to load more
    # individual transaction jsons, so this should not be set too high.
    # A reasonable value for a table with frequent updates (e.g. binned spikes) is 30m.
    file_cache_last_checkpoint_valid_duration: str | None = None

    def to_storage_options(self) -> dict[str, Any]:
        """Convert the cache options to a dictionary of storage options.

        Returns:
            dict[str, Any]: A dictionary of storage options that can be used with DeltaTable.
        """
        opts = {
            "file_cache_path": os.path.expanduser(self.file_cache_path),
        }
        if self.file_cache_last_checkpoint_valid_duration is not None:
            opts["file_cache_last_checkpoint_valid_duration"] = (
                self.file_cache_last_checkpoint_valid_duration
            )
        return opts

to_storage_options()

Convert the cache options to a dictionary of storage options.

Returns:

Type Description
dict[str, Any]

dict[str, Any]: A dictionary of storage options that can be used with DeltaTable.

Source code in src/datarepo/core/tables/deltalake_table.py
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def to_storage_options(self) -> dict[str, Any]:
    """Convert the cache options to a dictionary of storage options.

    Returns:
        dict[str, Any]: A dictionary of storage options that can be used with DeltaTable.
    """
    opts = {
        "file_cache_path": os.path.expanduser(self.file_cache_path),
    }
    if self.file_cache_last_checkpoint_valid_duration is not None:
        opts["file_cache_last_checkpoint_valid_duration"] = (
            self.file_cache_last_checkpoint_valid_duration
        )
    return opts

DeltalakeTable

Bases: TableProtocol

A table that is backed by a Delta Lake table.

Source code in src/datarepo/core/tables/deltalake_table.py
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
class DeltalakeTable(TableProtocol):
    """A table that is backed by a Delta Lake table."""

    def __init__(
        self,
        name: str,
        uri: str,
        schema: pa.Schema,
        description: str = "",
        docs_filters: list[Filter] = [],
        docs_columns: list[str] | None = None,
        roapi_opts: RoapiOptions | None = None,
        unique_columns: list[str] | None = None,
        table_metadata_args: dict[str, Any] | None = None,
        stats_cols: list[str] | None = None,
        extra_cols: list[tuple[pl.Expr, str]] | None = None,
        partition_columns: list[str] | None = None,
        schema_versions: dict[str, pa.Schema] | None = None,
        default_schema_version: str | None = None,
    ):
        """Initialize the DeltalakeTable.

        Args:
            name (str): table name, used as the table identifier in the DeltaTable
            uri (str): uri of the table, e.g. "s3://bucket/path/to/table". When
                ``schema_versions`` is set, this is the base URI and each version
                is read from ``{uri}/{version}``.
            schema (pa.Schema): schema of the table, used to define the table structure
            description (str, optional): description of the table, used for documentation. Defaults to "".
            docs_filters (list[Filter], optional): documentation filters, used to filter the table in the documentation. Defaults to [].
            docs_columns (list[str] | None, optional): documentation columns, used to define the columns in the documentation. Defaults to None.
            roapi_opts (RoapiOptions | None, optional): ROAPI options, used to configure the ROAPI for the table. Defaults to None.
            unique_columns (list[str] | None, optional): unique columns in the table, used to optimize the read performance. Defaults to None.
            table_metadata_args (dict[str, Any] | None, optional): table metadata arguments, used to configure the table metadata. Defaults to None.
            stats_cols (list[str] | None, optional): statistics columns, used to define the columns that have statistics. Defaults to None.
            extra_cols (list[tuple[pl.Expr, str]] | None, optional): extra columns to add to the table, where each tuple contains a Polars expression and its type annotation. Defaults to None.
            partition_columns (list[str] | None, optional): partition columns for the table. If provided, get_schema() will use these instead of querying the remote Delta table. This allows catalog export to work even when the table doesn't exist yet. Defaults to None.
            schema_versions (dict[str, pa.Schema] | None, optional): optional map of
                version labels (e.g. ``"v1"``, ``"v2"``) to schemas for tables stored
                at versioned URIs. Defaults to None.
            default_schema_version (str | None, optional): version label from
                ``schema_versions`` used for the default read URI and schema.
                Required when ``schema_versions`` is set. Defaults to None.
        """
        self.name = name
        self._base_uri = uri
        self.schema_versions = schema_versions
        self.unique_columns = unique_columns
        self.stats_cols = stats_cols or []
        self.extra_cols = extra_cols or []
        self.partition_columns = partition_columns
        self._schema_version: str | None = None

        if schema_versions:
            if default_schema_version is None:
                raise ValueError(
                    "default_schema_version is required when schema_versions is set"
                )
            if default_schema_version not in schema_versions:
                raise ValueError(
                    f"default_schema_version {default_schema_version!r} is not in "
                    f"schema_versions keys: {sorted(schema_versions)}"
                )
            self._schema_version = default_schema_version
            self.schema = schema_versions[self._schema_version]
        else:
            if default_schema_version is not None:
                raise ValueError(
                    "default_schema_version requires schema_versions to be set"
                )
            self.schema = schema

        self._table_metadata_args = table_metadata_args or {}

        self.table_metadata = TableMetadata(
            table_type="DELTA_LAKE",
            description=description,
            docs_args={"filters": docs_filters, "columns": docs_columns},
            roapi_opts=roapi_opts or DeltaRoapiOptions(),
            **self._table_metadata_args,
        )

    @property
    def uri(self) -> str:
        """Resolved table URI, including the schema version suffix when applicable."""
        if self._schema_version is not None and self.schema_versions is not None:
            return f"{self._base_uri.rstrip('/')}/{self._schema_version}"
        return self._base_uri

    @property
    def schema_version(self) -> str | None:
        """Active schema version label, or None when versioning is disabled."""
        return self._schema_version

    def at_version(self, version: str) -> DeltalakeTable:
        """Return a table bound to a specific schema version and URI suffix.

        Args:
            version (str): key from ``schema_versions`` (e.g. ``"v1"``).

        Returns:
            DeltalakeTable: a new table instance using the version's schema and URI.
        """
        if self.schema_versions is None:
            raise ValueError(
                "at_version() requires schema_versions to be set on DeltalakeTable"
            )
        if version not in self.schema_versions:
            raise KeyError(
                f"Unknown schema version {version!r}. "
                f"Known versions: {sorted(self.schema_versions)}"
            )

        return DeltalakeTable(
            name=self.name,
            uri=self._base_uri,
            schema=self.schema_versions[version],
            description=self.table_metadata.description,
            docs_filters=self.table_metadata.docs_args.get("filters", []),
            docs_columns=self.table_metadata.docs_args.get("columns"),
            roapi_opts=self.table_metadata.roapi_opts,
            unique_columns=self.unique_columns,
            table_metadata_args=self._table_metadata_args,
            stats_cols=self.stats_cols,
            extra_cols=self.extra_cols,
            partition_columns=self.partition_columns,
            schema_versions=self.schema_versions,
            default_schema_version=version,
        )

    def get_schema(
        self,
        boto3_session: boto3.Session | None = None,
        endpoint_url: str | None = None,
        timeout: str | None = None,
        cache_options: DeltaCacheOptions | None = None,
    ) -> TableSchema:
        """Generate and return the schema of the table, including partitions and columns.

        Returns:
            TableSchema: table schema containing partition and column information.
        """
        if self.partition_columns is not None:
            partition_cols = self.partition_columns
        else:
            storage_options = get_delta_storage_options(
                boto3_session=boto3_session,
                endpoint_url=endpoint_url,
                timeout=timeout,
            )
            dt = self.delta_table(storage_options=storage_options)
            partition_cols = dt.metadata().partition_columns

        schema = self.schema
        filters = {
            f.column: f
            for f in self.table_metadata.docs_args.get("filters", [])
            if isinstance(f, Filter)
        }
        partitions = [
            TablePartition(
                column_name=col,
                operator=filters[col].operator if col in filters else "=",
                type_annotation=str(schema.field(col).type),
                value=filters[col].value if col in filters else None,
            )
            for col in partition_cols
        ]
        columns = [
            TableColumn(
                name=name,
                type=str(schema.field(name).type),
                readonly=False,
                filter_only=False,
                has_stats=name in partition_cols or name in self.stats_cols,
            )
            for name in schema.names
        ]
        columns += [
            TableColumn(
                name=expr.meta.output_name(),
                type=expr_type,
                readonly=True,
                filter_only=False,
                has_stats=False,
            )
            for expr, expr_type in self.extra_cols
        ]
        return TableSchema(
            partitions=partitions,
            columns=columns,
        )

    def __call__(
        self,
        filters: DeltaInputFilters | None = None,
        columns: list[str] | None = None,
        boto3_session: boto3.Session | None = None,
        endpoint_url: str | None = None,
        timeout: str | None = None,
        cache_options: DeltaCacheOptions | None = None,
        schema_version: str | None = None,
        delta_version: int | None = None,
        **kwargs: Any,
    ) -> NlkDataFrame:
        """Fetch a dataframe from the Delta Lake table.

        Args:
            filters (DeltaInputFilters | None, optional): filters to apply to the table. Defaults to None.
            columns (list[str] | None, optional): columns to select from the table. Defaults to None.
            boto3_session (boto3.Session | None, optional): boto3 session to use for S3 access. Defaults to None.
            endpoint_url (str | None, optional): endpoint URL for S3 access. Defaults to None.
            timeout (str | None, optional): timeout for S3 access. Defaults to None.
            cache_options (DeltaCacheOptions | None, optional): cache options for the Delta Lake table. Defaults to None.
            schema_version (str | None, optional): when ``schema_versions`` is configured,
                read from ``{uri}/{schema_version}`` using that version's schema.
                Defaults to None (use the table's bound version).
            delta_version (int | None, optional): Delta Lake transaction log version for
                time travel within the resolved table URI. Defaults to None (latest).

        Returns:
            NlkDataFrame: a dataframe containing the data from the Delta Lake table, filtered and selected according to the provided parameters.
        """
        table = self.at_version(schema_version) if schema_version else self

        storage_options = get_delta_storage_options(
            boto3_session=boto3_session,
            endpoint_url=endpoint_url,
            timeout=timeout,
            cache_options=cache_options,
        )
        dt = table.delta_table(storage_options=storage_options, version=delta_version)

        return table.construct_df(dt=dt, filters=filters, columns=columns)

    def construct_df(
        self,
        dt: DeltaTable,
        filters: DeltaInputFilters | None = None,
        columns: list[str] | None = None,
    ) -> NlkDataFrame:
        """Construct a dataframe from the Delta Lake table.

        Args:
            dt (DeltaTable): The DeltaTable object representing the Delta Lake table.
            filters (DeltaInputFilters | None, optional): filters to apply to the table. Defaults to None.
            columns (list[str] | None, optional): columns to select from the table. Defaults to None.

        Returns:
            NlkDataFrame: a dataframe containing the data from the Delta Lake table, filtered and selected according to the provided parameters.
        """
        # Use schema defined on this table, the physical schema in deltalake metadata might be different
        schema = self.schema

        predicate_str = datafusion_predicate_from_filters(schema, filters)

        # These should not be read because they don't exist in the delta table
        extra_col_exprs = [expr for expr, _ in self.extra_cols]
        extra_column_names = set(expr.meta.output_name() for expr in extra_col_exprs)

        columns_to_read = None
        unique_column_names = set(self.unique_columns or [])
        if columns:
            columns_to_read = list(
                (set(columns) | unique_column_names) - extra_column_names
            )

        # TODO(peter): consider a sql builder for more complex queries?
        select_cols = (
            ", ".join([f'"{col}"' for col in columns_to_read])
            if columns_to_read
            else "*"
        )
        condition = f"WHERE {predicate_str}" if predicate_str else ""
        query_string = f"""
            SELECT {select_cols}
            FROM "{self.name}"
            {condition}
        """
        with warnings.catch_warnings():
            # Ignore ExperimentalWarning emitted from QueryBuilder
            warnings.filterwarnings("ignore", category=ExperimentalWarning)
            batches = (
                QueryBuilder().register(self.name, dt).execute(query_string).fetchall()
            )

        # Since we might cast unique string columns to categoricals, use a string cache to
        # improve performance when combining multiple dataframes
        with pl.StringCache():
            if batches:
                frame = pl.from_arrow(batches, rechunk=False)
                if isinstance(frame, pl.Series):
                    frame = frame.to_frame()
                frame = _normalize_df(frame, self.schema, columns=columns_to_read)
            else:
                # If dataset is empty, the returned dataframe will have no columns
                frame = _empty_normalized_df(schema)

            if self.extra_cols:
                frame = frame.with_columns(extra_col_exprs)

            if self.unique_columns:
                # Cast unique string columns to categoricals first
                # In some cases, this reduces peak memory usage up to 50%
                curr_schema = frame.schema
                cat_schema = {
                    col: pl.Categorical
                    for col in curr_schema
                    if curr_schema[col] == pl.String
                }
                # NOTE: Using .filter() + .is_last_distinct() instead of the
                # .unique() API as due to a polars bug that corrupts very large
                # dataframes, causing panics on some subsequent operations with
                # the dataframe (Specifically, `PanicException: polars' maximum
                # length reached.`). This is also faster.
                #
                # .is_last_distinct() was chosen such that later rows are kept
                # instead of earlier ones. This is useful for backfills
                # backfill where the underlying raw data may have been updated
                # due to a migration.
                frame = (
                    frame.cast(cat_schema)  # type: ignore[arg-type]
                    .filter(pl.struct(self.unique_columns).is_last_distinct())
                    .cast(curr_schema)  # type: ignore[arg-type]
                )

        if columns:
            frame = frame.select(columns)

        return frame.lazy()

    def delta_table(
        self, storage_options: dict[str, Any] | None = None, version: int | None = None
    ) -> DeltaTable:
        """Get the DeltaTable object for this table.

        Args:
            storage_options (dict[str, Any] | None, optional): Storage options for the DeltaTable, such as S3 access credentials. Defaults to None.
            version (int | None, optional): Version of the Delta table to read. If None, the latest version is used. Defaults to None.

        Returns:
            DeltaTable: The DeltaTable object representing the Delta Lake table.
        """
        return DeltaTable(
            table_uri=self.uri, storage_options=storage_options, version=version
        )

schema_version property

Active schema version label, or None when versioning is disabled.

uri property

Resolved table URI, including the schema version suffix when applicable.

__call__(filters=None, columns=None, boto3_session=None, endpoint_url=None, timeout=None, cache_options=None, schema_version=None, delta_version=None, **kwargs)

Fetch a dataframe from the Delta Lake table.

Parameters:

Name Type Description Default
filters DeltaInputFilters | None

filters to apply to the table. Defaults to None.

None
columns list[str] | None

columns to select from the table. Defaults to None.

None
boto3_session Session | None

boto3 session to use for S3 access. Defaults to None.

None
endpoint_url str | None

endpoint URL for S3 access. Defaults to None.

None
timeout str | None

timeout for S3 access. Defaults to None.

None
cache_options DeltaCacheOptions | None

cache options for the Delta Lake table. Defaults to None.

None
schema_version str | None

when schema_versions is configured, read from {uri}/{schema_version} using that version's schema. Defaults to None (use the table's bound version).

None
delta_version int | None

Delta Lake transaction log version for time travel within the resolved table URI. Defaults to None (latest).

None

Returns:

Name Type Description
NlkDataFrame NlkDataFrame

a dataframe containing the data from the Delta Lake table, filtered and selected according to the provided parameters.

Source code in src/datarepo/core/tables/deltalake_table.py
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
def __call__(
    self,
    filters: DeltaInputFilters | None = None,
    columns: list[str] | None = None,
    boto3_session: boto3.Session | None = None,
    endpoint_url: str | None = None,
    timeout: str | None = None,
    cache_options: DeltaCacheOptions | None = None,
    schema_version: str | None = None,
    delta_version: int | None = None,
    **kwargs: Any,
) -> NlkDataFrame:
    """Fetch a dataframe from the Delta Lake table.

    Args:
        filters (DeltaInputFilters | None, optional): filters to apply to the table. Defaults to None.
        columns (list[str] | None, optional): columns to select from the table. Defaults to None.
        boto3_session (boto3.Session | None, optional): boto3 session to use for S3 access. Defaults to None.
        endpoint_url (str | None, optional): endpoint URL for S3 access. Defaults to None.
        timeout (str | None, optional): timeout for S3 access. Defaults to None.
        cache_options (DeltaCacheOptions | None, optional): cache options for the Delta Lake table. Defaults to None.
        schema_version (str | None, optional): when ``schema_versions`` is configured,
            read from ``{uri}/{schema_version}`` using that version's schema.
            Defaults to None (use the table's bound version).
        delta_version (int | None, optional): Delta Lake transaction log version for
            time travel within the resolved table URI. Defaults to None (latest).

    Returns:
        NlkDataFrame: a dataframe containing the data from the Delta Lake table, filtered and selected according to the provided parameters.
    """
    table = self.at_version(schema_version) if schema_version else self

    storage_options = get_delta_storage_options(
        boto3_session=boto3_session,
        endpoint_url=endpoint_url,
        timeout=timeout,
        cache_options=cache_options,
    )
    dt = table.delta_table(storage_options=storage_options, version=delta_version)

    return table.construct_df(dt=dt, filters=filters, columns=columns)

__init__(name, uri, schema, description='', docs_filters=[], docs_columns=None, roapi_opts=None, unique_columns=None, table_metadata_args=None, stats_cols=None, extra_cols=None, partition_columns=None, schema_versions=None, default_schema_version=None)

Initialize the DeltalakeTable.

Parameters:

Name Type Description Default
name str

table name, used as the table identifier in the DeltaTable

required
uri str

uri of the table, e.g. "s3://bucket/path/to/table". When schema_versions is set, this is the base URI and each version is read from {uri}/{version}.

required
schema Schema

schema of the table, used to define the table structure

required
description str

description of the table, used for documentation. Defaults to "".

''
docs_filters list[Filter]

documentation filters, used to filter the table in the documentation. Defaults to [].

[]
docs_columns list[str] | None

documentation columns, used to define the columns in the documentation. Defaults to None.

None
roapi_opts RoapiOptions | None

ROAPI options, used to configure the ROAPI for the table. Defaults to None.

None
unique_columns list[str] | None

unique columns in the table, used to optimize the read performance. Defaults to None.

None
table_metadata_args dict[str, Any] | None

table metadata arguments, used to configure the table metadata. Defaults to None.

None
stats_cols list[str] | None

statistics columns, used to define the columns that have statistics. Defaults to None.

None
extra_cols list[tuple[Expr, str]] | None

extra columns to add to the table, where each tuple contains a Polars expression and its type annotation. Defaults to None.

None
partition_columns list[str] | None

partition columns for the table. If provided, get_schema() will use these instead of querying the remote Delta table. This allows catalog export to work even when the table doesn't exist yet. Defaults to None.

None
schema_versions dict[str, Schema] | None

optional map of version labels (e.g. "v1", "v2") to schemas for tables stored at versioned URIs. Defaults to None.

None
default_schema_version str | None

version label from schema_versions used for the default read URI and schema. Required when schema_versions is set. Defaults to None.

None
Source code in src/datarepo/core/tables/deltalake_table.py
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def __init__(
    self,
    name: str,
    uri: str,
    schema: pa.Schema,
    description: str = "",
    docs_filters: list[Filter] = [],
    docs_columns: list[str] | None = None,
    roapi_opts: RoapiOptions | None = None,
    unique_columns: list[str] | None = None,
    table_metadata_args: dict[str, Any] | None = None,
    stats_cols: list[str] | None = None,
    extra_cols: list[tuple[pl.Expr, str]] | None = None,
    partition_columns: list[str] | None = None,
    schema_versions: dict[str, pa.Schema] | None = None,
    default_schema_version: str | None = None,
):
    """Initialize the DeltalakeTable.

    Args:
        name (str): table name, used as the table identifier in the DeltaTable
        uri (str): uri of the table, e.g. "s3://bucket/path/to/table". When
            ``schema_versions`` is set, this is the base URI and each version
            is read from ``{uri}/{version}``.
        schema (pa.Schema): schema of the table, used to define the table structure
        description (str, optional): description of the table, used for documentation. Defaults to "".
        docs_filters (list[Filter], optional): documentation filters, used to filter the table in the documentation. Defaults to [].
        docs_columns (list[str] | None, optional): documentation columns, used to define the columns in the documentation. Defaults to None.
        roapi_opts (RoapiOptions | None, optional): ROAPI options, used to configure the ROAPI for the table. Defaults to None.
        unique_columns (list[str] | None, optional): unique columns in the table, used to optimize the read performance. Defaults to None.
        table_metadata_args (dict[str, Any] | None, optional): table metadata arguments, used to configure the table metadata. Defaults to None.
        stats_cols (list[str] | None, optional): statistics columns, used to define the columns that have statistics. Defaults to None.
        extra_cols (list[tuple[pl.Expr, str]] | None, optional): extra columns to add to the table, where each tuple contains a Polars expression and its type annotation. Defaults to None.
        partition_columns (list[str] | None, optional): partition columns for the table. If provided, get_schema() will use these instead of querying the remote Delta table. This allows catalog export to work even when the table doesn't exist yet. Defaults to None.
        schema_versions (dict[str, pa.Schema] | None, optional): optional map of
            version labels (e.g. ``"v1"``, ``"v2"``) to schemas for tables stored
            at versioned URIs. Defaults to None.
        default_schema_version (str | None, optional): version label from
            ``schema_versions`` used for the default read URI and schema.
            Required when ``schema_versions`` is set. Defaults to None.
    """
    self.name = name
    self._base_uri = uri
    self.schema_versions = schema_versions
    self.unique_columns = unique_columns
    self.stats_cols = stats_cols or []
    self.extra_cols = extra_cols or []
    self.partition_columns = partition_columns
    self._schema_version: str | None = None

    if schema_versions:
        if default_schema_version is None:
            raise ValueError(
                "default_schema_version is required when schema_versions is set"
            )
        if default_schema_version not in schema_versions:
            raise ValueError(
                f"default_schema_version {default_schema_version!r} is not in "
                f"schema_versions keys: {sorted(schema_versions)}"
            )
        self._schema_version = default_schema_version
        self.schema = schema_versions[self._schema_version]
    else:
        if default_schema_version is not None:
            raise ValueError(
                "default_schema_version requires schema_versions to be set"
            )
        self.schema = schema

    self._table_metadata_args = table_metadata_args or {}

    self.table_metadata = TableMetadata(
        table_type="DELTA_LAKE",
        description=description,
        docs_args={"filters": docs_filters, "columns": docs_columns},
        roapi_opts=roapi_opts or DeltaRoapiOptions(),
        **self._table_metadata_args,
    )

at_version(version)

Return a table bound to a specific schema version and URI suffix.

Parameters:

Name Type Description Default
version str

key from schema_versions (e.g. "v1").

required

Returns:

Name Type Description
DeltalakeTable DeltalakeTable

a new table instance using the version's schema and URI.

Source code in src/datarepo/core/tables/deltalake_table.py
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def at_version(self, version: str) -> DeltalakeTable:
    """Return a table bound to a specific schema version and URI suffix.

    Args:
        version (str): key from ``schema_versions`` (e.g. ``"v1"``).

    Returns:
        DeltalakeTable: a new table instance using the version's schema and URI.
    """
    if self.schema_versions is None:
        raise ValueError(
            "at_version() requires schema_versions to be set on DeltalakeTable"
        )
    if version not in self.schema_versions:
        raise KeyError(
            f"Unknown schema version {version!r}. "
            f"Known versions: {sorted(self.schema_versions)}"
        )

    return DeltalakeTable(
        name=self.name,
        uri=self._base_uri,
        schema=self.schema_versions[version],
        description=self.table_metadata.description,
        docs_filters=self.table_metadata.docs_args.get("filters", []),
        docs_columns=self.table_metadata.docs_args.get("columns"),
        roapi_opts=self.table_metadata.roapi_opts,
        unique_columns=self.unique_columns,
        table_metadata_args=self._table_metadata_args,
        stats_cols=self.stats_cols,
        extra_cols=self.extra_cols,
        partition_columns=self.partition_columns,
        schema_versions=self.schema_versions,
        default_schema_version=version,
    )

construct_df(dt, filters=None, columns=None)

Construct a dataframe from the Delta Lake table.

Parameters:

Name Type Description Default
dt DeltaTable

The DeltaTable object representing the Delta Lake table.

required
filters DeltaInputFilters | None

filters to apply to the table. Defaults to None.

None
columns list[str] | None

columns to select from the table. Defaults to None.

None

Returns:

Name Type Description
NlkDataFrame NlkDataFrame

a dataframe containing the data from the Delta Lake table, filtered and selected according to the provided parameters.

Source code in src/datarepo/core/tables/deltalake_table.py
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
def construct_df(
    self,
    dt: DeltaTable,
    filters: DeltaInputFilters | None = None,
    columns: list[str] | None = None,
) -> NlkDataFrame:
    """Construct a dataframe from the Delta Lake table.

    Args:
        dt (DeltaTable): The DeltaTable object representing the Delta Lake table.
        filters (DeltaInputFilters | None, optional): filters to apply to the table. Defaults to None.
        columns (list[str] | None, optional): columns to select from the table. Defaults to None.

    Returns:
        NlkDataFrame: a dataframe containing the data from the Delta Lake table, filtered and selected according to the provided parameters.
    """
    # Use schema defined on this table, the physical schema in deltalake metadata might be different
    schema = self.schema

    predicate_str = datafusion_predicate_from_filters(schema, filters)

    # These should not be read because they don't exist in the delta table
    extra_col_exprs = [expr for expr, _ in self.extra_cols]
    extra_column_names = set(expr.meta.output_name() for expr in extra_col_exprs)

    columns_to_read = None
    unique_column_names = set(self.unique_columns or [])
    if columns:
        columns_to_read = list(
            (set(columns) | unique_column_names) - extra_column_names
        )

    # TODO(peter): consider a sql builder for more complex queries?
    select_cols = (
        ", ".join([f'"{col}"' for col in columns_to_read])
        if columns_to_read
        else "*"
    )
    condition = f"WHERE {predicate_str}" if predicate_str else ""
    query_string = f"""
        SELECT {select_cols}
        FROM "{self.name}"
        {condition}
    """
    with warnings.catch_warnings():
        # Ignore ExperimentalWarning emitted from QueryBuilder
        warnings.filterwarnings("ignore", category=ExperimentalWarning)
        batches = (
            QueryBuilder().register(self.name, dt).execute(query_string).fetchall()
        )

    # Since we might cast unique string columns to categoricals, use a string cache to
    # improve performance when combining multiple dataframes
    with pl.StringCache():
        if batches:
            frame = pl.from_arrow(batches, rechunk=False)
            if isinstance(frame, pl.Series):
                frame = frame.to_frame()
            frame = _normalize_df(frame, self.schema, columns=columns_to_read)
        else:
            # If dataset is empty, the returned dataframe will have no columns
            frame = _empty_normalized_df(schema)

        if self.extra_cols:
            frame = frame.with_columns(extra_col_exprs)

        if self.unique_columns:
            # Cast unique string columns to categoricals first
            # In some cases, this reduces peak memory usage up to 50%
            curr_schema = frame.schema
            cat_schema = {
                col: pl.Categorical
                for col in curr_schema
                if curr_schema[col] == pl.String
            }
            # NOTE: Using .filter() + .is_last_distinct() instead of the
            # .unique() API as due to a polars bug that corrupts very large
            # dataframes, causing panics on some subsequent operations with
            # the dataframe (Specifically, `PanicException: polars' maximum
            # length reached.`). This is also faster.
            #
            # .is_last_distinct() was chosen such that later rows are kept
            # instead of earlier ones. This is useful for backfills
            # backfill where the underlying raw data may have been updated
            # due to a migration.
            frame = (
                frame.cast(cat_schema)  # type: ignore[arg-type]
                .filter(pl.struct(self.unique_columns).is_last_distinct())
                .cast(curr_schema)  # type: ignore[arg-type]
            )

    if columns:
        frame = frame.select(columns)

    return frame.lazy()

delta_table(storage_options=None, version=None)

Get the DeltaTable object for this table.

Parameters:

Name Type Description Default
storage_options dict[str, Any] | None

Storage options for the DeltaTable, such as S3 access credentials. Defaults to None.

None
version int | None

Version of the Delta table to read. If None, the latest version is used. Defaults to None.

None

Returns:

Name Type Description
DeltaTable DeltaTable

The DeltaTable object representing the Delta Lake table.

Source code in src/datarepo/core/tables/deltalake_table.py
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def delta_table(
    self, storage_options: dict[str, Any] | None = None, version: int | None = None
) -> DeltaTable:
    """Get the DeltaTable object for this table.

    Args:
        storage_options (dict[str, Any] | None, optional): Storage options for the DeltaTable, such as S3 access credentials. Defaults to None.
        version (int | None, optional): Version of the Delta table to read. If None, the latest version is used. Defaults to None.

    Returns:
        DeltaTable: The DeltaTable object representing the Delta Lake table.
    """
    return DeltaTable(
        table_uri=self.uri, storage_options=storage_options, version=version
    )

get_schema(boto3_session=None, endpoint_url=None, timeout=None, cache_options=None)

Generate and return the schema of the table, including partitions and columns.

Returns:

Name Type Description
TableSchema TableSchema

table schema containing partition and column information.

Source code in src/datarepo/core/tables/deltalake_table.py
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
def get_schema(
    self,
    boto3_session: boto3.Session | None = None,
    endpoint_url: str | None = None,
    timeout: str | None = None,
    cache_options: DeltaCacheOptions | None = None,
) -> TableSchema:
    """Generate and return the schema of the table, including partitions and columns.

    Returns:
        TableSchema: table schema containing partition and column information.
    """
    if self.partition_columns is not None:
        partition_cols = self.partition_columns
    else:
        storage_options = get_delta_storage_options(
            boto3_session=boto3_session,
            endpoint_url=endpoint_url,
            timeout=timeout,
        )
        dt = self.delta_table(storage_options=storage_options)
        partition_cols = dt.metadata().partition_columns

    schema = self.schema
    filters = {
        f.column: f
        for f in self.table_metadata.docs_args.get("filters", [])
        if isinstance(f, Filter)
    }
    partitions = [
        TablePartition(
            column_name=col,
            operator=filters[col].operator if col in filters else "=",
            type_annotation=str(schema.field(col).type),
            value=filters[col].value if col in filters else None,
        )
        for col in partition_cols
    ]
    columns = [
        TableColumn(
            name=name,
            type=str(schema.field(name).type),
            readonly=False,
            filter_only=False,
            has_stats=name in partition_cols or name in self.stats_cols,
        )
        for name in schema.names
    ]
    columns += [
        TableColumn(
            name=expr.meta.output_name(),
            type=expr_type,
            readonly=True,
            filter_only=False,
            has_stats=False,
        )
        for expr, expr_type in self.extra_cols
    ]
    return TableSchema(
        partitions=partitions,
        columns=columns,
    )

Filter

Bases: NamedTuple

Filter represents a condition to be applied to a column in a table.

Source code in src/datarepo/core/tables/filters.py
22
23
24
25
26
27
class Filter(NamedTuple):
    """Filter represents a condition to be applied to a column in a table."""

    column: str
    operator: FilterOperator
    value: Any

ParquetTable

Bases: TableProtocol

A table that is stored in Parquet format.

Source code in src/datarepo/core/tables/parquet_table.py
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
class ParquetTable(TableProtocol):
    """A table that is stored in Parquet format."""

    def __init__(
        self,
        name: str,
        uri: str,
        partitioning: list[Partition],
        partitioning_scheme: PartitioningScheme = PartitioningScheme.DIRECTORY,
        description: str = "",
        docs_filters: list[Filter] = [],
        docs_columns: list[str] | None = None,
        roapi_opts: RoapiOptions | None = None,
        parquet_file_name: str = "df.parquet",
        table_metadata_args: dict[str, Any] | None = None,
        schema: dict[str, pl.DataType] | None = None,
    ):
        """Initialize the ParquetTable.

        Args:
            name (str): name of the table, used for documentation and metadata.
            uri (str): uri of the table, typically an S3 bucket path.
            partitioning (list[Partition]): partitioning scheme for the table.
                This is a list of Partition objects, which define the columns and types used for partitioning
            partitioning_scheme (PartitioningScheme, optional): scheme used for partitioning.
                Defaults to PartitioningScheme.DIRECTORY.
            description (str, optional): description of the table, used for documentation.
                Defaults to "".
            docs_filters (list[Filter], optional): documentation filters for the table.
                These filters are used to generate documentation and are not applied to the data.
                Defaults to [].
            docs_columns (list[str] | None, optional): docsumentation columns for the table.
                These columns are used to generate documentation and are not applied to the data.
                Defaults to None.
            roapi_opts (RoapiOptions | None, optional): Read-only API options for the table.
                These options are used to configure the ROAPI endpoint for the table.
                Defaults to None.
            parquet_file_name (str, optional): parquet file name to use when building file fragments.
            table_metadata_args (dict[str, Any] | None, optional): additional metadata arguments for the table.
            schema (dict[str, pl.DataType] | None, optional): explicit schema for the table.
                When provided, allows scan_parquet to succeed even when the S3 path
                contains no files (returns an empty DataFrame with this schema).
                Defaults to None.

        Raises:
            ValueError: if the partitioning_scheme is not a valid PartitioningScheme.
        """
        if not isinstance(partitioning_scheme, PartitioningScheme):
            raise ValueError(f"Invalid partitioning scheme, got {partitioning_scheme}")

        self.name = name
        self.uri = uri
        self.partitioning = partitioning
        self.partitioning_scheme = partitioning_scheme
        self.schema = schema

        self.table_metadata = TableMetadata(
            table_type="PARQUET",
            description=description,
            docs_args={"filters": docs_filters, "columns": docs_columns},
            roapi_opts=roapi_opts,
            **(table_metadata_args or {}),
        )

        self.parquet_file_name = parquet_file_name

    def get_schema(self) -> TableSchema:
        """Generates the schema of the table, including partitions and columns.

        Returns:
            TableSchema: table schema containing partitions and columns.
        """
        partitions = [
            TablePartition(
                column_name=filter.column,
                operator=filter.operator,
                type_annotation=type(filter.value).__name__,
                value=filter.value,
            )
            for filter in self.table_metadata.docs_args.get("filters", [])
        ]

        # Pop the selected columns so that we still pull the full schema below
        docs_args = {**self.table_metadata.docs_args}
        docs_args.pop("columns")

        columns = []
        if docs_args or not partitions:
            table: NlkDataFrame = self(**docs_args)
            columns = [
                TableColumn(
                    name=key,
                    type=type.__str__(),
                    readonly=False,
                    filter_only=False,
                    has_stats=False,
                )
                for key, type in table.schema.items()
            ]

        return TableSchema(partitions=partitions, columns=columns)

    def __call__(
        self,
        filters: InputFilters | None = None,
        columns: Optional[list[str]] = None,
        boto3_session: boto3.Session | None = None,
        endpoint_url: str | None = None,
        **kwargs: Any,
    ) -> NlkDataFrame:
        """Fetches data from the Parquet table based on the provided filters and columns.

        Args:
            filters (InputFilters | None, optional): filters to apply to the data. Defaults to None.
            columns (Optional[list[str]], optional): columns to select from the data. Defaults to None.
            boto3_session (boto3.Session | None, optional): boto3 session to use for S3 access. Defaults to None.
            endpoint_url (str | None, optional): endpoint URL for S3 access. Defaults to None.

        Returns:
            NlkDataFrame: A DataFrame containing the filtered data from the Parquet table.
        """
        normalized_filters = normalize_filters(filters)
        (
            uri,
            remaining_partitions,
            remaining_filters,
            applied_filters,
        ) = self._build_uri_from_filters(normalized_filters)

        storage_options = get_storage_options(
            boto3_session=boto3_session,
            endpoint_url=endpoint_url,
        )

        hive_schema = (
            {partition.column: partition.col_type for partition in remaining_partitions}
            if remaining_partitions
            else None
        )

        df = pl.scan_parquet(
            uri,
            hive_partitioning=len(remaining_partitions) > 0,
            allow_missing_columns=True,
            storage_options=storage_options,
            hive_schema=hive_schema,
            schema=self.schema,
        )

        if applied_filters:
            # Add columns removed from partitions and added to uri
            df = df.with_columns(
                pl.lit(f.value)
                .cast(
                    next(
                        partition.col_type
                        for partition in self.partitioning
                        if partition.column == f.column
                    )
                )
                .alias(f.column)
                for f in applied_filters
            )

        if remaining_filters:
            filter_expr = _filters_to_expr(remaining_filters)
            if filter_expr is not None:
                df = df.filter(filter_expr)

        if columns:
            df = df.select(columns)

        return df

    def build_file_fragment(self, filters: list[Filter]) -> str:
        """
        Returns a file path from the base table URI with the given filters.
        This will raise an error if the filter does not specify all partitions.

        This is currently used to generate the file path used by ROAPI to infer schemas.
        """
        uri_with_prefix, partitions, _, _ = self._build_uri_from_filters(
            normalize_filters(filters), include_base_uri=False
        )
        if len(partitions) > 0:
            partition_names = [partition.column for partition in partitions]
            raise ValueError(
                f"Not enough partitions specified, missing: {partition_names}"
            )

        return path.join(uri_with_prefix, self.parquet_file_name)

    def _build_uri_from_filters(
        self,
        filters: NormalizedFilters,
        include_base_uri: bool = True,
    ) -> tuple[str, list[Partition], NormalizedFilters, list[Filter]]:
        """Attempts to build an S3 list prefix from the given filters.
        We do this because pyarrow runs an S3 List() query on the base URI
        before applying filters to possible files. This can be slow.

        We can force pyarrow to do more sophisticated prefix filtering
        if we pre-construct the URI before making the call. If a partition
        has exactly one filter that uses strict equality, we know that
        it will contain that filter string in the URI. We attempt to do this
        for each partition filter, in order, until we encounter a
        partition that:
            1. does not have a filter, or
            2. has a more than one filters, or
            3. has filters that are not strict equality checks.

        This gives us the longest URL that can be used as a prefix filter for
        all files returned by the S3 List() call.

        In a benchmark, this brought reading ~1 million rows of binned spike files
        from 12s to 1.5s.

        In the future, we can further optimize this by building a list of
        candidate URIs and doing separate prefix filtering with each of those, in parallel.

        Long-term we should push this logic into pyarrow and make an upstream commit.

        NOTE: Add trailing slash - this is important to ensure that,
        when partitioning only by implant ID, a future 5-digit implant beginning with the same 4 digits
        as a 4-digit implant is not included in a 4-digit implant's query.
        """
        uri = self.uri if include_base_uri else ""

        if not filters or not self.partitioning:
            return uri, self.partitioning, filters, []

        partitions = list(self.partitioning)
        filters = [list(filter_set) for filter_set in filters]
        applied_filters = []

        for i, partition in enumerate(self.partitioning):
            partition_filters = [
                exactly_one_equality_filter(partition, f) for f in filters
            ]

            if len(partition_filters) == 0:
                break

            # either 0 or multiple filters for this partition,
            # break and deal with the s3 list() query
            if any(partition_filter is None for partition_filter in partition_filters):
                break

            # Only move forward if all remaining filter sets have the same partition filter
            if not all(
                partition_filter == partition_filters[0]
                for partition_filter in partition_filters
            ):
                break

            partition_filter = partition_filters[0]

            if partition_filter is None:
                continue

            if self.partitioning_scheme == PartitioningScheme.DIRECTORY:
                partition_component = str(partition_filter.value)
            elif self.partitioning_scheme == PartitioningScheme.HIVE:
                partition_component = (
                    partition.column + "=" + str(partition_filter.value)
                )

            uri = path.join(uri, partition_component)

            # remove the partition and filter since it's been applied
            # technically might not need to remove the partitions,
            # but they are semantically meaningless as we already have
            # constructed the URI, so we can pop them
            partitions.remove(partition)
            for filter_set in filters:
                filter_set.remove(partition_filter)

            applied_filters.append(partition_filter)

            # An emptied branch is true within the selected partition. Keeping
            # the other OR branches would incorrectly restrict its rows, and
            # their predicates cannot justify pruning any deeper partitions.
            if any(not filter_set for filter_set in filters):
                filters = [[]]
                break

        uri = path.join(
            uri, ""
        )  # trailing slash prevents inclusion of partitions that are subsets of other partitions

        return (uri, partitions, filters, applied_filters)

__call__(filters=None, columns=None, boto3_session=None, endpoint_url=None, **kwargs)

Fetches data from the Parquet table based on the provided filters and columns.

Parameters:

Name Type Description Default
filters InputFilters | None

filters to apply to the data. Defaults to None.

None
columns Optional[list[str]]

columns to select from the data. Defaults to None.

None
boto3_session Session | None

boto3 session to use for S3 access. Defaults to None.

None
endpoint_url str | None

endpoint URL for S3 access. Defaults to None.

None

Returns:

Name Type Description
NlkDataFrame NlkDataFrame

A DataFrame containing the filtered data from the Parquet table.

Source code in src/datarepo/core/tables/parquet_table.py
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
def __call__(
    self,
    filters: InputFilters | None = None,
    columns: Optional[list[str]] = None,
    boto3_session: boto3.Session | None = None,
    endpoint_url: str | None = None,
    **kwargs: Any,
) -> NlkDataFrame:
    """Fetches data from the Parquet table based on the provided filters and columns.

    Args:
        filters (InputFilters | None, optional): filters to apply to the data. Defaults to None.
        columns (Optional[list[str]], optional): columns to select from the data. Defaults to None.
        boto3_session (boto3.Session | None, optional): boto3 session to use for S3 access. Defaults to None.
        endpoint_url (str | None, optional): endpoint URL for S3 access. Defaults to None.

    Returns:
        NlkDataFrame: A DataFrame containing the filtered data from the Parquet table.
    """
    normalized_filters = normalize_filters(filters)
    (
        uri,
        remaining_partitions,
        remaining_filters,
        applied_filters,
    ) = self._build_uri_from_filters(normalized_filters)

    storage_options = get_storage_options(
        boto3_session=boto3_session,
        endpoint_url=endpoint_url,
    )

    hive_schema = (
        {partition.column: partition.col_type for partition in remaining_partitions}
        if remaining_partitions
        else None
    )

    df = pl.scan_parquet(
        uri,
        hive_partitioning=len(remaining_partitions) > 0,
        allow_missing_columns=True,
        storage_options=storage_options,
        hive_schema=hive_schema,
        schema=self.schema,
    )

    if applied_filters:
        # Add columns removed from partitions and added to uri
        df = df.with_columns(
            pl.lit(f.value)
            .cast(
                next(
                    partition.col_type
                    for partition in self.partitioning
                    if partition.column == f.column
                )
            )
            .alias(f.column)
            for f in applied_filters
        )

    if remaining_filters:
        filter_expr = _filters_to_expr(remaining_filters)
        if filter_expr is not None:
            df = df.filter(filter_expr)

    if columns:
        df = df.select(columns)

    return df

__init__(name, uri, partitioning, partitioning_scheme=PartitioningScheme.DIRECTORY, description='', docs_filters=[], docs_columns=None, roapi_opts=None, parquet_file_name='df.parquet', table_metadata_args=None, schema=None)

Initialize the ParquetTable.

Parameters:

Name Type Description Default
name str

name of the table, used for documentation and metadata.

required
uri str

uri of the table, typically an S3 bucket path.

required
partitioning list[Partition]

partitioning scheme for the table. This is a list of Partition objects, which define the columns and types used for partitioning

required
partitioning_scheme PartitioningScheme

scheme used for partitioning. Defaults to PartitioningScheme.DIRECTORY.

DIRECTORY
description str

description of the table, used for documentation. Defaults to "".

''
docs_filters list[Filter]

documentation filters for the table. These filters are used to generate documentation and are not applied to the data. Defaults to [].

[]
docs_columns list[str] | None

docsumentation columns for the table. These columns are used to generate documentation and are not applied to the data. Defaults to None.

None
roapi_opts RoapiOptions | None

Read-only API options for the table. These options are used to configure the ROAPI endpoint for the table. Defaults to None.

None
parquet_file_name str

parquet file name to use when building file fragments.

'df.parquet'
table_metadata_args dict[str, Any] | None

additional metadata arguments for the table.

None
schema dict[str, DataType] | None

explicit schema for the table. When provided, allows scan_parquet to succeed even when the S3 path contains no files (returns an empty DataFrame with this schema). Defaults to None.

None

Raises:

Type Description
ValueError

if the partitioning_scheme is not a valid PartitioningScheme.

Source code in src/datarepo/core/tables/parquet_table.py
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def __init__(
    self,
    name: str,
    uri: str,
    partitioning: list[Partition],
    partitioning_scheme: PartitioningScheme = PartitioningScheme.DIRECTORY,
    description: str = "",
    docs_filters: list[Filter] = [],
    docs_columns: list[str] | None = None,
    roapi_opts: RoapiOptions | None = None,
    parquet_file_name: str = "df.parquet",
    table_metadata_args: dict[str, Any] | None = None,
    schema: dict[str, pl.DataType] | None = None,
):
    """Initialize the ParquetTable.

    Args:
        name (str): name of the table, used for documentation and metadata.
        uri (str): uri of the table, typically an S3 bucket path.
        partitioning (list[Partition]): partitioning scheme for the table.
            This is a list of Partition objects, which define the columns and types used for partitioning
        partitioning_scheme (PartitioningScheme, optional): scheme used for partitioning.
            Defaults to PartitioningScheme.DIRECTORY.
        description (str, optional): description of the table, used for documentation.
            Defaults to "".
        docs_filters (list[Filter], optional): documentation filters for the table.
            These filters are used to generate documentation and are not applied to the data.
            Defaults to [].
        docs_columns (list[str] | None, optional): docsumentation columns for the table.
            These columns are used to generate documentation and are not applied to the data.
            Defaults to None.
        roapi_opts (RoapiOptions | None, optional): Read-only API options for the table.
            These options are used to configure the ROAPI endpoint for the table.
            Defaults to None.
        parquet_file_name (str, optional): parquet file name to use when building file fragments.
        table_metadata_args (dict[str, Any] | None, optional): additional metadata arguments for the table.
        schema (dict[str, pl.DataType] | None, optional): explicit schema for the table.
            When provided, allows scan_parquet to succeed even when the S3 path
            contains no files (returns an empty DataFrame with this schema).
            Defaults to None.

    Raises:
        ValueError: if the partitioning_scheme is not a valid PartitioningScheme.
    """
    if not isinstance(partitioning_scheme, PartitioningScheme):
        raise ValueError(f"Invalid partitioning scheme, got {partitioning_scheme}")

    self.name = name
    self.uri = uri
    self.partitioning = partitioning
    self.partitioning_scheme = partitioning_scheme
    self.schema = schema

    self.table_metadata = TableMetadata(
        table_type="PARQUET",
        description=description,
        docs_args={"filters": docs_filters, "columns": docs_columns},
        roapi_opts=roapi_opts,
        **(table_metadata_args or {}),
    )

    self.parquet_file_name = parquet_file_name

build_file_fragment(filters)

Returns a file path from the base table URI with the given filters. This will raise an error if the filter does not specify all partitions.

This is currently used to generate the file path used by ROAPI to infer schemas.

Source code in src/datarepo/core/tables/parquet_table.py
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def build_file_fragment(self, filters: list[Filter]) -> str:
    """
    Returns a file path from the base table URI with the given filters.
    This will raise an error if the filter does not specify all partitions.

    This is currently used to generate the file path used by ROAPI to infer schemas.
    """
    uri_with_prefix, partitions, _, _ = self._build_uri_from_filters(
        normalize_filters(filters), include_base_uri=False
    )
    if len(partitions) > 0:
        partition_names = [partition.column for partition in partitions]
        raise ValueError(
            f"Not enough partitions specified, missing: {partition_names}"
        )

    return path.join(uri_with_prefix, self.parquet_file_name)

get_schema()

Generates the schema of the table, including partitions and columns.

Returns:

Name Type Description
TableSchema TableSchema

table schema containing partitions and columns.

Source code in src/datarepo/core/tables/parquet_table.py
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
def get_schema(self) -> TableSchema:
    """Generates the schema of the table, including partitions and columns.

    Returns:
        TableSchema: table schema containing partitions and columns.
    """
    partitions = [
        TablePartition(
            column_name=filter.column,
            operator=filter.operator,
            type_annotation=type(filter.value).__name__,
            value=filter.value,
        )
        for filter in self.table_metadata.docs_args.get("filters", [])
    ]

    # Pop the selected columns so that we still pull the full schema below
    docs_args = {**self.table_metadata.docs_args}
    docs_args.pop("columns")

    columns = []
    if docs_args or not partitions:
        table: NlkDataFrame = self(**docs_args)
        columns = [
            TableColumn(
                name=key,
                type=type.__str__(),
                readonly=False,
                filter_only=False,
                has_stats=False,
            )
            for key, type in table.schema.items()
        ]

    return TableSchema(partitions=partitions, columns=columns)

PartitioningScheme

Bases: Enum

Defines the partitioning scheme for the table.

DIRECTORY - e.g. s3://bucket/5956/2024-03-24 HIVE - e.g. s3://bucket/implant_id=5956/date=2024-03-24

Source code in src/datarepo/core/tables/util.py
31
32
33
34
35
36
37
38
39
40
class PartitioningScheme(Enum):
    """
    Defines the partitioning scheme for the table.

    DIRECTORY - e.g. s3://bucket/5956/2024-03-24
    HIVE - e.g. s3://bucket/implant_id=5956/date=2024-03-24
    """

    DIRECTORY = 1
    HIVE = 2

TableMetadata dataclass

Information about a table used for documentation generation.

Source code in src/datarepo/core/tables/metadata.py
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
@dataclass
class TableMetadata:
    """
    Information about a table used for documentation generation.
    """

    table_type: str
    description: str
    docs_args: Dict[str, Any]
    latency_info: str | None = None
    example_notebook: str | None = None
    data_input: str | None = None
    is_deprecated: bool = False
    roapi_opts: RoapiOptions | None = None

TableProtocol

Bases: Protocol

Source code in src/datarepo/core/tables/metadata.py
47
48
49
50
51
52
53
54
55
56
57
class TableProtocol(Protocol):
    # Properties used to generate the web catalog & roapi config
    table_metadata: TableMetadata

    def __call__(self, **kwargs: Dict[str, Any]) -> NlkDataFrame: ...

    def get_schema(self) -> TableSchema:
        """
        Returns the schema of the table, used to generate the web catalog.
        """
        ...

get_schema()

Returns the schema of the table, used to generate the web catalog.

Source code in src/datarepo/core/tables/metadata.py
53
54
55
56
57
def get_schema(self) -> TableSchema:
    """
    Returns the schema of the table, used to generate the web catalog.
    """
    ...

TableSchema dataclass

TableSchema represents the schema of a table, including partitions and columns.

Source code in src/datarepo/core/tables/metadata.py
39
40
41
42
43
44
@dataclass
class TableSchema:
    """TableSchema represents the schema of a table, including partitions and columns."""

    partitions: list[TablePartition]
    columns: list[TableColumn]

make_clickhouse_config(table_name=None, host=None, port=None, username=None, password=None, database=None)

Create a ClickHouse config from environment variables or explicit values.

Parameters:

Name Type Description Default
table_name str | None

The actual table name in ClickHouse. If not provided, queries will use the ClickHouseTable.name as the table name.

None

Environment variables (used if explicit values not provided): CLICKHOUSE_HOST: ClickHouse server hostname (default: localhost) CLICKHOUSE_PORT: ClickHouse HTTP port (default: 8123). Use 8443 with secure=True for TLS. CLICKHOUSE_USER: ClickHouse username (default: default) CLICKHOUSE_PASSWORD: ClickHouse password CLICKHOUSE_DATABASE: ClickHouse database (default: default)

Source code in src/datarepo/core/tables/clickhouse_table.py
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def make_clickhouse_config(
    table_name: str | None = None,
    host: str | None = None,
    port: int | None = None,
    username: str | None = None,
    password: str | None = None,
    database: str | None = None,
) -> ClickHouseTableConfig:
    """Create a ClickHouse config from environment variables or explicit values.

    Args:
        table_name: The actual table name in ClickHouse. If not provided,
            queries will use the ClickHouseTable.name as the table name.

    Environment variables (used if explicit values not provided):
        CLICKHOUSE_HOST: ClickHouse server hostname (default: localhost)
        CLICKHOUSE_PORT: ClickHouse HTTP port (default: 8123). Use 8443 with
            secure=True for TLS.
        CLICKHOUSE_USER: ClickHouse username (default: default)
        CLICKHOUSE_PASSWORD: ClickHouse password
        CLICKHOUSE_DATABASE: ClickHouse database (default: default)
    """
    return ClickHouseTableConfig(
        host=host or os.environ.get("CLICKHOUSE_HOST", "localhost"),
        port=port or int(os.environ.get("CLICKHOUSE_PORT", "8123")),
        username=username or os.environ.get("CLICKHOUSE_USER", "default"),
        password=password or os.environ.get("CLICKHOUSE_PASSWORD", ""),
        database=database or os.environ.get("CLICKHOUSE_DATABASE", "default"),
        table_name=table_name,
    )

table(*args, **kwargs)

Decorator to define a table using a function.

Example uage
@table(description="This is a sample table.")
def my_table_function(param1, param2):
    # Function logic to create a table
    return NlkDataFrame(...)

Returns:

Type Description
Callable[[U], U] | Callable[[Any], Callable[[U], U]]

Callable[[U], U] | Callable[[Any], Callable[[U], U]]: A decorator that wraps a function to create a table.

Source code in src/datarepo/core/tables/decorator.py
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def table(*args, **kwargs) -> Callable[[U], U] | Callable[[Any], Callable[[U], U]]:
    """Decorator to define a table using a function.

    Example uage:
        ``` py
        @table(description="This is a sample table.")
        def my_table_function(param1, param2):
            # Function logic to create a table
            return NlkDataFrame(...)
        ```

    Returns:
        Callable[[U], U] | Callable[[Any], Callable[[U], U]]: A decorator that wraps a function to create a table.
    """

    def wrapper(func):
        return FunctionTable(
            table_metadata=TableMetadata(
                table_type="FUNCTION",
                description=func.__doc__.strip() if func.__doc__ else "",
                docs_args=kwargs.get("docs_args", {}),
                latency_info=kwargs.get("latency_info"),
                example_notebook=kwargs.get("example_notebook"),
                data_input=kwargs.get("data_input"),
                is_deprecated=kwargs.get("is_deprecated", False),
            ),
            func=func,
        )

    if len(args) == 0:
        return wrapper
    else:
        return wrapper(args[0])